From 7fa80fc24dda8e120ef06ac543128b3283ec61f3 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Tue, 28 Jul 2026 19:52:30 +0000 Subject: [PATCH] [NNX] Delete Linen (1a/5): collapse dispatch in pre-train checkpoint tooling --- .../convert_gpt3_ckpt_from_paxml.py | 90 ++++------- src/maxtext/common/checkpointing.py | 33 ++-- src/maxtext/common/train_state_nnx.py | 6 +- .../utils/generate_param_only_checkpoint.py | 144 +++--------------- src/maxtext/utils/standalone_checkpointer.py | 32 ++-- ...generate_param_only_checkpoint_nnx_test.py | 3 - tests/unit/train_state_nnx_checkpoint_test.py | 52 ++----- 7 files changed, 88 insertions(+), 272 deletions(-) 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 2f3e26db7b..8278e6989a 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,7 +35,6 @@ """ import argparse -import functools import gc import os import sys @@ -47,11 +46,7 @@ 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 @@ -92,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, @@ -201,24 +188,18 @@ def init_state_fn(): "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None), } - if cfg.pure_nnx: - # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a - # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so - # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus - # `.value` for the Variable leaf, plus `[idx]` for the optax tuple: - # model params -> ['model'].value - # adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu'].value - # step -> ['optimizer']['step'].value - # opt count -> ['optimizer']['opt_state'][0]['count'].value - state_map = { - "['optimizer']['step'].value": ("step", None), - "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None), - } - else: - state_map = { - ".step": ("step", None), - ".opt_state.count": ("opt_states_0.no_prefix_0.count", None), - } + # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a + # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so + # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus + # `.value` for the Variable leaf, plus `[idx]` for the optax tuple: + # model params -> ['model'].value + # adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu'].value + # step -> ['optimizer']['step'].value + # opt count -> ['optimizer']['opt_state'][0]['count'].value + state_map = { + "['optimizer']['step'].value": ("step", None), + "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None), + } def get_layer_prefix(keystr_pax): # different path format between decoder_layer variable @@ -231,26 +212,15 @@ def get_layer_prefix(keystr_pax): for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items(): prefix_pax_opt_state = get_layer_prefix(keystr_pax) - if cfg.pure_nnx: - state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) - else: - state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f".opt_state.mu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f".opt_state.nu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) + state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn) + state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = ( + f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", + transform_fn, + ) + state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = ( + f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", + transform_fn, + ) def verify_fn(key_path, _): keystr = jax.tree_util.keystr(key_path) @@ -302,7 +272,7 @@ def map_fn(key_path, value): max_logging.log("converted state finished") max_utils.print_mem_stats("converted state finished") - step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step + step_value = int(converted_state.optimizer.step.value) if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state): max_logging.log(f"saved a checkpoint at step {step_value}") # Upon preemption, exit when and only when all ongoing saves are complete. diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 4827cbc1a4..b556465bf5 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -138,7 +138,7 @@ def _load_linen_checkpoint_into_nnx( use_ocdbt, use_zarr3, ): - """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 `_linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when @@ -273,7 +273,7 @@ def combine_sharding(sds, shardings): else: raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}") else: - # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. + # Checkpoints use the Linen on-disk layout; reshape it back into the NNX state. if isinstance(abstract_unboxed_pre_state, nnx.State): return _load_linen_checkpoint_into_nnx( path, @@ -730,12 +730,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. @@ -753,17 +749,16 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.") return - if config.pure_nnx: - # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. - if config.enable_diloco: - # DiLoCoTrainState: persist the synchronized global model (outer params). - # The per-replica inner optimizer / outer-momentum state is not checkpointed. - step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step - state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) - else: - # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout - # stream continues across resumes instead of resetting to a base key. - state = train_state_nnx.to_checkpoint_dict(state) + # Save in the Linen on-disk layout so NNX and Linen checkpoints are interchangeable. + if config.enable_diloco: + # DiLoCoTrainState: persist the synchronized global model (outer params). + # The per-replica inner optimizer / outer-momentum state is not checkpointed. + step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step + state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) + else: + # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout + # stream continues across resumes instead of resetting to a base key. + state = train_state_nnx.to_checkpoint_dict(state) try: checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save) diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index 7d73f45d4a..602b1aa3c4 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -58,9 +58,9 @@ 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 -# interchangeable. The NNX state pure dict differs from Linen's in three ways, all -# reshaped below at save time: +# NNX saves in the same on-disk layout as the old Linen format, so checkpoints +# stay interchangeable. The NNX state pure dict differs from that layout in three +# ways, all reshaped below at save time: # 1. top-level keys: {model, optimizer:{step, opt_state}} -> {params:{params:...}, step, opt_state} # 2. weights: model/... -> params/params/... (Linen `params` collection) # 3. opt_state: int-keyed dict (empty entries skipped) -> list with None for EmptyState, diff --git a/src/maxtext/utils/generate_param_only_checkpoint.py b/src/maxtext/utils/generate_param_only_checkpoint.py index 6e08166262..d38dc850e5 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,92 +161,35 @@ 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() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + 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) - 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) + 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) - 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): @@ -475,11 +374,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/standalone_checkpointer.py b/src/maxtext/utils/standalone_checkpointer.py index bd32f74b48..31cefa498d 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 from typing import Sequence @@ -51,23 +50,16 @@ 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() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - 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() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn) @@ -114,8 +106,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/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/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py index ef19ac427e..ff500b4fc9 100644 --- a/tests/unit/train_state_nnx_checkpoint_test.py +++ b/tests/unit/train_state_nnx_checkpoint_test.py @@ -338,12 +338,10 @@ class TestMaybeSaveCheckpointStepAlignment(unittest.TestCase): """Verify maybe_save_checkpoint's fallback step matches the last completed step. When the training loop's final save calls maybe_save_checkpoint without an - explicit `step`, it derives `actual_step` from the state: - - NNX: int(state.optimizer.step) - 1 - - Linen: int(state.step) - 1 - Both TrainStateNNX.apply_gradients (via nnx.Optimizer.update) and Linen - TrainState.apply_gradients increment the counter by 1 per call, so after N - gradient applications the counter is N and the "last completed step" is N-1. + explicit `step`, it derives `actual_step` from int(state.optimizer.step) - 1. + TrainStateNNX.apply_gradients (via nnx.Optimizer.update) increments the + counter by 1 per call, so after N gradient applications the counter is N and + the "last completed step" is N-1. """ N_STEPS = 5 @@ -384,20 +382,10 @@ def loss_fn(m): # (train_step returns nnx.state(new_state)). return nnx.state(state) - def _build_linen_state(self, num_steps): - """Build a Linen TrainState after num_steps gradient applications.""" - model = LinenMockModel() - variables = model.init(jax.random.key(0), jnp.ones((1, 2))) - state = train_state.TrainState.create(apply_fn=model.apply, params=variables["params"], tx=self.tx) - grads = jax.tree.map(jnp.ones_like, state.params) - for _ in range(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.reached_preemption.return_value = False @@ -415,30 +403,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, maybe_save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" + """maybe_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 - captured = self._invoke_maybe_save(state, pure_nnx=True) + captured = self._invoke_maybe_save(state) # save_checkpoint should receive a plain dict in Linen layout, not the nnx.State. self.assertIsInstance(captured["state"], dict) @@ -451,13 +424,6 @@ def test_nnx_state_is_saved_in_linen_layout(self): self.assertNotIn("optimizer", captured["state"]) self.assertIn("params", captured["state"]["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)