diff --git a/docs/tutorials/posttraining/lora.md b/docs/tutorials/posttraining/lora.md index c309078593..1b11583822 100644 --- a/docs/tutorials/posttraining/lora.md +++ b/docs/tutorials/posttraining/lora.md @@ -190,7 +190,7 @@ After completing the fine-tuning process, your LoRA weights are stored in MaxTex ```sh python3 -m maxtext.checkpoint_conversion.to_huggingface \ model_name="${MODEL_NAME?}" \ - lora.lora_restore_path="${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/checkpoints//model_params" \ + lora.lora_restore_path="${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/checkpoints//items" \ base_output_directory="${BASE_OUTPUT_DIRECTORY?}/hf_lora_adapter" ``` diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 854c1b3968..38c3c37e8d 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -150,7 +150,7 @@ def _raise_on_weight_mismatch(want, have, config=None): ) -def _linen_items_to_nnx(restored_linen, abstract_nnx_state): +def linen_items_to_nnx(restored_linen, abstract_nnx_state): """Reshapes a restored Linen-layout `items` dict into an NNX state. The inverse of `to_checkpoint_dict`, over the same `split_for_checkpoint` partition. The Linen @@ -184,7 +184,7 @@ def _load_linen_checkpoint_into_nnx( """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume). Restores a Linen-shape target that includes `nnx_aux`, then reshapes back via - `_linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when + `linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when present, else keep their fresh init value. A genuinely-missing weight raises. """ max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {path}") @@ -222,7 +222,7 @@ def _restored_linen_to_nnx(restored_linen, abstract_nnx_state, config=None): itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout. """ _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored_linen), config=config) - return _linen_items_to_nnx(restored_linen, abstract_nnx_state) + return linen_items_to_nnx(restored_linen, abstract_nnx_state) def _abstract_params(abstract_unboxed_pre_state): @@ -701,6 +701,25 @@ def setup_checkpoint_logger(config) -> Any | None: # pytype: disable=attribute- return orbax_cloud_logger +def _nnx_native_wrapper_key(ckptr, path): + """Returns the wrapper key an NNX-native checkpoint nests its weights under, if any. + + Post-training wrote DPO and RL checkpoints through `TunixMaxTextAdapter`, whose only child + module is `base`, so their weights sit one level deeper. SFT and the training engine save the + model directly and do not. + + Args: + ckptr: Checkpointer used to read the checkpoint metadata. + path: Path to the checkpoint. + + Returns: + "base" if the checkpoint nests its weights under it, else None. + """ + tree = ckptr.metadata(epath.Path(path)).item_metadata + tree = getattr(tree, "tree", tree) + return "base" if tree is not None and "base" in tree else None + + def load_params_from_path( load_parameters_from_path, abstract_unboxed_params, @@ -718,15 +737,16 @@ def load_params_from_path( is_nnx = isinstance(abstract_unboxed_params, nnx.State) want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params - # Determine the restore key based on the leaf directory name to support native and custom SFT - restore_key = os.path.basename(load_parameters_from_path) - if restore_key not in ("model_params", "model"): - restore_key = "params" - - if restore_key in ("model_params", "model"): - params_collection = want - else: - params_collection = {"params": want} if is_nnx else want + # A path ending in `model_params` (or `model`) holds an NNX state written straight from + # `nnx.state(model)`, rather than the Linen on-disk layout. Post-training wrote this before it + # moved to MaxText's layout, and the training engine still does. The tree is the whole + # checkpoint, not an item inside one. + is_nnx_native = os.path.basename(load_parameters_from_path) in ("model_params", "model") + if is_nnx_native and not is_nnx: + raise ValueError( + f"'{load_parameters_from_path}' holds an NNX state, which only restores into an NNX params " + "state. Point load_parameters_path at a checkpoint saved in the Linen on-disk layout instead." + ) # *_concurrent_gb should be set for large models, the default is 96. max_logging.log(f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}") @@ -743,18 +763,31 @@ def load_params_from_path( # Rather than pass the entire abstract state, which could unnecessarily restore opt_state and such and waste # memory, we instead specify here that we are just restoring the params field of the checkpoint # (which itself may be a dictionary containing a key named 'params' or 'model'). - restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection) - restored = ckptr.restore( - epath.Path(load_parameters_from_path), - item={restore_key: params_collection}, - transforms={}, - restore_args={restore_key: restore_args}, - ) - restored_collection = restored[restore_key] - - if restore_key in ("model_params", "model"): - restored_weights = restored_collection + if is_nnx_native: + wrapper_key = _nnx_native_wrapper_key(ckptr, load_parameters_from_path) + # Restore into the NNX state itself rather than a pure dict. Flax registers a Variable as a + # pytree holding its array under `value`, matching what `nnx.state(model)` wrote, so save and + # restore agree without reshaping either side by hand. + item = {wrapper_key: abstract_unboxed_params} if wrapper_key else abstract_unboxed_params + restored = ckptr.restore( + epath.Path(load_parameters_from_path), + item=item, + transforms={}, + restore_args=ocp.checkpoint_utils.construct_restore_args(item), + ) + # No `params` collection in this layout, so the weights are the whole collection. + restored_weights = nnx.to_pure_dict(restored[wrapper_key] if wrapper_key else restored) + restored_collection = restored_weights else: + params_collection = {"params": want} if is_nnx else want + restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection) + restored = ckptr.restore( + epath.Path(load_parameters_from_path), + item={"params": params_collection}, + transforms={}, + restore_args={"params": restore_args}, + ) + restored_collection = restored["params"] restored_weights = restored_collection["params"] if is_nnx else restored_collection # `transforms={}` lets Orbax return an unmaterialized leaf for a weight the checkpoint lacks, @@ -777,26 +810,64 @@ def save_params_to_path(checkpoint_dir, params, use_ocdbt=True, use_zarr3=True): print(f"Quantized params checkpoint saved at: {checkpoint_dir}") -def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]: - """Loads custom metadata from an Orbax checkpoint. +def checkpoint_custom_metadata(config) -> dict[str, Any]: + """Returns the metadata a checkpoint stores alongside its state. + + `verify_and_sync_scan_layers` and `lora_utils.sync_lora_metadata` read these back on load. + Post-training saves through its own manager, so it calls this too. Args: - checkpoint_dir_path: Path to the checkpoint directory. + config: The run's config, or None. Returns: - A dictionary containing custom metadata, or an empty dictionary if none is - present or loading fails. + The metadata dict, empty if there is no config to read it from. + """ + custom_metadata = {} + if config: + if hasattr(config, "scan_layers"): + custom_metadata["scan_layers"] = config.scan_layers + if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0: + custom_metadata["lora"] = config.lora.model_dump() + return custom_metadata + + +def _custom_metadata_at(checkpoint_dir: epath.Path) -> dict[str, Any]: + """Reads the custom metadata stored at exactly this directory. + + Args: + checkpoint_dir: Directory to read. + + Returns: + The metadata dict, empty if there is none or the read fails. """ - checkpoint_dir = epath.Path(checkpoint_dir_path) try: - ckptr = ocp.StandardCheckpointer() - metadata = ckptr.metadata(checkpoint_dir) + metadata = ocp.StandardCheckpointer().metadata(checkpoint_dir) return metadata.custom_metadata or {} except Exception as e: # pylint: disable=broad-except max_logging.log(f"Warning: Failed to load checkpoint metadata: {e}") return {} +def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]: + """Loads custom metadata from an Orbax checkpoint. + + The metadata belongs to the step, so it sits at `/` and not at `/items/`. Callers + pass `load_parameters_path`, which points at the item, so fall back to the parent directory. + + Args: + checkpoint_dir_path: Path to the checkpoint directory, item level or step level. + + Returns: + A dictionary containing custom metadata, or an empty dictionary if none is + present or loading fails. + """ + checkpoint_dir = epath.Path(checkpoint_dir_path) + metadata = _custom_metadata_at(checkpoint_dir) + if not metadata and checkpoint_dir.parent != checkpoint_dir: + metadata = _custom_metadata_at(checkpoint_dir.parent) + return metadata + + def _uses_local_checkpoint_period(config): return config.enable_emergency_checkpoint or config.enable_multi_tier_checkpointing @@ -1030,12 +1101,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= item=grain_iters_to_save ) # pyrefly: ignore[bad-assignment] - custom_metadata = {} - if config: - if hasattr(config, "scan_layers"): - custom_metadata["scan_layers"] = config.scan_layers - if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0: - custom_metadata["lora"] = config.lora.model_dump() + custom_metadata = checkpoint_custom_metadata(config) match (checkpoint_manager, config, data_iterator): case (checkpoint_manager, _, _) if isinstance( diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index 45dc386576..86b6c68654 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -120,6 +120,22 @@ def _as_chain_index(key): return None +def opt_state_to_linen(opt_state): + """Reshapes an optimizer state into the Linen on-disk layout. + + `to_checkpoint_dict` already does this for the state it is given. This is for a caller that + has unwrapped something the conversion could not see through -- `optax.inject_hyperparams` + hides mu and nu behind its own keys -- and needs the inner state converted on its own. + + Args: + opt_state: The optimizer state to reshape. + + Returns: + The same state in the Linen layout. + """ + return _opt_state_to_linen(opt_state) + + def _opt_state_to_linen(opt_state): """Reshapes the NNX opt_state to Linen's on-disk layout. diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index d8efa17359..8abb8631d0 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -55,6 +55,8 @@ save_checkpoint_on_start: true save_checkpoint_on_completion: true async_checkpointing: true checkpoint_period: 10_000 +# Post-training only: skip installing a checkpoint manager, so nothing is written back. +post_train_skip_checkpointing: false max_num_checkpoints_to_keep: None enable_continuous_checkpointing: false # enables one replica to read the ckpt then broadcast to the rest diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index db9a68d1cd..c942832979 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -348,6 +348,17 @@ class Checkpointing(BaseModel): load_checkpoint_only_once: bool = Field(False, description="If True, deep copy the reference model to the actor model.") async_checkpointing: bool = Field(True, description="If True, uses an asynchronous checkpointer for performance.") checkpoint_period: int = Field(10_000, description="The frequency (in steps) at which to save checkpoints.") + post_train_skip_checkpointing: bool = Field( + False, + description=( + "If True, post-training trainers do not install a checkpoint manager and write nothing." + " Base weights still load through load_parameters_path. Exists because enable_checkpointing" + " cannot express this: it is validated as required whenever load_parameters_path is set," + " so it cannot be turned off by a run that has to load base weights. Setting" + " checkpoint_period to 0 is not an alternative either, since step % checkpoint_period is" + " evaluated on the shared path and raises ZeroDivisionError." + ), + ) max_num_checkpoints_to_keep: int | None = Field(None, description="Maximum number of checkpoints to keep.") enable_single_replica_ckpt_restoring: bool = Field( False, description="One replica reads and broadcasts the checkpoint." diff --git a/src/maxtext/eval/README.md b/src/maxtext/eval/README.md index 7970ce63b2..5eb0a9d889 100644 --- a/src/maxtext/eval/README.md +++ b/src/maxtext/eval/README.md @@ -186,7 +186,7 @@ Example (Qwen3-30B-A3B, v6e-8): STEP=244 MODEL=qwen3-30b-a3b HF_PATH=Qwen/Qwen3-30B-A3B -CHECKPOINT=gs:///run/checkpoints/actor/${STEP}/model_params +CHECKPOINT=gs:///run/checkpoints/actor/${STEP}/items OUTPUT=gs:///eval/ python -m maxtext.eval.runner.run \ diff --git a/src/maxtext/examples/rl_llama3_demo.ipynb b/src/maxtext/examples/rl_llama3_demo.ipynb index 68a9ff95eb..4615c1207b 100644 --- a/src/maxtext/examples/rl_llama3_demo.ipynb +++ b/src/maxtext/examples/rl_llama3_demo.ipynb @@ -300,13 +300,13 @@ "# Define the output directory for the Hugging Face checkpoint\n", "hf_output_directory = epath.Path(BASE_OUTPUT_DIRECTORY) / \"hf_checkpoint\"\n", "\n", - "# Find the latest MaxText checkpoint\n", - "checkpoint_dir = epath.Path(config.checkpoint_dir) / 'actor'\n", + "# Find the latest MaxText checkpoint. RL checkpoints the actor under its own subdirectory.\n", + "checkpoint_dir = epath.Path(config.checkpoint_dir) / \"actor\"\n", "step_dirs = [d.name for d in checkpoint_dir.iterdir() if d.name.isdigit() and d.is_dir()]\n", "if not step_dirs:\n", " raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n", "latest_step = max(step_dirs, key=int)\n", - "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n", + "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"items\"\n", "\n", "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n", "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n", diff --git a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb index 216769ac48..2001622fe7 100644 --- a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb +++ b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb @@ -341,7 +341,7 @@ "if not step_dirs:\n", " raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n", "latest_step = max(step_dirs, key=int)\n", - "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n", + "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"items\"\n", "\n", "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n", "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n", diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py index 03514ce9cf..12be4a0b8e 100644 --- a/src/maxtext/inference/vllm_decode.py +++ b/src/maxtext/inference/vllm_decode.py @@ -110,6 +110,24 @@ def decode_with_vllm(config: Config) -> None: }, } + # vllm_additional_config exists in the config but was never passed through, so keys tpu_inference + # reads straight off additional_config -- skip_quantization among them -- had no way to be set. + # Merged at the top level so it cannot clobber the maxtext_config and sharding blocks built above. + if config.vllm_additional_config: + for key, value in dict(config.vllm_additional_config).items(): + if key in vllm_args["additional_config"]: + max_logging.log(f"Ignoring vllm_additional_config[{key!r}]: it would overwrite a generated block.") + else: + vllm_args["additional_config"][key] = value + + # The worker rebuilds its own MaxTextConfig from inference/vllm.yml, which pins attention to + # vllm_rpa, and nothing carried the caller's choice across. Some decoder blocks reject that + # outright -- deepseek4 raises "DeepSeek4 decoder block currently only supports dot_product + # attention" -- and there was no way to say otherwise. Only forwarded when explicitly set, so + # models that are happy on vllm_rpa keep it. + if config.attention != "autoselected": + vllm_args["additional_config"]["maxtext_config"]["attention"] = config.attention + if config.load_parameters_path: vllm_args["additional_config"]["maxtext_config"]["load_parameters_path"] = config.load_parameters_path else: diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index 56d9ef0498..89a4954202 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -33,10 +33,15 @@ def _get_pad_id(tokenizer): - if tokenizer.pad_token_id is not None: + """Returns the pad token id from the tokenizer, or -1 if not found.""" + if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None: pad_id = tokenizer.pad_token_id - elif tokenizer.unk_token_id is not None: + elif hasattr(tokenizer, "unk_token_id") and tokenizer.unk_token_id is not None: pad_id = tokenizer.unk_token_id + elif hasattr(tokenizer, "pad_id") and getattr(tokenizer, "pad_id", None) is not None: + pad_id = tokenizer.pad_id + elif hasattr(tokenizer, "unk_id") and getattr(tokenizer, "unk_id", None) is not None: + pad_id = tokenizer.unk_id else: pad_id = -1 return pad_id diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 8517fc4253..08b7327290 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -316,13 +316,25 @@ def apply_chat_template(example, tokenizer_model, data_column_name): def tokenization(example, hf_tokenizer, truncation, max_length, column_names): """Tokenize a HuggingFace dataset""" + is_hf = callable(hf_tokenizer) for column_name in column_names: if isinstance(example[column_name], list): - example[column_name] = [ - hf_tokenizer(x, truncation=truncation, max_length=max_length)["input_ids"] for x in example[column_name] - ] + if is_hf: + example[column_name] = [ + hf_tokenizer(x, truncation=truncation, max_length=max_length)["input_ids"] for x in example[column_name] + ] + else: + example[column_name] = [ + hf_tokenizer.encode(x)[:max_length] if truncation else hf_tokenizer.encode(x) for x in example[column_name] + ] elif isinstance(example[column_name], str): - example[column_name] = hf_tokenizer(example[column_name], truncation=truncation, max_length=max_length)["input_ids"] + if is_hf: + example[column_name] = hf_tokenizer(example[column_name], truncation=truncation, max_length=max_length)[ + "input_ids" + ] + else: + ids = hf_tokenizer.encode(example[column_name]) + example[column_name] = ids[:max_length] if truncation else ids return example diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py index 7c79eebfcb..4ed2f24b52 100644 --- a/src/maxtext/input_pipeline/synthetic_data_processing.py +++ b/src/maxtext/input_pipeline/synthetic_data_processing.py @@ -80,7 +80,9 @@ def __init__(self, config, mesh): self.config = config data_pspec_shardings = sharding.get_input_data_sharding(config, mesh) self.data_generator = jax.jit( - SyntheticDataIterator.raw_generate_synthetic_data, out_shardings=data_pspec_shardings, static_argnums=0 + SyntheticDataIterator.raw_generate_synthetic_data, + out_shardings=data_pspec_shardings, + static_argnums=1, ) tokens = jax.random.randint( @@ -104,6 +106,7 @@ def __init__(self, config, mesh): else: segmentation = jnp.ones((config.global_batch_size_to_load, config.max_target_length), dtype=jnp.int32) self.data = (tokens, batch_positions, segmentation) + self.num_diloco_replicas = config.num_diloco_replicas if config.enable_diloco else 0 def reset(self): pass # Synthetic data is stateless; nothing to reset. @@ -113,10 +116,10 @@ def __iter__(self): def __next__(self): with self.mesh: - return self.data_generator(self.config, self.data) # pylint: disable=not-callable + return self.data_generator(self.data, self.num_diloco_replicas) # pylint: disable=not-callable @staticmethod - def raw_generate_synthetic_data(config: pyconfig.HyperParameters, data): + def raw_generate_synthetic_data(data, num_diloco_replicas): """Generates a single batch of synthetic data""" tokens, positions, segmentation = data @@ -127,8 +130,8 @@ def raw_generate_synthetic_data(config: pyconfig.HyperParameters, data): output["targets"] = tokens[:, 1:] output["targets_position"] = positions[:, 1:] output["targets_segmentation"] = segmentation - if config.enable_diloco: - output = reshape_first_axis_with_diloco(config.num_diloco_replicas, output) + if num_diloco_replicas: + output = reshape_first_axis_with_diloco(num_diloco_replicas, output) return output diff --git a/src/maxtext/input_pipeline/tokenizer.py b/src/maxtext/input_pipeline/tokenizer.py index 10a528205d..502aab6d83 100644 --- a/src/maxtext/input_pipeline/tokenizer.py +++ b/src/maxtext/input_pipeline/tokenizer.py @@ -23,7 +23,53 @@ from sentencepiece import SentencePieceProcessor -class TikTokenTokenizer: +class ChatTemplateMixin: + """Mixin to provide Jinja2 chat template rendering for native tokenizers.""" + + def apply_chat_template(self, conversation, chat_template=None, add_generation_prompt=False, tokenize=True, **kwargs): + """Applies a Jinja2 chat template to a conversation.""" + if chat_template is None: + chat_template = getattr(self, "chat_template", None) + if chat_template is None: + raise ValueError("Cannot apply chat template because no chat template was provided or set.") + + import jinja2 # pylint: disable=import-outside-toplevel + + env = jinja2.Environment(autoescape=False) + + def raise_exception(message): + raise jinja2.exceptions.TemplateError(message) + + env.globals["raise_exception"] = raise_exception + + template = env.from_string(chat_template) + + bos_token = "" + eos_token = "" + if hasattr(self, "_tokenizer_model"): # SentencePiece + if self.bos_id is not None and self.bos_id >= 0: + bos_token = self._tokenizer_model.IdToPiece(self.bos_id) + if self.eos_id is not None and self.eos_id >= 0: + eos_token = self._tokenizer_model.IdToPiece(self.eos_id) + elif hasattr(self, "model"): # TikToken + if self.bos_id is not None and self.bos_id >= 0: + bos_token = self.decode([self.bos_id]) + if self.eos_id is not None and self.eos_id >= 0: + eos_token = self.decode([self.eos_id]) + + rendered = template.render( + messages=conversation, + add_generation_prompt=add_generation_prompt, + bos_token=bos_token, + eos_token=eos_token, + **kwargs, + ) + if tokenize: + return self.encode(rendered) + return rendered + + +class TikTokenTokenizer(ChatTemplateMixin): """ Tokenizing and encoding/decoding text using the Tiktoken tokenizer. """ @@ -180,7 +226,7 @@ def _split_whitespaces_or_nonwhitespaces(s: str, max_consecutive_slice_len: int) yield s[slice_start:] -class SentencePieceTokenizer: +class SentencePieceTokenizer(ChatTemplateMixin): """ Tokenizing and encoding/decoding text using the native sentencepiece library. Supports both local and GCS (gs://) model paths. @@ -245,6 +291,9 @@ def encode(self, s: str) -> list[int]: def decode(self, t: Sequence[int]) -> str: return self.tokenizer.decode(t) + def apply_chat_template(self, *args, **kwargs): + return self.tokenizer.apply_chat_template(*args, **kwargs) + def build_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token): """Loads the tokenizer at `tokenizer_path`""" diff --git a/src/maxtext/integration/tunix/weight_mapping/__init__.py b/src/maxtext/integration/tunix/weight_mapping/__init__.py index 39ab12ff8f..22577f1e93 100644 --- a/src/maxtext/integration/tunix/weight_mapping/__init__.py +++ b/src/maxtext/integration/tunix/weight_mapping/__init__.py @@ -30,7 +30,11 @@ class StandaloneVllmWeightMapping: """Mapping MaxText model weights to vLLM's model weights.""" def __getattr__(self, name): - if name.startswith("llama3.1"): + # "llama3", not "llama3.1": Llama 3 and 3.1 share this mapping -- the constant is + # LLAMA3_VLLM_MAPPING, not LLAMA31_ -- and their MaxText configs are identical apart from a + # comment. Matching only the 3.1 prefix left llama3-8b raising "vLLM weight mapping not found" + # from RL, after SFT and DPO had already run on it. + if name.startswith("llama3"): return LLAMA3_VLLM_MAPPING elif name.startswith("qwen2"): return QWEN2_VLLM_MAPPING diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index bd3b8a59f0..fc789ae938 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -48,6 +48,71 @@ ) from maxtext.integration.vllm.torchax_converter.gemma4_moe import Gemma4MaxTextToVLLMConverter +import vllm.config.vllm +import vllm.config.utils as vllm_config_utils + +# Monkey-patch VLLM's is_init_field to gracefully handle dynamically added +# fields (like sharding_config injected by tpu-inference) which would otherwise +# cause ValueError/AssertionError during vllm_config.with_hf_config replacement. +_orig_is_init_field = vllm_config_utils.is_init_field + + +def _engine_args_accepts(name): + """Whether the installed vLLM's EngineArgs still takes `name`. + + Fields come and go between vLLM releases, and passing one that has been + removed is a TypeError rather than something vLLM ignores. + """ + # pylint: disable=import-outside-toplevel + import dataclasses + + from vllm.engine.arg_utils import EngineArgs + + return any(field.name == name for field in dataclasses.fields(EngineArgs)) + + +def _patched_is_init_field(cls, name): + try: + return _orig_is_init_field(cls, name) + except ValueError: + return False + + +vllm_config_utils.is_init_field = _patched_is_init_field + + +_orig_with_hf_config = vllm.config.vllm.VllmConfig.with_hf_config + + +def _patched_with_hf_config(self, *args, **kwargs): + """ + Restore the original data_parallel_size which tpu_platform mutated, + so that the new VllmConfig passes the device_indexes length assertion. + tpu_inference deletes sharding_config before calling with_hf_config, + so we must reverse-engineer the data_parallel_size from device_indexes. + """ + if self.additional_config and "sharding" in self.additional_config: + sharding_strategy = self.additional_config["sharding"].get("sharding_strategy", {}) + device_indexes = sharding_strategy.get("device_indexes") + if device_indexes is not None: + pc = self.parallel_config + tp = sharding_strategy.get("tensor_parallelism") or pc.tensor_parallel_size + ep = sharding_strategy.get("expert_parallelism", 1) + sp = sharding_strategy.get("sequence_parallelism", 1) + attn_dp = sharding_strategy.get("attention_data_parallelism", 1) + attn_dp_ep = sharding_strategy.get("attention_data_expert_parallelism", 1) + dcp = pc.decode_context_parallel_size + + other_parallelism = tp * ep * sp * attn_dp * attn_dp_ep * dcp + if other_parallelism > 0: + self.parallel_config.data_parallel_size = len(device_indexes) // other_parallelism + + return _orig_with_hf_config(self, *args, **kwargs) + + +vllm.config.vllm.VllmConfig.with_hf_config = _patched_with_hf_config + + # Sentinel distinguishing "this model has no entry" from "this model has an # entry whose value is None", which means direct-sync-only. _NO_RULE_TABLE = object() @@ -571,6 +636,10 @@ def __init__( "max_logprobs": 1, "logprobs_mode": rollout_config.rollout_vllm_logprobs_mode, } + if _engine_args_accepts("swap_space"): + engine_kwargs["swap_space"] = getattr( + rollout_config, "rollout_vllm_swap_space_size_gb", maxtext_config.swap_space_vllm_gb + ) # Merge additional kwargs like dtype and hf_overrides provided by train_rl.py if hasattr(rollout_config, "rollout_vllm_kwargs") and rollout_config.rollout_vllm_kwargs: diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index 67e1f589ca..23d1bfa98f 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -166,6 +166,33 @@ def skip_update(): return optax.GradientTransformationExtraArgs(init_fn, update_fn) +def add_gradient_clipping(tx, clipping_threshold): + """Clips gradients by global norm ahead of `tx`, keeping `tx`'s optimizer state shape. + + `optax.chain(clip_by_global_norm(...), tx)` would nest tx's state under an extra chain level + even though the clip is stateless. Pre-training clips raw gradients in its train step instead, + so its checkpointed state has no such level and cannot restore one that does. + + Args: + tx: The optimizer to clip gradients for. + clipping_threshold: Global norm to clip to. + + Returns: + A transformation applying the same updates as the chained form, with tx's state tree. + """ + clip = optax.clip_by_global_norm(clipping_threshold) + inner = optax.with_extra_args_support(tx) + + def init_fn(params): + return inner.init(params) + + def update_fn(updates, state, params=None, **extra_args): + updates, _ = clip.update(updates, optax.EmptyState(), params) + return inner.update(updates, state, params, **extra_args) + + return optax.GradientTransformationExtraArgs(init_fn, update_fn) + + def get_optimizer(config, learning_rate_schedule, model=None): """Create optimizer.""" if config.opt_type == "adamw": diff --git a/src/maxtext/trainers/post_train/checkpointing.py b/src/maxtext/trainers/post_train/checkpointing.py new file mode 100644 index 0000000000..d40a3f6b1d --- /dev/null +++ b/src/maxtext/trainers/post_train/checkpointing.py @@ -0,0 +1,452 @@ +# Copyright 2023-2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpointing for the Tunix post-training trainers, in MaxText's on-disk layout. + +Lives here rather than in `maxtext.common.checkpointing` because the manager subclasses +Tunix's, and `maxtext.common.checkpointing` is imported by pre-training and inference, which +run without Tunix installed. +""" + +import os +from typing import Any, Sequence + +from flax import nnx +import jax +import jax.numpy as jnp +import orbax.checkpoint as ocp +from tunix.sft import checkpoint_manager as tunix_checkpoint_manager + +from maxtext.common import checkpointing +from maxtext.common import train_state_nnx +from maxtext.utils import max_logging + +# The item MaxText stores a checkpoint under, matching create_orbax_checkpoint_manager. +_ITEM_NAME = "items" + +# What Tunix stored a checkpoint under, kept registered so old checkpoints still restore. +_TUNIX_ITEM_NAMES = ("model_params", "optimizer_state") + +# The Tunix adapter's only child module. DPO and RL train through the adapter, so its state +# carries this extra level and a MaxText checkpoint must not. +_ADAPTER_CHILD = "base" + + +def unwrap_model(model: nnx.Module) -> nnx.Module: + """Returns the MaxText model, unwrapping the Tunix adapter if there is one. + + Matches on the child module rather than on `TunixMaxTextAdapter` itself, so any equivalent + wrapper unwraps the same way. + + Args: + model: The model a Tunix trainer holds. + + Returns: + The wrapped model, or `model` itself if it is not wrapped. + """ + base = getattr(model, _ADAPTER_CHILD, None) + if isinstance(base, nnx.Module): + return unwrap_model(base) + return model + + +def _drop_adapter_level(tree): + """Removes the adapter level wherever it wraps a weight-shaped subtree. + + The optimizer is built over the adapter, so its accumulators (mu, nu, acc_grads) are keyed by + the adapter's graph and carry the level even though the weights they shadow do not. + + Args: + tree: A pure dict, typically the optimizer state. + + Returns: + The same tree with every `{"base": subtree}` replaced by `subtree`. + """ + if isinstance(tree, dict): + if set(tree) == {_ADAPTER_CHILD}: + return _drop_adapter_level(tree[_ADAPTER_CHILD]) + return {k: _drop_adapter_level(v) for k, v in tree.items()} + if isinstance(tree, list): + return [_drop_adapter_level(v) for v in tree] + return tree + + +def _add_adapter_level(tree, guide): + """Inverse of `_drop_adapter_level`. + + Args: + tree: A pure dict with the adapter level removed. + guide: The same tree before removal, giving the positions to restore. + + Returns: + `tree` with the adapter level put back wherever `guide` carries it. + """ + if isinstance(guide, dict) and set(guide) == {_ADAPTER_CHILD}: + return {_ADAPTER_CHILD: _add_adapter_level(tree, guide[_ADAPTER_CHILD])} + if isinstance(guide, dict) and isinstance(tree, dict): + return {k: (_add_adapter_level(v, guide[k]) if k in guide else v) for k, v in tree.items()} + if isinstance(guide, list) and isinstance(tree, list) and len(guide) == len(tree): + return [_add_adapter_level(t, g) for t, g in zip(tree, guide)] + return tree + + +def _drop_inject_hyperparams(opt_state): + """Strips the `optax.inject_hyperparams` state wrapper if present. + + RL and distillation trainers wrap their optimizer in `inject_hyperparams`. To produce + a checkpoint fully compatible with pre-training, we strip the outer shell and only save + the inner state. + + Args: + opt_state: The optimizer state dict to inspect. + + Returns: + The inner state if `inject_hyperparams` was found, otherwise `opt_state`. + """ + if isinstance(opt_state, dict) and {"count", "hyperparams", "hyperparams_states", "inner_state"}.issubset( + opt_state.keys() + ): + return opt_state["inner_state"] + return opt_state + + +def _add_inject_hyperparams(restored_opt_state, guide, step): + """Restores the `optax.inject_hyperparams` wrapper state. + + Args: + restored_opt_state: The bare inner state loaded from disk. + guide: The currently initialized optimizer state dict, used as a structural guide. + step: The global step to restore into the wrapper's count. + + Returns: + The reconstructed full state dict. + """ + if isinstance(guide, dict) and {"count", "hyperparams", "hyperparams_states", "inner_state"}.issubset(guide.keys()): + + new_state = dict(guide) + new_state["inner_state"] = restored_opt_state + new_state["count"] = jnp.array(step, dtype=guide["count"].dtype) + return new_state + return restored_opt_state + + +class MaxTextLayoutCheckpointManager(tunix_checkpoint_manager.CheckpointManager): + """Tunix checkpoint manager that reads and writes MaxText's on-disk layout. + + Tunix stores `nnx.state(model)` verbatim under a `model_params` item. MaxText stores the Linen + layout under `items`: weights in `params/params`, the optimizer in `opt_state` and `step`, and + NNX-only state such as rngs in `nnx_aux`. Converting on the way out keeps post-training + checkpoints loadable by pre-training and everything else that reads MaxText checkpoints. + + Checkpoints written before this existed are still in the Tunix layout, so `maybe_restore` + falls back to the base class for those. + """ + + def __init__(self, root_directory=None, options=None, extra_item_handlers=None, config=None): + """Initializes the manager. + + Args: + root_directory: Directory to write checkpoints to. None disables checkpointing. + options: Orbax `CheckpointManagerOptions`. + extra_item_handlers: Handlers for items a subclass saves besides the state. + config: The run's config, read for the metadata the checkpoint stores. + """ + self._config = config + super().__init__(root_directory=root_directory, options=options) + # The base class built a manager over Tunix's item names. Close it before replacing it with + # one that knows MaxText's layout, or its open handles and threads outlive it. + # pylint: disable=access-member-before-definition + if getattr(self, "_checkpoint_manager", None) is not None: + self._checkpoint_manager.close() + # pylint: enable=access-member-before-definition + + if root_directory is not None: + # Pathways only supports the persistence APIs, so drop ocdbt/zarr3 there as Tunix does. + pathways = "proxy" in os.getenv("JAX_PLATFORMS", "") + + # Orbax otherwise materialises the whole tree on the host at once (its default concurrency + # is ~89GiB), which OOMKills the container the trainer runs in. MaxText already has a knob + # for this, checkpoint_storage_concurrent_gb, but it was never plumbed into this path. + concurrent_gb = getattr(config, "checkpoint_storage_concurrent_gb", None) if config is not None else None + + def pytree_handler(): + kwargs = {"use_ocdbt": not pathways, "use_zarr3": not pathways} + if concurrent_gb: + # Only the device-to-host budget: that is the one that decides how much of the tree is + # resident in host memory at once. Capping save_concurrent_gb/restore_concurrent_gb as + # well breaks reads of any single array larger than the cap, e.g. llama3.1-8b's + # mlp.wi_0.kernel at 3.75GiB ("Requested more bytes than we reserved space for"). + kwargs["save_device_host_concurrent_gb"] = concurrent_gb + return ocp.PyTreeCheckpointHandler(**kwargs) + + handlers = { + _ITEM_NAME: pytree_handler(), + # Tunix's item names stay registered so `maybe_restore` can fall back to checkpoints + # written before the layout change. + **{name: pytree_handler() for name in _TUNIX_ITEM_NAMES}, + "custom_metadata": ocp.JsonCheckpointHandler(), + **(extra_item_handlers or {}), + } + self._checkpoint_manager = ocp.CheckpointManager( + root_directory, + item_handlers=handlers, + options=options, + ) + else: + self._checkpoint_manager = None + + def wait_until_finished(self): + """Blocks until outstanding async checkpoint writes are complete.""" + if getattr(self, "_checkpoint_manager", None) is not None: + self._checkpoint_manager.wait_until_finished() + + def close(self): + """Closes the checkpoint manager.""" + if getattr(self, "_checkpoint_manager", None) is not None: + self._checkpoint_manager.close() + + def latest_step(self) -> int | None: + """Returns the latest step saved, reloading from storage if not cached.""" + if getattr(self, "_checkpoint_manager", None) is None: + return None + step = self._checkpoint_manager.latest_step() + if step is None: + steps = self.all_steps(read=True) + return steps[-1] if steps else None + return step + + def all_steps(self, read: bool = False) -> Sequence[int]: + """Returns all steps tracked by the manager.""" + if getattr(self, "_checkpoint_manager", None) is None: + return [] + return self._checkpoint_manager.all_steps(read=read) + + def model_to_checkpoint(self, model: nnx.Module) -> nnx.Module: + """Returns the module whose weights belong in the checkpoint. + + Args: + model: The model the trainer holds. + + Returns: + The module to checkpoint. Subclasses override this when it is not the trainer's model. + """ + return unwrap_model(model) + + def _train_state(self, model, optimizer): + """Returns the `{model, optimizer}` state to checkpoint. + + Args: + model: The model the trainer holds. + optimizer: The trainer's optimizer, or None to checkpoint weights only. + + Returns: + An `nnx.State` shaped like the one pre-training checkpoints. + """ + return nnx.state(train_state_nnx.TrainStateNNX(self.model_to_checkpoint(model), optimizer)) + + def _extra_save_args(self, step): + """Returns save args for items a subclass stores besides the state. + + Args: + step: The step being saved. + + Returns: + A dict of item name to Orbax save args. Empty by default. + """ + del step + return {} + + def save( # pylint: disable=too-many-positional-arguments + self, + step: int, + model: nnx.Module, + optimizer: nnx.Optimizer | None = None, + save_only_lora_params: bool = False, + force: bool = False, + custom_metadata: dict[str, Any] | None = None, + ) -> bool: + """Saves the model and optimizer in MaxText's on-disk layout. + + Args: + step: The step to save at. + model: The model the trainer holds. + optimizer: The trainer's optimizer, or None to save weights only. + save_only_lora_params: Whether to save only the LoRA params. + force: Whether to save regardless of the save decision policy. + custom_metadata: Metadata to store with the checkpoint. + + Returns: + Whether a checkpoint was written. + """ + if self._checkpoint_manager is None: + return False + if not force and not self._checkpoint_manager.should_save(step): + return False + + state = self._train_state(model, optimizer) + if save_only_lora_params: + state = nnx.split_state(state, nnx.LoRAParam, ...)[0] + items = train_state_nnx.to_checkpoint_dict(state) + if "opt_state" in items: + inner = _drop_inject_hyperparams(items["opt_state"]) + if inner is not items["opt_state"]: + # to_checkpoint_dict ran against the inject_hyperparams shell. It puts mu and nu into + # the Linen `params` collection by finding those keys at the top of the optimizer + # state, and behind the shell they are not there, so it left them bare. Convert what + # was behind it, or pre-training finds the accumulators one level short. + inner = train_state_nnx.opt_state_to_linen(inner) + items["opt_state"] = inner + if self.model_to_checkpoint(model) is not model: + items["opt_state"] = _drop_adapter_level(items["opt_state"]) + jax.block_until_ready(items) + + save_args = { + _ITEM_NAME: ocp.args.PyTreeSave(item=items, save_args=jax.tree.map(lambda _: ocp.SaveArgs(), items)), + **self._extra_save_args(step), + } + # The config-derived keys are the ones pre-training writes; a caller's own keys win. + metadata = checkpointing.checkpoint_custom_metadata(self._config) + metadata.update(custom_metadata or {}) + + if not force and step in self.all_steps(): + max_logging.log(f"Step {step} already exists in MaxText layout. Skipping save.") + return False + + try: + saved = self._checkpoint_manager.save( + step, + args=ocp.args.Composite(**save_args), + custom_metadata=metadata, + force=force, + ) + except Exception as e: # pylint: disable=broad-exception-caught + if "StepAlreadyExistsError" in type(e).__name__: + max_logging.log(f"Step {step} already exists. Skipping save.") + saved = False + else: + raise e + if saved: + max_logging.log(f"Saved post-training checkpoint at step {step} in MaxText's on-disk layout") + return saved + + def maybe_restore( + self, + model: nnx.Module, + optimizer: nnx.Optimizer | None = None, + step: int | None = None, + restore_only_lora_params: bool = False, + ) -> tuple[int, dict[str, Any]]: + """Restores the model and optimizer in place from the latest checkpoint. + + Args: + model: The model to restore into. + optimizer: The optimizer to restore into, or None to skip it. + step: The step to restore from. Defaults to the latest. + restore_only_lora_params: Whether to restore only the LoRA params. + + Returns: + A tuple of the restored step (0 if there is no checkpoint) and its custom metadata. + """ + if self._checkpoint_manager is None: + return 0, {} + if step is None: + step = self._checkpoint_manager.latest_step() + if step is None: + return 0, {} + + metadata = self._checkpoint_manager.metadata(step) + if _ITEM_NAME not in metadata.item_metadata: + max_logging.log(f"Step {step} predates MaxText-layout post-training checkpoints; restoring the Tunix layout") + return super().maybe_restore(model, optimizer, step=step, restore_only_lora_params=restore_only_lora_params) + + state = self._train_state(model, optimizer) + if restore_only_lora_params: + # save() narrows the state the same way, so a LoRA run writes only its adapter. Restoring + # against the full state asks for weights the checkpoint never held. + state = nnx.split_state(state, nnx.LoRAParam, ...)[0] + target = train_state_nnx.to_checkpoint_dict(state) + opt_state_guide = target.get("opt_state") + is_wrapped = self.model_to_checkpoint(model) is not model + if is_wrapped and opt_state_guide is not None: + target["opt_state"] = _drop_adapter_level(opt_state_guide) + + restored = self._checkpoint_manager.restore( + step, + args=ocp.args.Composite( + **{ + _ITEM_NAME: ocp.args.PyTreeRestore( + item=target, + restore_args=ocp.checkpoint_utils.construct_restore_args(target), + ) + } + ), + ) + + restored_items = dict(restored[_ITEM_NAME]) + if "opt_state" in restored_items and opt_state_guide is not None: + restored_items["opt_state"] = _add_inject_hyperparams(restored_items["opt_state"], opt_state_guide, step) + if is_wrapped: + restored_items["opt_state"] = _add_adapter_level(restored_items["opt_state"], opt_state_guide) + + new_state = checkpointing.linen_items_to_nnx(restored_items, state) + nnx.update(self.model_to_checkpoint(model), new_state["model"]) + if optimizer is not None and "optimizer" in new_state: + nnx.update(optimizer, new_state["optimizer"]) + + max_logging.log(f"Restored post-training checkpoint from step {step}") + return step, (metadata.custom_metadata if metadata else {}) or {} + + +def install(trainer, checkpoint_dir: str, config=None) -> None: + """Replaces a Tunix trainer's checkpoint manager with the MaxText-layout one and restores. + + `PeftTrainer.__init__` builds its own manager and restores from it, so callers pass a + `checkpoint_root_directory` of None and call this straight afterwards instead. + + Args: + trainer: A Tunix `PeftTrainer` or subclass. + checkpoint_dir: Directory to read and write checkpoints in. + config: The run's config, read for the metadata the checkpoint stores. + """ + # enable_checkpointing is a documented MaxText flag, and until now this path ignored it: the + # manager was installed regardless, so Tunix saved at the end of training no matter what the + # config said. That save is not free -- an 8B SFT writes 44.9 GiB and the transfer to host + # OOMKills the container -- so being able to turn it off is the difference between a smoke test + # that reports whether the trainer runs and one that cannot get past its first save. + # post_train_skip_checkpointing exists because neither existing flag can say this: + # enable_checkpointing is validated as required whenever load_parameters_path is set, and + # checkpoint_period=0 makes the shared checkpointing path raise ZeroDivisionError on + # `step % config.checkpoint_period`. + if config is not None and ( + not getattr(config, "enable_checkpointing", True) or getattr(config, "post_train_skip_checkpointing", False) + ): + max_logging.log("Checkpoint saving disabled: skipping post-train checkpoint manager install.") + return + + if trainer.checkpoint_manager is not None: + trainer.checkpoint_manager.close() + + trainer.checkpoint_manager = MaxTextLayoutCheckpointManager( + root_directory=checkpoint_dir, + options=trainer.config.checkpointing_options, + config=config, + ) + # pylint: disable=protected-access + trainer._train_steps, trainer._restored_custom_metadata = trainer.checkpoint_manager.maybe_restore( + trainer.model, + trainer.optimizer, + restore_only_lora_params=getattr(trainer, "_lora_enabled", False), + ) + trainer._iter_steps = trainer._train_steps * trainer.config.get_with_default("gradient_accumulation_steps", 1) + # pylint: enable=protected-access diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py index 132a808d5e..5c3926be70 100644 --- a/src/maxtext/trainers/post_train/distillation/distillation_utils.py +++ b/src/maxtext/trainers/post_train/distillation/distillation_utils.py @@ -22,17 +22,16 @@ from typing import Any, Callable, Iterator, List, Literal, Optional, Sequence import flax -from flax import nnx import jax import jax.numpy as jnp import numpy as np import optax -from orbax import checkpoint +import orbax.checkpoint as ocp from maxtext.utils import max_logging from maxtext.utils import maxtext_utils from maxtext.common import grain_utility -from tunix.sft import checkpoint_manager as tunix_checkpoint_manager +from maxtext.trainers.post_train import checkpointing as post_train_checkpointing from tunix.sft import peft_trainer @@ -86,13 +85,19 @@ class MaxTextToTunixIterator: Tunix expects an object with specific attributes (input_tokens, etc.). """ - def __init__(self, maxtext_iterator: Iterator): + def __init__(self, maxtext_iterator: Iterator, max_batches: int | None = None): """Initializes the adapter. Args: maxtext_iterator: The upstream iterator created by MaxText's input pipeline. + max_batches: Batches to yield before stopping, or None for as many as the upstream + iterator has. Distillation drives the training loop itself, which makes Tunix skip + its own `max_steps` check, so on an endless dataset the run only ends when the + batches do. """ self._iterator = maxtext_iterator + self._max_batches = max_batches + self._batches = 0 def __iter__(self): """Returns self as the iterator.""" @@ -105,9 +110,12 @@ def __next__(self) -> MaxTextTrainingInput: A MaxTextTrainingInput object containing the batch data. Raises: - StopIteration: If the upstream iterator is exhausted. + StopIteration: If the upstream iterator is exhausted, or the batch budget is spent. """ + if self._max_batches is not None and self._batches >= self._max_batches: + raise StopIteration batch = next(self._iterator) + self._batches += 1 # Ensure segmentation exists, default to ones if missing (standard non-packed) if "inputs_segmentation" in batch: @@ -647,7 +655,7 @@ def create_labels(self, targets, targets_segmentation=None, **kwargs): # ----------------------------------------------------------------------------- -class MaxTextCheckpointManager(tunix_checkpoint_manager.CheckpointManager): +class MaxTextCheckpointManager(post_train_checkpointing.MaxTextLayoutCheckpointManager): """Custom CheckpointManager that uses MaxText's native handlers. Model and optimizer are delegated to Tunix's v1 ``Checkpointer`` unchanged. @@ -660,95 +668,71 @@ def __init__( raw_iterator: Any | None, root_directory: str | None, student_config: Any, - options: checkpoint.CheckpointManagerOptions | None = None, + options: ocp.CheckpointManagerOptions | None = None, ): - super().__init__(root_directory=root_directory, options=options) + super().__init__( + root_directory=root_directory, + options=options, + # MaxText's Grain handler, so the input pipeline's position rides along with the state. + extra_item_handlers={"iter": grain_utility.GrainCheckpointHandler()}, + config=student_config, + ) self.student_config = student_config self._iterator = raw_iterator - def save( - self, - step, - model, - optimizer=None, - save_only_lora_params=False, - force=False, - custom_metadata=None, - ): - """Saves model, optimizer and the Grain input pipeline state.""" - if self._checkpointer is None: - return False - - # Standard Tunix Logic for Model/Optimizer. - # Accept either a ModelBundle (common path) or a plain nnx module. - target_model = getattr(model, "student_model", model) - if save_only_lora_params: - params = nnx.state(target_model, nnx.LoRAParam) - else: - params = nnx.state(target_model) - - checkpointables: dict[str, Any] = {"model_params": params} - # Exclude optimizer state when learn_to_init_mode is active. - exclude_opt = self.student_config.learn_to_init_mode - - if optimizer is not None and not exclude_opt: - checkpointables["optimizer_state"] = nnx.state(optimizer, nnx.optimizer.OptState) - - if self._iterator is not None: - # Follow MaxText's logic to handle multi-process saving - # Logic extracted from src/maxtext/common/checkpointing.py:save_checkpoint - data_iterator = self._iterator - if not isinstance(data_iterator, list): - data_iterator = [data_iterator] - - grain_iters_to_save = [] - process_count_total = jax.process_count() * len(data_iterator) - - for i, data_iter in enumerate(data_iterator): - process_index = jax.process_index() + i * jax.process_count() - # MaxText iterators (MultiHostDataLoadIterator) wrap the actual Grain iterator in .local_iterator - local_iter = data_iter.local_iterator if hasattr(data_iter, "local_iterator") else data_iter - grain_iters_to_save.append((local_iter, process_index, process_count_total)) - - checkpointables["iter"] = grain_utility.GrainCheckpointable( - save_args=grain_utility.GrainCheckpointSave(item=grain_iters_to_save) # pyrefly: ignore[bad-assignment] - ) + def model_to_checkpoint(self, model): + """Only the student is trained, so only the student is checkpointed.""" + return getattr(model, "student_model", model) + + def _train_state(self, model, optimizer): + # learn-to-init runs discard the optimizer state, so leave it out of the checkpoint. + if self.student_config.learn_to_init_mode: + optimizer = None + return super()._train_state(model, optimizer) + + def _extra_save_args(self, step): + """Saves the input pipeline's position alongside the state, when there is one to save.""" + del step + if self._iterator is None: + return {} + + # Follow MaxText's logic to handle multi-process saving. + # Logic extracted from src/maxtext/common/checkpointing.py:save_checkpoint + data_iterator = self._iterator + if not isinstance(data_iterator, list): + data_iterator = [data_iterator] + + grain_iters_to_save = [] + process_count_total = jax.process_count() * len(data_iterator) + for i, data_iter in enumerate(data_iterator): + process_index = jax.process_index() + i * jax.process_count() + # MaxText iterators (MultiHostDataLoadIterator) wrap the actual Grain iterator in .local_iterator + local_iter = data_iter.local_iterator if hasattr(data_iter, "local_iterator") else data_iter + grain_iters_to_save.append((local_iter, process_index, process_count_total)) - return self._save_checkpointables(step, checkpointables, force, custom_metadata) + return {"iter": grain_utility.GrainCheckpointSave(item=grain_iters_to_save)} def maybe_restore( # pyrefly: ignore[bad-override] self, model: Any, optimizer: Any = None, + step: int | None = None, restore_only_lora_params: bool = False, ) -> tuple[int, dict[str, Any]]: - """Restores model + optimizer by delegating to upstream Tunix. - - Unwraps `ModelBundle` if present (we only restore `student_model`). - - Returns: - (restored step, custom_metadata dict). Step is 0 if no checkpoint exists. - """ - if self._checkpointer is None: - return 0, {} - - target_model = getattr(model, "student_model", model) - + """Restores the student model and its optimizer from MaxText's on-disk layout.""" step, custom_metadata = super().maybe_restore( - model=target_model, # pyrefly: ignore[bad-argument-type] - optimizer=optimizer, + model, + optimizer, + step=step, restore_only_lora_params=restore_only_lora_params, ) - if step == 0: - return 0, {} - - max_logging.log(f"Restored from checkpoint step {step}.") - + if step: + max_logging.log(f"Restored from checkpoint step {step}.") return step, dict(custom_metadata or {}) def restore_iterator(self): """Restores the iterator using MaxText's logic.""" - if self._checkpointer is None or self._iterator is None: + if self._checkpoint_manager is None or self._iterator is None: return None step = self.latest_step() @@ -761,9 +745,11 @@ def restore_iterator(self): data_iter = self._iterator local_iter = data_iter.local_iterator if hasattr(data_iter, "local_iterator") else data_iter - self._checkpointer.load_checkpointables( + self._checkpoint_manager.restore( step, - {"iter": grain_utility.GrainCheckpointable(restore_args=grain_utility.GrainCheckpointRestore(item=local_iter))}, + args=ocp.args.Composite( + iter=grain_utility.GrainCheckpointRestore(item=local_iter), + ), ) # Since Grain restores in-place via set_state(), we return the original object return self._iterator @@ -771,8 +757,3 @@ def restore_iterator(self): except Exception as e: # pylint: disable=broad-exception-caught max_logging.log(f"Warning: Could not restore input pipeline: {e}") return None - - def wait_until_finished(self): - """Blocks until all outstanding checkpoint operations are complete.""" - if self._checkpointer is not None: - self._checkpointer.wait() diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index ab4f7bc5fa..8f824e827d 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/src/maxtext/trainers/post_train/distillation/train_distill.py @@ -43,6 +43,7 @@ from flax.linen import partitioning as nn_partitioning import jax import jax.numpy as jnp +import numpy as np import optax import re import os @@ -113,10 +114,7 @@ def optimizer_factory(learning_rate): # Apply Gradient Clipping if config.gradient_clipping_threshold > 0: - opt = optax.chain( - optax.clip_by_global_norm(max_norm=config.gradient_clipping_threshold), - opt, - ) + opt = optimizers.add_gradient_clipping(opt, config.gradient_clipping_threshold) return opt # 3. Create Injectable Optimizer @@ -206,6 +204,22 @@ def call_teacher(self, *args, **kwargs): return jax.lax.stop_gradient(self.teacher_model(*args, **kwargs)) # pyrefly: ignore[not-callable] +def _select_inputs(args, kwargs): + """Picks the training batch out of Tunix's train-step arguments. + + Tunix passes the gradient accumulator alongside the batch, and has swapped their order between + versions, so the batch cannot be taken from a fixed position. Identify the accumulator by type + and take the remaining argument. + """ + if "inputs" in kwargs: + return kwargs["inputs"] + accumulator_type = getattr(peft_trainer, "GradientAccumulator", ()) + candidates = [a for a in args if a is not None and not isinstance(a, accumulator_type)] + if not candidates: + raise ValueError(f"No training batch found in train-step arguments: {[type(a) for a in args]}") + return candidates[0] + + class MaxTextDistillationTrainer(peft_trainer.PeftTrainer): """Custom Trainer to preserve MaxText fields and log Teacher metrics. @@ -276,7 +290,7 @@ def wrt_filter(path, x): # Inherits _shard_optimizer from PeftTrainer. - def _train_step(self, model, optimizer, inputs, grad_accumulator=None, **kwargs): # pyrefly: ignore[bad-override] + def _train_step(self, model, optimizer, *args, **kwargs): # pyrefly: ignore[bad-override] """Overrides the main JIT block to natively handle ModelBundle module. Uses jax.value_and_grad with explicit split/merge to avoid nesting @@ -284,7 +298,15 @@ def _train_step(self, model, optimizer, inputs, grad_accumulator=None, **kwargs) conflicting outer_index values and raises: ValueError: The graph structure of a node added to cached_partial was mutated inside the transformation. + + Tunix has moved grad_accumulator either side of inputs across versions: older builds call + (model, optimizer, inputs, grad_accumulator), newer ones + (model, optimizer, grad_accumulator, inputs, is_update_step=...). Binding those positionally + hands the accumulator to gen_model_input_fn, which then fails with + "'GradientAccumulator' object has no attribute 'input_tokens'". Pick the batch out of the + positional arguments instead of trusting their order. """ + inputs = _select_inputs(args, kwargs) batch = self.gen_model_input_fn(inputs) student = model.student_model teacher = model.teacher_model @@ -530,16 +552,29 @@ def setup_checkpoint_manager_and_restore(self, raw_train_iter, config): # 3. Restore Model & Optimizer State correctly via MaxTextCheckpointManager. # Accessing protected variables of the base class IS allowed inside the subclass! - self._train_steps, self._restored_custom_metadata = self.checkpoint_manager.maybe_restore( - self.model, - self.optimizer, - restore_only_lora_params=getattr(self, "_lora_enabled", False), - ) + # + # post_train_skip_checkpointing covers restore as well as save. A run carrying the flag wants + # nothing to do with the checkpoint directory, and a partial checkpoint left there by an + # earlier run is not a state it can resume from: the 0816 rerun died in maybe_restore with + # "restore item and on-disk value metadata tree structures do not match", reading a step-1 + # checkpoint that a previous broken run had written into the same path. + if getattr(config, "post_train_skip_checkpointing", False): + max_logging.log("post_train_skip_checkpointing=True: starting from step 0, not restoring.") + # An empty dict rather than None: tunix/rl/trainer.py reads this with .get(), and while + # distillation does not take that path today, a None here would be a live trap for whoever + # does. + self._train_steps, self._restored_custom_metadata = 0, {} + else: + self._train_steps, self._restored_custom_metadata = self.checkpoint_manager.maybe_restore( + self.model, + self.optimizer, + restore_only_lora_params=getattr(self, "_lora_enabled", False), + ) grad_accum_steps = self.config.get_with_default("gradient_accumulation_steps", 1) self._iter_steps = self._train_steps * grad_accum_steps # 4. Restore input state (if applicable) - if enable_checkpointing: + if enable_checkpointing and not getattr(config, "post_train_skip_checkpointing", False): restored_iter = self.checkpoint_manager.restore_iterator() if restored_iter is not None: max_logging.log("Restored input pipeline state to match model step.") @@ -625,8 +660,25 @@ def build_training_components( # Prepare optimizer optimizer = get_distillation_optimizer(student_config, student_config.steps) + # post_train_skip_checkpointing is honoured here as well as in post_train/checkpointing.py's + # install(). Distillation does not go through install() -- it builds its own manager -- so a flag + # applied only there left this path saving anyway: the 0815 matrix passed the flag and every case + # still wrote a checkpoint at step 1, several dying in that save after training had already + # produced a loss. + # + # Widening the interval rather than removing the manager: Tunix's PeftTrainer.train calls + # self.checkpoint_manager.save() on every update step with no None check, so setting the manager + # to None killed all 15 cases with AttributeError before any of them reached a loss. An interval + # past the last step leaves that call in place and lets Orbax decline it. + save_interval = student_config.checkpoint_period + if getattr(student_config, "post_train_skip_checkpointing", False): + save_interval = student_config.steps + 1_000_000 + max_logging.log( + f"post_train_skip_checkpointing=True: save_interval_steps={save_interval}, no checkpoint will be written." + ) + checkpointing_options = checkpoint.CheckpointManagerOptions( - save_interval_steps=student_config.checkpoint_period, + save_interval_steps=save_interval, max_to_keep=student_config.max_num_checkpoints_to_keep, enable_async_checkpointing=student_config.async_checkpointing, create=True, @@ -788,7 +840,16 @@ def custom_gen_model_input_fn(batch): trainer = trainer.with_gen_model_input_fn(custom_gen_model_input_fn) # 7. Create Iterator Wrappers (Use Utils) - train_iter = distillation_utils.MaxTextToTunixIterator(raw_train_iter) + # The trainer is managed externally, so Tunix does not enforce max_steps for us. Bound the + # batches instead: one training step consumes gradient_accumulation_steps of them, and a + # resumed run has already spent some. + grad_accum = train_config.get_with_default("gradient_accumulation_steps", 1) + iter_steps = getattr(trainer, "_iter_steps", 0) + if not isinstance(iter_steps, (int, float, np.integer)): + iter_steps = 0 + batch_budget = max(0, student_config.steps * grad_accum - int(iter_steps)) # pylint: disable=protected-access + max_logging.log(f"Distillation will run at most {batch_budget} more batches ({student_config.steps} steps).") + train_iter = distillation_utils.MaxTextToTunixIterator(raw_train_iter, max_batches=batch_budget) eval_iter = None if raw_eval_iter is not None: @@ -808,7 +869,9 @@ def custom_gen_model_input_fn(batch): apply_lti_model_update(student_model, student_config) # 9. Final Save (Conditional) - if student_config.save_checkpoint_on_completion: + # The widened save_interval_steps above does not cover this block: it saves with force=True, + # which bypasses the interval entirely. Both have to be off for a run to write nothing. + if student_config.save_checkpoint_on_completion and not getattr(student_config, "post_train_skip_checkpointing", False): should_save = student_config.steps % student_config.checkpoint_period if should_save: diff --git a/src/maxtext/trainers/post_train/dpo/train_dpo.py b/src/maxtext/trainers/post_train/dpo/train_dpo.py index 35b407e4a5..d8fedc42b9 100644 --- a/src/maxtext/trainers/post_train/dpo/train_dpo.py +++ b/src/maxtext/trainers/post_train/dpo/train_dpo.py @@ -28,7 +28,6 @@ from absl import app import jax -import optax from orbax import checkpoint as ocp import pathwaysutils @@ -55,6 +54,7 @@ from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils +from maxtext.trainers.post_train import checkpointing as post_train_checkpointing from maxtext.utils import model_creation_utils @@ -94,8 +94,15 @@ def get_tunix_config(mt_config: MaxTextConfig) -> DPOTrainingConfig: return DPOTrainingConfig( eval_every_n_steps=mt_config.eval_interval, max_steps=mt_config.steps, - gradient_accumulation_steps=mt_config.gradient_accumulation_steps, - checkpoint_root_directory=mt_config.checkpoint_dir, + # None rather than 1: Tunix wraps the optimizer in optax.MultiSteps whenever this is set, + # and a 1-step wrap buys nothing while giving the optimizer state a shape pre-training + # can't resume from. Matches train_sft. + gradient_accumulation_steps=( + mt_config.gradient_accumulation_steps if mt_config.gradient_accumulation_steps > 1 else None + ), + # Checkpointing is handled by post_train.checkpointing, which writes MaxText's on-disk + # layout instead of Tunix's, so Tunix's own manager stays disabled. + checkpoint_root_directory=None, checkpointing_options=checkpointing_options, metrics_logging_options=metrics_logging_options, profiler_options=profiler_options, @@ -103,8 +110,9 @@ def get_tunix_config(mt_config: MaxTextConfig) -> DPOTrainingConfig: lambda_orpo=mt_config.dpo.orpo_lambda, beta=mt_config.dpo.dpo_beta, label_smoothing=mt_config.dpo.dpo_label_smoothing, - max_prompt_length=mt_config.dpo.max_prompt_length, - max_response_length=mt_config.max_target_length - mt_config.dpo.max_prompt_length, + max_prompt_length=mt_config.dpo.max_prompt_length or (mt_config.max_target_length // 2), + max_response_length=mt_config.max_target_length + - (mt_config.dpo.max_prompt_length or (mt_config.max_target_length // 2)), ) @@ -129,31 +137,29 @@ def setup_trainer_state(mt_config, goodput_recorder=None, test_only_training_hoo tokenizer_pad_id=tok.pad_id, ) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(mt_config.logical_axis_rules): learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(mt_config) # pass in model for muon optimizer = optimizers.get_optimizer(mt_config, learning_rate_schedule, model) if mt_config.gradient_clipping_threshold > 0: - optimizer = optax.chain( - optax.clip_by_global_norm(max_norm=mt_config.gradient_clipping_threshold), - optimizer, - ) + optimizer = optimizers.add_gradient_clipping(optimizer, mt_config.gradient_clipping_threshold) - # ORPO does not require a reference model. - ref_model = nnx.clone(model) if mt_config.dpo.algo == "dpo" else None + # ORPO does not require a reference model. + ref_model = nnx.clone(model) if mt_config.dpo.algo == "dpo" else None - with maybe_record_goodput(goodput_recorder, GoodputEvent.TRAINING_PREPARATION): - training_hooks_class = test_only_training_hooks_class or hooks.DPOTrainingHooks - training_hooks = training_hooks_class(mt_config, mesh, learning_rate_schedule, goodput_recorder) - data_hooks = hooks.DPODataHooks(mt_config, mesh, goodput_recorder) + with maybe_record_goodput(goodput_recorder, GoodputEvent.TRAINING_PREPARATION): + training_hooks_class = test_only_training_hooks_class or hooks.DPOTrainingHooks + training_hooks = training_hooks_class(mt_config, mesh, learning_rate_schedule, goodput_recorder) + data_hooks = hooks.DPODataHooks(mt_config, mesh, goodput_recorder) - # Provide rules context so logical axes (e.g. 'norm') are translated to mesh axes during maybe_restore - with nn_partitioning.axis_rules(mt_config.logical_axis_rules): + # Provide rules context so logical axes (e.g. 'norm') are translated to mesh axes during maybe_restore trainer = DPOTrainer( model=model, ref_model=ref_model, optimizer=optimizer, training_config=tunix_config, tokenizer=None ) trainer.with_training_hooks(training_hooks) trainer.with_data_hooks(data_hooks) + post_train_checkpointing.install(trainer, mt_config.checkpoint_dir, mt_config) return trainer, mesh diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 725d5cd48a..e1660276b0 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -73,6 +73,47 @@ from tunix.rl.grpo.grpo_learner import GrpoConfig, GrpoLearner from tunix.sft import metrics_logger, profiler import tunix.generate.utils as tunix_utils +from tunix.generate.tokenizer_adapter import TokenizerAdapter + +# Monkey-patch TokenizerAdapter to handle MaxText tokenizer properties +_old_eos_id = TokenizerAdapter.eos_id + + +def _patched_eos_id(self): + # pylint: disable=protected-access + if hasattr(self._tokenizer, "eos_id") and not callable(self._tokenizer.eos_id): + return self._tokenizer.eos_id + return _old_eos_id(self) + + +TokenizerAdapter.eos_id = _patched_eos_id + +_old_bos_id = TokenizerAdapter.bos_id + + +def _patched_bos_id(self): + # pylint: disable=protected-access + if hasattr(self._tokenizer, "bos_id") and not callable(self._tokenizer.bos_id): + return self._tokenizer.bos_id + return _old_bos_id(self) + + +TokenizerAdapter.bos_id = _patched_bos_id + +_old_pad_id = TokenizerAdapter.pad_id + + +def _patched_pad_id(self): + # pylint: disable=protected-access + if hasattr(self._tokenizer, "pad_id") and not callable(self._tokenizer.pad_id): + pad_id = self._tokenizer.pad_id + if pad_id is None or pad_id < 0: + return self.eos_id() + return pad_id + return _old_pad_id(self) + + +TokenizerAdapter.pad_id = _patched_pad_id @contextlib.contextmanager @@ -143,6 +184,7 @@ def _compat_unstack(src_val, tgt_val, key_path, scan_axis=None): from maxtext.trainers.post_train.rl.evaluate_rl import evaluate from maxtext.trainers.post_train.rl import utils_rl from maxtext.input_pipeline.instruction_data_processing import load_data_template_from_file +from maxtext.trainers.post_train import checkpointing as post_train_checkpointing from maxtext.utils import max_logging, max_utils, model_creation_utils @@ -358,7 +400,9 @@ def prepare_datasets( ) def _filter_long_prompts(x): - tokens = model_tokenizer.tokenize(x["prompts"]) + # tokenize() is a HuggingFace method. MaxText's own tokenizers only offer encode(), and the + # length being checked is the prefill length in token ids anyway. + tokens = model_tokenizer.encode(x["prompts"]) return len(tokens) <= trainer_config.max_prefill_predict_length train_dataset = train_dataset.filter(_filter_long_prompts) @@ -509,7 +553,9 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments rollout_micro_batch_size=rollout_micro_batch_size, metrics_logging_options=metrics_logging_options, profiler_options=profiler_options, - checkpoint_root_directory=checkpoint_dir, + # Checkpointing is handled by post_train.checkpointing, which writes MaxText's on-disk + # layout instead of Tunix's, so Tunix's own manager stays disabled. + checkpoint_root_directory=None, checkpointing_options=checkpointing_options, ), rollout_config=base_rollout.RolloutConfig( @@ -519,7 +565,7 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments temperature=trainer_config.decode_sampling_temperature, top_p=trainer_config.decode_sampling_nucleus_p, top_k=trainer_config.decode_sampling_top_k, - rollout_vllm_model_version=trainer_config.tokenizer_path, + rollout_vllm_model_version=trainer_config.vllm_hf_config_path or trainer_config.tokenizer_path, rollout_vllm_hbm_utilization=trainer_config.hbm_utilization_vllm, rollout_vllm_tpu_backend_type=getattr( trainer_config, @@ -539,6 +585,7 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments "hf_overrides": trainer_config.vllm_hf_overrides, "enable_expert_parallel": sampler_config.enable_expert_parallel, "enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config), + "trust_remote_code": True, # Ensures vLLM model initializes with correct dtype (not float32 default) "dtype": trainer_config.weight_dtype.value, }, @@ -577,6 +624,8 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments cluster_config=cluster_config, **rl_cluster_kwargs, ) + if checkpoint_dir is not None: + post_train_checkpointing.install(rl_cluster.actor_trainer, os.path.join(checkpoint_dir, "actor"), trainer_config) def make_reward_fn(fn): # pragma: no cover @@ -700,9 +749,14 @@ def _rl_train_impl(argv: Sequence[str], kwargs: dict): # adapter (used to synthesize segment_ids that mask pad positions from # attention — without this the trainer attends to pad tokens and produces # corrupted log-probs). - model_tokenizer = AutoTokenizer.from_pretrained( - trainer_config.tokenizer_path, - token=trainer_config.hf_access_token or None, + from maxtext.input_pipeline import tokenizer # pylint: disable=import-outside-toplevel + + model_tokenizer = tokenizer.build_tokenizer( + tokenizer_path=trainer_config.tokenizer_path, + tokenizer_type=trainer_config.tokenizer_type, + add_bos=False, + add_eos=False, + hf_access_token=trainer_config.hf_access_token, ) configure_tokenizer_chat_template(model_tokenizer, trainer_config) @@ -711,7 +765,7 @@ def _rl_train_impl(argv: Sequence[str], kwargs: dict): sampler_config, trainer_devices, sampler_devices, - tokenizer_pad_id=model_tokenizer.pad_token_id, + tokenizer_pad_id=model_tokenizer.pad_id, ) if not trainer_config.debug: diff --git a/src/maxtext/trainers/post_train/rl/utils_rl.py b/src/maxtext/trainers/post_train/rl/utils_rl.py index 708571ff7e..53c013cb4b 100644 --- a/src/maxtext/trainers/post_train/rl/utils_rl.py +++ b/src/maxtext/trainers/post_train/rl/utils_rl.py @@ -30,6 +30,7 @@ from tunix.rl.agentic.parser.chat_template_parser import parser as agentic_chat_template_parser +from maxtext.optimizers import optimizers from maxtext.trainers.post_train.rl.math_verify_pool import math_verify_pool, verify_math_worker from maxtext.utils import max_logging @@ -625,18 +626,15 @@ def get_optimizer(tmvp_config: Any) -> optax.GradientTransformation: # Grad clipping to prevent large gradients. We find this # important to keep KL divergence in check. def make_optimizer(learning_rate): - transforms = [] - if tmvp_config.gradient_clipping_threshold > 0: - transforms.append(optax.clip_by_global_norm(max_norm=tmvp_config.gradient_clipping_threshold)) - transforms.append( - optax.adamw( - learning_rate=learning_rate, - b1=tmvp_config.adam_b1, - b2=tmvp_config.adam_b2, - weight_decay=tmvp_config.adam_weight_decay, - ) + opt = optax.adamw( + learning_rate=learning_rate, + b1=tmvp_config.adam_b1, + b2=tmvp_config.adam_b2, + weight_decay=tmvp_config.adam_weight_decay, ) - return optax.chain(*transforms) + if tmvp_config.gradient_clipping_threshold > 0: + opt = optimizers.add_gradient_clipping(opt, tmvp_config.gradient_clipping_threshold) + return opt # Wrap the entire optimizer (including gradient clipping) with # inject_hyperparams so opt_state.hyperparams['learning_rate'] is at the diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py index c99b5f48b6..594bd882d5 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/src/maxtext/trainers/post_train/sft/train_sft.py @@ -70,6 +70,7 @@ from maxtext.utils import max_logging # Placeholder: internal from maxtext.utils import maxtext_utils +from maxtext.trainers.post_train import checkpointing as post_train_checkpointing from maxtext.utils import model_creation_utils @@ -233,7 +234,9 @@ def get_tunix_config(mt_config): gradient_accumulation_steps=( mt_config.gradient_accumulation_steps if mt_config.gradient_accumulation_steps > 1 else None ), - checkpoint_root_directory=mt_config.checkpoint_dir, + # Checkpointing is handled by post_train.checkpointing, which writes MaxText's on-disk + # layout instead of Tunix's, so Tunix's own manager stays disabled. + checkpoint_root_directory=None, checkpointing_options=checkpointing_options, metrics_logging_options=metrics_logging_options, profiler_options=profiler_options, @@ -304,10 +307,7 @@ def setup_trainer_state(mt_config, goodput_recorder=None): optimizer = optimizers.get_optimizer(mt_config, learning_rate_schedule, model) if mt_config.gradient_clipping_threshold > 0: - optimizer = optax.chain( - optax.clip_by_global_norm(max_norm=mt_config.gradient_clipping_threshold), - optimizer, - ) + optimizer = optimizers.add_gradient_clipping(optimizer, mt_config.gradient_clipping_threshold) with maybe_record_goodput(goodput_recorder, GoodputEvent.TRAINING_PREPARATION): training_hooks = hooks.SFTTrainingHooks(mt_config, mesh, learning_rate_schedule, goodput_recorder) @@ -320,6 +320,7 @@ def setup_trainer_state(mt_config, goodput_recorder=None): trainer.with_training_hooks(training_hooks) trainer.with_data_hooks(data_hooks) trainer = use_maxtext_loss_function(trainer, mt_config) + post_train_checkpointing.install(trainer, mt_config.checkpoint_dir, mt_config) return trainer, mesh diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 08de8ed8e0..4d5f561d68 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/src/maxtext/utils/lora_utils.py @@ -746,6 +746,12 @@ def _build_target_leaf(v): is_leaf=lambda n: isinstance(n, nnx.Variable), ) + # Match the nesting the checkpoint actually has, so the guided restore lines up instead of + # falling back to a shapeless read that the walk below then cannot follow. + layout_keys = _maxtext_layout_keys(lora_restore_path) + for key in reversed(layout_keys): + target_for_restore = {key: target_for_restore} + sharding_tree = jax.tree.map(lambda x: getattr(x, "sharding", None), target_for_restore) restore_args_tree = ocp.checkpoint_utils.construct_restore_args(target_for_restore, sharding_tree) @@ -763,6 +769,10 @@ def _build_target_leaf(v): max_logging.log(f"Guided restore failed: {e}. Falling back to basic restore.") restored_lora_params = ocp.PyTreeCheckpointer().restore(lora_restore_path) + for key in layout_keys: + if isinstance(restored_lora_params, dict) and key in restored_lora_params: + restored_lora_params = restored_lora_params[key] + # Post processing def _map_to_state(path, variable): if not isinstance(variable, nnx.Variable): @@ -821,6 +831,22 @@ def _map_to_state(path, variable): # is identical to the Linen path. +def _maxtext_layout_keys(restore_path): + """Returns the keys to descend through before reaching the model tree. + + An adapter saved in MaxText's layout sits under a `params` item, with Flax's `params` collection + inside it. One written straight from `nnx.state(model)` is already model-shaped and needs none. + """ + try: + tree = ocp.PyTreeCheckpointer().metadata(restore_path).item_metadata + except Exception: # pylint: disable=broad-exception-caught + return () + inner = tree.get("params") if tree is not None and "params" in tree else None + if inner is not None and "params" in inner: + return ("params", "params") + return () + + def _is_nnx_branch(x): """Return True if `x` should be recursed into as a sub-tree.""" return isinstance(x, Mapping) diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh index eed23bad03..c7f54c515e 100644 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh @@ -60,7 +60,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ # Step 3: Run inference on the checkpoint generated from the previous run python3 -m maxtext.inference.vllm_decode \ model_name=${MODEL_NAME} \ - load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \ + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \ vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ hbm_utilization_vllm=0.5 \ prompt='Suggest some famous landmarks in London.' \ diff --git a/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh b/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh index ba2e294e57..560ea2ce43 100644 --- a/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh +++ b/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh @@ -65,7 +65,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ # Step 3: Run inference on the checkpoint generated from the previous run python3 -m maxtext.inference.vllm_decode \ model_name=${MODEL_NAME} \ - load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \ + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \ vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ hbm_utilization_vllm=0.85 \ prompt='Suggest some famous landmarks in London.' \ diff --git a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh index e50fab88f0..1c1dad96eb 100644 --- a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh +++ b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh @@ -54,7 +54,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ # Step 3: Run inference on the checkpoint generated from the previous run python3 -m maxtext.inference.vllm_decode \ model_name=${MODEL_NAME} \ - load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \ + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \ tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \ vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ hbm_utilization_vllm=0.6 \ diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh index 3bd90eb519..a4916d8c24 100755 --- a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh +++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh @@ -95,7 +95,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ # Step 3: Run inference on the checkpoint produced by the RL run python3 -m maxtext.inference.vllm_decode \ model_name=${MODEL_NAME} \ - load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \ + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \ vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ hbm_utilization_vllm=0.85 \ prompt='Suggest some famous landmarks in London.' \ diff --git a/tests/post_training/integration/lora_e2e_nnx_test.py b/tests/post_training/integration/lora_e2e_nnx_test.py index 63a664cd42..438bb0249e 100644 --- a/tests/post_training/integration/lora_e2e_nnx_test.py +++ b/tests/post_training/integration/lora_e2e_nnx_test.py @@ -104,7 +104,7 @@ def _run_e2e_flow_sft(self, model_name, lora_weight_qtype=None, scan_layers=True base_ckpt_dir = os.path.join(self.test_dir, base_run_name, "checkpoints", "2") self.assertTrue(os.path.exists(base_ckpt_dir), f"Base checkpoint path does not exist: {base_ckpt_dir}") - base_ckpt_path = os.path.join(base_ckpt_dir, "model_params") + base_ckpt_path = os.path.join(base_ckpt_dir, "items") lora_config = {"enable_lora": True, "lora_rank": 4} if lora_weight_qtype: @@ -128,7 +128,7 @@ def _run_e2e_flow_sft(self, model_name, lora_weight_qtype=None, scan_layers=True lora_ckpt_dir = os.path.join(self.test_dir, lora_run_name, "checkpoints", "4") self.assertTrue(os.path.exists(lora_ckpt_dir), f"Saved LoRA checkpoint path does not exist: {lora_ckpt_dir}") - lora_ckpt_path = os.path.join(lora_ckpt_dir, "model_params") + lora_ckpt_path = os.path.join(lora_ckpt_dir, "items") # Step 3: Resume training under same run name (steps=6) config_step3 = _tiny_lora_pyconfig( diff --git a/tests/post_training/unit/distillation_checkpointing_test.py b/tests/post_training/unit/distillation_checkpointing_test.py index 940f3bbc37..372868a122 100644 --- a/tests/post_training/unit/distillation_checkpointing_test.py +++ b/tests/post_training/unit/distillation_checkpointing_test.py @@ -21,6 +21,10 @@ import json import os +from types import SimpleNamespace +from etils import epath +import jax.numpy as jnp +import optax import shutil import tempfile from unittest import mock @@ -30,6 +34,7 @@ import jax from flax import nnx import orbax.checkpoint as ocp +from maxtext.common import checkpointing from maxtext.trainers.post_train.distillation import distillation_utils @@ -88,6 +93,9 @@ def test_save_and_restore_iterator(self): # 2. Save Checkpoint mock_student_config = mock.Mock() mock_student_config.learn_to_init_mode = False + mock_student_config.scan_layers = True + mock_student_config.lora = None + mock_student_config.checkpoint_storage_concurrent_gb = None manager = distillation_utils.MaxTextCheckpointManager( raw_iterator=iterator, root_directory=self.test_dir, student_config=mock_student_config, options=self.options ) @@ -119,6 +127,9 @@ def test_save_and_restore_iterator(self): mock_student_config_restore = mock.Mock() mock_student_config_restore.learn_to_init_mode = False + mock_student_config_restore.scan_layers = True + mock_student_config_restore.lora = None + mock_student_config_restore.checkpoint_storage_concurrent_gb = None restore_manager = distillation_utils.MaxTextCheckpointManager( raw_iterator=new_iterator, root_directory=self.test_dir, @@ -137,6 +148,9 @@ def test_restore_returns_none_if_no_checkpoint(self): iterator = FakeGrainIterator() mock_student_config_restore = mock.Mock() mock_student_config_restore.learn_to_init_mode = False + mock_student_config_restore.scan_layers = True + mock_student_config_restore.lora = None + mock_student_config_restore.checkpoint_storage_concurrent_gb = None manager = distillation_utils.MaxTextCheckpointManager( raw_iterator=iterator, root_directory=self.test_dir, @@ -149,5 +163,107 @@ def test_restore_returns_none_if_no_checkpoint(self): self.assertIsNone(result) +class MaxTextCheckpointManagerLayoutTest(absltest.TestCase): + """Distillation checkpoints only the student, in MaxText's on-disk layout.""" + + class Bundle(nnx.Module): + """Stand-in for the teacher/student ModelBundle the trainer holds.""" + + def __init__(self, student, teacher): + self.student_model = student + self.teacher_model = teacher + + def setUp(self): + super().setUp() + self.test_dir = tempfile.mkdtemp() + self.options = ocp.CheckpointManagerOptions(max_to_keep=2, create=True) + + def tearDown(self): + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + super().tearDown() + + def _save(self, learn_to_init_mode=False): + """Saves a checkpoint and returns its on-disk leaf paths.""" + student, teacher = DummyModel(nnx.Rngs(0)), DummyModel(nnx.Rngs(1)) + bundle = self.Bundle(student, teacher) + optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param) + manager = distillation_utils.MaxTextCheckpointManager( + raw_iterator=None, + root_directory=self.test_dir, + student_config=SimpleNamespace(learn_to_init_mode=learn_to_init_mode, scan_layers=True, lora=None), + options=self.options, + ) + self.assertTrue(manager.save(1, bundle, optimizer, force=True)) + manager.wait_until_finished() + manager.close() + + metadata = ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).metadata(epath.Path(self.test_dir) / "1" / "items") + tree = getattr(metadata.item_metadata, "tree", metadata.item_metadata) + return list(ocp.tree.to_flat_dict(tree, sep="/")) + + def test_saves_the_student_in_maxtext_layout(self): + keys = self._save() + + self.assertTrue(any(k.startswith("params/params/layer/") for k in keys), keys) + self.assertTrue(any(k.startswith("opt_state/") for k in keys), keys) + # The bundle's wrapper level and the teacher stay out of the checkpoint. + self.assertEqual([k for k in keys if "student_model" in k.split("/") or "teacher_model" in k.split("/")], []) + + def test_the_students_scan_setting_is_recorded(self): + """Distillation checkpoints the student, so the metadata has to describe the student.""" + student, teacher = DummyModel(nnx.Rngs(0)), DummyModel(nnx.Rngs(1)) + bundle = self.Bundle(student, teacher) + optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param) + manager = distillation_utils.MaxTextCheckpointManager( + raw_iterator=None, + root_directory=self.test_dir, + student_config=SimpleNamespace(learn_to_init_mode=False, scan_layers=False, lora=None), + options=self.options, + ) + manager.save(1, bundle, optimizer, force=True) + manager.wait_until_finished() + manager.close() + + metadata = checkpointing.load_checkpoint_metadata(os.path.join(self.test_dir, "1", "items")) + self.assertIs(metadata.get("scan_layers"), False) + + def test_learn_to_init_mode_leaves_the_optimizer_out(self): + keys = self._save(learn_to_init_mode=True) + + self.assertTrue(any(k.startswith("params/params/layer/") for k in keys), keys) + self.assertEqual([k for k in keys if k.startswith("opt_state")], []) + + def test_restores_the_student(self): + student, teacher = DummyModel(nnx.Rngs(0)), DummyModel(nnx.Rngs(1)) + bundle = self.Bundle(student, teacher) + optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param) + optimizer.update(student, jax.tree.map(jnp.ones_like, nnx.state(student, nnx.Param))) + trained = jnp.asarray(student.layer.kernel[...]) + + def manager(): + return distillation_utils.MaxTextCheckpointManager( + raw_iterator=None, + root_directory=self.test_dir, + student_config=SimpleNamespace(learn_to_init_mode=False), + options=self.options, + ) + + saver = manager() + saver.save(1, bundle, optimizer, force=True) + saver.wait_until_finished() + saver.close() + + fresh_student = DummyModel(nnx.Rngs(0)) + fresh_bundle = self.Bundle(fresh_student, teacher) + fresh_optimizer = nnx.Optimizer(fresh_student, optax.adamw(1e-3), wrt=nnx.Param) + restorer = manager() + step, _ = restorer.maybe_restore(fresh_bundle, fresh_optimizer) + restorer.close() + + self.assertEqual(step, 1) + self.assertTrue(jnp.array_equal(trained, fresh_student.layer.kernel[...])) + + if __name__ == "__main__": absltest.main() diff --git a/tests/post_training/unit/post_train_checkpointing_test.py b/tests/post_training/unit/post_train_checkpointing_test.py new file mode 100644 index 0000000000..d5d0eaa45f --- /dev/null +++ b/tests/post_training/unit/post_train_checkpointing_test.py @@ -0,0 +1,576 @@ +# Copyright 2023-2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for post-training checkpointing in MaxText's on-disk layout.""" + +import os +import tempfile +from types import SimpleNamespace +import unittest +from unittest import mock + +from etils import epath +from flax import nnx +import jax +import jax.numpy as jnp +import optax +import orbax.checkpoint as ocp +import pytest +from tunix.sft import checkpoint_manager as tunix_checkpoint_manager + +from maxtext.common import checkpointing +from maxtext.trainers.post_train import checkpointing as post_train_checkpointing + +pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only] + + +class _Model(nnx.Module): + """Tiny stand-in for a MaxText Transformer.""" + + def __init__(self, rngs: nnx.Rngs): + self.linear = nnx.Linear(2, 3, rngs=rngs) + self.dropout = nnx.Dropout(rate=0.1, rngs=rngs) + + +class _Adapter(nnx.Module): + """Stand-in for TunixMaxTextAdapter: holds the model as its only child.""" + + def __init__(self, base): + self.base = base + + +class _ScannedModel(nnx.Module): + """scan_layers=True: every decoder layer stacked under one `layers` key.""" + + def __init__(self, rngs: nnx.Rngs, num_layers=3): + self.layers = nnx.Param(jnp.zeros((num_layers, 2, 3))) + self.decoder_norm = nnx.Param(jnp.ones((3,))) + self.dropout = nnx.Dropout(rate=0.1, rngs=rngs) + + +class _UnscannedModel(nnx.Module): + """scan_layers=False: one key per decoder layer, each without the stacking axis.""" + + def __init__(self, rngs: nnx.Rngs, num_layers=3): + for i in range(num_layers): + setattr(self, f"layers_{i}", nnx.Linear(2, 3, rngs=rngs)) + self.decoder_norm = nnx.Param(jnp.ones((3,))) + self.dropout = nnx.Dropout(rate=0.1, rngs=rngs) + + +def _build(wrapped): + model = _Model(nnx.Rngs(0)) + outer = _Adapter(model) if wrapped else model + optimizer = nnx.Optimizer(outer, optax.adamw(1e-3), wrt=nnx.Param) + return outer, optimizer + + +def _on_disk_keys(directory, step=1): + """Returns the checkpoint's leaf paths, slash-separated. + + Args: + directory: Checkpoint root directory. + step: Step to read. + + Returns: + A list of leaf paths. + """ + metadata = ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).metadata(epath.Path(directory) / str(step) / "items") + tree = getattr(metadata.item_metadata, "tree", metadata.item_metadata) + return list(ocp.tree.to_flat_dict(tree, sep="/")) + + +def _train_a_step(model, optimizer): + """Moves the weights and fills opt_state, so a restore has something to prove.""" + target = model.base if isinstance(model, _Adapter) else model + grads = jax.tree.map(lambda p: jnp.full_like(p, 0.1), nnx.state(model, nnx.Param)) + optimizer.update(model, grads) + return jnp.asarray(target.linear.kernel[...]) + + +class PostTrainCheckpointInjectHyperparamsTest(unittest.TestCase): + """RL and distillation wrap the optimizer in inject_hyperparams; the layout must not notice.""" + + def test_opt_state_matches_an_unwrapped_optimizer(self): + """mu and nu belong under the Linen `params` collection either way. + + The conversion finds them by name at the top of the optimizer state. Behind the + inject_hyperparams shell they are not there, so stripping it afterwards leaves them + unwrapped and pre-training cannot line the optimizer up. + """ + + def mu_keys(directory, optimizer): + model = _Model(nnx.Rngs(0)) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=directory, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + self.assertTrue(manager.save(1, model, optimizer, force=True)) + manager.close() + keys = _on_disk_keys(directory) + return sorted({k.split("/")[3] for k in keys if k.startswith("opt_state/0/mu/")}) + + with tempfile.TemporaryDirectory() as plain_dir: # pylint: disable=consider-using-with + plain_model = _Model(nnx.Rngs(0)) + plain = mu_keys(plain_dir, nnx.Optimizer(plain_model, optax.adamw(1e-3), wrt=nnx.Param)) + + with tempfile.TemporaryDirectory() as injected_dir: # pylint: disable=consider-using-with + injected_model = _Model(nnx.Rngs(0)) + # A schedule, as RL and distillation use. A plain float produces a different shell -- + # optax only adds hyperparams_states for callables -- and would not reproduce this. + schedule = optax.constant_schedule(1e-3) + tx = optax.inject_hyperparams(optax.adamw)(learning_rate=schedule) + injected = mu_keys(injected_dir, nnx.Optimizer(injected_model, tx, wrt=nnx.Param)) + + self.assertEqual(plain, ["params"], "an unwrapped optimizer should already be under params") + self.assertEqual(injected, plain, "inject_hyperparams changed where mu landed on disk") + + +class PostTrainCheckpointLayoutTest(unittest.TestCase): + """The on-disk layout has to be MaxText's, so pre-training can read what post-training wrote.""" + + def _save(self, directory, wrapped): + model, optimizer = _build(wrapped) + trained = _train_a_step(model, optimizer) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=directory, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + self.assertTrue(manager.save(1, model, optimizer, force=True)) + manager.close() + return trained + + def test_saves_maxtext_layout_not_the_tunix_one(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + self._save(d, wrapped=False) + self.assertEqual(sorted(os.listdir(os.path.join(d, "1"))), ["_CHECKPOINT_METADATA", "items"]) + keys = _on_disk_keys(d) + + self.assertTrue(any(k.startswith("params/params/linear/") for k in keys), keys) + self.assertTrue(any(k.startswith("opt_state/") for k in keys), keys) + self.assertIn("step", keys) + # rngs are NNX-only, so they belong in nnx_aux rather than the Linen collections. + self.assertTrue(any(k.startswith("nnx_aux/") for k in keys), keys) + + def test_adapter_level_is_stripped_from_weights_and_optimizer(self): + """DPO and RL train through the adapter; its `base` level must not reach the checkpoint. + + Pre-training builds its params and opt_state from the bare model, so a stray `base` level + puts every weight at a path it will not look for. + """ + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + self._save(d, wrapped=True) + keys = _on_disk_keys(d) + + self.assertTrue(keys) + self.assertEqual([k for k in keys if "base" in k.split("/")], []) + + def test_restores_weights_and_optimizer_it_saved(self): + for wrapped in (False, True): + with self.subTest(wrapped=wrapped): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + trained = self._save(d, wrapped=wrapped) + + model, optimizer = _build(wrapped) + target = model.base if wrapped else model + self.assertFalse(jnp.array_equal(trained, target.linear.kernel[...]), "fresh model should differ") + + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + step, _ = manager.maybe_restore(model, optimizer) + manager.close() + + self.assertEqual(step, 1) + self.assertTrue(jnp.array_equal(trained, target.linear.kernel[...]), "weights were not restored") + # A resume that dropped opt_state would silently restart the optimizer's moments. + opt_leaves = jax.tree.leaves(nnx.state(optimizer, nnx.optimizer.OptState)) + self.assertTrue(any(jnp.any(jnp.asarray(leaf) != 0) for leaf in opt_leaves), "opt_state came back empty") + + def test_no_checkpoint_yet_reports_step_zero(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + step, metadata = manager.maybe_restore(model, optimizer) + manager.close() + + self.assertEqual(step, 0) + self.assertEqual(metadata, {}) + + def test_custom_metadata_survives_the_round_trip(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + manager.save(1, model, optimizer, force=True, custom_metadata={"run": "abc"}) + manager.wait_until_finished() + _, metadata = manager.maybe_restore(*_build(wrapped=False)) + manager.close() + + self.assertEqual(metadata.get("run"), "abc") + + def test_config_metadata_is_stamped_like_pre_training_does(self): + """`scan_layers` and the LoRA settings have to ride along for the loaders to read them back.""" + config = SimpleNamespace(scan_layers=False, lora=SimpleNamespace(lora_rank=8, model_dump=lambda: {"lora_rank": 8})) + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1), config=config + ) + manager.save(1, model, optimizer, force=True) + manager.wait_until_finished() + _, metadata = manager.maybe_restore(*_build(wrapped=False)) + manager.close() + + self.assertIs(metadata.get("scan_layers"), False) + self.assertEqual(metadata.get("lora"), {"lora_rank": 8}) + + def test_a_caller_key_wins_over_the_config_derived_one(self): + config = SimpleNamespace(scan_layers=True, lora=None) + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1), config=config + ) + manager.save(1, model, optimizer, force=True, custom_metadata={"scan_layers": False, "run": "abc"}) + manager.wait_until_finished() + _, metadata = manager.maybe_restore(*_build(wrapped=False)) + manager.close() + + self.assertIs(metadata.get("scan_layers"), False) + self.assertEqual(metadata.get("run"), "abc") + + def test_no_config_still_saves(self): + """The config is optional; a manager built without one just writes no sidecar metadata.""" + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + self.assertTrue(manager.save(1, model, optimizer, force=True)) + manager.wait_until_finished() + _, metadata = manager.maybe_restore(*_build(wrapped=False)) + manager.close() + + self.assertEqual(metadata, {}) + + def test_weights_only_when_there_is_no_optimizer(self): + """Some callers checkpoint the model alone; opt_state must simply be absent.""" + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, _ = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + self.assertTrue(manager.save(1, model, optimizer=None, force=True)) + manager.close() + keys = _on_disk_keys(d) + + self.assertTrue(any(k.startswith("params/params/linear/") for k in keys), keys) + self.assertEqual([k for k in keys if k.startswith("opt_state")], []) + + +class PostTrainCheckpointScanLayoutTest(unittest.TestCase): + """Both scan settings have to survive the layout conversion. + + Scanned stacks the decoder layers under one `layers` key; unscanned splits them into + `layers_0 … layers_N`. The RL scripts ship both and vLLM requires unscanned. + """ + + def _round_trip(self, build): + """Saves a trained model and restores it into a fresh one. + + Args: + build: Callable taking rngs and returning the model to checkpoint. + + Returns: + A tuple of the on-disk leaf paths, the restored step, and the params before and after. + """ + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model = build(nnx.Rngs(0)) + optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param) + grads = jax.tree.map(lambda p: jnp.full_like(p, 0.1), nnx.state(model, nnx.Param)) + optimizer.update(model, grads) + trained = jax.tree.leaves(nnx.state(model, nnx.Param)) + + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + self.assertTrue(manager.save(1, model, optimizer, force=True)) + manager.wait_until_finished() + keys = _on_disk_keys(d) + + restored_model = build(nnx.Rngs(0)) + restored_optimizer = nnx.Optimizer(restored_model, optax.adamw(1e-3), wrt=nnx.Param) + step, _ = manager.maybe_restore(restored_model, restored_optimizer) + manager.close() + + restored = jax.tree.leaves(nnx.state(restored_model, nnx.Param)) + return keys, step, trained, restored + + def test_scanned_layers_round_trip(self): + keys, step, trained, restored = self._round_trip(_ScannedModel) + self.assertEqual(step, 1) + self.assertIn("params/params/layers", keys) + self.assertEqual([k for k in keys if k.startswith("params/params/layers_")], []) + for want, got in zip(trained, restored): + self.assertTrue(jnp.array_equal(want, got)) + + def test_unscanned_layers_round_trip(self): + keys, step, trained, restored = self._round_trip(_UnscannedModel) + self.assertEqual(step, 1) + layer_keys = {k.split("/")[2] for k in keys if k.startswith("params/params/layers_")} + self.assertEqual(layer_keys, {"layers_0", "layers_1", "layers_2"}) + for want, got in zip(trained, restored): + self.assertTrue(jnp.array_equal(want, got)) + + def test_the_two_layouts_are_actually_different_on_disk(self): + """Guards the test itself: if both models wrote the same keys, neither case would prove much.""" + scanned, _, _, _ = self._round_trip(_ScannedModel) + unscanned, _, _, _ = self._round_trip(_UnscannedModel) + self.assertNotEqual(sorted(scanned), sorted(unscanned)) + + +class PostTrainCheckpointMetadataReaderTest(unittest.TestCase): + """The metadata has to come back through the function the loaders actually call.""" + + def test_load_checkpoint_metadata_reads_what_the_manager_wrote(self): + config = SimpleNamespace( + scan_layers=False, + lora=SimpleNamespace(lora_rank=8, model_dump=lambda: {"lora_rank": 8, "lora_alpha": 16.0}), + ) + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1), config=config + ) + manager.save(1, model, optimizer, force=True) + manager.wait_until_finished() + manager.close() + + # `verify_and_sync_scan_layers` and `sync_lora_metadata` both read a checkpoint this way. + metadata = checkpointing.load_checkpoint_metadata(os.path.join(d, "1", "items")) + + self.assertIs(metadata.get("scan_layers"), False) + self.assertEqual(metadata.get("lora"), {"lora_rank": 8, "lora_alpha": 16.0}) + + +class PostTrainCheckpointSaveDecisionTest(unittest.TestCase): + """Saving has to honour the same enable/interval rules the Tunix manager applied.""" + + def test_disabled_when_there_is_no_root_directory(self): + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(root_directory=None) + + self.assertFalse(manager.save(1, model, optimizer, force=True)) + self.assertEqual(manager.maybe_restore(model, optimizer), (0, {})) + + def test_declines_when_the_policy_says_not_to_save(self): + """`force` bypasses the policy; without it the manager's decision is honoured.""" + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + with mock.patch.object(manager._checkpoint_manager, "should_save", return_value=False): # pylint: disable=protected-access + declined = manager.save(3, model, optimizer) + forced = manager.save(3, model, optimizer, force=True) + manager.close() + + self.assertFalse(declined) + self.assertTrue(forced) + self.assertEqual(sorted(x for x in os.listdir(d) if x.isdigit()), ["3"]) + + +class PostTrainCheckpointLegacyLayoutTest(unittest.TestCase): + """Checkpoints written before the layout change are still in Tunix's, and must still restore.""" + + def test_falls_back_to_the_tunix_layout(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + trained = _train_a_step(model, optimizer) + + legacy = tunix_checkpoint_manager.CheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + legacy.save(1, model, optimizer, force=True) + legacy.close() + + fresh_model, fresh_optimizer = _build(wrapped=False) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + step, _ = manager.maybe_restore(fresh_model, fresh_optimizer) + manager.close() + + self.assertEqual(step, 1) + self.assertTrue(jnp.array_equal(trained, fresh_model.linear.kernel[...])) + + +class PostTrainCheckpointSubclassHookTest(unittest.TestCase): + """Distillation checkpoints a sub-module and an extra item, through these hooks.""" + + class _Bundle(nnx.Module): + + def __init__(self, student): + self.student_model = student + + class _Manager(post_train_checkpointing.MaxTextLayoutCheckpointManager): + """Stand-in for the distillation manager: checkpoints a sub-module plus an extra item.""" + + def __init__(self, root_directory, options): + super().__init__( + root_directory=root_directory, + options=options, + extra_item_handlers={"note": ocp.JsonCheckpointHandler()}, + ) + + def model_to_checkpoint(self, model): + return getattr(model, "student_model", model) + + def _extra_save_args(self, step): + del step + return {"note": ocp.args.JsonSave({"hello": "world"})} + + def test_hooks_pick_the_submodule_and_add_the_extra_item(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + student = _Model(nnx.Rngs(0)) + bundle = self._Bundle(student) + optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param) + manager = self._Manager(d, ocp.CheckpointManagerOptions(save_interval_steps=1)) + self.assertTrue(manager.save(1, bundle, optimizer, force=True)) + manager.close() + + self.assertIn("note", os.listdir(os.path.join(d, "1"))) + metadata = ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).metadata(epath.Path(d) / "1" / "items") + tree = getattr(metadata.item_metadata, "tree", metadata.item_metadata) + keys = list(ocp.tree.to_flat_dict(tree, sep="/")) + + # The student's weights, not the bundle's wrapper level. + self.assertTrue(any(k.startswith("params/params/linear/") for k in keys), keys) + self.assertEqual([k for k in keys if "student_model" in k.split("/")], []) + + +class PostTrainCheckpointBaseManagerTest(unittest.TestCase): + """The base class builds a manager over Tunix's item names before we replace it.""" + + def test_closes_the_base_class_manager_it_replaces(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + mock_base_cm = mock.MagicMock() + + def fake_base_init(self, root_directory=None, options=None): + del root_directory, options + self._checkpoint_manager = mock_base_cm + + with mock.patch.object(tunix_checkpoint_manager.CheckpointManager, "__init__", fake_base_init): + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, + options=ocp.CheckpointManagerOptions(save_interval_steps=1), + ) + mock_base_cm.close.assert_called_once() + # pylint: disable=protected-access + self.assertIsNotNone(manager._checkpoint_manager) + self.assertIsNot(manager._checkpoint_manager, mock_base_cm) + # pylint: enable=protected-access + manager.close() + + +class InstallTest(unittest.TestCase): + """`install` swaps in the MaxText-layout manager and restores what it finds.""" + + class _FakeConfig: + + def __init__(self): + self.checkpointing_options = ocp.CheckpointManagerOptions(save_interval_steps=1) + + def get_with_default(self, key, default): + del key + return default + + class _FakeTrainer: + + def __init__(self, model, optimizer, checkpoint_manager): + self.model = model + self.optimizer = optimizer + self.checkpoint_manager = checkpoint_manager + self.config = InstallTest._FakeConfig() + self._train_steps = 0 + self._iter_steps = 0 + + def test_replaces_the_manager_and_restores_the_step(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + trained = _train_a_step(model, optimizer) + manager = post_train_checkpointing.MaxTextLayoutCheckpointManager( + root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1) + ) + manager.save(4, model, optimizer, force=True) + manager.close() + + fresh_model, fresh_optimizer = _build(wrapped=False) + trainer = self._FakeTrainer(fresh_model, fresh_optimizer, checkpoint_manager=None) + post_train_checkpointing.install(trainer, d) + trainer.checkpoint_manager.close() + + self.assertIsInstance(trainer.checkpoint_manager, post_train_checkpointing.MaxTextLayoutCheckpointManager) + self.assertEqual(trainer._train_steps, 4) # pylint: disable=protected-access + self.assertEqual(trainer._iter_steps, 4) # pylint: disable=protected-access + self.assertTrue(jnp.array_equal(trained, fresh_model.linear.kernel[...])) + + def test_forwards_the_run_config_to_the_manager(self): + """The manager reads the config for the metadata it stamps, so `install` has to pass it on.""" + config = SimpleNamespace(scan_layers=False, lora=None) + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + trainer = self._FakeTrainer(*_build(wrapped=False), checkpoint_manager=None) + post_train_checkpointing.install(trainer, d, config) + trainer.checkpoint_manager.save(1, trainer.model, trainer.optimizer, force=True) + trainer.checkpoint_manager.wait_until_finished() + _, metadata = trainer.checkpoint_manager.maybe_restore(*_build(wrapped=False)) + trainer.checkpoint_manager.close() + + self.assertIs(metadata.get("scan_layers"), False) + + def test_closes_the_manager_it_replaces(self): + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + model, optimizer = _build(wrapped=False) + replaced = tunix_checkpoint_manager.CheckpointManager(root_directory=None) + closed = [] + replaced.close = lambda: closed.append(True) + + trainer = self._FakeTrainer(model, optimizer, checkpoint_manager=replaced) + post_train_checkpointing.install(trainer, d) + trainer.checkpoint_manager.close() + + self.assertEqual(closed, [True]) + + +class UnwrapModelTest(unittest.TestCase): + """The adapter is matched on its child module, not on its class.""" + + def test_unwraps_a_wrapper(self): + model = _Model(nnx.Rngs(0)) + self.assertIs(post_train_checkpointing.unwrap_model(_Adapter(model)), model) + + def test_leaves_a_bare_model_alone(self): + model = _Model(nnx.Rngs(0)) + self.assertIs(post_train_checkpointing.unwrap_model(model), model) + + def test_ignores_a_base_attribute_that_is_not_a_module(self): + model = _Model(nnx.Rngs(0)) + setattr(model, "base", "not a module") + self.assertIs(post_train_checkpointing.unwrap_model(model), model) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py index 845014db32..ed77fc4fbc 100644 --- a/tests/post_training/unit/train_distill_test.py +++ b/tests/post_training/unit/train_distill_test.py @@ -95,6 +95,49 @@ def test_maxtext_to_tunix_iterator(self): expected_mask = dummy_batch["inputs_segmentation"] != 0 np.testing.assert_array_equal(tunix_input.input_mask, expected_mask) + def test_maxtext_to_tunix_iterator_stops_at_the_batch_budget(self): + """The trainer is managed externally, so this bound is what ends an endless dataset.""" + + def endless(): + while True: + yield { + "inputs": np.array([[10, 11]]), + "inputs_position": np.array([[0, 1]]), + "inputs_segmentation": np.array([[1, 1]]), + "targets": np.array([[11, 12]]), + } + + adapter = distillation_utils.MaxTextToTunixIterator(endless(), max_batches=3) + self.assertEqual(len(list(adapter)), 3) + + def test_maxtext_to_tunix_iterator_is_unbounded_without_a_budget(self): + """A finite upstream iterator still ends on its own.""" + batches = [ + { + "inputs": np.array([[10, 11]]), + "inputs_position": np.array([[0, 1]]), + "inputs_segmentation": np.array([[1, 1]]), + "targets": np.array([[11, 12]]), + } + ] * 4 + adapter = distillation_utils.MaxTextToTunixIterator(iter(batches)) + self.assertEqual(len(list(adapter)), 4) + + def test_maxtext_to_tunix_iterator_spends_nothing_on_a_zero_budget(self): + """A resume that is already at its step count must not train further.""" + + def endless(): + while True: + yield { + "inputs": np.array([[10, 11]]), + "inputs_position": np.array([[0, 1]]), + "inputs_segmentation": np.array([[1, 1]]), + "targets": np.array([[11, 12]]), + } + + adapter = distillation_utils.MaxTextToTunixIterator(endless(), max_batches=0) + self.assertEqual(len(list(adapter)), 0) + def test_maxtext_to_tunix_iterator_sft(self): """Verifies SFT-related fields are handled correctly.""" # 1. Create a dummy batch with SFT fields @@ -589,6 +632,8 @@ def test_setup_pipeline_grain_enabled(self): mock_iter.save = mock.Mock() # Has save method config = mock.Mock() + config.post_train_skip_checkpointing = False + config.checkpoint_storage_concurrent_gb = None config.dataset_type = "grain" config.checkpoint_dir = self.test_dir @@ -634,6 +679,8 @@ def test_setup_pipeline_restored(self): mock_iter = mock.Mock() mock_iter.save = mock.Mock() config = mock.Mock() + config.post_train_skip_checkpointing = False + config.checkpoint_storage_concurrent_gb = None config.dataset_type = "grain" config.checkpoint_dir = self.test_dir @@ -721,6 +768,8 @@ def test_setup_pipeline_disabled(self): mock_iter = object() # No save method config = mock.Mock() + config.post_train_skip_checkpointing = False + config.checkpoint_storage_concurrent_gb = None config.dataset_type = "synthetic" # Not grain config.checkpoint_dir = self.test_dir @@ -960,6 +1009,12 @@ def __call__(self, input_tokens, **kwargs): config.checkpoint_dir = self.test_dir config.dataset_type = "synthetic" config.lora_enabled = False + # The checkpoint manager reads these for the metadata it stamps, and a bare Mock puts a Mock + # where a bool belongs. + config.scan_layers = True + config.lora = None + config.post_train_skip_checkpointing = False + config.checkpoint_storage_concurrent_gb = None # pylint: disable=import-outside-toplevel from tunix.sft import peft_trainer @@ -967,7 +1022,12 @@ def __call__(self, input_tokens, **kwargs): train_config = peft_trainer.TrainingConfig( max_steps=2, eval_every_n_steps=0, - checkpointing_options=ocp.CheckpointManagerOptions(save_interval_steps=1, max_to_keep=2, create=True), + checkpointing_options=ocp.CheckpointManagerOptions( + save_interval_steps=1, + max_to_keep=2, + create=True, + enable_async_checkpointing=False, + ), gradient_accumulation_steps=1, ) diff --git a/tests/post_training/unit/train_dpo_test.py b/tests/post_training/unit/train_dpo_test.py index 489d3ff9ed..9bf51467ee 100644 --- a/tests/post_training/unit/train_dpo_test.py +++ b/tests/post_training/unit/train_dpo_test.py @@ -51,5 +51,44 @@ def test_validate_config_invalid_vocab_tiling(self): train_dpo.validate_config(config) +class TrainDPOTunixConfigTest(unittest.TestCase): + """The Tunix config decides who checkpoints and whether the optimizer gets wrapped.""" + + def _mt_config(self, grad_accum=1): + return SimpleNamespace( + checkpoint_period=5, + async_checkpointing=False, + tensorboard_dir="/tmp/tb", + profiler="", + eval_interval=1, + steps=10, + checkpoint_dir="/tmp/ckpt", + data_sharding=["data"], + gradient_accumulation_steps=grad_accum, + max_target_length=128, + dpo=SimpleNamespace( + algo="dpo", + orpo_lambda=1.0, + dpo_beta=0.1, + dpo_label_smoothing=0.0, + max_prompt_length=32, + ), + ) + + @pytest.mark.cpu_only + def test_tunix_checkpointing_is_disabled(self): + """post_train.checkpointing owns checkpointing, so Tunix's own manager must stay off.""" + self.assertIsNone(train_dpo.get_tunix_config(self._mt_config()).checkpoint_root_directory) + + @pytest.mark.cpu_only + def test_single_step_accumulation_is_not_passed_through(self): + """Tunix wraps the optimizer in MultiSteps whenever this is set, changing the state shape.""" + self.assertIsNone(train_dpo.get_tunix_config(self._mt_config(grad_accum=1)).gradient_accumulation_steps) + + @pytest.mark.cpu_only + def test_real_accumulation_is_passed_through(self): + self.assertEqual(train_dpo.get_tunix_config(self._mt_config(grad_accum=4)).gradient_accumulation_steps, 4) + + if __name__ == "__main__": unittest.main() diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py index 41daec72e9..bec71c192e 100644 --- a/tests/post_training/unit/train_rl_test.py +++ b/tests/post_training/unit/train_rl_test.py @@ -20,8 +20,11 @@ import pytest from types import SimpleNamespace import jax +import jax.numpy as jnp +import optax from maxtext.trainers.post_train.rl import train_rl +from maxtext.trainers.post_train.rl import utils_rl pytestmark = [pytest.mark.post_training] from maxtext.configs import types @@ -342,13 +345,13 @@ def test_prompt_filtering(self): mock_tokenizer = mock.MagicMock() # Define tokenizer side effect - def tokenize_side_effect(text): + def encode_side_effect(text): if text == "short": return [0] * 5 else: return [0] * 15 - mock_tokenizer.tokenize.side_effect = tokenize_side_effect + mock_tokenizer.encode.side_effect = encode_side_effect # Define dataset mock data train_data = [ @@ -662,5 +665,43 @@ def apply_chat_template(self, conversation, tokenize=False): self.assertEqual(rendered, "Hello!") +class RLOptimizerClippingTest(unittest.TestCase): + """RL clips gradients without giving the optimizer state an extra chain level.""" + + def _config(self, threshold): + return SimpleNamespace( + learning_rate=1e-3, + learning_rate_schedule_steps=-1, + steps=-1, + train_steps=10, + # No warmup, so the learning rate at step 0 is non-zero and updates are comparable. + warmup_steps_fraction=0.0, + adam_b1=0.9, + adam_b2=0.95, + adam_weight_decay=0.1, + gradient_clipping_threshold=threshold, + ) + + def _opt_state_structure(self, threshold): + params = {"w": jnp.array([3.0, 4.0])} + return jax.tree_util.tree_structure(utils_rl.get_optimizer(self._config(threshold)).init(params)) + + @pytest.mark.cpu_only + def test_clipping_does_not_change_the_optimizer_state_shape(self): + self.assertEqual(self._opt_state_structure(1.0), self._opt_state_structure(0.0)) + + @pytest.mark.cpu_only + def test_gradients_are_clipped(self): + params = {"w": jnp.array([3.0, 4.0])} # global norm 5.0 + grads = {"w": jnp.array([3.0, 4.0])} + + clipped = utils_rl.get_optimizer(self._config(0.1)) + unclipped = utils_rl.get_optimizer(self._config(0.0)) + clipped_updates, _ = clipped.update(grads, clipped.init(params), params) + unclipped_updates, _ = unclipped.update(grads, unclipped.init(params), params) + + self.assertLess(float(optax.tree.norm(clipped_updates)), float(optax.tree.norm(unclipped_updates))) + + if __name__ == "__main__": unittest.main() diff --git a/tests/post_training/unit/train_sft_test.py b/tests/post_training/unit/train_sft_test.py index 3e71ba6273..72d4d07e4f 100644 --- a/tests/post_training/unit/train_sft_test.py +++ b/tests/post_training/unit/train_sft_test.py @@ -104,5 +104,35 @@ def test_maxtext_peft_trainer_train_step_signature(self): self.assertEqual(params, ["model", "optimizer", "grad_accumulator", "inputs", "is_update_step"]) +class TrainSFTTunixConfigTest(unittest.TestCase): + """The Tunix config decides who checkpoints and whether the optimizer gets wrapped.""" + + def _mt_config(self, grad_accum=1): + return SimpleNamespace( + checkpoint_period=5, + async_checkpointing=False, + tensorboard_dir="/tmp/tb", + profiler="", + eval_interval=1, + steps=10, + checkpoint_dir="/tmp/ckpt", + data_sharding=["data"], + gradient_accumulation_steps=grad_accum, + ) + + @pytest.mark.cpu_only + def test_tunix_checkpointing_is_disabled(self): + """post_train.checkpointing owns checkpointing, so Tunix's own manager must stay off.""" + self.assertIsNone(train_sft.get_tunix_config(self._mt_config()).checkpoint_root_directory) + + @pytest.mark.cpu_only + def test_single_step_accumulation_is_not_passed_through(self): + self.assertIsNone(train_sft.get_tunix_config(self._mt_config(grad_accum=1)).gradient_accumulation_steps) + + @pytest.mark.cpu_only + def test_real_accumulation_is_passed_through(self): + self.assertEqual(train_sft.get_tunix_config(self._mt_config(grad_accum=4)).gradient_accumulation_steps, 4) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py index 0c2dae51ff..976d82fc5a 100644 --- a/tests/unit/checkpointing_nnx_load_test.py +++ b/tests/unit/checkpointing_nnx_load_test.py @@ -525,7 +525,7 @@ def __init__(self, rngs): self.assertIn("table", ckpt["nnx_aux"]["model"]) # the custom variable persists here instead self.assertNotIn("cache", ckpt["nnx_aux"]["model"]) # caches are recomputed, never stored - restored = checkpointing._linen_items_to_nnx(ckpt, nnx.eval_shape(lambda: state)) # pylint: disable=protected-access + restored = checkpointing.linen_items_to_nnx(ckpt, nnx.eval_shape(lambda: state)) pure = restored.to_pure_dict() self.assertTrue(jnp.array_equal(pure["model"]["table"], model.table.value)) self.assertIsInstance(pure["model"]["cache"], jax.ShapeDtypeStruct) # left for the init to fill @@ -595,11 +595,11 @@ def test_no_nnx_aux_when_state_has_none(self): class TestLinenItemsToNnx(unittest.TestCase): - """checkpointing._linen_items_to_nnx reshapes restored items into the NNX-layout overlay.""" + """checkpointing.linen_items_to_nnx reshapes restored items into the NNX-layout overlay.""" def _to_nnx(self, restored): """Reshape `restored` against the `_ModelDropout` abstract, as the restore paths do.""" - state = checkpointing._linen_items_to_nnx(restored, _abstract_dropout_state()) # pylint: disable=protected-access + state = checkpointing.linen_items_to_nnx(restored, _abstract_dropout_state()) return state.to_pure_dict() def test_materialized_aux_is_kept(self): @@ -723,5 +723,79 @@ def test_expected_and_restored_params_splits_by_param_type(self): checkpointing._raise_on_weight_mismatch(want, have) # pylint: disable=protected-access +class TestLoadNnxNativeParams(unittest.TestCase): + """Weight-only load of an NNX-native checkpoint -- the layout post-training writes. + + Tunix and the training engine save `nnx.state(model)` under a `model_params` item, so each + leaf lands in a `value` box and the tree is the whole checkpoint. DPO and RL train through + `TunixMaxTextAdapter`, adding a `base` level on top. + """ + + def _save(self, directory, weights, wrapper_key=None): + """Writes `weights` the way `nnx.state(model)` serializes, under a `model_params` dir.""" + boxed = jax.tree.map(lambda leaf: {"value": leaf}, weights) + path = os.path.join(directory, "model_params") + ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save( + epath.Path(path), + {wrapper_key: boxed} if wrapper_key else boxed, + force=True, + ) + return path + + def _weights(self): + return { + "linear": { + "kernel": jnp.arange(2, dtype=jnp.float32).reshape(2, 1), + "bias": jnp.array([5.0]), + } + } + + def _restore(self, path): + # The real caller hands over an abstract params state, as load_state_if_possible does. + params_abstract = nnx.eval_shape(lambda: nnx.split(_Model(nnx.Rngs(0)), nnx.Param, ...)[1]) + return checkpointing.load_params_from_path(path, params_abstract, 8) + + def test_restores_weights_saved_by_a_bare_model(self): + """SFT, distillation and the training engine save the model directly, with no wrapper level.""" + weights = self._weights() + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + restored = self._restore(self._save(d, weights)) + + pure = restored.to_pure_dict() + self.assertTrue(jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"])) + self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"])) + + def test_restores_weights_saved_through_the_tunix_adapter(self): + """DPO and RL nest the whole tree under `base`; the weights still have to come back.""" + weights = self._weights() + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + restored = self._restore(self._save(d, weights, wrapper_key="base")) + + pure = restored.to_pure_dict() + self.assertTrue(jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"])) + self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"])) + + def test_linen_target_against_an_nnx_checkpoint_raises(self): + """This layout only restores into an NNX state, so say so instead of failing inside Orbax.""" + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + path = self._save(d, self._weights()) + linen_params = jax.tree.map(jnp.zeros_like, {"params": self._weights()}) + with self.assertRaises(ValueError) as ctx: + checkpointing.load_params_from_path(path, linen_params, 8) + + self.assertIn("NNX", str(ctx.exception)) + + def test_weight_missing_from_the_checkpoint_raises(self): + """A params-only load has no init state to fall back on, so a gap must not pass silently.""" + weights = self._weights() + del weights["linear"]["bias"] + with tempfile.TemporaryDirectory() as d: # pylint: disable=consider-using-with + path = self._save(d, weights) + with self.assertRaises(ValueError) as ctx: + self._restore(path) + + self.assertIn("linear/bias", str(ctx.exception)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index 4ef2ab30b7..3f1a213914 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -17,6 +17,7 @@ import asyncio import json import os +from types import SimpleNamespace from unittest import mock from absl.testing import absltest @@ -319,6 +320,32 @@ def test_load_checkpoint_metadata(self, mock_checkpointer_cls): self.assertEqual(loaded_metadata.get("lora"), {"lora_rank": 8, "lora_alpha": 16.0}) mock_ckptr.metadata.assert_called_once() + @mock.patch.object(checkpointing.ocp, "StandardCheckpointer") + def test_load_checkpoint_metadata_falls_back_to_the_step_directory(self, mock_checkpointer_cls): + """`load_parameters_path` points at the item, but the metadata belongs to the step above it.""" + mock_ckptr = mock_checkpointer_cls.return_value + + def metadata_for(path): + result = mock.MagicMock() + result.custom_metadata = {"scan_layers": False} if str(path).endswith("/0") else None + return result + + mock_ckptr.metadata.side_effect = metadata_for + + self.assertEqual(checkpointing.load_checkpoint_metadata("/ckpt/0/items"), {"scan_layers": False}) + self.assertEqual([str(c.args[0]) for c in mock_ckptr.metadata.call_args_list], ["/ckpt/0/items", "/ckpt/0"]) + + @mock.patch.object(checkpointing.ocp, "StandardCheckpointer") + def test_load_checkpoint_metadata_stops_at_the_path_it_was_given(self, mock_checkpointer_cls): + """A hit at the given path is used as is; no reason to look at the parent.""" + mock_ckptr = mock_checkpointer_cls.return_value + metadata = mock.MagicMock() + metadata.custom_metadata = {"scan_layers": True} + mock_ckptr.metadata.return_value = metadata + + self.assertEqual(checkpointing.load_checkpoint_metadata("/ckpt/0"), {"scan_layers": True}) + mock_ckptr.metadata.assert_called_once() + @mock.patch.object(checkpointing.ocp, "StandardCheckpointer") def test_load_checkpoint_metadata_handles_exceptions(self, mock_checkpointer_cls): mock_ckptr = mock_checkpointer_cls.return_value @@ -326,7 +353,44 @@ def test_load_checkpoint_metadata_handles_exceptions(self, mock_checkpointer_cls loaded_metadata = checkpointing.load_checkpoint_metadata("corrupt/path") self.assertEqual(loaded_metadata, {}) - mock_ckptr.metadata.assert_called_once() + # A failed read falls through to the step directory above, which fails the same way. + self.assertEqual(mock_ckptr.metadata.call_count, 2) + + +class CheckpointCustomMetadataTest(parameterized.TestCase): + """What a checkpoint records about the run that wrote it. + + Loaders read this back to fill in a value the run left at its default, or to reject one that + contradicts the checkpoint. Post-training saves through its own manager and calls the same + builder, so the two paths cannot drift apart. + """ + + def _config(self, scan_layers=True, lora_rank=0): + lora = SimpleNamespace( + lora_rank=lora_rank, + lora_alpha=16.0, + model_dump=lambda: {"lora_rank": lora_rank, "lora_alpha": 16.0}, + ) + return SimpleNamespace(scan_layers=scan_layers, lora=lora) + + @parameterized.parameters(True, False) + def test_scan_layers_is_recorded_either_way(self, scan_layers): + metadata = checkpointing.checkpoint_custom_metadata(self._config(scan_layers=scan_layers)) + self.assertIs(metadata["scan_layers"], scan_layers) + + def test_lora_is_recorded_once_there_is_a_rank(self): + metadata = checkpointing.checkpoint_custom_metadata(self._config(lora_rank=8)) + self.assertEqual(metadata["lora"], {"lora_rank": 8, "lora_alpha": 16.0}) + + def test_no_lora_key_without_a_rank(self): + """A rank of 0 means LoRA is off; recording it would make `sync_lora_metadata` sync a zero.""" + self.assertNotIn("lora", checkpointing.checkpoint_custom_metadata(self._config(lora_rank=0))) + + def test_no_lora_key_when_the_config_has_no_lora_section(self): + self.assertNotIn("lora", checkpointing.checkpoint_custom_metadata(SimpleNamespace(scan_layers=True))) + + def test_no_config_records_nothing(self): + self.assertEqual(checkpointing.checkpoint_custom_metadata(None), {}) class GrainCheckpointableEquivalenceTest(parameterized.TestCase): diff --git a/tests/unit/optimizers_test.py b/tests/unit/optimizers_test.py index 6f43c420cd..bf45130f56 100644 --- a/tests/unit/optimizers_test.py +++ b/tests/unit/optimizers_test.py @@ -825,5 +825,85 @@ def test_muon_newton_schulz_config(self): self.assertEqual(kwargs["ns_coeffs"], (3.4445, -4.7750, 2.0315)) +class AddGradientClippingTest(parameterized.TestCase): + """Clipping must stay identical to the chained form while leaving the state tree alone.""" + + THRESHOLD = 1.0 + + def _params(self): + return {"w": jnp.array([3.0, 4.0]), "b": jnp.array([12.0])} + + def _grads(self, scale): + # Global norm of [3, 4, 12] is 13, so scale=1.0 is well over the threshold. + return jax.tree.map(lambda x: x * scale, self._params()) + + @parameterized.named_parameters(("over_threshold", 1.0), ("under_threshold", 0.01)) + def test_updates_match_the_chained_form(self, scale): + params, grads = self._params(), self._grads(scale) + + chained = optax.chain(optax.clip_by_global_norm(self.THRESHOLD), optax.adamw(1e-2)) + chained_updates, _ = chained.update(grads, chained.init(params), params) + + inline = optimizers.add_gradient_clipping(optax.adamw(1e-2), self.THRESHOLD) + inline_updates, _ = inline.update(grads, inline.init(params), params) + + jax.tree.map(lambda a, b: self.assertTrue(jnp.allclose(a, b), f"{a} != {b}"), chained_updates, inline_updates) + + def test_gradients_over_the_threshold_are_clipped(self): + """Guards against the wrapper quietly becoming a passthrough.""" + params, grads = self._params(), self._grads(1.0) + + inline = optimizers.add_gradient_clipping(optax.sgd(1.0), self.THRESHOLD) + inline_updates, _ = inline.update(grads, inline.init(params), params) + unclipped = optax.sgd(1.0) + unclipped_updates, _ = unclipped.update(grads, unclipped.init(params), params) + + inline_norm = optax.tree.norm(inline_updates) + self.assertAlmostEqual(float(inline_norm), self.THRESHOLD, places=5) + self.assertGreater(float(optax.tree.norm(unclipped_updates)), float(inline_norm)) + + def test_state_tree_matches_the_unclipped_optimizer(self): + """A checkpointed opt_state must not gain a level from clipping. + + Pre-training clips in its train step, so its optimizer state is the bare tx state. A + post-training checkpoint one chain level deeper cannot be resumed by it. + """ + params = self._params() + inner = optax.adamw(1e-2) + + bare = jax.tree_util.tree_structure(inner.init(params)) + inline = jax.tree_util.tree_structure(optimizers.add_gradient_clipping(inner, self.THRESHOLD).init(params)) + chained = jax.tree_util.tree_structure(optax.chain(optax.clip_by_global_norm(self.THRESHOLD), inner).init(params)) + + self.assertEqual(bare, inline) + self.assertNotEqual(bare, chained) + + def test_extra_args_reach_the_inner_optimizer(self): + """`get_optimizer` can return a transform taking extra args, as skip_step_on_spikes does.""" + params, grads = self._params(), self._grads(1.0) + seen = {} + + def update_fn(updates, state, params=None, **extra_args): + del params + seen.update(extra_args) + return updates, state + + inner = optax.GradientTransformationExtraArgs(lambda p: optax.EmptyState(), update_fn) + wrapped = optimizers.add_gradient_clipping(inner, self.THRESHOLD) + wrapped.update(grads, wrapped.init(params), params, loss=jnp.array(1.0)) + + self.assertEqual(list(seen), ["loss"]) + + def test_clipping_composes_with_skip_step_on_spikes(self): + """get_optimizer can return a spike-skipping transform, which takes loss/grad_norm.""" + params, grads = self._params(), self._grads(1.0) + + inner = optimizers.skip_step_on_spikes(optax.adamw(1e-2), interval=4, scaling_factor=2.0) + wrapped = optimizers.add_gradient_clipping(inner, self.THRESHOLD) + updates, _ = wrapped.update(grads, wrapped.init(params), params, loss=jnp.array(1.0), grad_norm=jnp.array(1.0)) + + self.assertEqual(jax.tree_util.tree_structure(updates), jax.tree_util.tree_structure(params)) + + if __name__ == "__main__": unittest.main()