From efc5353e1166e3a30e6e6ddf79f0c9d32194affe Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Wed, 19 Aug 2026 15:36:52 +0000 Subject: [PATCH] [NNX] Delete Linen (2/5): collapse dispatch in trainers (pre-train, DiLoCo, GRPO) --- src/maxtext/experimental/rl/grpo_trainer.py | 309 ++------ src/maxtext/trainers/pre_train/train.py | 702 ++++++------------ .../trainers/pre_train/train_compile.py | 90 +-- .../golden_logits/golden_dpo_correctness.json | 18 +- .../generate_grpo_golden_logits.py | 93 +-- tests/integration/diloco_test.py | 163 ++-- .../integration/dpo_correctness_base.py | 3 - .../integration/grpo_correctness.py | 91 +-- .../grpo_trainer_correctness_test.py | 54 +- .../sft_trainer_correctness_test.py | 41 +- tests/unit/grpo_nnx_test.py | 72 +- tests/unit/pre_train_loss_mask_test.py | 61 +- tests/unit/router_replay_test.py | 59 +- 13 files changed, 490 insertions(+), 1266 deletions(-) diff --git a/src/maxtext/experimental/rl/grpo_trainer.py b/src/maxtext/experimental/rl/grpo_trainer.py index d8661e9814..6d68133822 100644 --- a/src/maxtext/experimental/rl/grpo_trainer.py +++ b/src/maxtext/experimental/rl/grpo_trainer.py @@ -38,7 +38,6 @@ import datetime import time import os -import functools import threading from typing import Sequence, Callable, Iterator @@ -109,43 +108,6 @@ # ----------------------------------------------------------------------------- -def _split_grpo_state(state): - """Splits the reference parameters from the main training state. - - This is a utility function to separate the reference model's parameters, - which are kept frozen, from the policy model's parameters that are actively - being trained. - - Args: - state: The combined training state, expected to contain a 'reference_params' - key within its `params` attribute. - - Returns: - A tuple containing: - - new_state: The training state with 'reference_params' removed. - - reference_params: The extracted reference parameters. - """ - reference_params = state.params["reference_params"] - new_state = state.replace(params={k: v for k, v in state.params.items() if k != "reference_params"}) - return new_state, reference_params - - -def _merge_grpo_state(state, reference_params): - """Merges the reference parameters back into the training state. - - This is the inverse operation of `_split_grpo_state`, used to reconstruct - the full state object after a training step. - - Args: - state: The training state, without 'reference_params'. - reference_params: The frozen reference parameters to be added back. - - Returns: - A new state object with the 'reference_params' re-integrated. - """ - return state.replace(params=dict(state.params, reference_params=reference_params)) - - @struct.dataclass class LossAux: """A dataclass to hold auxiliary outputs from the GRPO loss function. @@ -609,120 +571,22 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat """Run one GRPO training step. Computes the GRPO loss and gradients and applies the update to the policy - parameters; the reference parameters are held constant. The Linen and NNX - paths share this entry point: on the NNX path `model` is an NNX - `GraphDef` and `state` is the matching flat `nnx.State` of a - `TrainStateNNX`. On the Linen path they are the usual `nn.Module` and - `TrainState`. + parameters; the reference parameters are held constant. `model` is an NNX + `GraphDef` and `state` is the matching flat `nnx.State` of a `TrainStateNNX`. Args: - model: Linen `nn.Module` or NNX `GraphDef`, depending on `config.pure_nnx`. + model: NNX `GraphDef` of the `TrainStateNNX`. config: Training configuration object. state_mesh_shardings: Pytree of shardings matching `state`. - params_shardings: Param-only shardings, used for gradient accumulation - on the Linen path. Ignored on NNX. - state: Linen `TrainState` or NNX `nnx.State` matching `model`. + params_shardings: Param-only shardings. + state: NNX `nnx.State` matching `model`. data: A batch dict produced by the GRPO input pipeline. - dropout_rng: PRNG key for dropout (Linen only). + dropout_rng: PRNG key for dropout. Returns: A tuple `(new_state, metrics)`. """ - if config.pure_nnx: - return _train_step_nnx(model, config, state_mesh_shardings, state, data) - - state, reference_params = _split_grpo_state(state) - state_mesh_shardings, reference_params_sharding = _split_grpo_state(state_mesh_shardings) - extra_grpo_args = [reference_params] - _loss_fn = grpo_loss_fn - - if config.gradient_accumulation_steps > 1: - - def accumulate_gradient(acc_grad_and_loss, data): - grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True) - (_, aux), cur_batch_gradient = grad_func( - model, config, data, dropout_rng, state.params, *extra_grpo_args, is_train=True - ) - acc_grad_and_loss["loss"] += aux["total_loss"] - acc_grad_and_loss["moe_lb_loss"] += aux["moe_lb_loss"] - acc_grad_and_loss["grad"] = jax.tree_util.tree_map( - lambda x, y: x * aux["total_weights"] + y, cur_batch_gradient, acc_grad_and_loss["grad"] - ) - acc_grad_and_loss["total_weights"] += aux["total_weights"] - return acc_grad_and_loss, aux - - def reshape_to_microbatch_accumulations(batch_arr): - """Reshape global batch to microbatches, assuming batch axis is leading.""" - microbatches = config.gradient_accumulation_steps - microbatch_shape = (microbatches, batch_arr.shape[0] // microbatches) + batch_arr.shape[1:] - return jnp.reshape(batch_arr, microbatch_shape) - - data = jax.tree_util.tree_map(reshape_to_microbatch_accumulations, data) - init_grad = jax.tree_util.tree_map(jnp.zeros_like, state.params) - init_grad_and_loss = {"loss": 0.0, "grad": init_grad, "total_weights": 0, "moe_lb_loss": 0.0} - - grad_and_loss, aux = jax.lax.scan( - accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps - ) - loss = ( - grad_and_loss["loss"] / grad_and_loss["total_weights"] - + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps - ) - raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], grad_and_loss["grad"]) - aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) - else: - if config.optimizer_memory_host_offload: - cast_params = jax.device_put(state.params, max_utils.with_memory_kind(state_mesh_shardings.params, "device")) - cast_params = max_utils.cast_to_bf16(cast_params) - state = state.replace(params=cast_params) - if config.use_grpo: - reference_params = jax.device_put( - reference_params, max_utils.with_memory_kind(reference_params_sharding, "device") - ) - reference_params = max_utils.cast_to_bf16(reference_params) - extra_grpo_args = [reference_params] - grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True) - (loss, aux), raw_grads = grad_func(model, config, data, dropout_rng, state.params, *extra_grpo_args, is_train=True) - - total_weights = aux.total_weights - moe_lb_loss = aux.moe_lb_loss - - if config.gradient_clipping_threshold > 0: - grads = maxtext_utils.apply_gradient_clipping(raw_grads, state, config.gradient_clipping_threshold) - else: - grads = raw_grads - if config.optimizer_memory_host_offload: - state = state.replace( - opt_state=jax.device_put( - state.opt_state, - jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="device"), state_mesh_shardings.opt_state), - ) - ) - new_state = state.apply_gradients(grads=grads) - - scalar_metrics = { - "learning/loss": loss, - "learning/avg_reward": aux.avg_reward, - "learning/avg_reward_std": aux.avg_reward_std, - "learning/avg_advantage": aux.avg_advantage, - "learning/avg_kl": aux.avg_kl, - "learning/completion_length": aux.completion_length, - "learning/moe_lb_loss": moe_lb_loss, - "learning/total_weights": total_weights, - } - if not config.optimizer_memory_host_offload: - scalar_metrics["learning/grad_norm"] = max_utils.l2norm_pytree(grads) - scalar_metrics["learning/raw_grad_norm"] = max_utils.l2norm_pytree(raw_grads) - scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(new_state.params) - scalar_metrics["learning/avg_reward"] = aux.avg_reward - metrics = { - "scalar": scalar_metrics, - "scalars": {}, - } - - new_state = _merge_grpo_state(new_state, reference_params) - - return new_state, metrics + return _train_step_nnx(model, config, state_mesh_shardings, state, data) def eval_step(model, config, state, data, dropout_rng): @@ -741,31 +605,7 @@ def eval_step(model, config, state, data, dropout_rng): Returns: A dictionary of evaluation metrics. """ - if config.pure_nnx: - return _eval_step_nnx(model, config, state, data) - - reference_params, extra_grpo_args, _loss_fn = [], [], grpo_loss_fn - state, reference_params = _split_grpo_state(state) - extra_grpo_args = [reference_params] - _loss_fn = grpo_loss_fn - - eval_loss_fn = functools.partial(_loss_fn, model, config, data, dropout_rng, is_train=False) - loss, aux = eval_loss_fn(state.params, *extra_grpo_args) - total_loss = aux["total_loss"] - total_weights = aux["total_weights"] - moe_lb_loss = aux["moe_lb_loss"] - metrics = { - "scalar": { - "evaluation/loss": loss, - "evaluation/total_loss": total_loss, - "evaluation/total_weights": total_weights, - "evaluation/moe_lb_loss": moe_lb_loss, - }, - } - if config.use_dpo: - metrics["scalar"]["evaluation/grpo_reward_accuracy"] = aux["reward_accuracy"] - - return metrics + return _eval_step_nnx(model, config, state, data) def setup_train_loop( @@ -812,54 +652,41 @@ def setup_train_loop( - eval_data_iterator: The iterator for the evaluation dataset (or None). - state: The initialized training state. """ - if config.pure_nnx != config_inference.pure_nnx: - raise ValueError( - f"config.pure_nnx ({config.pure_nnx}) and config_inference.pure_nnx " f"({config_inference.pure_nnx}) must agree." - ) with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT): max_logging.log("Training mesh used for the workload") num_inference_devices = config.inference_devices_per_replica * config.inference_replicas training_devices = jax.devices()[num_inference_devices:] init_rng = jax.random.PRNGKey(config.init_weights_seed) - if config.pure_nnx: - training_mesh = maxtext_utils.get_mesh_from_config(config, devices=training_devices) - training_rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) - model = mt.from_config(config, devices=training_devices, mesh=training_mesh, rngs=training_rngs) - else: - model = mt.from_config(config, devices=training_devices) + training_mesh = maxtext_utils.get_mesh_from_config(config, devices=training_devices) + training_rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) + model = mt.from_config(config, devices=training_devices, mesh=training_mesh, rngs=training_rngs) mesh = model.mesh max_logging.log("Inference mesh used for the workload") inference_devices = jax.devices()[:num_inference_devices] - if config_inference.pure_nnx: - inference_mesh_obj = maxtext_utils.get_mesh_from_config(config_inference, devices=inference_devices) - inference_rngs = maxtext_utils_nnx.create_nnx_rngs(config_inference, rng_key=init_rng) - inference_model = mt.from_config( - config_inference, devices=inference_devices, mesh=inference_mesh_obj, rngs=inference_rngs - ) - else: - inference_model = mt.from_config(config_inference, devices=inference_devices) + inference_mesh_obj = maxtext_utils.get_mesh_from_config(config_inference, devices=inference_devices) + inference_rngs = maxtext_utils_nnx.create_nnx_rngs(config_inference, rng_key=init_rng) + inference_model = mt.from_config( + config_inference, devices=inference_devices, mesh=inference_mesh_obj, rngs=inference_rngs + ) inference_mesh = inference_model.mesh learning_rate_schedule, tx = train_utils.create_training_optimizer(config, model) - if config.pure_nnx: - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh, devices=training_devices) - - def init_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - # Reference uses the same init seed so it starts identical to the policy. - reference_model = _create_model_partial() - # TrainStateNNX only takes (model, optimizer); reference_model is an NNX - # sibling attribute set after construction (nnx.Module is mutable). - state = train_state_nnx.TrainStateNNX(nnx_model, optimizer) - state.reference_model = reference_model - return state - - else: - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, init_rng) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh, devices=training_devices) + + def init_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + # Reference uses the same init seed so it starts identical to the policy. + reference_model = _create_model_partial() + # TrainStateNNX only takes (model, optimizer); reference_model is an NNX + # sibling attribute set after construction (nnx.Module is mutable). + state = train_state_nnx.TrainStateNNX(nnx_model, optimizer) + state.reference_model = reference_model + return state + checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn) with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION): @@ -868,29 +695,21 @@ def init_state_fn(): data_iterator, config, mesh, checkpoint_manager, init_state_fn ) - if config_inference.pure_nnx: - _create_inference_partial, _ = model_creation_utils.create_nnx_abstract_model( - config_inference, inference_mesh, devices=inference_devices - ) + _create_inference_partial, _ = model_creation_utils.create_nnx_abstract_model( + config_inference, inference_mesh, devices=inference_devices + ) - def init_inference_state_fn(): - inference_nnx_model = _create_inference_partial() - return train_state_nnx.TrainStateNNX(inference_nnx_model, None) + def init_inference_state_fn(): + inference_nnx_model = _create_inference_partial() + return train_state_nnx.TrainStateNNX(inference_nnx_model, None) - else: - init_inference_state_fn = functools.partial( - maxtext_utils.init_initial_state, inference_model, tx, config_inference, False, init_rng - ) inference_state_mesh_shardings = maxtext_utils.get_abstract_state( config_inference, inference_mesh, init_inference_state_fn, is_training=False )[2] if not config.using_pipeline_parallelism: # The vocab tensor(s) of shape [vocab, embed] (and transpose) are not sharded by stage - if config.pure_nnx: - params_for_check = nnx.state(state.model, nnx.Param) - sharding.assert_params_sufficiently_sharded(params_for_check, mesh, config.sharding_tolerance) - else: - sharding.assert_params_sufficiently_sharded(state.params, mesh, config.sharding_tolerance) + params_for_check = nnx.state(state.model, nnx.Param) + sharding.assert_params_sufficiently_sharded(params_for_check, mesh, config.sharding_tolerance) return ( init_rng, @@ -1003,16 +822,10 @@ def train_loop(config, config_inference, recorder, state=None): token=config.hf_access_token, ) - if config.pure_nnx: - # `reference_model` is a sibling field on TrainStateNNX, populated by - # init_state_fn. Nothing to merge here; just verify it is present. - if not hasattr(state, "reference_model"): - raise RuntimeError("NNX GRPO state is missing reference_model; check setup_train_loop.") - else: - if "reference_params" not in state.params: - reference_params = jax.tree.map(jnp.copy, state.params["params"]) - state = _merge_grpo_state(state, reference_params) - state_mesh_shardings = _merge_grpo_state(state_mesh_shardings, state_mesh_shardings.params["params"]) + # `reference_model` is a sibling field on TrainStateNNX, populated by + # init_state_fn. Nothing to merge here; just verify it is present. + if not hasattr(state, "reference_model"): + raise RuntimeError("NNX GRPO state is missing reference_model; check setup_train_loop.") p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( config, model, mesh, state, state_mesh_shardings, train_step, eval_step, eval_data_iterator @@ -1037,11 +850,8 @@ def train_loop(config, config_inference, recorder, state=None): metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) # Write train config params, num model params, and XLA flags to tensorboard - if config.pure_nnx: - params_for_metrics = nnx.state(state.model, nnx.Param) - metric_logger.write_setup_info_to_tensorboard(params_for_metrics) - else: - metric_logger.write_setup_info_to_tensorboard(state.params["params"]) + params_for_metrics = nnx.state(state.model, nnx.Param) + metric_logger.write_setup_info_to_tensorboard(params_for_metrics) def generation_worker_fn( worker_inference_engine, @@ -1165,32 +975,21 @@ def generation_worker_fn( state, metrics = p_train_step(state, example_batch, train_rng) with jax.profiler.StepTraceAnnotation("transfer data", step_num=step): if step != 0 and step % config.inference_rollouts == 0: - if config.pure_nnx: - grpo_utils.pathways_reshard_nnx( - config_inference, - inference_engine, - state.model, - state_mesh_shardings.model, - inference_state_mesh_shardings.model, - ) - else: - grpo_utils.pathways_reshard( - config_inference, - inference_engine, - {"params": state.params["params"]}, - {"params": state_mesh_shardings.params["params"]}, - mesh, - {"params": inference_state_mesh_shardings.params["params"]}, - ) + grpo_utils.pathways_reshard_nnx( + config_inference, + inference_engine, + state.model, + state_mesh_shardings.model, + inference_state_mesh_shardings.model, + ) with data_buffer_lock: data_buffer.clear() step_time_delta = datetime.datetime.now() - last_step_completion - # On the Linen path, the reference is embedded in `state.params` and is - # stripped before saving. On the NNX path, the reference is a sibling - # field on TrainStateNNX, so the whole state can be saved as-is. - state_to_save = state if config.pure_nnx else _split_grpo_state(state)[0] + # The reference is a sibling field on TrainStateNNX, so the whole state + # can be saved as-is. + state_to_save = state checkpointing.maybe_save_checkpoint(checkpoint_manager, state_to_save, config, data_iterator, step) if config.dump_hlo and step == start_step: @@ -1234,7 +1033,7 @@ def generation_worker_fn( metric_logger.buffer_and_write_metrics(metrics, step, step_time_delta) if config.save_checkpoint_on_completion: - state_to_save = state if config.pure_nnx else _split_grpo_state(state)[0] + state_to_save = state checkpointing.maybe_save_checkpoint(checkpoint_manager, state_to_save, config, data_iterator) elif checkpoint_manager is not None: # in case the last checkpoint_period checkpoint is still in progress diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index c13f8bba19..8777b928c2 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -27,7 +27,6 @@ from absl import app import numpy as np -import optax import pathwaysutils # pylint: disable=unused-import @@ -42,7 +41,7 @@ import jax.numpy as jnp from jax.sharding import NamedSharding -from flax import linen as nn, nnx, traverse_util +from flax import nnx, traverse_util from flax.linen import partitioning as nn_partitioning from flax.nnx import variablelib @@ -80,7 +79,7 @@ from maxtext.utils import maxtext_utils_nnx from maxtext.utils import train_utils from maxtext.utils.gradient_accumulation import gradient_accumulation_loss_and_grad -from maxtext.utils.vocabulary_tiling import vocab_tiling_linen_loss, vocab_tiling_nnx_loss +from maxtext.utils.vocabulary_tiling import vocab_tiling_nnx_loss class EncoderKwargs(TypedDict, total=False): @@ -97,8 +96,7 @@ class EncoderKwargs(TypedDict, total=False): def get_first_step(model, state): - if isinstance(model, nn.Module): - return int(state.step) + del model # NNX-only; kept for call-site signature parity if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo) step_val = state.step.get_value() if hasattr(state.step, "get_value") else state.step return int(step_val) @@ -114,17 +112,18 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr """loss_fn for both train and eval. Args: - model: A nn.Module (Linen) or nnx.Module (NNX). + model: An NNX model. config: Config of parameters data: Batch of data to apply to the model - dropout_rng: A key to use to generate rng for dropout (Linen); unused for NNX. - params: Model params (Linen); unused for NNX (params are part of the model). + dropout_rng: Unused for NNX (kept for signature parity). + params: Unused for NNX; params are part of the model. is_train: True for train_step and False for eval_step Returns: loss: average loss aux: a dictionary including intermediate_outputs, xent_sum, and total_weights """ + del dropout_rng, params, sparsity_state # unused for NNX (kept for signature parity) is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion" if getattr(config, "attention_type", "global") == "block_diffusion" and not is_block_diffusion: raise ValueError( @@ -177,175 +176,86 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr "encoder_images": encoder_images, "encoder_image_masks": encoder_image_masks, } - mutable_collections = ["intermediates"] - if config.mtp_num_layers > 0 and is_train: - # The single model.apply call now triggers the entire chain if MTP is enabled: - # Decoder runs -> returns hidden_state -> MTPBlock uses it -> MTPBlock sows losses -> we reap them here. - mutable_collections.append("mtp_losses") - - # During evaluation, if the acceptance rate test is enabled, we must - # make its specific collection mutable so the MTPBlock can sow into it. - if config.mtp_eval_target_module > 0 and not is_train: - mutable_collections.append("mtp_acceptance") - if config.use_indexer and is_train: - mutable_collections.append("indexer_losses") - - sparsity_enabled = is_train and config.weight_sparsity_n and config.weight_sparsity_m - if sparsity_enabled: - mutable_collections.append("batch_stats") - if isinstance(model, nn.Module): - # inputs, targets, segments, positions = apply_args - if dropout_rng is not None: - rng1, aqt_rng = jax.random.split(dropout_rng) - else: - rng1, aqt_rng = None, None - - # Flax Linen model - if sparsity_enabled: - model_vars = {"params": params} - if sparsity_state: - model_vars["batch_stats"] = sparsity_state - else: - model_vars = params - logits, intermediate_outputs = model.apply( - model_vars, - data["inputs"], - data["inputs_position"], - decoder_segment_ids=data["inputs_segmentation"], - **encoder_kwargs, - enable_dropout=config.enable_dropout if is_train else False, - rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] - mutable=mutable_collections, - decoder_target_tokens=data["targets"], - decoder_target_mask=data["targets_segmentation"], - **forced_routing_kwargs, - ) - - if (config.use_indexer and not config.indexer_sparse_training) and is_train: - # In Dense Warm-up stage, we skip main model loss calculation for efficiency. - # The main model parameters are frozen and only the indexer is trained via KL divergence. - xent_sum = 0.0 - total_z_loss = 0.0 - elif config.num_vocab_tiling > 1: - hidden_state_key = ("intermediates", "decoder", "hidden_states") - hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] - xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train) - else: - if is_block_diffusion: - logits = block_diffusion_target_alignment.align_logits_to_targets( - logits, - config.block_diffusion_logit_alignment, - target_positions, - data["targets_segmentation"] != 0, - ) - one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) - xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) - - xent = sharding.maybe_shard_with_logical( - xent, - ("activation_embed_and_logits_batch", "activation_length"), - model.mesh, - config.shard_mode, - debug_sharding=config.debug_sharding, - ) - z_loss = sharding.maybe_shard_with_logical( - z_loss, - ("activation_embed_and_logits_batch", "activation_length"), - model.mesh, - config.shard_mode, - debug_sharding=config.debug_sharding, - ) - - if is_block_diffusion: - xent = xent * targets_loss_mask - z_loss = z_loss * targets_loss_mask - else: - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) - - xent_sum = jnp.sum(xent) - total_z_loss = jnp.sum(z_loss) + # Flax NNX model: forward pass, then pop Intermediates sown during it. + logits = model( + decoder_input_tokens=data["inputs"], + decoder_positions=data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + **encoder_kwargs, + enable_dropout=config.enable_dropout if is_train else False, + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + **forced_routing_kwargs, + ) + # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match + # subclasses. Pop them before the generic Intermediate pop below, which would otherwise + # take them too and leave the MTP loss silently reading as 0. + mtp_losses_state, mtp_acceptance_state = None, None + if config.mtp_num_layers > 0: + mtp_losses_state = nnx.pop(model, mtp_losses) + mtp_acceptance_state = nnx.pop(model, mtp_acceptance) + + indexer_losses_state = None + if config.use_indexer: + # Pop dedicated indexer_losses to harvest auxiliary KL loss and prevent model state PyTree mismatches. + indexer_losses_state = nnx.pop(model, indexer_losses) + + intermediates = nnx.pop(model, nnx.Intermediate) + intermediate_outputs = intermediates.to_pure_dict() + + # Store them under the collection name so calculate_mtp_loss and + # calculate_mtp_acceptance_rate find them at the path they expect. + if mtp_losses_state is not None and mtp_acceptance_state is not None: + intermediate_outputs["mtp_losses"] = mtp_losses_state.to_pure_dict() + intermediate_outputs["mtp_acceptance"] = mtp_acceptance_state.to_pure_dict() + + if indexer_losses_state is not None: + intermediate_outputs["indexer_losses"] = indexer_losses_state.to_pure_dict() + + if (config.use_indexer and not config.indexer_sparse_training) and is_train: + # In Dense Warm-up stage, we skip main model loss calculation for efficiency. + # The main model parameters are frozen and only the indexer is trained via KL divergence. + xent_sum = 0.0 + total_z_loss = 0.0 + elif config.num_vocab_tiling > 1: + hidden_state_key = ("decoder", "hidden_states") + hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] + xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train) else: - # Flax NNX model: forward pass, then pop Intermediates sown during it. - logits = model( - decoder_input_tokens=data["inputs"], - decoder_positions=data["inputs_position"], - decoder_segment_ids=data["inputs_segmentation"], - **encoder_kwargs, - enable_dropout=config.enable_dropout if is_train else False, - decoder_target_tokens=data["targets"], - decoder_target_mask=data["targets_segmentation"], - **forced_routing_kwargs, - ) - # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match - # subclasses. Pop them before the generic Intermediate pop below, which would otherwise - # take them too and leave the MTP loss silently reading as 0. - mtp_losses_state, mtp_acceptance_state = None, None - if config.mtp_num_layers > 0: - mtp_losses_state = nnx.pop(model, mtp_losses) - mtp_acceptance_state = nnx.pop(model, mtp_acceptance) - - indexer_losses_state = None - if config.use_indexer: - # Pop dedicated indexer_losses to harvest auxiliary KL loss and prevent model state PyTree mismatches. - indexer_losses_state = nnx.pop(model, indexer_losses) - - intermediates = nnx.pop(model, nnx.Intermediate) - intermediate_outputs = intermediates.to_pure_dict() - - # Store them under the collection name so calculate_mtp_loss and - # calculate_mtp_acceptance_rate find them at the same path as the Linen collections. - if mtp_losses_state is not None and mtp_acceptance_state is not None: - intermediate_outputs["mtp_losses"] = mtp_losses_state.to_pure_dict() - intermediate_outputs["mtp_acceptance"] = mtp_acceptance_state.to_pure_dict() - - if indexer_losses_state is not None: - intermediate_outputs["indexer_losses"] = indexer_losses_state.to_pure_dict() - - if (config.use_indexer and not config.indexer_sparse_training) and is_train: - # In Dense Warm-up stage, we skip main model loss calculation for efficiency. - # The main model parameters are frozen and only the indexer is trained via KL divergence. - xent_sum = 0.0 - total_z_loss = 0.0 - elif config.num_vocab_tiling > 1: - hidden_state_key = ("decoder", "hidden_states") - hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] - xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train) - else: - if is_block_diffusion: - logits = block_diffusion_target_alignment.align_logits_to_targets( - logits, - config.block_diffusion_logit_alignment, - target_positions, - data["targets_segmentation"] != 0, - ) - one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) - xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) - - xent = sharding.maybe_shard_with_logical( - xent, - ("activation_embed_and_logits_batch", "activation_length"), - model.mesh, - config.shard_mode, - debug_sharding=config.debug_sharding, - ) - z_loss = sharding.maybe_shard_with_logical( - z_loss, - ("activation_embed_and_logits_batch", "activation_length"), - model.mesh, - config.shard_mode, - debug_sharding=config.debug_sharding, + if is_block_diffusion: + logits = block_diffusion_target_alignment.align_logits_to_targets( + logits, + config.block_diffusion_logit_alignment, + target_positions, + data["targets_segmentation"] != 0, ) + one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) + xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) + + xent = sharding.maybe_shard_with_logical( + xent, + ("activation_embed_and_logits_batch", "activation_length"), + model.mesh, + config.shard_mode, + debug_sharding=config.debug_sharding, + ) + z_loss = sharding.maybe_shard_with_logical( + z_loss, + ("activation_embed_and_logits_batch", "activation_length"), + model.mesh, + config.shard_mode, + debug_sharding=config.debug_sharding, + ) - if is_block_diffusion: - xent = xent * targets_loss_mask - z_loss = z_loss * targets_loss_mask - else: - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + if is_block_diffusion: + xent = xent * targets_loss_mask + z_loss = z_loss * targets_loss_mask + else: + xent = xent * (data["targets_segmentation"] != 0) + z_loss = z_loss * (data["targets_segmentation"] != 0) - xent_sum = jnp.sum(xent) - total_z_loss = jnp.sum(z_loss) + xent_sum = jnp.sum(xent) + total_z_loss = jnp.sum(z_loss) if is_block_diffusion: assert targets_loss_mask is not None @@ -403,24 +313,20 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr moe_bias_updates = None mtp_moe_bias_updates = None if config.routed_bias and config.routed_bias_update_rate > 0.0: - if isinstance(model, nn.Module): - nested_key = ("intermediates", "decoder", "moe_layers", "moe_bias_updates") - moe_bias_updates = maxtext_utils.get_nested_value(intermediate_outputs, nested_key, None) - else: - # NNX intermediates are model-rooted (no "intermediates" prefix), - # so match by suffix instead. Unlike collect_intermediates_by_suffix - # we must not ravel: the decoder update is a 2-D matrix that's - # transposed and MTP update is 1-D matrix. - for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs): - keys = tuple(k.key for k in path if hasattr(k, "key")) - if not keys or keys[-1] != "moe_bias_updates": - continue - if "decoder" in keys: - moe_bias_updates = (val,) - elif "mtp_block" in keys: - if mtp_moe_bias_updates is None: - mtp_moe_bias_updates = [] - mtp_moe_bias_updates.append(val) + # NNX intermediates are model-rooted (no "intermediates" prefix), + # so match by suffix instead. Unlike collect_intermediates_by_suffix + # we must not ravel: the decoder update is a 2-D matrix that's + # transposed and MTP update is 1-D matrix. + for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs): + keys = tuple(k.key for k in path if hasattr(k, "key")) + if not keys or keys[-1] != "moe_bias_updates": + continue + if "decoder" in keys: + moe_bias_updates = (val,) + elif "mtp_block" in keys: + if mtp_moe_bias_updates is None: + mtp_moe_bias_updates = [] + mtp_moe_bias_updates.append(val) # Add the model's primary output to the intermediates dict so it can be used # by the acceptance rate calculation in eval_step. @@ -474,29 +380,26 @@ def _find_gate_bias(module: nnx.Module | None) -> nnx.Variable | None: def train_step(model, config, state_mesh_shardings, params_shardings, state, data, dropout_rng=None): - """Training step for both Linen and NNX models. + """Training step for the NNX model. Args: - model: A nn.Module (Linen) or nnx.GraphDef of the TrainStateNNX (NNX). + model: An nnx.GraphDef of the TrainStateNNX. config: Hyperparameters. state_mesh_shardings: PyTree of PartitionSpecs for the train state. params_shardings: PyTree of PartitionSpecs for model parameters, used for gradient accumulation. - state: Linen TrainState or NNX pure State. + state: NNX pure State. data: Training data batch. - dropout_rng: A key to use to generate rng for dropout (Linen); unused for NNX. + dropout_rng: Unused for NNX (kept for jit signature parity). Returns: - new_state: Updated Linen TrainState or NNX pure State. + new_state: Updated NNX pure State. metrics: Dictionary of model metrics such as loss, training rate, etc. """ + del dropout_rng # unused for NNX (kept for jit signature parity) # pylint: disable=too-many-nested-blocks # --- Per-path initialization --- - if isinstance(model, nn.Module): - params = state.params - loss_model, loss_params, loss_rng = model, params, dropout_rng - else: - state = nnx.merge(model, state) # reconstruct TrainStateNNX - loss_model, loss_params, loss_rng = state.model, None, None + state = nnx.merge(model, state) # reconstruct TrainStateNNX + loss_model, loss_params, loss_rng = state.model, None, None # --- Gradient computation --- if config.gradient_accumulation_steps > 1: @@ -510,86 +413,62 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat loss_rng, ) else: - if isinstance(model, nn.Module): - if config.shard_optimizer_over_data: - params = jax.tree.map( - functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode), - params, # pyrefly: ignore[unbound-name] - params_shardings, - ) - sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - pure_params = params["params"] if sparsity_enabled else params # pyrefly: ignore[unbound-name] - batch_stats = params.get("batch_stats", {}) - - grad_func = jax.value_and_grad(loss_fn, argnums=4, has_aux=True) - (loss, aux), raw_grads = grad_func( - model, - config, - data, - dropout_rng, - pure_params, - sparsity_state=batch_stats, - is_train=True, + owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True) + custom_param_filter = nnx.Any(owg_type) + train_param_type = ( + getattr(nnx, "LoRAParam", nnx.Param) + if getattr(getattr(config, "lora", None), "enable_lora", False) + else nnx.Param + ) + nnx.pop(state.model, nnx.Intermediate) + model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, train_param_type, custom_param_filter, ...) + if config.parameter_memory_host_offload: + # Params are kept on host (pinned_host) in in_shardings. Move only Param + # variables to device before the forward/backward pass so that all dot_general + # operands share the same memory space (XLA on GPU requires this). + # Using params_shardings (Param-only) avoids Shardy rank mismatches that + # occur when applying PartitionSpec() (rank-0 in SDY) to rank-1 RNG key tensors. + device_param_shardings = jax.tree_util.tree_map_with_path( + maxtext_utils_nnx.move_memory_to_device, + params_shardings, + is_leaf=lambda x: isinstance(x, NamedSharding), ) - else: - owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True) - custom_param_filter = nnx.Any(owg_type) - train_param_type = ( - getattr(nnx, "LoRAParam", nnx.Param) - if getattr(getattr(config, "lora", None), "enable_lora", False) - else nnx.Param + curr_params = jax.device_put(curr_params, device_param_shardings) + nnx.update(state.model, curr_params) # ensure state.model has device params for optimizer update + if config.shard_optimizer_over_data: + param_sharding_lookup = {} + for p, s in jax.tree_util.tree_leaves_with_path( + params_shardings, is_leaf=lambda x: isinstance(x, (nnx.Variable, NamedSharding, jax.sharding.Sharding)) + ): + param_sharding_lookup[p] = s.get_value() if isinstance(s, nnx.Variable) else s + + def _maybe_shard_param(path, var): + if path in param_sharding_lookup: + return sharding.maybe_shard_with_name(var, param_sharding_lookup[path], shard_mode=config.shard_mode) + return var + + curr_params = jax.tree_util.tree_map_with_path( + _maybe_shard_param, + curr_params, + is_leaf=lambda x: isinstance(x, nnx.Variable), ) - nnx.pop(state.model, nnx.Intermediate) - model_graphdef, curr_params, custom_params, rest = nnx.split( - state.model, train_param_type, custom_param_filter, ... + nnx.update(state.model, curr_params) + + def diff_wrapper(curr_params, custom_params, rest, config, data): + local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True) + loss, aux = loss_fn(local_model, config, data, None, None, is_train=True) + # Exclude parameters, custom-gradient state, and intermediates. Custom- + # gradient state is updated separately, so stale values must not overwrite + # `custom_grads`. + non_param_rest = nnx.state( + local_model, + nnx.Not(nnx.Any(nnx.Param, custom_param_filter, nnx.Intermediate)), ) - if config.parameter_memory_host_offload: - # Params are kept on host (pinned_host) in in_shardings. Move only Param - # variables to device before the forward/backward pass so that all dot_general - # operands share the same memory space (XLA on GPU requires this). - # Using params_shardings (Param-only) avoids Shardy rank mismatches that - # occur when applying PartitionSpec() (rank-0 in SDY) to rank-1 RNG key tensors. - device_param_shardings = jax.tree_util.tree_map_with_path( - maxtext_utils_nnx.move_memory_to_device, - params_shardings, - is_leaf=lambda x: isinstance(x, NamedSharding), - ) - curr_params = jax.device_put(curr_params, device_param_shardings) - nnx.update(state.model, curr_params) # ensure state.model has device params for optimizer update - if config.shard_optimizer_over_data: - param_sharding_lookup = {} - for p, s in jax.tree_util.tree_leaves_with_path( - params_shardings, is_leaf=lambda x: isinstance(x, (nnx.Variable, NamedSharding, jax.sharding.Sharding)) - ): - param_sharding_lookup[p] = s.get_value() if isinstance(s, nnx.Variable) else s - - def _maybe_shard_param(path, var): - if path in param_sharding_lookup: - return sharding.maybe_shard_with_name(var, param_sharding_lookup[path], shard_mode=config.shard_mode) - return var - - curr_params = jax.tree_util.tree_map_with_path( - _maybe_shard_param, - curr_params, - is_leaf=lambda x: isinstance(x, nnx.Variable), - ) - nnx.update(state.model, curr_params) - - def diff_wrapper(curr_params, custom_params, rest, config, data): - local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True) - loss, aux = loss_fn(local_model, config, data, None, None, is_train=True) - # Exclude parameters, custom-gradient state, and intermediates. Custom- - # gradient state is updated separately, so stale values must not overwrite - # `custom_grads`. - non_param_rest = nnx.state( - local_model, - nnx.Not(nnx.Any(nnx.Param, custom_param_filter, nnx.Intermediate)), - ) - return loss, (aux, non_param_rest) + return loss, (aux, non_param_rest) - grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True) - (loss, (aux, non_param_rest)), (raw_grads, custom_grads) = grad_func(curr_params, custom_params, rest, config, data) - nnx.update(state.model, nnx.State.merge(custom_grads, non_param_rest)) + grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True) + (loss, (aux, non_param_rest)), (raw_grads, custom_grads) = grad_func(curr_params, custom_params, rest, config, data) + nnx.update(state.model, nnx.State.merge(custom_grads, non_param_rest)) raw_grads = jax.tree_util.tree_map( lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x, @@ -615,140 +494,78 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): new_opt_state = None bias_metrics = {} - if isinstance(model, nn.Module): - if config.gradient_clipping_threshold > 0: - grads = maxtext_utils.apply_gradient_clipping(raw_grads, state, config.gradient_clipping_threshold) - else: - grads = raw_grads - if config.optimizer_memory_host_offload: - state = state.replace( - opt_state=jax.device_put( - state.opt_state, - jax.tree_util.tree_map( - lambda x: x.with_memory_kind(kind="device"), - state_mesh_shardings.opt_state, - ), - ) - ) - # Move all parameters to device before optimizer update - if config.parameter_memory_host_offload: - max_logging.log("\nMoving all parameters to device before optimizer update") - - def move(path, value): - max_logging.log(f"train.py: Moving f{path} to device") - return value.with_memory_kind(kind="device") - - state = state.replace( - params=jax.device_put( - state.params, - jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params), - ) - ) - # Re-wrap grads to match state.params structure if it's a dict of collections - # (when weight_sparsity is enabled, params has both 'params' and 'batch_stats' keys). - sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - if sparsity_enabled: - full_grads = {"params": grads} - if "batch_stats" in state.params: - batch_stats_grads = jax.tree_util.tree_map(jnp.zeros_like, state.params.get("batch_stats", {})) - full_grads["batch_stats"] = batch_stats_grads - full_grads = max_utils.unbox_logicallypartioned(full_grads) - else: - full_grads = grads - - if getattr(config, "skip_step_on_spikes", False): - grad_norm = max_utils.l2norm_pytree(grads) - # TrainState.apply_gradients doesn't pass **kwargs to tx.update, so we unpack it manually. - updates, new_opt_state = state.tx.update(grads, state.opt_state, state.params, loss=loss, grad_norm=grad_norm) - new_params = optax.apply_updates(state.params, updates) - - new_state = state.replace( - step=state.step + 1, - params=new_params, - opt_state=new_opt_state, - ) - else: - new_state = state.apply_gradients(grads=full_grads) - - # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family - if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: - target_path = ("params", "decoder", "moe_layers", "DeepSeekMoeBlock_0", "MoeBlock_0", "gate", "bias") - # Updates the shape to be aligned with state. - moe_bias_updates = jnp.array(moe_bias_updates[0]).transpose() - new_state = maxtext_utils.update_state_param(new_state, target_path, moe_bias_updates) + if config.gradient_clipping_threshold > 0: + grads = maxtext_utils.apply_gradient_clipping(raw_grads, None, config.gradient_clipping_threshold) else: - if config.gradient_clipping_threshold > 0: - grads = maxtext_utils.apply_gradient_clipping(raw_grads, None, config.gradient_clipping_threshold) - else: - grads = raw_grads - if config.optimizer_memory_host_offload: - # state.optimizer is an NNX Optimizer module; state_mesh_shardings.optimizer - # is an NNX State. Use nnx.state() to get a compatible State for device_put. - device_opt_shardings = jax.tree_util.tree_map_with_path( - maxtext_utils_nnx.move_memory_to_device, - state_mesh_shardings.optimizer, - is_leaf=lambda x: isinstance(x, NamedSharding), - ) - opt_state = nnx.state(state.optimizer) - new_opt_state = jax.device_put(opt_state, device_opt_shardings) - nnx.update(state.optimizer, new_opt_state) - if config.skip_step_on_spikes: - # The skip-step optimizer is a GradientTransformationExtraArgs that reads - # loss/grad_norm to decide whether to zero the update on a spike. nnx - # Optimizer.update forwards these kwargs to tx.update. - grad_norm = max_utils.l2norm_pytree(grads) - state.apply_gradients(grads, loss=loss, grad_norm=grad_norm) + grads = raw_grads + if config.optimizer_memory_host_offload: + # state.optimizer is an NNX Optimizer module; state_mesh_shardings.optimizer + # is an NNX State. Use nnx.state() to get a compatible State for device_put. + device_opt_shardings = jax.tree_util.tree_map_with_path( + maxtext_utils_nnx.move_memory_to_device, + state_mesh_shardings.optimizer, + is_leaf=lambda x: isinstance(x, NamedSharding), + ) + opt_state = nnx.state(state.optimizer) + new_opt_state = jax.device_put(opt_state, device_opt_shardings) + nnx.update(state.optimizer, new_opt_state) + if config.skip_step_on_spikes: + # The skip-step optimizer is a GradientTransformationExtraArgs that reads + # loss/grad_norm to decide whether to zero the update on a spike. nnx + # Optimizer.update forwards these kwargs to tx.update. + grad_norm = max_utils.l2norm_pytree(grads) + state.apply_gradients(grads, loss=loss, grad_norm=grad_norm) + else: + state.apply_gradients(grads) + new_state = state + + # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family + # pylint: disable=too-many-nested-blocks + if config.routed_bias and config.routed_bias_update_rate > 0.0: + if getattr(config, "model_name", "").startswith("deepseek4"): + max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.") + flat_intermediates = traverse_util.flatten_dict(aux.get("intermediate_outputs", {})) + for path, update in flat_intermediates.items(): + if path[-1] != "moe_bias_updates": + continue + target = new_state.model + prefix = path[1:-1] if path[0] == "intermediates" else path[:-1] + for key in prefix: + if hasattr(target, key): + target = getattr(target, key) + elif isinstance(target, dict) and key in target: + target = target[key] + else: + target = None + break + if target is None: + continue + for _, node in nnx.iter_graph(target): + if type(node).__name__ == "GateLogit" and hasattr(node, "bias") and node.bias is not None: + update_val = update[0] if isinstance(update, (tuple, list)) else update + name_prefix = "-".join(map(str, prefix)) + if getattr(config, "log_moe_bias_norms", False): + bias_metrics[f"learning/moe_bias_before_norm_{name_prefix}"] = jnp.linalg.norm(node.bias.value) + node.bias.value = node.bias.value + jnp.array(update_val) + if getattr(config, "log_moe_bias_norms", False): + bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val)) else: - state.apply_gradients(grads) - new_state = state - - # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family - # pylint: disable=too-many-nested-blocks - if config.routed_bias and config.routed_bias_update_rate > 0.0: - if getattr(config, "model_name", "").startswith("deepseek4"): - max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.") - flat_intermediates = traverse_util.flatten_dict(aux.get("intermediate_outputs", {})) - for path, update in flat_intermediates.items(): - if path[-1] != "moe_bias_updates": - continue - target = new_state.model - prefix = path[1:-1] if path[0] == "intermediates" else path[:-1] - for key in prefix: - if hasattr(target, key): - target = getattr(target, key) - elif isinstance(target, dict) and key in target: - target = target[key] - else: - target = None - break - if target is None: - continue - for _, node in nnx.iter_graph(target): - if type(node).__name__ == "GateLogit" and hasattr(node, "bias") and node.bias is not None: - update_val = update[0] if isinstance(update, (tuple, list)) else update - name_prefix = "-".join(map(str, prefix)) - if getattr(config, "log_moe_bias_norms", False): - bias_metrics[f"learning/moe_bias_before_norm_{name_prefix}"] = jnp.linalg.norm(node.bias.value) - node.bias.value = node.bias.value + jnp.array(update_val) - if getattr(config, "log_moe_bias_norms", False): - bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val)) - else: - # 1. Update main decoder scanned MoE layers. - # The update from the scan is (num_moe_layers, num_experts) and must be transposed. - decoder_layer = getattr(new_state.model.decoder, "moe_layers", new_state.model.decoder) - decoder_bias = _find_gate_bias(decoder_layer) - if decoder_bias is not None: - decoder_bias.value = decoder_bias.value + jnp.array(moe_bias_updates[0]) - - # 2. Update auxiliary MTP MoE layers (if enabled). - # Unlike the main decoder, each MTP layer is an individual un-scanned layer - # with a 1D bias of shape (num_experts,). - if mtp_moe_bias_updates is not None and hasattr(new_state.model, "mtp_block"): - for i, update in enumerate(mtp_moe_bias_updates): - mtp_layer = getattr(new_state.model.mtp_block, f"mtp_layer_{i + 1}", None) - mtp_bias = _find_gate_bias(mtp_layer) - if mtp_bias is not None: - mtp_bias.value = mtp_bias.value + jnp.array(update) + # 1. Update main decoder scanned MoE layers. + # The update from the scan is (num_moe_layers, num_experts) and must be transposed. + decoder_layer = getattr(new_state.model.decoder, "moe_layers", new_state.model.decoder) + decoder_bias = _find_gate_bias(decoder_layer) + if decoder_bias is not None: + decoder_bias.value = decoder_bias.value + jnp.array(moe_bias_updates[0]) + + # 2. Update auxiliary MTP MoE layers (if enabled). + # Unlike the main decoder, each MTP layer is an individual un-scanned layer + # with a 1D bias of shape (num_experts,). + if mtp_moe_bias_updates is not None and hasattr(new_state.model, "mtp_block"): + for i, update in enumerate(mtp_moe_bias_updates): + mtp_layer = getattr(new_state.model.mtp_block, f"mtp_layer_{i + 1}", None) + mtp_bias = _find_gate_bias(mtp_layer) + if mtp_bias is not None: + mtp_bias.value = mtp_bias.value + jnp.array(update) lm_loss = xent_sum / (total_weights + EPS) scalar_metrics = { @@ -771,10 +588,7 @@ def move(path, value): ) scalar_metrics.update(bias_metrics) if config.use_qk_clip: - if isinstance(model, nn.Module): - new_state = qk_clip_utils.apply_qk_clip(new_state, intermediate_outputs, config) - else: - new_state = qk_clip_utils.apply_qk_clip_nnx(new_state, intermediate_outputs, config) + new_state = qk_clip_utils.apply_qk_clip_nnx(new_state, intermediate_outputs, config) global_max_logit = qk_clip_utils.calculate_max_logit_metric(intermediate_outputs) if global_max_logit is not None: @@ -783,21 +597,14 @@ def move(path, value): if not config.optimizer_memory_host_offload: scalar_metrics["learning/grad_norm"] = max_utils.l2norm_pytree(grads) scalar_metrics["learning/raw_grad_norm"] = max_utils.l2norm_pytree(raw_grads) - if isinstance(model, nn.Module): - scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(new_state.params) - else: - model_params = nnx.state(new_state.model, nnx.Param) - scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(model_params) + model_params = nnx.state(new_state.model, nnx.Param) + scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(model_params) # Surface skip-step rejections as a TB metric. The skip-step optimizer stores - # is_skipped in its opt_state: the Linen path gets it from the tx.update return, - # the NNX path reads it back off the optimizer it just updated in place. + # is_skipped in its opt_state; read it back off the optimizer just updated in place. if config.skip_step_on_spikes: - if isinstance(model, nn.Module): - is_skipped = new_opt_state.get("is_skipped") if isinstance(new_opt_state, dict) else None - else: - opt_state = nnx.to_pure_dict(nnx.state(new_state.optimizer)).get("opt_state", {}) - is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None + opt_state = nnx.to_pure_dict(nnx.state(new_state.optimizer)).get("opt_state", {}) + is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None if is_skipped is not None: scalar_metrics["optim/step_skipped"] = is_skipped.astype(jnp.float32) metrics = { @@ -809,8 +616,6 @@ def move(path, value): if getattr(config, "record_internal_nn_metrics", False): record_activation_metrics(metrics, intermediate_outputs, config) - if isinstance(model, nn.Module): - return new_state, metrics # Drop Intermediates (e.g. sowed max_logits for QK-Clip) and the MTP sown # vars (mtp_losses/mtp_acceptance) before returning. They're absent from # state_mesh_shardings and would cause a leaf-count / structure mismatch. @@ -819,16 +624,9 @@ def move(path, value): def eval_step(model, config, state, data, dropout_rng=None): """eval_step no backprop and new state compared with train_step.""" - if isinstance(model, nn.Module): - sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - pure_params = state.params["params"] if sparsity_enabled else state.params - batch_stats = state.params.get("batch_stats", {}) - - eval_loss_fn = functools.partial(loss_fn, model, config, data, dropout_rng, is_train=False) - loss, aux = eval_loss_fn(pure_params, sparsity_state=batch_stats) - else: - state = nnx.merge(model, state) # reconstruct TrainStateNNX - loss, aux = loss_fn(state.model, config, data, None, None, is_train=False) + del dropout_rng # unused for NNX (kept for jit signature parity) + state = nnx.merge(model, state) # reconstruct TrainStateNNX + loss, aux = loss_fn(state.model, config, data, None, None, is_train=False) mtp_acceptance_rate = 0.0 if config.mtp_eval_target_module > 0: @@ -867,11 +665,9 @@ def training_loop_iteration( state = jax_device_state["state"] init_rng = jax_device_state["init_rng"] mesh = jax_device_state["mesh"] - state_mesh_shardings = jax_device_state["state_mesh_shardings"] p_train_step = jax_device_state["p_train_step"] p_train_step_dropless = jax_device_state.get("p_train_step_dropless", None) p_eval_step = jax_device_state["p_eval_step"] - model = jax_device_state["model"] # Unpack python_vars step = python_vars["step"] @@ -889,8 +685,6 @@ def training_loop_iteration( config = immutable_data["config"] # for helpers logical_axis_rules_for_train = immutable_data["logical_axis_rules_for_train"] logical_axis_rules_for_eval = immutable_data["logical_axis_rules_for_eval"] - shard_optimizer_over_data = immutable_data["shard_optimizer_over_data"] - shard_mode = immutable_data["shard_mode"] eval_interval = immutable_data["eval_interval"] eval_steps = immutable_data["eval_steps"] start_step = immutable_data["start_step"] @@ -911,16 +705,14 @@ def training_loop_iteration( with jax.profiler.StepTraceAnnotation("train", step_num=step): example_batch = data_loader.load_next_batch(rampup_manager=rampup_manager) - # DiLoCo's inner step takes the rng like the Linen step does. - if isinstance(model, nn.Module) or config.enable_diloco: + # DiLoCo's inner step takes the rng like the inner NNX step. + if config.enable_diloco: # pylint: disable=not-callable step_rng_args = (jax.jit(jax.random.fold_in)(init_rng, step),) else: step_rng_args = () with maybe_record_goodput(recorder, GoodputEvent.STEP, step): with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_train): - if shard_optimizer_over_data and isinstance(model, nn.Module): - state = sharding.maybe_shard_with_name(state, state_mesh_shardings, shard_mode) if config.retry_when_tokens_dropped and p_train_step_dropless is not None: candidate_state, metrics = p_train_step(state, example_batch, *step_rng_args) if bool(metrics.get("has_moe_overflow")): @@ -1026,9 +818,7 @@ def train_loop(config, recorder, state=None): jit_model_dropless = None - if isinstance(model, nn.Module): - jit_model = model - elif config.enable_diloco: + if config.enable_diloco: # state is the DiLoCoTrainState; `model` is already the TrainStateNNX graphdef the inner step needs. jit_model = model else: @@ -1046,7 +836,7 @@ def train_loop(config, recorder, state=None): del reconstructed, _ gc.collect() - if config.pure_nnx and config.enable_diloco: + if config.enable_diloco: # DiLoCoTrainState.params already holds the param shardings the inner step needs; # the Zero-1 opt overlay doesn't apply through the diloco wrapper. params_shardings = state_mesh_shardings.params @@ -1085,13 +875,11 @@ def train_loop(config, recorder, state=None): with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): data_sharding = sharding.get_input_data_sharding(config, mesh) shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding) - if config.shard_optimizer_over_data and isinstance(model, nn.Module): - state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode) - elif config.shard_optimizer_over_data: + if config.shard_optimizer_over_data: # NNX: reshard state so params match the data-sharded in_shardings (Zero-1 layout) state = jax.device_put(state, state_mesh_shardings) - if isinstance(model, nn.Module) or config.enable_diloco: - # The DiLoCo train step takes (state, batch, rng), like the Linen step. + if config.enable_diloco: + # The DiLoCo train step takes (state, batch, rng), like the inner NNX step. lower_args = (state, shaped_batch, init_rng) else: lower_args = (state, shaped_batch) @@ -1110,9 +898,7 @@ def train_loop(config, recorder, state=None): metric_logger_instance = metric_logger.MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) # Write train config params, num model params, and XLA flags to tensorboard - if isinstance(model, nn.Module): - setup_params = state.params - elif config.enable_diloco: + if config.enable_diloco: setup_params = state.params # DiLoCoTrainState.params: the outer (global) params else: _, setup_params, _ = nnx.split(state.model, nnx.Param, ...) diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py index b705bef562..d87ab1a991 100644 --- a/src/maxtext/trainers/pre_train/train_compile.py +++ b/src/maxtext/trainers/pre_train/train_compile.py @@ -38,9 +38,8 @@ import jax.numpy as jnp from jax.sharding import AxisType, Mesh from maxtext.common import train_state_nnx -from maxtext.common.common_types import MODEL_MODE_TRAIN, ShardMode +from maxtext.common.common_types import ShardMode from maxtext.configs import pyconfig -from maxtext.layers import quantizations from maxtext.models import models from maxtext.optimizers import optimizers from maxtext.trainers.diloco import diloco @@ -130,68 +129,49 @@ def _nnx_forward(decoder_input_tokens, decoder_positions, decoder_segment_ids): def get_shaped_inputs(topology_mesh, config): """Get shaped abstractions of inputs to train_step: state, batch and rng""" # Construct the model and optimizer to get shaped versions of the state - quant = quantizations.configure_quantization(config) - if config.pure_nnx: - _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, topology_mesh) - else: - model = Transformer(config, topology_mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, topology_mesh) # The learning_rate_schedule is baked into the compiled object. learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) # pass in model for muon tx = optimizers.get_optimizer(config, learning_rate_schedule, model) - # Shaped RNG keys - _, example_rng = jax.random.split(jax.random.PRNGKey(0), 2) - shaped_rng = jax.ShapeDtypeStruct(example_rng.shape, example_rng.dtype) - - if config.pure_nnx: - - def create_train_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) + def create_train_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) - init_state_fn = create_train_state_fn - else: - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, example_rng) + init_state_fn = create_train_state_fn # Shaped state abstract_state, _, state_mesh_shardings = maxtext_utils.get_abstract_state(config, topology_mesh, init_state_fn, True) - if config.pure_nnx: - # NNX doesn't use Linen logical annotations; derive PartitionSpecs from the physical shardings. - logical_annotations = maxtext_utils_nnx.get_partition_spec_nnx(state_mesh_shardings) - # For NNX, get_functional_train_with_signature expects the graphdef (static structure), - # not the raw model — mirroring how the training loop does nnx.split(train_state). - with nn_partitioning.axis_rules(config.logical_axis_rules): - abs_train_state = nnx.eval_shape(init_state_fn) - graphdef, _ = nnx.split(abs_train_state) - model = graphdef - else: - # unsharded logical annotations - logical_annotations = maxtext_utils.get_logical_annotations(config, topology_mesh, init_state_fn) + # NNX doesn't use Linen logical annotations; derive PartitionSpecs from the physical shardings. + logical_annotations = maxtext_utils_nnx.get_partition_spec_nnx(state_mesh_shardings) + # For NNX, get_functional_train_with_signature expects the graphdef (static structure), + # not the raw model — mirroring how the training loop does nnx.split(train_state). + with nn_partitioning.axis_rules(config.logical_axis_rules): + abs_train_state = nnx.eval_shape(init_state_fn) + graphdef, _ = nnx.split(abs_train_state) + model = graphdef # Shaped batch data_sharding = sharding.get_input_data_sharding(config, topology_mesh) shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding) - if config.pure_nnx: - shaped_train_args = ( - abstract_state, - shaped_batch, - ) # NNX doesn't use dropout_rng - else: - shaped_train_args = (abstract_state, shaped_batch, shaped_rng) + shaped_train_args = ( + abstract_state, + shaped_batch, + ) # NNX doesn't use dropout_rng shaped_train_kwargs = {} # Collect NNX activation shardings via an abstract forward pass (must run # after get_abstract_state, which only traces __init__). - if config.debug_sharding and config.pure_nnx: + if config.debug_sharding: _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh) # pyrefly: ignore[unbound-name] return ( @@ -392,20 +372,12 @@ def main(argv: Sequence[str]) -> None: # print weights sharding info under debug sharding mode if config.debug_sharding: max_utils.print_non_trivial_mesh_axis(topology_mesh) - if config.pure_nnx: - maxtext_utils.print_shardings_params( - shaped_train_args[0], - state_mesh_shardings, - topology_mesh, - logical_annotations, - ) - else: - maxtext_utils.print_shardings_params( - shaped_train_args[0].params, - state_mesh_shardings.params, - topology_mesh, - logical_annotations.params, - ) + maxtext_utils.print_shardings_params( + shaped_train_args[0], + state_mesh_shardings, + topology_mesh, + logical_annotations, + ) # Compile print("Jitting and compiling train step...", flush=True) diff --git a/tests/assets/golden_logits/golden_dpo_correctness.json b/tests/assets/golden_logits/golden_dpo_correctness.json index 5552ebde26..7b3dced630 100644 --- a/tests/assets/golden_logits/golden_dpo_correctness.json +++ b/tests/assets/golden_logits/golden_dpo_correctness.json @@ -2,17 +2,17 @@ "explicit_prompt_len_3_column": { "loss_step_1": 0.6931471824645996, "margin_step_1": 0.0, - "loss": 0.6573728919029236, - "margin": 0.0728759765625, - "chosen_logps": -854.4208984375, - "rejected_logps": -530.0099487304688 + "loss": 0.690365, + "margin": 0.005573, + "chosen_logps": -882.7361450195312, + "rejected_logps": -587.572693 }, "default_prompt_len_2_column": { "loss_step_1": 0.6931471824645996, "margin_step_1": 0.0, - "loss": 0.6954630613327026, - "margin": -0.004626465495675802, - "chosen_logps": -875.7848510742188, - "rejected_logps": -502.052978515625 + "loss": 0.620875, + "margin": 0.150177, + "chosen_logps": -907.599609375, + "rejected_logps": -578.226013 } -} \ No newline at end of file +} diff --git a/tests/assets/logits_generation/generate_grpo_golden_logits.py b/tests/assets/logits_generation/generate_grpo_golden_logits.py index a7d87a1c3f..193030b8a6 100644 --- a/tests/assets/logits_generation/generate_grpo_golden_logits.py +++ b/tests/assets/logits_generation/generate_grpo_golden_logits.py @@ -25,7 +25,6 @@ import unittest from datasets import load_dataset -from flax import linen as nn from flax import nnx import jax import jax.numpy as jnp @@ -33,11 +32,10 @@ import jsonlines from maxtext.configs import pyconfig, types from maxtext.utils.globals import MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT -from maxtext.common.common_types import Array, MODEL_MODE_TRAIN -from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, generate_completions, grpo_loss_fn, grpo_loss_fn_nnx -from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx +from maxtext.common.common_types import Array +from maxtext.experimental.rl.grpo_trainer import generate_completions, grpo_loss_fn_nnx +from maxtext.experimental.rl.grpo_utils import compute_log_probs_nnx from maxtext.inference.maxengine import maxengine -from maxtext.models import models from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils from tests.post_training.integration.grpo_trainer_correctness_test import prepare_maxtext_inputs @@ -49,40 +47,28 @@ def _setup_model(config, mesh, rng): - """Builds the model, and for NNX a frozen reference clone, dispatching on pure_nnx. + """Builds the model and a frozen reference clone. - Returns (model, reference_model, state). For NNX the model carries its own params - (from_pretrained loads the checkpoint or inits) and state is None; for Linen the - model is a ToLinen module with a separate decode state. + Returns (model, reference_model). The model carries its own params (from_pretrained + loads the checkpoint or inits) and the reference is a clone of the policy. """ - if config.pure_nnx: - model = model_creation_utils.from_pretrained(config, mesh=mesh, rng_key=rng) - return model, nnx.clone(model), None - model = models.transformer_as_linen(config=config, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, config, False, rng) - state, state_mesh_annotations = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn) - return model, None, (state, state_mesh_annotations) + model = model_creation_utils.from_pretrained(config, mesh=mesh, rng_key=rng) + return model, nnx.clone(model) -def _logps(config, model, state, ids, pos, seg, comp_seg): - """Policy per-token log-probs, dispatching between NNX and Linen.""" - if config.pure_nnx: - return compute_log_probs_nnx(model, ids, pos, seg, comp_seg, config, is_train=False) - return compute_log_probs(model, state.params, ids, pos, seg, comp_seg, config, is_train=False) +def _logps(config, model, ids, pos, seg, comp_seg): + """Policy per-token log-probs.""" + return compute_log_probs_nnx(model, ids, pos, seg, comp_seg, config, is_train=False) -def _reference_logps(config, model, reference_model, reference_params, ids, pos, seg, comp_seg): - """Reference per-token log-probs. NNX uses the cloned reference model; Linen uses the saved params.""" - if config.pure_nnx: - return compute_log_probs_nnx(reference_model, ids, pos, seg, comp_seg, config, is_train=False) - return compute_log_probs(model, {"params": reference_params}, ids, pos, seg, comp_seg, config, is_train=False) +def _reference_logps(config, reference_model, ids, pos, seg, comp_seg): + """Reference per-token log-probs, using the cloned reference model.""" + return compute_log_probs_nnx(reference_model, ids, pos, seg, comp_seg, config, is_train=False) -def _grpo_loss(config, model, reference_model, state, reference_params, data, rng): - """GRPO loss, dispatching between NNX (reference model) and Linen (reference params).""" - if config.pure_nnx: - return grpo_loss_fn_nnx(model, config, data, rng, None, reference_model) - return grpo_loss_fn(model, config, data, rng, state.params, reference_params) +def _grpo_loss(config, model, reference_model, data, rng): + """GRPO loss, using the cloned reference model.""" + return grpo_loss_fn_nnx(model, config, data, rng, None, reference_model) class GRPOTest(unittest.TestCase): @@ -116,19 +102,12 @@ def setUp(self): mesh = Mesh(devices_array, self.cfg.mesh_axes) self.mesh = mesh # With checkpoint - self.model, self.reference_model, linen_state = _setup_model(self.cfg, mesh, self.rng) - if self.cfg.pure_nnx: - self.state = None - self.state_mesh_shardings = None # NNX param shardings are derived in the generation step. - else: - self.state, state_mesh_annotations = linen_state - self.state_mesh_shardings = nn.logical_to_mesh_sharding(state_mesh_annotations, mesh, self.cfg.logical_axis_rules) + self.model, self.reference_model = _setup_model(self.cfg, mesh, self.rng) self.data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None)) # Without checkpoint - self.model_no_ckpt_loading, self.reference_model_no_ckpt_loading, linen_state_no_ckpt = _setup_model( + self.model_no_ckpt_loading, self.reference_model_no_ckpt_loading = _setup_model( self.cfg_no_ckpt_loading, mesh, self.rng ) - self.state_no_ckpt_loading = None if self.cfg_no_ckpt_loading.pure_nnx else linen_state_no_ckpt[0] self.tokenizer_model = transformers.AutoTokenizer.from_pretrained( "meta-llama/Llama-3.1-8B", @@ -217,29 +196,16 @@ def test_w_trl_and_write_golden_data(self): self.cfg.prompt, self.tokenizer_model ) maxtext_per_token_logps, _ = _logps( - self.cfg, self.model, self.state, input_ids, input_position, input_segmentation, completion_segmentation + self.cfg, self.model, input_ids, input_position, input_segmentation, completion_segmentation ) - # The reference is a frozen copy of the step-0 policy. NNX holds it as a cloned - # model (built in setUp); Linen snapshots the params and merges them into the state. - reference_params = None - reference_params_no_ckpt_loading = None - if not self.cfg.pure_nnx: - reference_params = jax.tree.map(jnp.copy, self.state.params["params"]) - self.state = _merge_grpo_state(self.state, reference_params) - if not self.cfg_no_ckpt_loading.pure_nnx: - reference_params_no_ckpt_loading = jax.tree.map(jnp.copy, self.state_no_ckpt_loading.params["params"]) - self.state_no_ckpt_loading = _merge_grpo_state(self.state_no_ckpt_loading, reference_params_no_ckpt_loading) - data = { "prompt_completions": input_ids, "prompt_completions_position": input_position, "prompt_completions_segmentation": input_segmentation, "ar_completions_segmentation": completion_segmentation, } - maxtext_loss, aux = _grpo_loss( - self.cfg, self.model, self.reference_model, self.state, reference_params, data, self.rng - ) + maxtext_loss, aux = _grpo_loss(self.cfg, self.model, self.reference_model, data, self.rng) # pylint: disable=protected-access self.assertEqual(self.trainer._metrics["train"]["kl"][0], aux.avg_kl.tolist()) self.assertEqual(hf_loss.item(), maxtext_loss.tolist()) @@ -247,13 +213,11 @@ def test_w_trl_and_write_golden_data(self): self.assertEqual(aux.avg_advantage.tolist(), 0.0) # since we are at step 0 maxtext_per_token_logps, _ = _logps( - self.cfg, self.model, self.state, input_ids, input_position, input_segmentation, completion_segmentation + self.cfg, self.model, input_ids, input_position, input_segmentation, completion_segmentation ) maxtext_per_token_logps_ref, _ = _reference_logps( self.cfg, - self.model, self.reference_model, - reference_params, input_ids, input_position, input_segmentation, @@ -274,7 +238,6 @@ def test_w_trl_and_write_golden_data(self): maxtext_per_token_logps_no_ckpt_loading, _ = _logps( self.cfg_no_ckpt_loading, self.model_no_ckpt_loading, - self.state_no_ckpt_loading, input_ids, input_position, input_segmentation, @@ -285,8 +248,6 @@ def test_w_trl_and_write_golden_data(self): self.cfg_no_ckpt_loading, self.model_no_ckpt_loading, self.reference_model_no_ckpt_loading, - self.state_no_ckpt_loading, - reference_params_no_ckpt_loading, data, self.rng, ) @@ -301,13 +262,9 @@ def test_w_trl_and_write_golden_data(self): ) prompt_true_length = jnp.array([len(prompt_tokens)] * 4) engine_data = {"prompt": prompt, "prompt_true_length": prompt_true_length} - if self.cfg_no_ckpt_loading.pure_nnx: - # NNX params live on the model; the inference engine is NNX-aware (config.pure_nnx). - gen_params = nnx.state(self.model_no_ckpt_loading, nnx.Param) - gen_param_shardings = jax.tree.map(lambda _: jax.NamedSharding(self.mesh, jax.sharding.PartitionSpec()), gen_params) - else: - gen_params = {"params": self.state_no_ckpt_loading.params["params"]} - gen_param_shardings = self.state_mesh_shardings.params + # Params live on the model; the inference engine reads them directly. + gen_params = nnx.state(self.model_no_ckpt_loading, nnx.Param) + gen_param_shardings = jax.tree.map(lambda _: jax.NamedSharding(self.mesh, jax.sharding.PartitionSpec()), gen_params) p_generate_completions: Callable[[dict, dict, Array], Array] = jax.jit( functools.partial(generate_completions, self.cfg, self.tokenizer_model, engine), in_shardings=(self.data_sharding, gen_param_shardings, None), diff --git a/tests/integration/diloco_test.py b/tests/integration/diloco_test.py index 2aeef5e408..9002176825 100644 --- a/tests/integration/diloco_test.py +++ b/tests/integration/diloco_test.py @@ -23,7 +23,6 @@ import chex from flax.experimental import nnx -from flax.training import train_state import jax import jax.numpy as jnp import jax.sharding @@ -88,72 +87,36 @@ def test_diloco_training_simulation_with_mesh(self): tx = optax.sgd(learning_rate=0.1) rngs = nnx.Rngs(params=jax.random.key(seed=42)) model = SimpleNNXModel(rngs=rngs) - graphdef, params = nnx.split(model) - if test_config.pure_nnx: - optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) - # diloco_test_state expects a TrainStateNNX instance when pure_nnx is True. - initial_test_state = TrainStateNNX(model, optimizer) + optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) + # diloco_test_state expects a TrainStateNNX instance. + initial_test_state = TrainStateNNX(model, optimizer) - # For NNX, train_step needs to take the TrainStateNNX and mutate it + # train_step takes the TrainStateNNX and mutates it. - def _test_train_step(state, batch, prng_key: diloco.PRNGKey): - del prng_key + def _test_train_step(state, batch, prng_key: diloco.PRNGKey): + del prng_key - def loss_fn(model, batch): - inputs, labels = batch - logits = jax.vmap(model)(inputs) - residual = logits - labels - return jnp.mean(jnp.square(residual)) + def loss_fn(model, batch): + inputs, labels = batch + logits = jax.vmap(model)(inputs) + residual = logits - labels + return jnp.mean(jnp.square(residual)) - loss, grads = nnx.value_and_grad(loss_fn)(state.model, batch) - state.optimizer.update(state.model, grads) - return state, loss - - else: - - def nnx_apply_fn(params, inputs): - model_replica = nnx.merge(graphdef, params) - return model_replica(inputs) - - # 2. Vmap this new wrapper function - vmapped_apply = jax.vmap(nnx_apply_fn, in_axes=(None, 0)) - - def _test_train_step(state: train_state.TrainState, batch, prng_key: diloco.PRNGKey): - """A simple MSE loss train step to enable numerics testing.""" - del prng_key - - def loss_fn(params, batch): - inputs, labels = batch - logits = vmapped_apply(params, inputs) - residual = logits - labels - sq_residual = jnp.square(residual) - msq_residual = jnp.mean(sq_residual) - return msq_residual - - loss, grad = jax.value_and_grad(loss_fn)(state.params, batch) - return state.apply_gradients(grads=grad), loss - - initial_test_state = train_state.TrainState.create( - apply_fn=vmapped_apply, - params=params, - tx=tx, - ) + loss, grads = nnx.value_and_grad(loss_fn)(state.model, batch) + state.optimizer.update(state.model, grads) + return state, loss diloco_test_state, _ = diloco.build_diloco_state(test_config, lambda: initial_test_state) chex.assert_equal(diloco_test_state.step, 0) - if test_config.pure_nnx: - _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...) - - # diloco_test_state.params might contain nnx.Variables instead of pure arrays. - # We need to unwrap them if they do. - diloco_params_pure = jax.tree_util.tree_map( - lambda x: x.value if hasattr(x, "value") else x, - diloco_test_state.params, - ) - chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict()) - else: - chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params) + _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...) + + # diloco_test_state.params might contain nnx.Variables instead of pure arrays. + # We need to unwrap them if they do. + diloco_params_pure = jax.tree_util.tree_map( + lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params + ) + chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict()) diloco_train_step = diloco.build_diloco_train_step(test_config, _test_train_step) inputs = jnp.array( @@ -201,18 +164,14 @@ def loss_fn(params, batch): chex.assert_equal(diloco_test_state.step, 1.0) chex.assert_equal(loss, 1.0) # Assert no updates to the global model yet (no synchronization) - if test_config.pure_nnx: - _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...) - - # diloco_test_state.params might contain nnx.Variables instead of pure arrays. - # We need to unwrap them if they do. - diloco_params_pure = jax.tree_util.tree_map( - lambda x: x.value if hasattr(x, "value") else x, - diloco_test_state.params, - ) - chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict()) - else: - chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params) + _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...) + + # diloco_test_state.params might contain nnx.Variables instead of pure arrays. + # We need to unwrap them if they do. + diloco_params_pure = jax.tree_util.tree_map( + lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params + ) + chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict()) # Run the second step (no synchronization). # Replica 0: @@ -242,18 +201,14 @@ def loss_fn(params, batch): chex.assert_equal(diloco_test_state.step, 2.0) chex.assert_trees_all_close(loss, 0.49, rtol=1e-2, atol=1e-2) # Assert no updates to the global model yet (no synchronization) - if test_config.pure_nnx: - _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...) - - # diloco_test_state.params might contain nnx.Variables instead of pure arrays. - # We need to unwrap them if they do. - diloco_params_pure = jax.tree_util.tree_map( - lambda x: x.value if hasattr(x, "value") else x, - diloco_test_state.params, - ) - chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict()) - else: - chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params) + _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...) + + # diloco_test_state.params might contain nnx.Variables instead of pure arrays. + # We need to unwrap them if they do. + diloco_params_pure = jax.tree_util.tree_map( + lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params + ) + chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict()) # Run the third step, which synchronizes afterwards. # Replica 0: @@ -288,33 +243,21 @@ def loss_fn(params, batch): chex.assert_trees_all_close(loss, 0.2401, rtol=1e-2, atol=1e-2) # Assert that inner and outer parameters are all equal now that # synchronization has happened. - if test_config.pure_nnx: - _, inner_params, _ = nnx.split(diloco_test_state.inner_state.model, nnx.Param, ...) - inner_params_pure = jax.tree_util.tree_map( - lambda x: x.value if hasattr(x, "value") else x, - inner_params.to_pure_dict(), - ) - diloco_params_pure_3 = jax.tree_util.tree_map( - lambda x: x.value if hasattr(x, "value") else x, - diloco_test_state.params, - ) - chex.assert_trees_all_equal( - diloco_params_pure_3, - jax.tree.map(lambda arr: arr[0, ...], inner_params_pure), - ) - chex.assert_trees_all_equal( - diloco_params_pure_3, - jax.tree.map(lambda arr: arr[1, ...], inner_params_pure), - ) - else: - chex.assert_trees_all_equal( - diloco_test_state.params, - jax.tree.map(lambda arr: arr[0, ...], diloco_test_state.inner_state.params), - ) - chex.assert_trees_all_equal( - diloco_test_state.params, - jax.tree.map(lambda arr: arr[1, ...], diloco_test_state.inner_state.params), - ) + _, inner_params, _ = nnx.split(diloco_test_state.inner_state.model, nnx.Param, ...) + inner_params_pure = jax.tree_util.tree_map( + lambda x: x.value if hasattr(x, "value") else x, inner_params.to_pure_dict() + ) + diloco_params_pure_3 = jax.tree_util.tree_map( + lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params + ) + chex.assert_trees_all_equal( + diloco_params_pure_3, + jax.tree.map(lambda arr: arr[0, ...], inner_params_pure), + ) + chex.assert_trees_all_equal( + diloco_params_pure_3, + jax.tree.map(lambda arr: arr[1, ...], inner_params_pure), + ) # Run the fourth step (no synchronization). # Replica 0: diff --git a/tests/post_training/integration/dpo_correctness_base.py b/tests/post_training/integration/dpo_correctness_base.py index eab78cae1e..11346f065b 100644 --- a/tests/post_training/integration/dpo_correctness_base.py +++ b/tests/post_training/integration/dpo_correctness_base.py @@ -202,9 +202,6 @@ def build_tiny_qwen2_jax_config( "per_device_batch_size=1", f"max_target_length={max_target_length}", "skip_jax_distributed_system=True", - "enable_nnx=True", - "pure_nnx=True", - "pure_nnx_decoder=False", "remat_policy=full", "log_config=0", # Tiny architecture specifications. diff --git a/tests/post_training/integration/grpo_correctness.py b/tests/post_training/integration/grpo_correctness.py index d2783ea372..98a568d08c 100644 --- a/tests/post_training/integration/grpo_correctness.py +++ b/tests/post_training/integration/grpo_correctness.py @@ -13,7 +13,6 @@ # limitations under the License. """GRPO correctness tests""" -import functools import os import unittest @@ -23,11 +22,9 @@ import jax.numpy as jnp from jax.sharding import Mesh from maxtext.configs import pyconfig, types -from maxtext.common.common_types import MODEL_MODE_TRAIN -from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, grpo_loss_fn, grpo_loss_fn_nnx -from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx +from maxtext.experimental.rl.grpo_trainer import grpo_loss_fn_nnx +from maxtext.experimental.rl.grpo_utils import compute_log_probs_nnx from maxtext.utils.globals import MAXTEXT_PKG_DIR -from maxtext.models import models from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils import numpy as np @@ -68,17 +65,11 @@ def setUp(self): self.rng = jax.random.PRNGKey(42) devices_array = maxtext_utils.create_device_mesh(self.cfg) mesh = Mesh(devices_array, self.cfg.mesh_axes) - if self.cfg.pure_nnx: - # NNX: from_pretrained loads the checkpoint (or inits) into the model, which - # carries its own params. The frozen reference is a clone of the policy. - self.model = model_creation_utils.from_pretrained(self.cfg, mesh=mesh, rng_key=self.rng) - self.reference_model = nnx.clone(self.model) - self.state = None - else: - self.model = models.transformer_as_linen(config=self.cfg, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.cfg, False, self.rng) - self.reference_model = None - self.state, _ = maxtext_utils.setup_decode_state(self.cfg, mesh, None, init_state_fn) + # NNX: from_pretrained loads the checkpoint (or inits) into the model, which + # carries its own params. The frozen reference is a clone of the policy. + self.model = model_creation_utils.from_pretrained(self.cfg, mesh=mesh, rng_key=self.rng) + self.reference_model = nnx.clone(self.model) + self.state = None tokenizer_path = ensure_tokenizer_downloaded("llama3.1-tokenizer", skip_test_on_failure=True) self.tokenizer_model = transformers.AutoTokenizer.from_pretrained( tokenizer_path, @@ -152,57 +143,25 @@ def _prepare_trl_inputs(self): return input_ids, attention_mask, logits_to_keep def _maxtext_logits(self, inputs, inputs_position, inputs_segmentation): - """Forward pass logits, dispatching between the NNX and Linen models.""" - if self.cfg.pure_nnx: - return self.model( - decoder_input_tokens=inputs, - decoder_positions=inputs_position, - decoder_segment_ids=inputs_segmentation, - enable_dropout=False, - ) - logits, _ = self.model.apply( - self.state.params, - inputs, - inputs_position, + """NNX forward pass logits.""" + return self.model( + decoder_input_tokens=inputs, + decoder_positions=inputs_position, decoder_segment_ids=inputs_segmentation, enable_dropout=False, - rngs=self.rng, - mutable="intermediates", ) - return logits def _policy_logps(self, input_ids, input_position, input_segmentation, completion_segmentation): - """Policy per-token log-probs, dispatching between NNX and Linen.""" - if self.cfg.pure_nnx: - return compute_log_probs_nnx( - self.model, input_ids, input_position, input_segmentation, completion_segmentation, self.cfg, is_train=False - ) - return compute_log_probs( - self.model, - self.state.params, - input_ids, - input_position, - input_segmentation, - completion_segmentation, - self.cfg, - is_train=False, + """Policy per-token log-probs.""" + return compute_log_probs_nnx( + self.model, input_ids, input_position, input_segmentation, completion_segmentation, self.cfg, is_train=False ) def _reference_logps(self, input_ids, input_position, input_segmentation, completion_segmentation, reference_params): - """Reference per-token log-probs. NNX uses the cloned reference model; Linen uses the saved params.""" - if self.cfg.pure_nnx: - return compute_log_probs_nnx( - self.reference_model, - input_ids, - input_position, - input_segmentation, - completion_segmentation, - self.cfg, - is_train=False, - ) - return compute_log_probs( - self.model, - {"params": reference_params}, + """Reference per-token log-probs from the cloned reference model.""" + del reference_params # NNX reads the cloned reference model directly + return compute_log_probs_nnx( + self.reference_model, input_ids, input_position, input_segmentation, @@ -212,10 +171,9 @@ def _reference_logps(self, input_ids, input_position, input_segmentation, comple ) def _grpo_loss(self, data, reference_params): - """GRPO loss, dispatching between NNX (reference model) and Linen (reference params).""" - if self.cfg.pure_nnx: - return grpo_loss_fn_nnx(self.model, self.cfg, data, self.rng, None, self.reference_model) - return grpo_loss_fn(self.model, self.cfg, data, self.rng, self.state.params, reference_params) + """GRPO loss from the NNX policy and cloned reference model.""" + del reference_params # NNX reads the cloned reference model directly + return grpo_loss_fn_nnx(self.model, self.cfg, data, self.rng, None, self.reference_model) def test_logits(self): def _prepare_inputs(): @@ -311,12 +269,9 @@ def test_loss_kl_div(self): input_ids, input_position, input_segmentation, completion_segmentation ) - # The reference is a frozen copy of the step-0 policy. NNX holds it as a cloned - # model (built in setUp); Linen snapshots the params and merges them into the state. + # The reference is a frozen copy of the step-0 policy, held as a cloned model + # built in setUp; the NNX loss path reads it directly. reference_params = None - if not self.cfg.pure_nnx: - reference_params = jax.tree.map(jnp.copy, self.state.params["params"]) - self.state = _merge_grpo_state(self.state, reference_params) data = { "prompt_completions": input_ids, "prompt_completions_position": input_position, diff --git a/tests/post_training/integration/grpo_trainer_correctness_test.py b/tests/post_training/integration/grpo_trainer_correctness_test.py index 95bfd28190..725e19a455 100644 --- a/tests/post_training/integration/grpo_trainer_correctness_test.py +++ b/tests/post_training/integration/grpo_trainer_correctness_test.py @@ -25,11 +25,9 @@ pytest tests/post_training/integration/grpo_trainer_correctness_test.py """ -import functools import os import sys import unittest -from flax import linen as nn import jax import jax.numpy as jnp from jax.sharding import Mesh @@ -37,17 +35,14 @@ import maxtext as mt from maxtext.configs import pyconfig, types from maxtext.utils.globals import MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT -from maxtext.common.common_types import MODEL_MODE_TRAIN from flax import nnx from maxtext.experimental.rl import grpo_utils -from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, grpo_loss_fn, grpo_loss_fn_nnx, setup_train_loop -from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx +from maxtext.experimental.rl.grpo_trainer import grpo_loss_fn_nnx, setup_train_loop +from maxtext.experimental.rl.grpo_utils import compute_log_probs_nnx from tests.utils.test_helpers import ensure_tokenizer_downloaded from maxtext.inference import offline_engine from maxtext.inference.maxengine import maxengine from maxtext.inference.offline_engine import InputData -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils import numpy as np @@ -73,49 +68,28 @@ def get_golden_data(config): def setup_maxtext_model(config, mesh): """Sets up the MaxText model. - Returns (model, state, reference, init_rng, state_mesh_shardings, data_sharding). On - the NNX path the model carries its own params (state is None) and reference is a - cloned frozen model; on Linen, state is the decode state and reference is a copy of - the params merged into it. + Returns (model, state, reference, init_rng, state_mesh_shardings, data_sharding). The + model carries its own params (state is None) and reference is a cloned frozen model. """ init_rng = jax.random.PRNGKey(config.init_weights_seed) data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None)) - if config.pure_nnx: - maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, rng_key=init_rng) - reference = nnx.clone(maxtext_model) - # state_mesh_shardings is unused by the live correctness test on the NNX path. - return maxtext_model, None, reference, init_rng, None, data_sharding - - quant = quantizations.configure_quantization(config) - maxtext_model = models.transformer_as_linen(config=config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, init_rng) - state, state_mesh_annotations = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn) - state_mesh_shardings = nn.logical_to_mesh_sharding(state_mesh_annotations, mesh, config.logical_axis_rules) - reference_params = jax.tree.map(jnp.copy, state.params["params"]) - state = _merge_grpo_state(state, reference_params) - return ( - maxtext_model, - state, - reference_params, - init_rng, - state_mesh_shardings, - data_sharding, - ) + maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, rng_key=init_rng) + reference = nnx.clone(maxtext_model) + # state_mesh_shardings is unused by the live correctness test on the NNX path. + return maxtext_model, None, reference, init_rng, None, data_sharding def _logps(config, model, state, ids, pos, seg, comp_seg, rngs=None): - """Per-token log-probs, dispatching between the NNX and Linen models.""" - if config.pure_nnx: - return compute_log_probs_nnx(model, ids, pos, seg, comp_seg, config, is_train=False) - return compute_log_probs(model, state.params, ids, pos, seg, comp_seg, config, is_train=False, rngs=rngs) + """Per-token log-probs from the NNX model.""" + del state, rngs # unused on the NNX path + return compute_log_probs_nnx(model, ids, pos, seg, comp_seg, config, is_train=False) def _grpo_loss(config, model, state, reference, data, rng): - """GRPO loss. On NNX `reference` is the frozen reference model; on Linen it is the reference params.""" - if config.pure_nnx: - return grpo_loss_fn_nnx(model, config, data, rng, None, reference) - return grpo_loss_fn(model, config, data, rng, state.params, reference) + """GRPO loss with the NNX policy and frozen reference model.""" + del state # unused on the NNX path + return grpo_loss_fn_nnx(model, config, data, rng, None, reference) def prepare_maxtext_inputs(input_str, tokenizer_model): diff --git a/tests/post_training/integration/sft_trainer_correctness_test.py b/tests/post_training/integration/sft_trainer_correctness_test.py index 253ba7ade9..edcfba2fe5 100644 --- a/tests/post_training/integration/sft_trainer_correctness_test.py +++ b/tests/post_training/integration/sft_trainer_correctness_test.py @@ -24,7 +24,6 @@ pytest tests/post_training/integration/sft_trainer_correctness_test.py """ -import functools import os.path import sys import unittest @@ -34,11 +33,8 @@ import jax.numpy as jnp from jax.sharding import Mesh import jsonlines -from maxtext.common.common_types import MODEL_MODE_TRAIN from maxtext.configs import pyconfig from maxtext.input_pipeline import input_pipeline_utils -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx from maxtext.utils import model_creation_utils @@ -118,40 +114,23 @@ def setup_maxtext_model(config): init_rng = jax.random.PRNGKey(config.init_weights_seed) devices_array = maxtext_utils.create_device_mesh(config) mesh = Mesh(devices_array, config.mesh_axes) - if config.pure_nnx: - # NNX model: params live on the module, so there is no separate train state. - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) - with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): - maxtext_model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) - return maxtext_model, None, init_rng - quant = quantizations.configure_quantization(config) - maxtext_model = models.transformer_as_linen(config=config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, init_rng) - state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn) - return maxtext_model, state, init_rng + # NNX model: params live on the module, so there is no separate train state. + rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + maxtext_model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) + return maxtext_model, None, init_rng def get_maxtext_logits(config, maxtext_data): """Get logits generated by MaxText.""" - maxtext_model, state, rng = setup_maxtext_model(config) - if config.pure_nnx: - # NNX forward: the model carries its own params and rng state. - return maxtext_model( - decoder_input_tokens=maxtext_data["inputs"], - decoder_positions=maxtext_data["inputs_position"], - decoder_segment_ids=maxtext_data["inputs_segmentation"], - enable_dropout=False, - ) - maxtext_logits, _ = maxtext_model.apply( - state.params, - maxtext_data["inputs"], - maxtext_data["inputs_position"], + maxtext_model, _, _ = setup_maxtext_model(config) + # NNX forward: the model carries its own params and rng state. + return maxtext_model( + decoder_input_tokens=maxtext_data["inputs"], + decoder_positions=maxtext_data["inputs_position"], decoder_segment_ids=maxtext_data["inputs_segmentation"], enable_dropout=False, - rngs=rng, - mutable="intermediates", ) - return maxtext_logits def get_token_log_probs(logits, inputs): diff --git a/tests/unit/grpo_nnx_test.py b/tests/unit/grpo_nnx_test.py index ddc2de5da0..0ccf30f7a0 100644 --- a/tests/unit/grpo_nnx_test.py +++ b/tests/unit/grpo_nnx_test.py @@ -12,14 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for `grpo_loss_fn_nnx`, `compute_log_probs_nnx`, plus a small -Linen-path regression block (the repo's existing Linen GRPO integration test -is TPU-only).""" +"""Unit tests for `grpo_loss_fn_nnx` and `compute_log_probs_nnx`.""" import types import unittest -import flax.linen as nn import jax import jax.numpy as jnp import numpy as np @@ -95,7 +92,7 @@ def setUp(self): self.data = _make_grpo_batch() def test_aux_structure_matches_linen(self): - """`grpo_loss_fn_nnx` returns the same `LossAux` dataclass shape as `grpo_loss_fn`.""" + """`grpo_loss_fn_nnx` returns a `LossAux` dataclass with the expected fields.""" loss, aux = grpo_trainer.grpo_loss_fn_nnx( self.policy, self.config, self.data, None, None, self.reference, is_train=True ) @@ -295,70 +292,5 @@ def test_host_offload_matches_no_offload(self): np.testing.assert_allclose(np.asarray(a), np.asarray(b), rtol=1e-6, atol=1e-6) -# --------------------------------------------------------------------------- -# Linen-path regression smoke tests -# --------------------------------------------------------------------------- - - -class _MockLinenTransformer(nn.Module): - """Tiny Linen module with the `model.apply(...)` signature that Linen `compute_log_probs` expects.""" - - vocab_size: int - embed_dim: int - - @nn.compact - def __call__(self, inputs, positions, decoder_segment_ids=None, enable_dropout=False): - del positions, decoder_segment_ids, enable_dropout - embed = nn.Embed(num_embeddings=self.vocab_size, features=self.embed_dim, name="embed")(inputs) - return nn.Dense(features=self.vocab_size, name="proj")(embed) - - -class TestLinenGrpoRegression(unittest.TestCase): - """Smoke tests that the Linen `grpo_loss_fn` and `compute_log_probs` still run on Linen-shaped inputs.""" - - def setUp(self): - self.config = _make_grpo_config() - self.config.pure_nnx = False # Force the Linen dispatch branch. - self.config.gradient_accumulation_steps = 1 - self.data = _make_grpo_batch() - self.model = _MockLinenTransformer(vocab_size=8, embed_dim=4) - rng = jax.random.key(0) - inputs = self.data["prompt_completions"] - self.params = self.model.init(rng, inputs, inputs, decoder_segment_ids=jnp.ones_like(inputs), enable_dropout=False) - self.reference_params = jax.tree_util.tree_map(jnp.copy, self.params) - - def test_linen_grpo_loss_fn_still_runs(self): - """Linen `grpo_loss_fn` returns a finite loss + a `LossAux`.""" - loss, aux = grpo_trainer.grpo_loss_fn( - self.model, - self.config, - self.data, - jax.random.key(1), - self.params, - self.reference_params["params"], # On Linen, reference_params is the inner subtree. - is_train=True, - ) - self.assertTrue(jnp.isfinite(loss)) - self.assertTrue(hasattr(aux, "total_loss")) - self.assertTrue(hasattr(aux, "moe_lb_loss")) - self.assertTrue(hasattr(aux, "total_weights")) - - def test_linen_compute_log_probs_still_runs(self): - """Linen `compute_log_probs` produces shape `[B, S-1]`.""" - log_probs, _ = grpo_utils.compute_log_probs( - self.model, - self.params, - self.data["prompt_completions"], - self.data["prompt_completions_position"], - self.data["prompt_completions_segmentation"], - self.data["ar_completions_segmentation"], - self.config, - is_train=False, - rngs={"dropout": jax.random.key(2), "params": jax.random.key(3)}, - ) - S = self.data["prompt_completions"].shape[1] - self.assertEqual(log_probs.shape, (self.data["prompt_completions"].shape[0], S - 1)) - - if __name__ == "__main__": unittest.main() diff --git a/tests/unit/pre_train_loss_mask_test.py b/tests/unit/pre_train_loss_mask_test.py index 549d134562..fe05a7c52a 100644 --- a/tests/unit/pre_train_loss_mask_test.py +++ b/tests/unit/pre_train_loss_mask_test.py @@ -20,7 +20,6 @@ import unittest from unittest import mock -from flax import linen as nn from flax import nnx import jax import jax.numpy as jnp @@ -85,29 +84,6 @@ def __call__( return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) -class _UniformLinenDecoder(nn.Module): - """Returns uniform logits through the Linen call contract.""" - - vocab_size: int - mesh: object - - @nn.compact - def __call__( - self, - decoder_input_tokens, - decoder_positions, - decoder_segment_ids=None, - encoder_images=None, - encoder_image_masks=None, - enable_dropout=False, - decoder_target_tokens=None, - decoder_target_mask=None, - ): - del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks - del enable_dropout, decoder_target_tokens, decoder_target_mask - return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) - - def _make_data(include_loss_mask=True): """Builds a batch whose explicit loss mask differs from segmentation.""" data = { @@ -124,7 +100,7 @@ def _make_data(include_loss_mask=True): class PreTrainLossMaskTest(unittest.TestCase): - """Checks explicit masks in both Linen and NNX loss branches.""" + """Checks explicit target-loss masking in the pre-training loss.""" def setUp(self): super().setUp() @@ -143,20 +119,6 @@ def _use_block_diffusion(self): self.config.attention_type = "block_diffusion" self.config.training_objective = "block_diffusion" - def _linen_model_and_variables(self, data): - """Initializes the test Linen decoder for the supplied batch.""" - mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) - model = _UniformLinenDecoder(vocab_size=self.config.vocab_size, mesh=mesh) - variables = model.init( - jax.random.key(0), - data["inputs"], - data["inputs_position"], - decoder_segment_ids=data["inputs_segmentation"], - decoder_target_tokens=data["targets"], - decoder_target_mask=data["targets_segmentation"], - ) - return model, variables - def _assert_explicit_mask_result(self, loss, aux): expected_mask = _make_data()["targets_loss_mask"] != 0 expected_xent = jnp.sum(self.per_token_xent * expected_mask) @@ -174,27 +136,6 @@ def test_nnx_loss_uses_targets_loss_mask(self): self._assert_explicit_mask_result(loss, aux) - def test_linen_loss_uses_targets_loss_mask(self): - self._use_block_diffusion() - data = _make_data() - model, variables = self._linen_model_and_variables(data) - with self._cross_entropy_patch(): - loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True) - - self._assert_explicit_mask_result(loss, aux) - - def test_linen_causal_loss_skips_diffusion_alignment(self): - data = _make_data(include_loss_mask=False) - model, variables = self._linen_model_and_variables(data) - with self._cross_entropy_patch(): - loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True) - - expected_mask = data["targets_segmentation"] != 0 - expected_xent = jnp.sum(self.per_token_xent * expected_mask) - self.assertEqual(int(aux["total_weights"]), 7) - self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) - self.assertAlmostEqual(float(loss), float(expected_xent / 7.0)) - def test_causal_fallback_uses_targets_segmentation(self): data = _make_data(include_loss_mask=False) model = _UniformNnxDecoder(self.config.vocab_size) diff --git a/tests/unit/router_replay_test.py b/tests/unit/router_replay_test.py index 6344c7fd80..b0c8b52d19 100644 --- a/tests/unit/router_replay_test.py +++ b/tests/unit/router_replay_test.py @@ -41,6 +41,8 @@ from maxtext.models import models from maxtext.trainers.pre_train import train from maxtext.utils import maxtext_utils +from maxtext.utils import maxtext_utils_nnx +from maxtext.utils import model_creation_utils from tests.utils.test_helpers import get_test_config_path import pytest @@ -495,24 +497,17 @@ def _assert_forced_routing_loss_finite(self, cfg, seq_len, batch_size, forced_ex "forced_routed_experts": forced_experts, } - model = models.transformer_as_linen(config=cfg, mesh=mesh, quant=None, model_mode="train") - init_params_rng, init_dropout_rng = jax.random.split(rng) - params = model.init( - {"params": init_params_rng, "dropout": init_dropout_rng}, - inputs, - positions, - segmentation, - enable_dropout=False, - ) - - loss, aux = train.loss_fn( - model, - cfg, - data_batch, - dropout_rng=init_dropout_rng, - params=params, - is_train=True, - ) + rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=rng) + with nn_partitioning.axis_rules(cfg.logical_axis_rules): + model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs, model_mode=ctypes.MODEL_MODE_TRAIN) + loss, aux = train.loss_fn( + model, + cfg, + data_batch, + dropout_rng=None, + params=None, + is_train=True, + ) self.assertIsNotNone(loss) self.assertFalse(jnp.isnan(loss), "Loss must not be NaN") @@ -530,15 +525,7 @@ def _loss_for_routing(self, cfg, seq_len, batch_size, forced_experts): ) segmentation = jnp.ones((batch_size, seq_len), dtype=jnp.int32) - model = models.transformer_as_linen(config=cfg, mesh=mesh, quant=None, model_mode="train") - params_rng, dropout_rng = jax.random.split(jax.random.PRNGKey(42)) - params = model.init( - {"params": params_rng, "dropout": dropout_rng}, - inputs, - positions, - segmentation, - enable_dropout=False, - ) + rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=jax.random.PRNGKey(42)) data_batch = { "inputs": inputs, @@ -549,14 +536,16 @@ def _loss_for_routing(self, cfg, seq_len, batch_size, forced_experts): } if forced_experts is not None: data_batch["forced_routed_experts"] = forced_experts - loss, _ = train.loss_fn( - model, - cfg, - data_batch, - dropout_rng=dropout_rng, - params=params, - is_train=True, - ) + with nn_partitioning.axis_rules(cfg.logical_axis_rules): + model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs, model_mode=ctypes.MODEL_MODE_TRAIN) + loss, _ = train.loss_fn( + model, + cfg, + data_batch, + dropout_rng=None, + params=None, + is_train=True, + ) return float(loss) def _assert_routing_is_load_bearing(self, cfg, seq_len, batch_size, forced_a, forced_b, label):