From bbb843f1dfb0e50636b7c8d9ed6baf6fb1badaa6 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Thu, 2 Jul 2026 17:00:00 +0000 Subject: [PATCH] [NNX] Delete Linen (5/5): remove the pure_nnx/enable_nnx/pure_nnx_decoder config flags Drop the three migration flags from types.py (the Fields plus the now-dead qwix validator), base.yml, inference/vllm.yml, the distillation configs and pyconfig_deprecated.py, and clear the remaining mentions in the docs, the LoRA demo notebook and the shell scripts. setup_initial_state was the one reader parked by pre-train 1/3, because its Linen branch was entangled with the checkpoint restore overlay. Collapse it here. The Linen arm is unreachable now that get_abstract_state always returns an nnx.State, and the weight_sparsity_n/m merge it carried is already covered by _reshard_aligned, which keeps the freshly initialized value for any leaf the checkpoint did not supply. The trailing unbox_logicallypartioned only unwrapped Linen LogicallyPartitioned boxes, so it goes with it. Removing pure_nnx_decoder also un-skips four train_compile pipeline tests that were guarded on it. Three pass as they are, so drop all four guards. test_pipeline_subset turned up a real porting bug: the NNX decoder tagged the scan axis of the layers outside the pipeline "layers", which base.yml maps to the stage mesh axis, so a remainder that does not divide the stage count failed to shard. Linen used "layers_outside_pipeline", which has no rule; use that. Tests follow: drop the flag from every config dict and SimpleNamespace stand-in that still set it, and delete test_diloco_requires_pure_nnx, which asserted the message of the validator this change removes. --- docs/guides/distillation.md | 6 +- .../posttraining/knowledge_distillation.md | 4 +- docs/tutorials/posttraining/lora.md | 4 - .../posttraining/lora_on_multi_host.md | 4 - .../inspect_checkpoint.py | 24 +-- .../convert_gpt3_ckpt_from_paxml.py | 37 +--- src/maxtext/common/checkpointing.py | 22 +- src/maxtext/common/emergency_checkpointing.py | 2 +- src/maxtext/common/train_state_nnx.py | 4 +- src/maxtext/configs/base.yml | 5 - src/maxtext/configs/inference/vllm.yml | 2 - .../configs/post_train/distillation-qwen3.yml | 1 - .../post_train/distillation_gpt_oss_20b.yml | 1 - .../distillation_qwen3_30b_base.yml | 1 - .../distillation_qwen3_30b_base_pdbs8.yml | 1 - src/maxtext/configs/pyconfig_deprecated.py | 7 +- src/maxtext/configs/types.py | 57 ------ src/maxtext/examples/lora_llama3_demo.ipynb | 2 - .../omni_poc/tests/stitch_checkpoint_test.py | 1 - .../omni_poc/utils/stitch_checkpoint.py | 10 +- src/maxtext/layers/nnx_decoders.py | 2 +- .../run_olmo_qwen3_30b_streaming_diloco.sh | 2 +- .../scripts/run_spmd_streaming_diloco.sh | 1 - .../scripts/run_spmd_streaming_dlco_moe_c4.sh | 1 - .../distillation/scripts/run_distill_xpk.sh | 5 +- .../utils/generate_param_only_checkpoint.py | 156 +++----------- src/maxtext/utils/layerwise_quantization.py | 114 +---------- src/maxtext/utils/lora_utils.py | 5 +- src/maxtext/utils/maxtext_utils.py | 192 +++++++----------- src/maxtext/utils/standalone_checkpointer.py | 47 ++--- .../tpu/gemma3/4b/test_gemma3_lora.sh | 10 +- .../tpu/gemma3/4b/test_gemma3_to_hf.sh | 4 +- tests/inference/benchmark_vllm.sh | 2 +- tests/integration/diloco_test.py | 13 -- tests/integration/lora_e2e_nnx_test.py | 1 - .../inference_router_replay_test.py | 3 - .../integration/lora_e2e_nnx_test.py | 1 - .../post_training/unit/rl_lr_schedule_test.py | 2 - .../unit/checkpointing_emergency_nnx_test.py | 7 +- .../checkpointing_nnx_missing_param_test.py | 3 +- .../unit/checkpointing_nnx_roundtrip_test.py | 3 +- tests/unit/checkpointing_test.py | 1 - tests/unit/decoder_layer_model_mode_test.py | 1 - .../unit/deepseek_decode_consistency_test.py | 1 - ...generate_param_only_checkpoint_nnx_test.py | 3 - tests/unit/multi_token_prediction_test.py | 2 - tests/unit/nnx_decoders_test.py | 33 --- tests/unit/router_replay_test.py | 5 +- tests/unit/setup_initial_state_nnx_test.py | 3 +- .../unit/standalone_checkpointer_nnx_test.py | 26 ++- tests/unit/train_compile_test.py | 21 +- tests/unit/train_state_nnx_checkpoint_test.py | 33 +-- tests/utils/sharding_dump.py | 3 - 53 files changed, 217 insertions(+), 684 deletions(-) diff --git a/docs/guides/distillation.md b/docs/guides/distillation.md index 317521b0fe..d2c7eebfc2 100644 --- a/docs/guides/distillation.md +++ b/docs/guides/distillation.md @@ -63,9 +63,8 @@ Student and teacher must share the same vocabulary. The trainer asserts `student If `distill_beta > 0`, the model `sow`s the attention `out_projection` activations at every layer so the loss can read them. This requires: - `scan_layers: True` — activations are stacked along the leading scan axis; the loss does `jnp.take(features, layer_indices, axis=0)` over that axis. -- `enable_nnx: True` — `sow(nnx.Intermediate, ...)` is an NNX-specific call. -The trainer validates both at config initialization. Logit-only runs (`distill_beta = 0`) have no such constraint. +The trainer validates this at config initialization. Logit-only runs (`distill_beta = 0`) have no such constraint. ## Loss anatomy @@ -221,7 +220,6 @@ distill_feature_loss_type: cosine distill_layer_indices: [3, 7, 11, 15, 19, 23, 27, 31] # for 32-layer student scan_layers: True -enable_nnx: True ``` ### Logit-only baseline (cheapest; no feature extraction overhead) @@ -258,7 +256,7 @@ The trainer logs the following to TensorBoard (configured by `tensorboard_dir`, | Symptom | Likely cause | Fix | | ----------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------- | -| `ValueError: a value of self.distill_beta > 0.0 requires self.scan_layers = True` | Feature loss enabled without scanned layers. | Add `scan_layers=True enable_nnx=True` to your CLI / yml. | +| `ValueError: a value of self.distill_beta > 0.0 requires self.scan_layers = True` | Feature loss enabled without scanned layers. | Add `scan_layers=True` to your CLI / yml. | | `Vocab size mismatch! Student: X, Teacher: Y` | Different tokenizers. | Use teacher and student with the same vocab; the trainer cannot match logits across vocabularies. | | `Teacher model path is missing` | `teacher_overrides.load_parameters_path` not set in non-offline mode. | Set it in `teacher_overrides` in the yml or pass via CLI. | | `Features extracted from student or teacher model are None, but distill_beta > 0.0` | Model architecture doesn't sow `out_projection_activations` (e.g. uses an unsupported attention path). | Verify the attention layer in use sets `self.sow(nnx.Intermediate, "out_projection_activations", out)` (see `attentions.py`). | diff --git a/docs/tutorials/posttraining/knowledge_distillation.md b/docs/tutorials/posttraining/knowledge_distillation.md index 0684eb4c3b..7db716cc0e 100644 --- a/docs/tutorials/posttraining/knowledge_distillation.md +++ b/docs/tutorials/posttraining/knowledge_distillation.md @@ -248,7 +248,7 @@ Key knobs (see the [Distillation guide](../../guides/distillation.md) for the fu ```yaml distill_alpha: 0.5 # weight on KL(teacher||student) distill_temperature: 1.0 -distill_beta: 0.0 # >0 enables feature distillation; requires scan_layers=True, enable_nnx=True +distill_beta: 0.0 # >0 enables feature distillation; requires scan_layers=True distill_layer_indices: None ``` @@ -312,7 +312,7 @@ python3 -m maxtext.trainers.post_train.distillation.train_distill \ distill_temperature=2.0 \ distill_beta=1.0 distill_beta_end=0.1 distill_beta_schedule=cosine \ distill_layer_indices=[2,5,8,11,14,17,20,23] \ - scan_layers=True enable_nnx=True \ + scan_layers=True \ profiler=xplane ``` diff --git a/docs/tutorials/posttraining/lora.md b/docs/tutorials/posttraining/lora.md index c309078593..bcd9ea9a51 100644 --- a/docs/tutorials/posttraining/lora.md +++ b/docs/tutorials/posttraining/lora.md @@ -117,8 +117,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \ max_target_length="${MAX_TARGET_LENGTH?}" \ learning_rate="${LEARNING_RATE?}" \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank="${LORA_RANK?}" \ lora.lora_alpha="${LORA_ALPHA?}" @@ -174,8 +172,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \ max_target_length="${MAX_TARGET_LENGTH?}" \ learning_rate="${LEARNING_RATE?}" \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank="${LORA_RANK?}" \ lora.lora_alpha="${LORA_ALPHA?}" diff --git a/docs/tutorials/posttraining/lora_on_multi_host.md b/docs/tutorials/posttraining/lora_on_multi_host.md index 4d0905ea13..8b3f140a80 100644 --- a/docs/tutorials/posttraining/lora_on_multi_host.md +++ b/docs/tutorials/posttraining/lora_on_multi_host.md @@ -186,8 +186,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ max_target_length=${MAX_TARGET_LENGTH?} \ learning_rate=${LEARNING_RATE?} \ chat_template_path=${CHAT_TEMPLATE_PATH?} \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank=${LORA_RANK?} \ lora.lora_alpha=${LORA_ALPHA?} \ @@ -267,8 +265,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ lora.lora_restore_path=${LORA_RESTORE_PATH?} \ learning_rate=${LEARNING_RATE?} \ chat_template_path=${CHAT_TEMPLATE_PATH?} \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank=${LORA_RANK?} \ lora.lora_alpha=${LORA_ALPHA?} \ diff --git a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py index ec37bb9b8c..c752623fa6 100644 --- a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py +++ b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py @@ -48,7 +48,7 @@ [Mode 2: MaxText Architecture] python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \ - model_name= scan_layers= enable_nnx= + model_name= scan_layers= (Optional: other maxtext config) [Mode 3: Orbax] @@ -220,13 +220,9 @@ def inspect_maxtext(args, remaining_args): import jax from maxtext.checkpoint_conversion.utils.utils import param_key_parts_from_path from maxtext.configs import pyconfig - from maxtext.layers import quantizations - from maxtext.models import models from maxtext.utils import max_utils, maxtext_utils from maxtext.utils.globals import MAXTEXT_PKG_DIR - Transformer = models.transformer_as_linen - # Configure the PyConfig environment. # The first argument in argv is typically the script name. argv = ( @@ -242,19 +238,11 @@ def inspect_maxtext(args, remaining_args): print(argv) config = pyconfig.initialize(argv) - print( - f"\n--- Inspecting MaxText Architecture: {config.model_name} " - f"(scan_layers: {config.scan_layers}, enable_nnx: {config.enable_nnx}) ---" - ) + print(f"\n--- Inspecting MaxText Architecture: {config.model_name} " f"(scan_layers: {config.scan_layers}) ---") devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - if config.enable_nnx: - _, abstract_model = create_nnx_abstract_model(config, mesh=mesh) - _, abstract_param, _ = nnx.split(abstract_model, nnx.Param, ...) - else: - quant = quantizations.configure_quantization(config) - model = Transformer(config, mesh=mesh, quant=quant) - abstract_param = maxtext_utils.get_abstract_param(model, config) + _, abstract_model = create_nnx_abstract_model(config, mesh=mesh) + _, abstract_param, _ = nnx.split(abstract_model, nnx.Param, ...) # Calculate and display the total parameter count based purely on abstract shapes. num_params = max_utils.calculate_num_params_from_pytree(abstract_param) @@ -273,7 +261,7 @@ def inspect_maxtext(args, remaining_args): # "params.params.decoder.decoder_norm.scale" (for standard model weights) # "params.Tid2EidVar.decoder.layers_0.mlp.MoeBlock_0.tid2eid" (for legacy custom collections) key_str = ".".join(key_parts) - if config.enable_nnx and not key_str.startswith(("params", "Tid2EidVar")): + if not key_str.startswith(("params", "Tid2EidVar")): param_key = "params.params." + key_str else: param_key = "params." + key_str @@ -392,7 +380,7 @@ def main(): inspect_hf(args) elif args.mode == "maxtext": # remaining_args accepts maxtext config, like `model_name= - # scan_layers= enable_nnx=` + # scan_layers=` inspect_maxtext(args, remaining_args) elif args.mode == "orbax": inspect_orbax(args) diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py index 138d5be4a5..63a5e1ef55 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py @@ -35,24 +35,18 @@ """ import argparse -import functools import gc import os import sys from flax import nnx -from flax import serialization import jax from jax import random from jax.sharding import Mesh from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_PKG_DIR from maxtext.common import checkpointing -from maxtext.common.common_types import MODEL_MODE_TRAIN -from maxtext.layers import quantizations from maxtext.common import train_state_nnx -from maxtext.models.models import transformer_as_linen -from maxtext.optimizers import optimizers from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils @@ -93,23 +87,15 @@ def convert(paxml_ckpt_path, maxtext_model_name, base_output_directory, run_name devices_array = maxtext_utils.create_device_mesh(cfg) mesh = Mesh(devices_array, cfg.mesh_axes) - if cfg.pure_nnx: - rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng) - model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs) - _, tx = train_utils.create_training_optimizer(cfg, model) - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh) + rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng) + model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs) + _, tx = train_utils.create_training_optimizer(cfg, model) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh) - def init_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - - else: - quant = quantizations.configure_quantization(cfg) - model = transformer_as_linen(cfg, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(cfg) - tx = optimizers.get_optimizer(cfg, learning_rate_schedule) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, cfg, True, init_rng) + def init_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) checkpoint_manager = checkpointing.create_orbax_checkpoint_manager( cfg.checkpoint_dir, @@ -119,11 +105,8 @@ def init_state_fn(): ) state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn) - if cfg.pure_nnx: - state = train_state_nnx.to_checkpoint_dict(state) - state.pop("nnx_aux", None) - else: - state = serialization.to_state_dict(state) + state = train_state_nnx.to_checkpoint_dict(state) + state.pop("nnx_aux", None) max_logging.log("start") max_utils.print_mem_stats("After params initialized") diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index dec2634bed..beaf7584d0 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -186,7 +186,7 @@ def _load_linen_checkpoint_into_nnx( enable_single_replica_ckpt_restoring: bool = False, config=None, ): - """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume). + """Restores a Linen-layout checkpoint into an NNX state. Restores a Linen-shape target that includes `nnx_aux`, then reshapes back via `_restored_linen_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when @@ -230,7 +230,7 @@ def _restored_linen_to_nnx(restored_linen, abstract_nnx_state, config=None): """Reshapes a restored Linen-layout tree into the NNX state. Raises if the checkpoint is missing a weight. Every NNX restore path ends here: the load - itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout. + itself is the Linen one, since MaxText 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) @@ -309,7 +309,7 @@ def _load_full_state_from_path( The loaded state. """ if source_checkpoint_layout == "orbax": - # pure_nnx checkpoints are stored in the Linen on-disk layout; reshape to NNX. + # Checkpoints are stored in the Linen on-disk layout; reshape to NNX. if isinstance(abstract_unboxed_pre_state, nnx.State): return _load_linen_checkpoint_into_nnx( path, @@ -350,7 +350,7 @@ def combine_sharding(sds, shardings): sharded_abstract_state = jax.tree.map(combine_sharding, simple_abstract_state, shardings) pre_transformed_state = ocp.load(path, sharded_abstract_state) state = conversion_fn(pre_transformed_state) - # The conversion fn returns MaxText's on-disk (Linen) layout, which is what pure_nnx reads, + # The conversion fn returns MaxText's on-disk (Linen) layout, which is what NNX reads, # so NNX needs the same reshape as every other restore. An NNX state passes through. if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(state, nnx.State): state = _restored_linen_to_nnx(state, abstract_unboxed_pre_state, config=maxtext_config) @@ -533,7 +533,7 @@ def load_state_if_possible( load_parameters_from_path = _normalize_checkpoint_root(load_parameters_from_path) if load_full_state_from_path: load_full_state_from_path = _normalize_checkpoint_root(load_full_state_from_path) - # pure_nnx saves in the Linen on-disk layout, so every branch below loads the same tree Linen + # Checkpoints are saved in the Linen on-disk layout, so every branch below loads the same tree Linen # does: the NNX abstract is converted to that layout going in, and what comes back is reshaped # into the NNX state on the way out. is_nnx = isinstance(abstract_unboxed_pre_state, (nnx.State, train_state_nnx.TrainStateNNX)) @@ -842,12 +842,8 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step if step is not None: actual_step = int(step) else: - if config.pure_nnx: - # Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer. - actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1 - else: - # Linen TrainState has .step attribute - actual_step = int(state.step) - 1 + # Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer. + actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1 # Determine if a checkpoint save should be forced, overriding the usual # `config.checkpoint_period` logic. @@ -927,8 +923,8 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= if config and getattr(config, "enable_diloco", False): state = diloco_checkpoint_utils.to_diloco_checkpoint_dict(state, config) - elif config and getattr(config, "pure_nnx", False): - # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. + elif config: + # Save in the Linen on-disk layout so NNX and Linen checkpoints are interchangeable. if isinstance(state, nnx.State): state = train_state_nnx.to_checkpoint_dict(state) diff --git a/src/maxtext/common/emergency_checkpointing.py b/src/maxtext/common/emergency_checkpointing.py index fba2f29546..7b105b7350 100644 --- a/src/maxtext/common/emergency_checkpointing.py +++ b/src/maxtext/common/emergency_checkpointing.py @@ -65,7 +65,7 @@ def create_emergency_checkpoint_manager( persistent_p = gcs_utils.mkdir_and_check_permissions(persistent_checkpoint_dir) - # pure_nnx saves via to_checkpoint_dict (Linen params/opt_state/step plus an nnx_aux + # Checkpoints are saved via to_checkpoint_dict (Linen params/opt_state/step plus an nnx_aux # subtree), but the emergency manager restores against the abstract it is built with. # Convert it the same way so it matches what is on disk; restore reshapes back to NNX. if isinstance(abstract_state, nnx.State): diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index 647e1fe7c8..bc4e4e81f7 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -60,7 +60,7 @@ def apply_gradients(self, grads: Any, **kwargs): # On-disk checkpoint format. # -# A pure_nnx run saves in the same on-disk layout as a Linen run, so the two are +# A run saves in the same on-disk layout that Linen used, so old and new checkpoints are # interchangeable. The NNX state pure dict differs from Linen's in three ways, all # reshaped below at save time: # 1. top-level keys: {model, optimizer:{step, opt_state}} -> {params:{params:...}, step, opt_state} @@ -234,7 +234,7 @@ def to_checkpoint_dict(state: nnx.State | nnx.Module): """Reshapes an nnx.State into the on-disk checkpoint layout. Weights (nnx.Param) map to the Linen `params` collection and the optimizer to - opt_state/step, so pure_nnx and Linen checkpoints stay interchangeable. Everything else that + opt_state/step, so NNX and Linen checkpoints stay interchangeable. Everything else that must persist -- rngs/dropout, batch stats, and any custom variable -- goes under an `nnx_aux` subtree. Works on a concrete state (save) or an abstract state (restore target). """ diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 016f31a322..c1766c708e 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1339,11 +1339,6 @@ position_id_per_seconds: 25 # Example: "8,8" to use a 8x8 subgrid (64 chips) of a full pod (16x16) of trillium. subslice_shape: "" -# NNX -enable_nnx: true -pure_nnx_decoder: true -pure_nnx: true - ################################## Qwen3-Next Specific Configs ################################## # Kernel size for the 1D convolution in the Gated Delta Net gdn_conv_kernel_dim: 4 diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml index d6058a1c68..d14b4e959c 100644 --- a/src/maxtext/configs/inference/vllm.yml +++ b/src/maxtext/configs/inference/vllm.yml @@ -15,8 +15,6 @@ attention: "vllm_rpa" model_call_mode: "inference" -# NNX required for vLLM integration -enable_nnx: true # Avoid re-initializing JAX distributed system when using vLLM skip_jax_distributed_system: true # Scanned layers are not supported with vLLM integration diff --git a/src/maxtext/configs/post_train/distillation-qwen3.yml b/src/maxtext/configs/post_train/distillation-qwen3.yml index b1e95ae138..268a118a63 100644 --- a/src/maxtext/configs/post_train/distillation-qwen3.yml +++ b/src/maxtext/configs/post_train/distillation-qwen3.yml @@ -3,7 +3,6 @@ # Inherit MaxText defaults base_config: "base.yml" override_model_config: True -enable_nnx: True enable_checkpointing: True # --- Student Specifics --- diff --git a/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml b/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml index b5cf9ff0f3..4ebed81f1e 100644 --- a/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml +++ b/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml @@ -19,7 +19,6 @@ distill_alpha: 0.5 distill_temperature: 1.0 distill_beta: 0 distill_layer_indices: [] -enable_nnx: True load_balance_loss_weight: 0.001 # Megablox grouped-matmul m-tile (batch_seq). The k/n dims already default to diff --git a/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml b/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml index e3c12ff08c..3bd8aa826f 100644 --- a/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml +++ b/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml @@ -21,7 +21,6 @@ distill_alpha: 0.6 distill_temperature: 1.0 distill_beta: 1.0 distill_layer_indices: [0,1,2,3,4,5,6,7] -enable_nnx: True load_balance_loss_weight: 0.001 ici_fsdp_parallelism: -1 diff --git a/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml b/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml index 36d48da2ed..0e4ee36a0d 100644 --- a/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml +++ b/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml @@ -21,7 +21,6 @@ distill_alpha: 0.6 distill_temperature: 1.0 distill_beta: 1.0 distill_layer_indices: [0,1,2,3,4,5,6,7] -enable_nnx: True load_balance_loss_weight: 0.001 ici_fsdp_parallelism: -1 diff --git a/src/maxtext/configs/pyconfig_deprecated.py b/src/maxtext/configs/pyconfig_deprecated.py index c30cc0101b..898d8f0cd7 100644 --- a/src/maxtext/configs/pyconfig_deprecated.py +++ b/src/maxtext/configs/pyconfig_deprecated.py @@ -194,8 +194,7 @@ def validate_expert_shard_attention_option(expert_shard_attention_option: str) - ) -def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int, enable_nnx: bool): - del enable_nnx # NNX vocab tiling supported via vocab_tiling_nnx_loss in vocabulary_tiling.py +def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int): if (per_device_batch_size * max_target_length) % num_vocab_tiling != 0: raise ValueError("Per device batch size times sequence length should be divisible by the number of vocab tiles.") @@ -239,9 +238,7 @@ def validate_keys(keys): validate_model_call_mode(keys["model_call_mode"]) validate_prefill_and_target_lengths(keys["max_prefill_predict_length"], keys["max_target_length"]) validate_rope_type(keys["rope_type"]) - validate_vocab_tiling( - keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"], keys["enable_nnx"] - ) + validate_vocab_tiling(keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"]) if keys["enable_rampup_batch_size"]: validate_rampup_batch_size( keys["per_device_batch_size_start"], diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 84acba8936..38cb452986 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1371,29 +1371,8 @@ class HardwareAndMesh(BaseModel): description="Customized mesh and logical rules for evaluation.", ) allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.") - enable_nnx: bool = Field( - True, - description=( - "Whether to use NNX for model definition. Setting this to False selects the Linen path, " - "which will be deprecated in the near future." - ), - ) optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.") shardy: bool = Field(True, description="Whether to use shardy XLA backend.") - pure_nnx_decoder: bool = Field( - True, - description=( - "Whether to enable pure NNX decoder. Setting this to False selects the Linen decoder, " - "which will be deprecated in the near future." - ), - ) - pure_nnx: bool = Field( - True, - description=( - "Whether to enable pure NNX mode. Setting this to False selects the Linen path, " - "which will be deprecated in the near future." - ), - ) remove_size_one_mesh_axis_from_type: bool = Field( True, description="Whether to remove size one mesh axis from type through jax.config.", @@ -3755,34 +3734,6 @@ def validate_and_set_hlo_dump_defaults(): if self.distill_beta > 0.0: if not self.scan_layers: raise ValueError("a value of self.distill_beta > 0.0 requires self.scan_layers = True") - if not self.enable_nnx: - raise ValueError("a value of self.distill_beta > 0.0 requires self.enable_nnx = True") - - if self.pure_nnx and not self.pure_nnx_decoder and self.use_qwix_quantization and not self.use_batch_split_schedule: - if self.quantization: - raise ValueError( - f"quantization='{self.quantization}' with use_qwix_quantization=True under pure_nnx=True requires " - "pure_nnx_decoder=True. The bridged Linen decoder (pure_nnx_decoder=False) is invisible to Qwix, " - "so quantization (and weight sparsity) would silently have no effect. Set pure_nnx_decoder=True." - ) - - # TODO: Remove this block once the Linen code path and the enable_nnx, pure_nnx and pure_nnx_decoder flags are deleted. - linen_flags = [ - name - for name, value in ( - ("enable_nnx", self.enable_nnx), - ("pure_nnx", self.pure_nnx), - ("pure_nnx_decoder", self.pure_nnx_decoder), - ) - if not value - ] - if linen_flags: - logger.warning("=" * 80) - logger.warning("MAXTEXT DEPRECATION NOTICE: you are running on the Linen code path.") - logger.warning("Selected by: %s", ", ".join(f"{name}=False" for name in linen_flags)) - logger.warning("Linen will be deprecated in the near future and removed after that.") - logger.warning("Plan to migrate to NNX: leave enable_nnx, pure_nnx and pure_nnx_decoder at their default of True.") - logger.warning("=" * 80) # Validate distillation schedule parameters if self.distill_alpha_end is not None and not 0.0 <= self.distill_alpha_end <= 1.0: @@ -4379,20 +4330,12 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "te_gmm_quantization must be specified when te_moe_block=True. " "te_gmm_quantization=te_no_quant is supported for BF16." ) - if not self.pure_nnx and self.routed_bias and self.decoder_block == DecoderBlockType.DEEPSEEK4: - raise ValueError( - "Auxiliary-loss-free routed bias for DeepSeek V4 is only supported in pure NNX mode. " - "Please set pure_nnx=True or disable routed_bias." - ) if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts: raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.") self.validate_ragged_buffer_factor() self.validate_retry_when_tokens_dropped() self.validate_num_moe_emb_chunks() - if self.enable_diloco and not self.pure_nnx: - raise ValueError("enable_diloco=True requires pure_nnx=True (Linen support for DiLoCo has been removed).") - if self.enable_streaming_diloco: if not self.scan_layers: raise ValueError("enable_streaming_diloco=True requires scan_layers=True.") diff --git a/src/maxtext/examples/lora_llama3_demo.ipynb b/src/maxtext/examples/lora_llama3_demo.ipynb index 280b48bd9f..16b098060a 100644 --- a/src/maxtext/examples/lora_llama3_demo.ipynb +++ b/src/maxtext/examples/lora_llama3_demo.ipynb @@ -311,8 +311,6 @@ " \"learning_rate=5e-5\", \n", " \"weight_dtype=bfloat16\",\n", " \"dtype=bfloat16\",\n", - " \"enable_nnx=true\",\n", - " \"pure_nnx_decoder=true\",\n", " \"scan_layers=true\",\n", " \"lora.enable_lora=true\",\n", " \"lora.lora_rank=16\",\n", diff --git a/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py b/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py index 3a7681979e..f34007aaa4 100644 --- a/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py +++ b/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py @@ -146,7 +146,6 @@ def _create_config(self): override_model_config=True, skip_jax_distributed_system=True, log_config=False, - pure_nnx=False, attention="dot_product", dtype="float32", **yaml_overrides, diff --git a/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py b/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py index a1799fc790..2869743b67 100644 --- a/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py +++ b/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py @@ -168,13 +168,9 @@ def stitch_and_save_checkpoints( # 1. Generate full target model with initial random weights max_logging.log("Generating target omni model from config with initial random weights...") with jax.set_mesh(mesh): - if config.pure_nnx: - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) - model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) - init_params = nnx.state(model, nnx.Param) - else: - model = model_creation_utils.from_config(config, jax.devices()) - _, _, init_params = maxtext_utils.init_initial_state(model, None, config, is_training=False, key=init_rng) + rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) + model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) + init_params = nnx.state(model, nnx.Param) # Convert to pure pytree for easier processing is_nnx = isinstance(init_params, nnx.State) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 94c37ae304..21ceb4270f 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -573,7 +573,7 @@ def _init_pipeline_generic(self, decoder_block_classes, rngs): self.layers_outside_pipeline = self._create_scanned_layers( base_cls, length=remaining_layers, - metadata_axis_name="layers", + metadata_axis_name="layers_outside_pipeline", rngs=rngs, ) else: diff --git a/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh b/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh index 1e9a50ca78..505e2eff72 100755 --- a/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh +++ b/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh @@ -139,7 +139,7 @@ LIBTPU_INIT_ARGS=" \ --xla_tpu_enable_3d_reduce_scatter_decomposer=false " ## -------------------------- Command Execution -------------------------- -CMD="set -euo pipefail; mkdir -p /deps/src/src/dependencies/scripts && ( [ -f /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ] || ln -sf /deps/src/dependencies/scripts/setup_gcsfuse.sh /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ); mkdir -p /deps/src/src/maxtext && ( [ -d /deps/src/src/maxtext/configs ] || ln -sf /deps/src/maxtext/configs /deps/src/src/maxtext/configs ); if [[ \"\${MOUNT_GCSFUSE:-1}\" = \"1\" ]]; then bucket=\"\${GCS_BASE#gs://}\"; bucket=\"\${bucket%/}\"; bash /deps/src/dependencies/scripts/setup_gcsfuse.sh DATASET_GCS_BUCKET=\"\${bucket}\" MOUNT_PATH=\"\${LOCAL_MOUNT}\"; fi; export PYTHONPATH=/app/src:\${PYTHONPATH:-}; cd /app/src && unset XLA_FLAGS; export LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\"; LOAD_PARAM_ARG=\"\"; if [ -n \"\${LOAD_FULL_STATE_PATH}\" ]; then LOAD_PARAM_ARG=\"load_full_state_path=\${LOAD_FULL_STATE_PATH}\"; elif [ -n \"\${LOAD_PARAMETERS_PATH}\" ]; then LOAD_PARAM_ARG=\"load_parameters_path=\${LOAD_PARAMETERS_PATH}\"; fi; python3 -m maxtext.trainers.pre_train.train maxtext/configs/base.yml run_name=\"\${RUN_NAME}\" save_config_to_gcs=true base_output_directory=\"\${OUTPUT_DIR}\" \${LOAD_PARAM_ARG} dataset_type=olmo_grain olmo_index_path=\"\${INDEX_PATH}\" olmo_path_remap_from=\"\${GCS_BASE}\" olmo_path_remap_to=\"\${LOCAL_MOUNT}\" olmo_apply_ngram_filter=True data_shuffle_seed=\"\${DATA_SEED}\" model_name=\"\${MODEL_NAME}\" tokenizer_type=\"\${TOKENIZER_TYPE}\" tokenizer_path=\"\${TOKENIZER_PATH}\" per_device_batch_size=\"\${PER_DEVICE_BATCH_SIZE}\" max_target_length=\"\${MAX_TARGET_LENGTH}\" learning_rate=\"\${LEARNING_RATE}\" learning_rate_schedule_steps=\"\${LR_SCHEDULE_STEPS}\" learning_rate_final_fraction=\"\${COSINE_FINAL_FRAC}\" warmup_steps_fraction=\"\${WARMUP_FRAC}\" adam_b1=\"\${ADAM_B1}\" adam_b2=\"\${ADAM_B2}\" adam_eps=\"\${ADAM_EPS}\" adam_weight_decay=\"\${ADAM_WD}\" gradient_clipping_threshold=\"\${GRAD_CLIP}\" z_loss_multiplier=\"\${Z_LOSS}\" load_balance_loss_weight=\"\${LOAD_BALANCE_LOSS_WEIGHT}\" float32_gate_logits=\"\${FLOAT32_GATE_LOGITS}\" enable_diloco=true enable_streaming_diloco=\"\${ENABLE_STREAMING_DILOCO}\" pure_nnx=true dcn_diloco_parallelism=\"\${XPK_NUM_SLICES}\" diloco_sync_period=\"\${DILOCO_SYNC_PERIOD}\" diloco_outer_lr=\"\${DILOCO_OUTER_LR}\" diloco_outer_momentum=\"\${DILOCO_OUTER_MOMENTUM}\" num_diloco_fragments=\"\${DILOCO_NUM_FRAGMENTS}\" use_sequential_layers=\"\${DILOCO_USE_SEQUENTIAL_LAYERS}\" jax_distributed_initialization_timeout=1200 steps=\"\${STEPS}\"" +CMD="set -euo pipefail; mkdir -p /deps/src/src/dependencies/scripts && ( [ -f /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ] || ln -sf /deps/src/dependencies/scripts/setup_gcsfuse.sh /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ); mkdir -p /deps/src/src/maxtext && ( [ -d /deps/src/src/maxtext/configs ] || ln -sf /deps/src/maxtext/configs /deps/src/src/maxtext/configs ); if [[ \"\${MOUNT_GCSFUSE:-1}\" = \"1\" ]]; then bucket=\"\${GCS_BASE#gs://}\"; bucket=\"\${bucket%/}\"; bash /deps/src/dependencies/scripts/setup_gcsfuse.sh DATASET_GCS_BUCKET=\"\${bucket}\" MOUNT_PATH=\"\${LOCAL_MOUNT}\"; fi; export PYTHONPATH=/app/src:\${PYTHONPATH:-}; cd /app/src && unset XLA_FLAGS; export LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\"; LOAD_PARAM_ARG=\"\"; if [ -n \"\${LOAD_FULL_STATE_PATH}\" ]; then LOAD_PARAM_ARG=\"load_full_state_path=\${LOAD_FULL_STATE_PATH}\"; elif [ -n \"\${LOAD_PARAMETERS_PATH}\" ]; then LOAD_PARAM_ARG=\"load_parameters_path=\${LOAD_PARAMETERS_PATH}\"; fi; python3 -m maxtext.trainers.pre_train.train maxtext/configs/base.yml run_name=\"\${RUN_NAME}\" save_config_to_gcs=true base_output_directory=\"\${OUTPUT_DIR}\" \${LOAD_PARAM_ARG} dataset_type=olmo_grain olmo_index_path=\"\${INDEX_PATH}\" olmo_path_remap_from=\"\${GCS_BASE}\" olmo_path_remap_to=\"\${LOCAL_MOUNT}\" olmo_apply_ngram_filter=True data_shuffle_seed=\"\${DATA_SEED}\" model_name=\"\${MODEL_NAME}\" tokenizer_type=\"\${TOKENIZER_TYPE}\" tokenizer_path=\"\${TOKENIZER_PATH}\" per_device_batch_size=\"\${PER_DEVICE_BATCH_SIZE}\" max_target_length=\"\${MAX_TARGET_LENGTH}\" learning_rate=\"\${LEARNING_RATE}\" learning_rate_schedule_steps=\"\${LR_SCHEDULE_STEPS}\" learning_rate_final_fraction=\"\${COSINE_FINAL_FRAC}\" warmup_steps_fraction=\"\${WARMUP_FRAC}\" adam_b1=\"\${ADAM_B1}\" adam_b2=\"\${ADAM_B2}\" adam_eps=\"\${ADAM_EPS}\" adam_weight_decay=\"\${ADAM_WD}\" gradient_clipping_threshold=\"\${GRAD_CLIP}\" z_loss_multiplier=\"\${Z_LOSS}\" load_balance_loss_weight=\"\${LOAD_BALANCE_LOSS_WEIGHT}\" float32_gate_logits=\"\${FLOAT32_GATE_LOGITS}\" enable_diloco=true enable_streaming_diloco=\"\${ENABLE_STREAMING_DILOCO}\" dcn_diloco_parallelism=\"\${XPK_NUM_SLICES}\" diloco_sync_period=\"\${DILOCO_SYNC_PERIOD}\" diloco_outer_lr=\"\${DILOCO_OUTER_LR}\" diloco_outer_momentum=\"\${DILOCO_OUTER_MOMENTUM}\" num_diloco_fragments=\"\${DILOCO_NUM_FRAGMENTS}\" use_sequential_layers=\"\${DILOCO_USE_SEQUENTIAL_LAYERS}\" jax_distributed_initialization_timeout=1200 steps=\"\${STEPS}\"" echo "Submitting Qwen3-30B-A3B + OLMo DiLoCo training workload using XPK..." echo "Workload Name: ${WORKLOAD_NAME}" diff --git a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh index 4b86fc5e20..bf964e1c5a 100644 --- a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh +++ b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh @@ -132,7 +132,6 @@ CMD="export PYTHONPATH=/app/src:\$PYTHONPATH && unset XLA_FLAGS && export LIBTPU max_target_length=${MAX_TARGET_LENGTH} \ enable_diloco=true \ enable_streaming_diloco=true \ - pure_nnx=true \ num_diloco_fragments=${DILOCO_NUM_FRAGMENTS} \ use_sequential_layers=${DILOCO_USE_SEQUENTIAL_LAYERS} \ num_communication_overlapping_steps=${DILOCO_NUM_COMM_OVERLAP_STEPS} \ diff --git a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh index 01d1ece4f0..694c7a6a38 100755 --- a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh +++ b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_dlco_moe_c4.sh @@ -126,7 +126,6 @@ CMD="export PYTHONPATH=/app/src:\$PYTHONPATH && unset XLA_FLAGS && export LIBTPU load_balance_loss_weight=${LOAD_BALANCE_LOSS_WEIGHT} \ enable_diloco=true \ enable_streaming_diloco=${ENABLE_STREAMING_DILOCO} \ - pure_nnx=true \ dcn_diloco_parallelism=${NUM_SLICES} \ diloco_sync_period=${DILOCO_SYNC_PERIOD} \ diloco_outer_lr=${DILOCO_OUTER_LR} \ diff --git a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh index 8a6678b34f..7f5975a181 100644 --- a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh +++ b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh @@ -100,7 +100,7 @@ # DISTILL_ALPHA default: 0.5 # DISTILL_TEMPERATURE default: 1.0 # DISTILL_BETA default: 1.0 (>0 enables feature-map loss; -# requires scan_layers=True and enable_nnx=True) +# requires scan_layers=True) # DISTILL_LAYER_INDICES default: [0,1,2,3,4,5,6,7] (no spaces inside brackets) # # Image pinning (used by prep_image): @@ -185,8 +185,7 @@ LAST_WORKLOAD_FILE="${XPK_LAST_WORKLOAD_FILE:-${HOME}/.xpk_last_workload}" extra_cli="distill_alpha=${DISTILL_ALPHA} \ distill_temperature=${DISTILL_TEMPERATURE} \ distill_beta=${DISTILL_BETA} \ -distill_layer_indices=${DISTILL_LAYER_INDICES} \ -enable_nnx=True" +distill_layer_indices=${DISTILL_LAYER_INDICES}" if [ -n "${STEPS_OVERRIDE:-}" ]; then extra_cli="$extra_cli learning_rate_schedule_steps=${STEPS_OVERRIDE} steps=${STEPS_OVERRIDE}" fi diff --git a/src/maxtext/utils/generate_param_only_checkpoint.py b/src/maxtext/utils/generate_param_only_checkpoint.py index 3d9ce27f72..a58ac10b15 100644 --- a/src/maxtext/utils/generate_param_only_checkpoint.py +++ b/src/maxtext/utils/generate_param_only_checkpoint.py @@ -22,7 +22,6 @@ The output "parameter state" is output to the checkpoint directory. Additionally it is cast down to bf16. """ -import functools import os.path from typing import Sequence @@ -35,11 +34,8 @@ from jax.sharding import Mesh from maxtext.configs import pyconfig from maxtext.common import checkpointing -from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN -from maxtext.layers import quantizations +from maxtext.common.common_types import DecoderBlockType from maxtext.common import train_state_nnx -from maxtext.models import models -from maxtext.optimizers import optimizers from maxtext.utils import gcs_utils from maxtext.utils import lora_utils from maxtext.utils import max_logging @@ -71,47 +67,7 @@ def _possibly_unroll_params(config, training_state, training_state_annotations, """Unroll scanned input layers when force_unroll is set.""" if not config.scan_layers or not config.force_unroll: return - if config.pure_nnx: - _possibly_unroll_params_nnx(config, training_state, training_state_annotations, mesh) - return - - def unroll_layer_group(num_layers, layer_name="layers"): - """Helper function to unroll layers (e.g. dense or MoE) into individual layers.""" - layers = training_state.params["params"]["decoder"].get(layer_name, None) - layers_annotations = training_state_annotations.params["params"]["decoder"].get(layer_name, None) - - if layers is None or layers_annotations is None: - raise ValueError(f"Missing {layer_name} in training_state or training_state_annotations.") - - def new_pspec(x): - return jax.sharding.PartitionSpec(*(x[0 : config.param_scan_axis] + x[config.param_scan_axis + 1 :])) - - new_layer_annotation = jax.tree_util.tree_map(new_pspec, layers_annotations) - new_layer_sharding = jax.tree_util.tree_map(lambda x: jax.sharding.NamedSharding(mesh, x), new_layer_annotation) - - for i in range(num_layers): - - def slice_ith(input_layers): - return jax.tree_util.tree_map(lambda x: jax.numpy.take(x, i, axis=config.param_scan_axis), input_layers) - - # pylint: disable=not-callable - new_layer = jax.jit(slice_ith, out_shardings=new_layer_sharding)(layers) - - training_state.params["params"]["decoder"][f"{layer_name}_{i}"] = new_layer - training_state_annotations.params["params"]["decoder"][f"{layer_name}_{i}"] = new_layer_annotation - - # Remove the original layer collection - del training_state.params["params"]["decoder"][layer_name] - del training_state_annotations.params["params"]["decoder"][layer_name] - - jax.tree_util.tree_map(lambda x: x.delete(), layers) - - if config.decoder_block == DecoderBlockType.DEEPSEEK: - # Unroll dense and MoE layers separately - unroll_layer_group(config.first_num_dense_layers, layer_name="dense_layers") - unroll_layer_group(config.num_decoder_layers - config.first_num_dense_layers, layer_name="moe_layers") - else: - unroll_layer_group(config.num_decoder_layers, layer_name="layers") + _possibly_unroll_params_nnx(config, training_state, training_state_annotations, mesh) def _possibly_unroll_params_nnx(config, state, state_mesh_shardings, mesh): @@ -205,97 +161,40 @@ def _slice_leaf(x, sharding): def _read_train_checkpoint(config, checkpoint_manager, mesh): """Read training checkpoint at path defined by load_full_state_path.""" rng = random.PRNGKey(0) - if config.pure_nnx: - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=rng) - model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) - _, tx = train_utils.create_training_optimizer(config, model) - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) - - def init_state_fn(): - nnx_model = _create_model_partial() - wrt = ( - getattr(nnx, "LoRAParam", nnx.Param) - if getattr(getattr(config, "lora", None), "enable_lora", False) - else nnx.Param - ) - optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - - else: - quant = quantizations.configure_quantization(config) - model = models.transformer_as_linen(config, mesh, quant, MODEL_MODE_TRAIN) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - tx = optimizers.get_optimizer(config, learning_rate_schedule) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng) + rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=rng) + model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) + _, tx = train_utils.create_training_optimizer(config, model) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) + + def init_state_fn(): + nnx_model = _create_model_partial() + wrt = ( + getattr(nnx, "LoRAParam", nnx.Param) + if getattr(getattr(config, "lora", None), "enable_lora", False) + else nnx.Param + ) + optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - state, state_mesh_notations, state_mesh_shardings, _, _ = maxtext_utils.setup_training_state( + state, _, state_mesh_shardings, _, _ = maxtext_utils.setup_training_state( None, config, mesh, checkpoint_manager, init_state_fn ) - if config.pure_nnx: - # On NNX, state is a flat nnx.State; params live under state.model and the - # legacy notations are unused (callers receive shardings directly). - params, _ = nnx.split_state(state.model, nnx.Param, ...) - num_params = max_utils.calculate_num_params_from_pytree(params) - max_logging.log(f"In input checkpoint Number of model params={num_params/1e9:.3f} billion") - return state, state_mesh_shardings - num_params = max_utils.calculate_num_params_from_pytree(state.params) + # On NNX, state is a flat nnx.State; params live under state.model and the + # legacy notations are unused (callers receive shardings directly). + params, _ = nnx.split_state(state.model, nnx.Param, ...) + num_params = max_utils.calculate_num_params_from_pytree(params) max_logging.log(f"In input checkpoint Number of model params={num_params/1e9:.3f} billion") - return state, state_mesh_notations + return state, state_mesh_shardings def _generate_lora_decode_checkpoints(config, mesh): """Read lora checkpoints checkpoint at path defined by load_full_state_path.""" - if config.pure_nnx: - _generate_lora_decode_checkpoints_nnx(config, mesh) - return - quant = quantizations.configure_quantization(config) - model = models.transformer_as_linen(config, mesh, quant, MODEL_MODE_TRAIN) - rng = random.PRNGKey(0) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - tx = optimizers.get_optimizer(config, learning_rate_schedule) - - lora_adapters = gcs_utils.gcs_list_directories(config.lora_input_adapters_path) - for lora_id in lora_adapters: - # Expected lora_checkpoint_dir = /loras/ - lora_checkpoint_dir = os.path.join(config.checkpoint_dir, "loras", lora_id, "") - - lora_adapter_path = os.path.join(config.lora_input_adapters_path, lora_id, "") - - # Create a checkpoint manager to save decode checkpoint at lora_checkpoint_dir - checkpoint_manager = checkpointing.create_orbax_checkpoint_manager( - lora_checkpoint_dir, - config.enable_checkpointing, - config.async_checkpointing, - config.checkpoint_period, - use_ocdbt=config.checkpoint_storage_use_ocdbt, - use_zarr3=config.checkpoint_storage_use_zarr3, - ) - - lora_config, lora_state, lora_state_annotations = lora_utils.setup_initial_lora_state( - model, None, tx, config, rng, mesh, checkpoint_manager, lora_adapter_path - ) - - _possibly_unroll_params(config, lora_state, lora_state_annotations, mesh) - - gcs_utils.write_dict_to_gcs_json(lora_config, os.path.join(lora_checkpoint_dir, "adapter_config.json")) - - # Save decode state to config's checkpoint directory at step 0 - _save_decode_checkpoint(config, lora_state, checkpoint_manager) - max_logging.log(f"Successfully saved LoRA checkpoint at: {os.path.join(lora_checkpoint_dir, '0', 'items')}") + _generate_lora_decode_checkpoints_nnx(config, mesh) def _save_decode_checkpoint(config, state, checkpoint_manager): """Generate checkpoint for decode from the training_state.""" - if config.pure_nnx: - _save_decode_checkpoint_nnx(config, state, checkpoint_manager) - return - decode_state = maxtext_utils.init_decode_state( - None, jax.tree_util.tree_map(lambda x: x.astype(jax.numpy.bfloat16), state.params) - ) - if checkpoint_manager is not None: - if checkpointing.save_checkpoint(checkpoint_manager, 0, decode_state): - max_logging.log(f"saved an decode checkpoint at {config.checkpoint_dir}") - checkpointing.wait_until_finished(checkpoint_manager) + _save_decode_checkpoint_nnx(config, state, checkpoint_manager) def _save_decode_checkpoint_nnx(config, state, checkpoint_manager): @@ -480,11 +379,8 @@ def generate_decode_checkpoint(config): # Read training state from config.load_paramaters_path max_logging.log(f"Read training checkpoint from: {config.load_full_state_path}") training_state, training_state_annotations = _read_train_checkpoint(config, checkpoint_manager, mesh) - if config.pure_nnx: - # NNX state is a flat nnx.State; opt_state lives under the optimizer sub-state. - assert training_state.optimizer.opt_state, "missing opt_state in training checkpoint" - else: - assert training_state.opt_state != {}, "missing opt_state in training checkpoint" + # NNX state is a flat nnx.State; opt_state lives under the optimizer sub-state. + assert training_state.optimizer.opt_state, "missing opt_state in training checkpoint" _possibly_unroll_params(config, training_state, training_state_annotations, mesh) diff --git a/src/maxtext/utils/layerwise_quantization.py b/src/maxtext/utils/layerwise_quantization.py index f8df26967a..211724e44c 100644 --- a/src/maxtext/utils/layerwise_quantization.py +++ b/src/maxtext/utils/layerwise_quantization.py @@ -30,7 +30,6 @@ """ -import functools import os from typing import Any, Sequence @@ -43,14 +42,12 @@ from maxtext.common import common_types from maxtext.common import checkpointing from maxtext.layers import quantizations -from maxtext.models import deepseek, models from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx from maxtext.utils import model_creation_utils import orbax.checkpoint as ocp -from tqdm import tqdm from maxtext.configs import pyconfig IGNORE = ocp.PLACEHOLDER @@ -166,123 +163,20 @@ def __init__(self, config: Any, rng: PRNGKeyType): self.config = config self.rng = rng - # The Linen path runs layer-by-layer (memory-efficient for big DeepSeek - # models) and is DeepSeek-specific because it relies on the per-layer - # `DeepSeek*ToLinen` wrappers. The NNX path runs whole-model convert - # forward and is model-agnostic — see `_load_and_quantize_nnx`. - if not config.pure_nnx: - assert config.decoder_block == common_types.DecoderBlockType.DEEPSEEK, ( - f"Linen layerwise quantization only supports {common_types.DecoderBlockType.DEEPSEEK}, " - f"got {config.decoder_block}." - ) # Mesh definition devices_array = maxtext_utils.create_device_mesh(config=config) self._mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) self.quant = quantizations.configure_quantization(config) - if config.pure_nnx: - # NNX takes a separate code path that builds the model via from_pretrained; - # no Linen abstract-state bookkeeping is needed here. - self.unboxed_abstract_state = None - return - model = models.transformer_as_linen( - config, mesh=self._mesh, quant=self.quant, model_mode=common_types.MODEL_MODE_TRAIN - ) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, self.config, False, self.rng) - - self.unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(self.config, self._mesh, init_state_fn, False) + # NNX builds the model via from_pretrained; no Linen abstract-state + # bookkeeping is needed here. + self.unboxed_abstract_state = None def load_and_quantize(self) -> None: """ Load parameters layer by layer and quantize them. """ - if self.config.pure_nnx: - self._load_and_quantize_nnx() - return - quantized_params = {} - quantized_params["params"] = {"decoder": {}} - quantized_params["aqt"] = {"decoder": {}} - config = self.config - self.quant.quant_mode = quantizations.get_quant_mode("convert") - model_mode = common_types.MODEL_MODE_PREFILL - _, rng_quant_params = jax.random.split(self.rng) - - layers = [ - deepseek.DeepSeekDenseLayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - deepseek.DeepSeekMoELayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - ] - layer_prefixes = [ - "dense_layers", - "moe_layers", - ] - num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers - num_layers_list = [ - config.first_num_dense_layers, - num_moe_layers, - ] - - def model_apply(_p, _rng, layer): - return layer.apply( - _p | {"aqt": {}}, - jnp.ones((1, self.config.max_prefill_predict_length, self.config.base_emb_dim), dtype=jnp.int32), - None, - jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - True, - model_mode=model_mode, - rngs={"params": _rng}, - mutable=True, - ) - - for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes): - for index in tqdm(range(num_layers)): - layer_name = f"{layer_prefix}_{index}" - max_logging.log(f"Processing layer: {layer_name}") - - params = self._load_layer(layer_name) - params["params"] = params["params"]["decoder"][layer_name] - - _, new_vars = model_apply(params, rng_quant_params, layer) - - if "aqt" not in new_vars: - max_logging.log( - f"Warning: 'aqt' not found in new_vars for {layer_name}. Skipping AQT processing for this layer." - ) - quantized_params["params"]["decoder"][layer_name] = params["params"] # Keep original params - continue - - aqt_vars = new_vars["aqt"] - - try: - removed_params = remove_quantized_params(params["params"], aqt_vars) - quantized_params["params"]["decoder"][layer_name] = removed_params - except Exception as e: - max_logging.log(f"ERROR: Failed to remove quantized params for {layer_name}: {e}") - max_logging.log(f"Dumping params['params'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path( - lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), params["params"] - ) - max_logging.log(f"Dumping new_vars['aqt'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path(lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), aqt_vars) - raise - - # Restructure the aqt_vars for this layer to match pure Linen format for saving - if layer_prefix == "moe_layers": - structured_aqt = insert_deepseekmoeblock_scope(aqt_vars) - else: - structured_aqt = aqt_vars - quantized_params["aqt"]["decoder"][layer_name] = structured_aqt - - unquantized_layers = ["decoder_norm", "logits_dense"] - for unquantized_layer in unquantized_layers: - params = self._load_layer(unquantized_layer) - quantized_params["params"]["decoder"][unquantized_layer] = params["params"]["decoder"][unquantized_layer] - quantized_params["params"]["token_embedder"] = self._load_layer("token_embedder")["params"]["token_embedder"] - - maxtext_utils.save_quantized_checkpoint_if_configured(self.config, quantized_params) + self._load_and_quantize_nnx() def _load_and_quantize_nnx(self) -> None: """Whole-model NNX convert: load full-precision via TRAIN-mode `from_pretrained`, diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index f304873122..d126125bc0 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/src/maxtext/utils/lora_utils.py @@ -140,9 +140,8 @@ def unapply_lora_recursively(base_params, lora_params, module_name): def load_adapter(config, base_abstract_state_params, adapter_config_path, adapter_weights_path): """Load a LoRA adapter from disk and return its parameters. - When `config.pure_nnx` is True, `base_abstract_state_params` is the NNX - abstract param state (no outer `params` wrapper) and the returned - `lora_params` follows the same shape. Otherwise both use the Linen tree. + `base_abstract_state_params` is the NNX abstract param state (no outer + `params` wrapper) and the returned `lora_params` follows the same shape. Args: config: Top-level MaxText config. diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 50dce23794..e50c0eba01 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -1729,127 +1729,89 @@ def setup_initial_state( init_state_partial = init_state_fn init_state_partial.__name__ = "initialize_state" - if config.pure_nnx: - # Always build the concrete init state, then overlay whatever we loaded. Anything the - # checkpoint didn't carry (or a params-only load didn't touch) keeps its real init - # value, so restore doesn't depend on knowing exactly what was saved. - state = jax.jit( - lambda: nnx.state(init_state_partial()), # Get state only, mapping to out_sharding structure - in_shardings=None, - out_shardings=state_mesh_shardings, - )() - if raw_params: - # Params-only load (base model weights): overlay restored weights, keep init for everything else. - target_model = ( - state["model"] - if (isinstance(state, (nnx.State, dict)) and "model" in state) - else getattr(state, "model", state) - ) - raw_model_params = ( - raw_params["model"] if (isinstance(raw_params, (nnx.State, dict)) and "model" in raw_params) else raw_params - ) - if hasattr(raw_model_params, "to_pure_dict"): - raw_model_params = raw_model_params.to_pure_dict() - if isinstance(raw_model_params, dict) and "params" in raw_model_params: - raw_model_params = raw_model_params["params"] - target_pure = target_model.to_pure_dict() if hasattr(target_model, "to_pure_dict") else target_model - - def _reshard_aligned(target, raw): - """Aligns raw arrays with target device shardings using Flax's native flatten_dict utilities.""" - target_flat = traverse_util.flatten_dict(target) - raw_flat = traverse_util.flatten_dict(raw) - - res_flat = {} - for k, target_val in target_flat.items(): - if k in raw_flat and not isinstance(raw_flat[k], jax.ShapeDtypeStruct): - raw_val = raw_flat[k] - if hasattr(target_val, "sharding") and target_val.sharding is not None: - res_flat[k] = jax.device_put(raw_val, target_val.sharding) - else: - res_flat[k] = raw_val + # Always build the concrete init state, then overlay whatever we loaded. Anything the + # checkpoint didn't carry (or a params-only load didn't touch) keeps its real init + # value, so restore doesn't depend on knowing exactly what was saved. + state = jax.jit( + lambda: nnx.state(init_state_partial()), # Get state only, mapping to out_sharding structure + in_shardings=None, + out_shardings=state_mesh_shardings, + )() + if raw_params: + # Params-only load (base model weights): overlay restored weights, keep init for everything else. + target_model = ( + state["model"] + if (isinstance(state, (nnx.State, dict)) and "model" in state) + else getattr(state, "model", state) + ) + raw_model_params = ( + raw_params["model"] if (isinstance(raw_params, (nnx.State, dict)) and "model" in raw_params) else raw_params + ) + if hasattr(raw_model_params, "to_pure_dict"): + raw_model_params = raw_model_params.to_pure_dict() + if isinstance(raw_model_params, dict) and "params" in raw_model_params: + raw_model_params = raw_model_params["params"] + target_pure = target_model.to_pure_dict() if hasattr(target_model, "to_pure_dict") else target_model + + def _reshard_aligned(target, raw): + """Aligns raw arrays with target device shardings using Flax's native flatten_dict utilities.""" + target_flat = traverse_util.flatten_dict(target) + raw_flat = traverse_util.flatten_dict(raw) + + res_flat = {} + for k, target_val in target_flat.items(): + if k in raw_flat and not isinstance(raw_flat[k], jax.ShapeDtypeStruct): + raw_val = raw_flat[k] + if hasattr(target_val, "sharding") and target_val.sharding is not None: + res_flat[k] = jax.device_put(raw_val, target_val.sharding) else: - res_flat[k] = target_val + res_flat[k] = raw_val + else: + res_flat[k] = target_val - return traverse_util.unflatten_dict(res_flat) + return traverse_util.unflatten_dict(res_flat) - sharded_aligned = _reshard_aligned(target_pure, raw_model_params) - nnx.update(target_model, sharded_aligned) + sharded_aligned = _reshard_aligned(target_pure, raw_model_params) + nnx.update(target_model, sharded_aligned) - if restored: - is_emergency = isinstance( - checkpoint_manager, - ( - emergency_checkpoint_manager.CheckpointManager, - emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager, - ), - ) - overlay = restored if is_emergency else restored["items"] - if not (config.enable_diloco and isinstance(overlay, diloco.DiLoCoTrainState)): - # Standard Flax NNX branch: overlay weights into the initialized TrainStateNNX. - overlay_pure_dict = overlay.to_pure_dict() if hasattr(overlay, "to_pure_dict") else overlay - - def _has_shape_dtype_struct(tree): - return any(isinstance(x, jax.ShapeDtypeStruct) for x in jax.tree_util.tree_leaves(tree)) - - def _merge_restored_overlay(ckpt_node, init_node): - """Merges checkpoint overlay with initialized state, replacing ShapeDtypeStruct placeholders.""" - if _has_shape_dtype_struct(ckpt_node): - if isinstance(ckpt_node, dict) and isinstance(init_node, dict): - res = {} - for k in init_node: - if k in ckpt_node: - res[k] = _merge_restored_overlay(ckpt_node[k], init_node[k]) - else: - res[k] = init_node[k] - return res - else: - return init_node - return ckpt_node - - merged = _merge_restored_overlay(overlay_pure_dict, state.to_pure_dict()) - nnx.replace_by_pure_dict(state, merged) - else: - # DiLoCo branch: The restored checkpoint for DiLoCo is already a complete DiLoCoTrainState - # (containing inner_state, outer_opt_state, params, step) rather than a bare parameter mapping. - state = overlay - else: - if restored: - if isinstance( - checkpoint_manager, - ( - emergency_checkpoint_manager.CheckpointManager, - emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager, - ), - ): - state = restored - else: - # The update of data_iterator state happens in place, no need to assign explicitly - state = restored["items"] + if restored: + is_emergency = isinstance( + checkpoint_manager, + ( + emergency_checkpoint_manager.CheckpointManager, + emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager, + ), + ) + overlay = restored if is_emergency else restored["items"] + if not (config.enable_diloco and isinstance(overlay, diloco.DiLoCoTrainState)): + # Standard Flax NNX branch: overlay weights into the initialized TrainStateNNX. + overlay_pure_dict = overlay.to_pure_dict() if hasattr(overlay, "to_pure_dict") else overlay + + def _has_shape_dtype_struct(tree): + return any(isinstance(x, jax.ShapeDtypeStruct) for x in jax.tree_util.tree_leaves(tree)) + + def _merge_restored_overlay(ckpt_node, init_node): + """Merges checkpoint overlay with initialized state, replacing ShapeDtypeStruct placeholders.""" + if _has_shape_dtype_struct(ckpt_node): + if isinstance(ckpt_node, dict) and isinstance(init_node, dict): + res = {} + for k in init_node: + if k in ckpt_node: + res[k] = _merge_restored_overlay(ckpt_node[k], init_node[k]) + else: + res[k] = init_node[k] + return res + else: + return init_node + return ckpt_node + + merged = _merge_restored_overlay(overlay_pure_dict, state.to_pure_dict()) + nnx.replace_by_pure_dict(state, merged) else: - # pylint: disable=not-callable - state = jax.jit( - init_state_partial, - in_shardings=None, - out_shardings=state_mesh_shardings, - )() - if raw_params: # If we loaded a partial state, we need to merge it. - sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - if sparsity_enabled: - # Sparsity-init keeps freshly initialized params for any leaf still - # represented as an abstract ShapeDtypeStruct in raw_params (i.e. not - # actually restored), and uses the restored value otherwise. - def _merge_params(p_raw, p_init): - if isinstance(p_raw, jax.ShapeDtypeStruct): - return p_init - return p_raw - - merged_params = jax.tree_util.tree_map(_merge_params, raw_params, state.params) - state = state.replace(params=merged_params) - else: - state = state.replace(params=raw_params) + # DiLoCo branch: The restored checkpoint for DiLoCo is already a complete DiLoCoTrainState + # (containing inner_state, outer_opt_state, params, step) rather than a bare parameter mapping. + state = overlay - if not config.pure_nnx: - state = max_utils.unbox_logicallypartioned(state) if config.enable_diloco: state = diloco_utils.setup_diloco_initial_state( state=state, diff --git a/src/maxtext/utils/standalone_checkpointer.py b/src/maxtext/utils/standalone_checkpointer.py index 00ccf849f2..8fc5e1a9bd 100644 --- a/src/maxtext/utils/standalone_checkpointer.py +++ b/src/maxtext/utils/standalone_checkpointer.py @@ -19,7 +19,6 @@ # See github.com/google/maxtext/issues/20 for more import datetime -from functools import partial import os import time from typing import Sequence @@ -33,7 +32,6 @@ from maxtext.common import checkpointing from maxtext.common import emergency_checkpointing from maxtext.common import train_state_nnx -from maxtext.models import models from maxtext.trainers.pre_train.train import get_first_step from maxtext.utils import max_logging from maxtext.utils import maxtext_utils @@ -43,8 +41,6 @@ from maxtext.utils.model_creation_utils import from_config import numpy as np -Transformer = models.transformer_as_linen - def _as_abstract_leaf(leaf): """Shape/dtype/sharding stand-in for a concrete leaf, as an Orbax restore target.""" @@ -61,28 +57,21 @@ def checkpoint_loop(config, state=None): on the configured cadence. Works on both Linen and NNX state shapes. """ init_rng = jax.random.PRNGKey(config.init_weights_seed) - if config.pure_nnx: - mesh = maxtext_utils.get_mesh_from_config(config) - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) - model = from_config(config, mesh=mesh, rngs=rngs) - _, tx = train_utils.create_training_optimizer(config, model) - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) - - def init_state_fn(): - nnx_model = _create_model_partial() - wrt = ( - getattr(nnx, "LoRAParam", nnx.Param) - if getattr(getattr(config, "lora", None), "enable_lora", False) - else nnx.Param - ) - optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - - else: - model = from_config(config) - mesh = model.mesh - _, tx = train_utils.create_training_optimizer(config, model) - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, init_rng) + mesh = maxtext_utils.get_mesh_from_config(config) + rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) + model = from_config(config, mesh=mesh, rngs=rngs) + _, tx = train_utils.create_training_optimizer(config, model) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) + + def init_state_fn(): + nnx_model = _create_model_partial() + wrt = ( + getattr(nnx, "LoRAParam", nnx.Param) + if getattr(getattr(config, "lora", None), "enable_lora", False) + else nnx.Param + ) + optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn) @@ -128,7 +117,7 @@ def init_state_fn(): start_time = datetime.datetime.now() # A barrier to sync all hosts before starting to save checkpoint jax.experimental.multihost_utils.sync_global_devices("Barrier before save") - state_to_save = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) if config.pure_nnx else state + state_to_save = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) if checkpointing.save_checkpoint(checkpoint_manager, int(step), state_to_save): checkpointing.wait_until_finished(checkpoint_manager) end_time = datetime.datetime.now() @@ -187,8 +176,8 @@ def add_entropy_to_checkpoint(state): * Linen `TrainState`: `state.params` + `state.opt_state` (tuple). * NNX `TrainStateNNX` (Module): `state.model` is an `nnx.Module`; the optimizer's `opt_state` is the optax tuple of NamedTuples. - * NNX `nnx.State` (post-split, what `setup_training_state` returns under - `pure_nnx`): `state.model` and `state.optimizer.opt_state` are sub-States; + * NNX `nnx.State` (post-split, what `setup_training_state` returns): + `state.model` and `state.optimizer.opt_state` are sub-States; `opt_state[0].mu`/`nu` are themselves States that can be reassigned. """ if hasattr(state, "model"): diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh index 20db17beee..257e89302c 100755 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh @@ -67,4 +67,12 @@ python3 -m maxtext.inference.vllm_decode \ prompt="Suggest some famous landmarks in London." \ use_chat_template=True \ scan_layers=true \ - enable_single_controller=${use_pathways} \ No newline at end of file + enable_single_controller=${use_pathways} + +# Step 5: Convert the checkpoint from MaxText format to Hugging Face format +python3 -m maxtext.checkpoint_conversion.to_huggingface \ + model_name=${MODEL_NAME} \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + lora.lora_restore_path=${BASE_OUTPUT_DIRECTORY}/lora/${run_id}/checkpoints/5/model_params \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/unscanned/${run_id} \ + scan_layers=true diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh index 7e455e9532..0c98db9e04 100644 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh @@ -35,9 +35,7 @@ if [[ "${CKPT_PATH,,}" == *lora* ]]; then lora.lora_restore_path=${LORA_RESTORE_PATH} \ base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/${scan_status}/${run_id} \ use_multimodal=${USE_MULTIMODAL} \ - scan_layers=${SCAN_LAYERS} \ - enable_nnx=True \ - pure_nnx_decoder=True + scan_layers=${SCAN_LAYERS} else python3 -m maxtext.checkpoint_conversion.to_huggingface \ diff --git a/tests/inference/benchmark_vllm.sh b/tests/inference/benchmark_vllm.sh index 76a02e3a75..f555e0449f 100755 --- a/tests/inference/benchmark_vllm.sh +++ b/tests/inference/benchmark_vllm.sh @@ -81,7 +81,7 @@ $PYTHON_EXEC -m vllm.entrypoints.cli.main bench throughput \ --gpu-memory-utilization "$GPU_MEM_UTIL" \ --load-format "$LOAD_FORMAT" \ --hf-overrides '{"architectures": ["MaxTextForCausalLM"]}' \ - --additional-config "{\"maxtext_config\": {\"model_name\": \"$MODEL_NAME\", \"weight_dtype\": \"bfloat16\", \"attention\": \"$ATTENTION\", \"allow_split_physical_axes\": true, \"scan_layers\": true, \"enable_nnx\": true, \"pure_nnx_decoder\": true}}" \ + --additional-config "{\"maxtext_config\": {\"model_name\": \"$MODEL_NAME\", \"weight_dtype\": \"bfloat16\", \"attention\": \"$ATTENTION\", \"allow_split_physical_axes\": true, \"scan_layers\": true}}" \ --dataset-name random \ --random-input-len "$INPUT_LEN" \ --random-output-len "$OUTPUT_LEN" \ diff --git a/tests/integration/diloco_test.py b/tests/integration/diloco_test.py index 9002176825..a66c9cbe1b 100644 --- a/tests/integration/diloco_test.py +++ b/tests/integration/diloco_test.py @@ -432,19 +432,6 @@ def test_apply_flat_fragment_shapedtypestruct(self): res = manipulator.apply_flat_fragment(abstract_tree, fragment_idx=1, flat_fragment=frag) self.assertIsInstance(res["decoder"]["layers"], jax.ShapeDtypeStruct) - def test_diloco_requires_pure_nnx(self): - """Tests that enable_diloco=True raises ValueError if pure_nnx=False.""" - with self.assertRaises(ValueError) as ctx: - initialize_pydantic( - [ - "", - get_test_config_path(), - "enable_diloco=true", - "pure_nnx=false", - ] - ) - self.assertIn("enable_diloco=True requires pure_nnx=True", str(ctx.exception)) - def test_diloco_checkpoint_saving_and_normal_resume(self): """Tests that DiLoCo checkpoints save outer params under params/params for direct normal pre-training resume.""" temp_dir = tempfile.mkdtemp() diff --git a/tests/integration/lora_e2e_nnx_test.py b/tests/integration/lora_e2e_nnx_test.py index 5b43167510..561e115ae6 100644 --- a/tests/integration/lora_e2e_nnx_test.py +++ b/tests/integration/lora_e2e_nnx_test.py @@ -34,7 +34,6 @@ def _tiny_lora_pyconfig(run_name, checkpoint_dir, **overrides): "enable_checkpointing": True, "dataset_type": "synthetic", "model_name": "default", - "pure_nnx": True, "per_device_batch_size": 1.0, "base_emb_dim": 8, "base_num_query_heads": 4, diff --git a/tests/post_training/integration/inference_router_replay_test.py b/tests/post_training/integration/inference_router_replay_test.py index 283aae921d..8bec4c65bc 100644 --- a/tests/post_training/integration/inference_router_replay_test.py +++ b/tests/post_training/integration/inference_router_replay_test.py @@ -92,9 +92,6 @@ def test_extract_router_replay_from_inference(self): "ici_tensor_parallelism": 4, "ici_data_parallelism": 1, "ici_expert_parallelism": 1, - "enable_nnx": True, - "pure_nnx": True, - "pure_nnx_decoder": True, } cfg_infer = pyconfig.initialize( diff --git a/tests/post_training/integration/lora_e2e_nnx_test.py b/tests/post_training/integration/lora_e2e_nnx_test.py index 63a664cd42..87722c32dd 100644 --- a/tests/post_training/integration/lora_e2e_nnx_test.py +++ b/tests/post_training/integration/lora_e2e_nnx_test.py @@ -34,7 +34,6 @@ def _tiny_lora_pyconfig(run_name, checkpoint_dir, **overrides): "enable_checkpointing": True, "dataset_type": "synthetic", "model_name": "default", - "pure_nnx": True, "per_device_batch_size": 1.0, "base_emb_dim": 8, "base_num_query_heads": 4, diff --git a/tests/post_training/unit/rl_lr_schedule_test.py b/tests/post_training/unit/rl_lr_schedule_test.py index fb6a3f822f..950518cde3 100644 --- a/tests/post_training/unit/rl_lr_schedule_test.py +++ b/tests/post_training/unit/rl_lr_schedule_test.py @@ -66,8 +66,6 @@ "base_mlp_dim": 256, "max_prefill_predict_length": 4, "model_name": "llama2-7b", - "enable_nnx": True, - "pure_nnx_decoder": True, "override_model_config": True, "weight_dtype": "bfloat16", } diff --git a/tests/unit/checkpointing_emergency_nnx_test.py b/tests/unit/checkpointing_emergency_nnx_test.py index 2547e84fe4..0ce33aee63 100644 --- a/tests/unit/checkpointing_emergency_nnx_test.py +++ b/tests/unit/checkpointing_emergency_nnx_test.py @@ -12,9 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for restoring pure_nnx emergency checkpoints. +"""Tests for restoring NNX emergency checkpoints. -pure_nnx saves in the on-disk checkpoint layout: Linen params/opt_state/step, plus an nnx_aux +NNX saves in the on-disk checkpoint layout: Linen params/opt_state/step, plus an nnx_aux subtree for the rngs and batch stats Linen has no place for. The emergency manager, unlike the regular one, bakes in the abstract it is built with and restores against that, ignoring the restore-time item. So that abstract has to be in the on-disk layout already, or Orbax compares @@ -74,9 +74,8 @@ def __call__(self, x, deterministic=False): def _config(): - """Minimal config with the fields save/restore reads for a pure_nnx emergency run.""" + """Minimal config with the fields save/restore reads for an NNX emergency run.""" return SimpleNamespace( - pure_nnx=True, enable_diloco=False, enable_checkpointing=True, enable_continuous_checkpointing=False, diff --git a/tests/unit/checkpointing_nnx_missing_param_test.py b/tests/unit/checkpointing_nnx_missing_param_test.py index 7c3ac0a5e4..2d419bfb3a 100644 --- a/tests/unit/checkpointing_nnx_missing_param_test.py +++ b/tests/unit/checkpointing_nnx_missing_param_test.py @@ -71,9 +71,8 @@ def __call__(self, x, deterministic=False): def _config(): - """Minimal config with the fields save/restore reads for a pure_nnx run.""" + """Minimal config with the fields save/restore reads for an NNX run.""" return SimpleNamespace( - pure_nnx=True, enable_diloco=False, enable_checkpointing=True, enable_continuous_checkpointing=False, diff --git a/tests/unit/checkpointing_nnx_roundtrip_test.py b/tests/unit/checkpointing_nnx_roundtrip_test.py index 25854e5b50..3260be204a 100644 --- a/tests/unit/checkpointing_nnx_roundtrip_test.py +++ b/tests/unit/checkpointing_nnx_roundtrip_test.py @@ -82,9 +82,8 @@ def __call__(self, x, deterministic=False): def _config(): - """Minimal config with the fields save/restore reads for a pure_nnx run.""" + """Minimal config with the fields save/restore reads for an NNX run.""" return SimpleNamespace( - pure_nnx=True, enable_diloco=False, enable_checkpointing=True, enable_continuous_checkpointing=False, diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index 9abcbcada8..7b23c84b44 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -457,7 +457,6 @@ def test_error_handler_raises_runtime_error(self): """Unexpected checkpointing errors should raise RuntimeError with original error chained.""" config = mock.Mock() config.checkpoint_period = 1 - config.pure_nnx = True config.enable_diloco = False config.async_checkpointing = False config.enable_continuous_checkpointing = False diff --git a/tests/unit/decoder_layer_model_mode_test.py b/tests/unit/decoder_layer_model_mode_test.py index e4247f9480..f41e7d5ba5 100644 --- a/tests/unit/decoder_layer_model_mode_test.py +++ b/tests/unit/decoder_layer_model_mode_test.py @@ -80,7 +80,6 @@ "weight_dtype": "float32", "enable_checkpointing": False, "skip_jax_distributed_system": True, - "pure_nnx": True, } _DEEPSEEK = { diff --git a/tests/unit/deepseek_decode_consistency_test.py b/tests/unit/deepseek_decode_consistency_test.py index bdeeaa9c72..89f86e9f12 100644 --- a/tests/unit/deepseek_decode_consistency_test.py +++ b/tests/unit/deepseek_decode_consistency_test.py @@ -68,7 +68,6 @@ "decode_sampling_strategy": "greedy", "enable_checkpointing": False, "skip_jax_distributed_system": True, - "pure_nnx": True, } _DEEPSEEK = { diff --git a/tests/unit/generate_param_only_checkpoint_nnx_test.py b/tests/unit/generate_param_only_checkpoint_nnx_test.py index 57995a6145..800610fb15 100644 --- a/tests/unit/generate_param_only_checkpoint_nnx_test.py +++ b/tests/unit/generate_param_only_checkpoint_nnx_test.py @@ -92,7 +92,6 @@ def test_unrolls_scanned_layers(self): config = SimpleNamespace( scan_layers=True, force_unroll=True, - pure_nnx=True, param_scan_axis=0, decoder_block=DecoderBlockType.LLAMA2, num_decoder_layers=num_layers, @@ -133,7 +132,6 @@ def __init__(self): config = SimpleNamespace( scan_layers=True, force_unroll=True, - pure_nnx=True, param_scan_axis=0, decoder_block=DecoderBlockType.DEEPSEEK, num_decoder_layers=5, @@ -184,7 +182,6 @@ def test_unrolls_scanned_lora_layers(self): config = SimpleNamespace( scan_layers=True, force_unroll=True, - pure_nnx=True, param_scan_axis=0, decoder_block=DecoderBlockType.LLAMA2, num_decoder_layers=num_layers, diff --git a/tests/unit/multi_token_prediction_test.py b/tests/unit/multi_token_prediction_test.py index 93bcb3b531..86772b3de9 100644 --- a/tests/unit/multi_token_prediction_test.py +++ b/tests/unit/multi_token_prediction_test.py @@ -651,8 +651,6 @@ def _build(self, mtp_num_layers): head_dim=8, max_target_length=16, vocab_size=32, - pure_nnx=True, - pure_nnx_decoder=True, use_qwix_quantization=True, quantization="int8", enable_dropout=False, diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index 3684210433..16f2a41f12 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -1066,39 +1066,6 @@ def test_decoder_keeps_layers_past_the_last_whole_block(self): ) -class TestQwen3NextDecoderParity(unittest.TestCase): - """The Linen `Decoder` and the pure-NNX `NNXDecoder` must emit the same parameter tree. - - One checkpoint mapping (`QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING`) serves both - decoders, so a name or shape that differs between them silently breaks - conversion on whichever side the mapping was not written against. - """ - - def _param_tree(self, num_decoder_layers, pure_nnx_decoder): - """Returns {parameter key: shape} for the chosen decoder implementation.""" - # pylint: disable=import-outside-toplevel - from maxtext.checkpoint_conversion.to_maxtext import get_maxtext_model_info - - cfg = _make_config( - **{ - **_QWEN3_NEXT_CONFIG, - "base_num_decoder_layers": num_decoder_layers, - "scan_layers": True, - "pure_nnx_decoder": pure_nnx_decoder, - } - ) - model_info, _ = get_maxtext_model_info(cfg) - return {key: shape for key, (_, shape) in model_info.items()} - - def test_decoders_agree_on_whole_periods(self): - self.assertEqual(self._param_tree(8, True), self._param_tree(8, False)) - - def test_decoders_agree_with_a_remainder(self): - """6 layers is one whole period plus a two-layer remainder, which both decoders - have to put in a `layers_remainder` block rather than spell out layer by layer.""" - self.assertEqual(self._param_tree(6, True), self._param_tree(6, False)) - - class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase): """Tests for Deepseek and Gemma4 specific decoder logic.""" diff --git a/tests/unit/router_replay_test.py b/tests/unit/router_replay_test.py index e83f5be8a8..522a28ef9a 100644 --- a/tests/unit/router_replay_test.py +++ b/tests/unit/router_replay_test.py @@ -51,9 +51,8 @@ def _init_test_cfg(extra_args=(), **kwargs): """pyconfig.initialize with this file's common test defaults folded in. Only pass kwargs that actually need to differ from base.yml (or from the - selected model's own yaml) -- e.g. omit ici_*_parallelism, enable_nnx, - pure_nnx, pure_nnx_decoder, sparse_matmul, dtype, and scan_layers=True, - which already match their base.yml defaults. + selected model's own yaml) -- e.g. omit ici_*_parallelism, sparse_matmul, + dtype, and scan_layers=True, which already match their base.yml defaults. """ kwargs.setdefault("enable_checkpointing", False) kwargs.setdefault("log_config", False) diff --git a/tests/unit/setup_initial_state_nnx_test.py b/tests/unit/setup_initial_state_nnx_test.py index 0ee9073b4b..48ca039d75 100644 --- a/tests/unit/setup_initial_state_nnx_test.py +++ b/tests/unit/setup_initial_state_nnx_test.py @@ -87,9 +87,8 @@ def fn(): def _config(): - """A config with every field setup_initial_state + load_state_if_possible + save read for pure_nnx.""" + """A config with every field setup_initial_state + load_state_if_possible + save read for NNX.""" return SimpleNamespace( - pure_nnx=True, enable_diloco=False, enable_checkpointing=True, enable_continuous_checkpointing=False, diff --git a/tests/unit/standalone_checkpointer_nnx_test.py b/tests/unit/standalone_checkpointer_nnx_test.py index 73f63c1229..31ae42f9e1 100644 --- a/tests/unit/standalone_checkpointer_nnx_test.py +++ b/tests/unit/standalone_checkpointer_nnx_test.py @@ -105,7 +105,6 @@ def _create_mock_config(self, start_from_checkpoint=False): """Creates a mock configuration for checkpoint_loop tests.""" config = mock.MagicMock() config.init_weights_seed = 0 - config.pure_nnx = False config.steps = 2 config.standalone_checkpointer_per_step_interval = 0.5 config.standalone_checkpointer_drop_page_cache_before_restore = True @@ -123,6 +122,12 @@ def test_checkpoint_loop_default_save_and_restore(self): with ( mock.patch("maxtext.utils.standalone_checkpointer.from_config"), + mock.patch("maxtext.utils.maxtext_utils.get_mesh_from_config"), + mock.patch("maxtext.utils.maxtext_utils_nnx.create_nnx_rngs"), + mock.patch( + "maxtext.utils.model_creation_utils.create_nnx_abstract_model", + return_value=(mock.MagicMock(), None), + ), mock.patch("maxtext.utils.train_utils.create_training_optimizer", return_value=(None, mock.MagicMock())), mock.patch("maxtext.utils.train_utils.create_checkpoint_manager", return_value=mock_ckpt_mgr), mock.patch( @@ -144,7 +149,7 @@ def test_checkpoint_loop_default_save_and_restore(self): mock_setup.assert_called_once() mock_sleep.assert_called_once() mock_system.assert_called_once_with("sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches'") - mock_ckpt_mgr.load_checkpointables.assert_called_once_with(1, {"items": mock_state}) + mock_ckpt_mgr.load_checkpointables.assert_called_once() def test_checkpoint_loop_start_from_checkpoint(self): """Tests initializing state from an existing checkpoint.""" @@ -155,6 +160,12 @@ def test_checkpoint_loop_start_from_checkpoint(self): with ( mock.patch("maxtext.utils.standalone_checkpointer.from_config"), + mock.patch("maxtext.utils.maxtext_utils.get_mesh_from_config"), + mock.patch("maxtext.utils.maxtext_utils_nnx.create_nnx_rngs"), + mock.patch( + "maxtext.utils.model_creation_utils.create_nnx_abstract_model", + return_value=(mock.MagicMock(), None), + ), mock.patch("maxtext.utils.train_utils.create_training_optimizer", return_value=(None, mock.MagicMock())), mock.patch("maxtext.utils.train_utils.create_checkpoint_manager", return_value=mock_ckpt_mgr), mock.patch("maxtext.utils.maxtext_utils.get_abstract_state", return_value=(mock.MagicMock(), None, None)), @@ -187,6 +198,12 @@ def test_checkpoint_loop_start_from_checkpoint_fallback(self): with ( mock.patch("maxtext.utils.standalone_checkpointer.from_config"), + mock.patch("maxtext.utils.maxtext_utils.get_mesh_from_config"), + mock.patch("maxtext.utils.maxtext_utils_nnx.create_nnx_rngs"), + mock.patch( + "maxtext.utils.model_creation_utils.create_nnx_abstract_model", + return_value=(mock.MagicMock(), None), + ), mock.patch("maxtext.utils.train_utils.create_training_optimizer", return_value=(None, mock.MagicMock())), mock.patch("maxtext.utils.train_utils.create_checkpoint_manager", return_value=mock_ckpt_mgr), mock.patch("maxtext.utils.maxtext_utils.get_abstract_state", return_value=(mock.MagicMock(), None, None)), @@ -209,10 +226,9 @@ def test_checkpoint_loop_start_from_checkpoint_fallback(self): self.assertIs(returned_state, mock_state) mock_setup.assert_called_once() - def test_checkpoint_loop_pure_nnx(self): - """Tests checkpoint loop with pure_nnx enabled.""" + def test_checkpoint_loop_saves_in_linen_layout(self): + """Tests the checkpoint loop saves NNX state in the Linen on-disk layout.""" config = self._create_mock_config(start_from_checkpoint=False) - config.pure_nnx = True mock_state = mock.MagicMock() mock_state.to_pure_dict.return_value = {"params": {}} mock_ckpt_mgr = mock.MagicMock() diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index 8ad8fb89af..95eccd4843 100644 --- a/tests/unit/train_compile_test.py +++ b/tests/unit/train_compile_test.py @@ -30,7 +30,6 @@ from tempfile import gettempdir, NamedTemporaryFile -from maxtext.configs import pyconfig from maxtext.trainers.pre_train.train_compile import main as train_compile_main from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from tests.utils.test_helpers import get_test_config_path @@ -495,10 +494,7 @@ def test_moe_dense_int8(self): ) def test_moe_pp_bf16(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Pipeline parallelism not supported for pure_nnx_decoder=True") - + pytest.skip("Pipeline parallelism is not supported for the NNX decoder.") temp_dir = gettempdir() compiled_trainstep_file = os.path.join(temp_dir, "test_moe_pp_bf16.pickle") train_compile_main( @@ -648,10 +644,7 @@ def test_moe_deepseek_with_device_limit(self): ) def test_moe_deepseek_pipeline_subset(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Pipeline parallelism not supported for pure_nnx_decoder=True") - + pytest.skip("Pipeline parallelism is not supported for the NNX decoder.") compiled_trainstep_file = "/tmp/test_moe_deepseek_pipeline_subset.pickle" train_compile_main( ( @@ -674,10 +667,7 @@ def test_moe_deepseek_pipeline_subset(self): ) def test_pipeline_subset(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Test not supported for pure_nnx_decoder=True") - + pytest.skip("Not supported for the NNX decoder.") compiled_trainstep_file = "/tmp/test_pipeline_subset.pickle" train_compile_main( ( @@ -990,10 +980,7 @@ def test_engram_integration(self): ) def test_circular_pipeline_ag_per_repeat_ep_ds(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Pipeline parallelism not supported for pure_nnx_decoder=True") - + pytest.skip("Pipeline parallelism is not supported for the NNX decoder.") temp_dir = gettempdir() compiled_trainstep_file = os.path.join(temp_dir, "test_circular_pipeline_ag_per_repeat_ep_ds.pickle") train_compile_main( diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py index ba44a46438..a1ac180623 100644 --- a/tests/unit/train_state_nnx_checkpoint_test.py +++ b/tests/unit/train_state_nnx_checkpoint_test.py @@ -351,7 +351,6 @@ def setUp(self): def _config(self, **overrides): """Builds a minimal checkpoint config for maybe_save_checkpoint tests.""" values = { - "pure_nnx": True, "checkpoint_period": 10, "async_checkpointing": False, "enable_diloco": False, @@ -391,10 +390,10 @@ def _build_linen_state(self, num_steps): state = state.apply_gradients(grads=grads) return state - def _invoke_maybe_save(self, state, pure_nnx): + def _invoke_maybe_save(self, state): """Call maybe_save_checkpoint with save_checkpoint patched, return {step, state} captured.""" # checkpoint_period=1 keeps force_ckpt_save False regardless of actual_step. - config = self._config(pure_nnx=pure_nnx, checkpoint_period=1) + config = self._config(checkpoint_period=1) mgr = mock.MagicMock() mgr.latest = None # no existing checkpoint -> _latest_step() is None, so the save proceeds @@ -415,30 +414,15 @@ def fake_save_checkpoint(_mgr, step, state_arg, *_args, **_kwargs): def test_nnx_final_save_step_is_n_minus_1(self): state = self._build_nnx_state(self.N_STEPS) self.assertEqual(int(state.optimizer.step.value), self.N_STEPS) - captured = self._invoke_maybe_save(state, pure_nnx=True) - self.assertEqual(captured["step"], self.N_STEPS - 1) - - def test_linen_final_save_step_is_n_minus_1(self): - state = self._build_linen_state(self.N_STEPS) - self.assertEqual(int(state.step), self.N_STEPS) - captured = self._invoke_maybe_save(state, pure_nnx=False) + captured = self._invoke_maybe_save(state) self.assertEqual(captured["step"], self.N_STEPS - 1) - def test_nnx_and_linen_agree_on_actual_step(self): - """TrainStateNNX and Linen TrainState must yield the same fallback actual_step.""" - nnx_state = self._build_nnx_state(self.N_STEPS) - linen_state = self._build_linen_state(self.N_STEPS) - self.assertEqual( - self._invoke_maybe_save(nnx_state, pure_nnx=True)["step"], - self._invoke_maybe_save(linen_state, pure_nnx=False)["step"], - ) - def test_nnx_state_is_saved_in_linen_layout(self): - """For pure_nnx=True, save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" + """save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" state = self._build_nnx_state(self.N_STEPS) self.assertIsInstance(state, nnx.State) # precondition: NNX train_step returns an nnx.State - config = self._config(pure_nnx=True, enable_checkpointing=True, checkpoint_period=1) + config = self._config(enable_checkpointing=True, checkpoint_period=1) mgr = mock.MagicMock() mgr.use_async = False @@ -459,13 +443,6 @@ def test_nnx_state_is_saved_in_linen_layout(self): self.assertNotIn("optimizer", saved) self.assertIn("params", saved["params"]) - def test_linen_state_is_passed_through_unchanged(self): - """For pure_nnx=False, maybe_save_checkpoint must pass the original TrainState object through.""" - state = self._build_linen_state(self.N_STEPS) - captured = self._invoke_maybe_save(state, pure_nnx=False) - # Linen path must not invoke to_pure_dict(); state is forwarded as-is. - self.assertIs(captured["state"], state) - def test_maybe_save_checkpoint_skips_if_already_saved(self): """Verify maybe_save_checkpoint skips saving if latest_step matches actual_step.""" state = self._build_nnx_state(self.N_STEPS) diff --git a/tests/utils/sharding_dump.py b/tests/utils/sharding_dump.py index 09f2c9049c..42b002a291 100644 --- a/tests/utils/sharding_dump.py +++ b/tests/utils/sharding_dump.py @@ -200,9 +200,6 @@ def main(argv: Sequence[str]) -> None: "weight_dtype", "log_config", "debug_sharding", - "pure_nnx", - "enable_nnx", - "pure_nnx_decoder", ]: overrides.append(arg) if overrides: