From 9ec47abc708b8f7e2f9d260a34e69da0295c59cd Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Sun, 23 Aug 2026 22:51:15 +0000 Subject: [PATCH] Write post-training checkpoints in MaxText's on-disk layout Tunix writes its own checkpoint shape -- model_params/ and optimizer_state/ beside a custom_metadata -- while pre-training expects items/ holding params under the Linen params collection, opt_state, step and nnx_aux. A model trained by SFT, DPO, RL or distillation therefore could not be handed back to pre-training without converting it by hand. MaxTextLayoutCheckpointManager subclasses Tunix's manager and reshapes the tree on the way out: it restores the Linen collection level, strips the inject_hyperparams shell that optax wraps a scheduled optimizer in -- otherwise mu and nu land a level short of where pre-training looks for them -- and drops the adapter level a LoRA wrapper adds. install() swaps it in for the manager the base class built, closing that one so it does not leak its writer thread. Gradient clipping moves out of optax.chain for the same reason: chaining a stateless clip in front of the optimizer nests its state under an extra level that pre-training cannot read. RL keeps its actor/ level, since GRPO checkpoints an actor and a reference separately. The end-to-end scripts, the demo notebooks and the eval README follow the path from model_params to items. Three things had to work before that code was reachable. Post-training reached for transformers' AutoTokenizer directly, so a run had to name a HuggingFace repo even when the model ships a sentencepiece or tiktoken asset in this repo. The tokenizers built by build_tokenizer now serve post-training too, and the two that are ours gain a chat template renderer, since RL applies one and only HuggingFace tokenizers carried that method. The vLLM rollout had three call-site assumptions that no longer hold, each of which stops RL before a single rollout: EngineArgs dropped swap_space, and passing a field it no longer takes is a TypeError rather than something vLLM ignores; tpu_inference mutates data_parallel_size and then deletes the sharding config before calling with_hf_config, leaving the rebuilt config to fail its own device-count assertion; is_init_field raises on the fields tpu_inference injects dynamically. Distillation then ran past its step count, because Tunix stops on max_steps only when it owns the checkpoint manager and installing ours sets is_managed_externally. The iterator carries a batch budget of steps times gradient accumulation, minus whatever a resumed run already consumed, and raises StopIteration once spent. post_train_skip_checkpointing exists for runs that only want to load base weights and write nothing back. --- docs/tutorials/posttraining/lora.md | 2 +- src/maxtext/common/checkpointing.py | 140 +++-- src/maxtext/common/train_state_nnx.py | 16 + src/maxtext/configs/base.yml | 2 + src/maxtext/configs/types.py | 11 + src/maxtext/eval/README.md | 2 +- src/maxtext/examples/rl_llama3_demo.ipynb | 6 +- .../examples/sft_llama3_demo_tpu.ipynb | 2 +- src/maxtext/inference/vllm_decode.py | 18 + .../input_pipeline/hf_data_processing.py | 9 +- .../input_pipeline/input_pipeline_utils.py | 20 +- .../synthetic_data_processing.py | 13 +- src/maxtext/input_pipeline/tokenizer.py | 53 +- .../tunix/weight_mapping/__init__.py | 6 +- .../integration/vllm/maxtext_vllm_rollout.py | 69 +++ src/maxtext/optimizers/optimizers.py | 27 + .../trainers/post_train/checkpointing.py | 452 ++++++++++++++ .../distillation/distillation_utils.py | 147 ++--- .../post_train/distillation/train_distill.py | 91 ++- .../trainers/post_train/dpo/train_dpo.py | 40 +- .../trainers/post_train/rl/train_rl.py | 68 ++- .../trainers/post_train/rl/utils_rl.py | 20 +- .../trainers/post_train/sft/train_sft.py | 11 +- src/maxtext/utils/lora_utils.py | 26 + .../tpu/gemma3/4b/test_gemma3_rl.sh | 2 +- .../tpu/gemma4/26b/test_gemma4_rl.sh | 2 +- .../tpu/llama3.1/70b/test_llama3.1_70b_rl.sh | 2 +- .../end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh | 2 +- .../integration/lora_e2e_nnx_test.py | 4 +- .../unit/distillation_checkpointing_test.py | 116 ++++ .../unit/post_train_checkpointing_test.py | 576 ++++++++++++++++++ .../post_training/unit/train_distill_test.py | 62 +- tests/post_training/unit/train_dpo_test.py | 39 ++ tests/post_training/unit/train_rl_test.py | 45 +- tests/post_training/unit/train_sft_test.py | 30 + tests/unit/checkpointing_nnx_load_test.py | 80 ++- tests/unit/checkpointing_test.py | 66 +- tests/unit/optimizers_test.py | 80 +++ 38 files changed, 2150 insertions(+), 207 deletions(-) create mode 100644 src/maxtext/trainers/post_train/checkpointing.py create mode 100644 tests/post_training/unit/post_train_checkpointing_test.py 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()