From cf8744e7cb7b3aff21dbeb0df02acc2ed8cc1b78 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Thu, 4 Jun 2026 18:43:53 +0000 Subject: [PATCH 1/4] [NNX] Delete Linen (1/4): collapse pure_nnx/enable_nnx/isinstance dispatch to NNX-only Across the core training/utils/inference/RL/checkpoint-conversion code, statically collapse every pure_nnx / enable_nnx / isinstance(model, nn.Module) branch to the NNX path (the model is always NNX now). No flag reads remain in these files. --- .../convert_gpt3_ckpt_from_paxml.py | 90 +-- src/maxtext/common/checkpointing.py | 25 +- src/maxtext/experimental/rl/grpo_trainer.py | 309 ++-------- src/maxtext/inference/kvcache.py | 195 ++----- src/maxtext/inference/maxengine/maxengine.py | 538 ++++-------------- src/maxtext/inference/vllm_decode.py | 34 +- src/maxtext/layers/quantizations.py | 35 +- src/maxtext/trainers/diloco/diloco.py | 128 ++--- src/maxtext/trainers/pre_train/train.py | 536 ++++++----------- .../trainers/pre_train/train_compile.py | 80 +-- .../utils/generate_param_only_checkpoint.py | 146 +---- src/maxtext/utils/layerwise_quantization.py | 253 +------- src/maxtext/utils/lora_utils.py | 39 +- src/maxtext/utils/maxtext_utils.py | 110 +--- src/maxtext/utils/model_creation_utils.py | 13 +- src/maxtext/utils/muon_utils.py | 44 +- src/maxtext/utils/sharding.py | 45 +- src/maxtext/utils/standalone_checkpointer.py | 29 +- src/maxtext/utils/train_utils.py | 87 +-- 19 files changed, 634 insertions(+), 2102 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py index b8a513c1df..e5bcfc302e 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py @@ -35,7 +35,6 @@ """ import argparse -import functools import gc import os import sys @@ -47,11 +46,7 @@ from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_PKG_DIR from maxtext.common import checkpointing -from maxtext.common.common_types import MODEL_MODE_TRAIN -from maxtext.layers import quantizations from maxtext.common import train_state_nnx -from maxtext.models.models import transformer_as_linen -from maxtext.optimizers import optimizers from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils @@ -92,23 +87,15 @@ def convert(paxml_ckpt_path, maxtext_model_name, base_output_directory, run_name devices_array = maxtext_utils.create_device_mesh(cfg) mesh = Mesh(devices_array, cfg.mesh_axes) - if cfg.pure_nnx: - rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng) - model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs) - _, tx = train_utils.create_training_optimizer(cfg, model) - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh) + rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng) + model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs) + _, tx = train_utils.create_training_optimizer(cfg, model) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh) - def init_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - - else: - quant = quantizations.configure_quantization(cfg) - model = transformer_as_linen(cfg, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(cfg) - tx = optimizers.get_optimizer(cfg, learning_rate_schedule) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, cfg, True, init_rng) + def init_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) checkpoint_manager = checkpointing.create_orbax_checkpoint_manager( cfg.checkpoint_dir, @@ -201,24 +188,18 @@ def init_state_fn(): "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None), } - if cfg.pure_nnx: - # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a - # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so - # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus - # `.value` for the Variable leaf, plus `[idx]` for the optax tuple: - # model params -> ['model'].value - # adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu'].value - # step -> ['optimizer']['step'].value - # opt count -> ['optimizer']['opt_state'][0]['count'].value - state_map = { - "['optimizer']['step'].value": ("step", None), - "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None), - } - else: - state_map = { - ".step": ("step", None), - ".opt_state.count": ("opt_states_0.no_prefix_0.count", None), - } + # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a + # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so + # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus + # `.value` for the Variable leaf, plus `[idx]` for the optax tuple: + # model params -> ['model'].value + # adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu'].value + # step -> ['optimizer']['step'].value + # opt count -> ['optimizer']['opt_state'][0]['count'].value + state_map = { + "['optimizer']['step'].value": ("step", None), + "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None), + } def get_layer_prefix(keystr_pax): # different path format between decoder_layer variable @@ -231,26 +212,15 @@ def get_layer_prefix(keystr_pax): for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items(): prefix_pax_opt_state = get_layer_prefix(keystr_pax) - if cfg.pure_nnx: - state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) - else: - state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f".opt_state.mu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f".opt_state.nu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) + state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn) + state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = ( + f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", + transform_fn, + ) + state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = ( + f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", + transform_fn, + ) def verify_fn(key_path, _): keystr = jax.tree_util.keystr(key_path) @@ -302,7 +272,7 @@ def map_fn(key_path, value): max_logging.log("converted state finished") max_utils.print_mem_stats("converted state finished") - step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step + step_value = int(converted_state.optimizer.step.value) if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state): max_logging.log(f"saved a checkpoint at step {step_value}") # Upon preemption, exit when and only when all ongoing saves are complete. diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 0316cd29b5..5c515f567f 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -1093,26 +1093,21 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step if step is not None: actual_step = int(step) else: - if config.pure_nnx: - # Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer. - actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1 - else: - # Linen TrainState has .step attribute - actual_step = int(state.step) - 1 + # Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer. + actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1 if checkpoint_manager.latest_step() == actual_step: max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.") return - if config.pure_nnx: - # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. - if config.enable_diloco: - # DiLoCoTrainState: persist the synchronized global model (outer params). - # The per-replica inner optimizer / outer-momentum state is not checkpointed. - step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step - state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) - else: - state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) + # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. + if config.enable_diloco: + # DiLoCoTrainState: persist the synchronized global model (outer params). + # The per-replica inner optimizer / outer-momentum state is not checkpointed. + step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step + state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) + else: + state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) # Determine if a checkpoint save should be forced, overriding the usual `config.checkpoint_period` logic. # This occurs if this function was called: diff --git a/src/maxtext/experimental/rl/grpo_trainer.py b/src/maxtext/experimental/rl/grpo_trainer.py index 9ea0a68f82..20089be819 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/inference/kvcache.py b/src/maxtext/inference/kvcache.py index ba2266060f..f1e8867bd6 100644 --- a/src/maxtext/inference/kvcache.py +++ b/src/maxtext/inference/kvcache.py @@ -26,9 +26,6 @@ from aqt.jax.v2.aqt_tensor import QTensor as KVTensor from aqt.jax.v2.flax import aqt_flax -from maxtext.layers import nnx_wrappers -from maxtext.layers.initializers import variable_to_logically_partitioned - from maxtext.common.common_types import Array, AxisNames, AxisIdxes, Config, CACHE_BATCH_PREFILL, DType, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, MODEL_MODE_AUTOREGRESSIVE, CACHE_HEADS_NONE, DECODING_ACTIVE_SEQUENCE_INDICATOR from maxtext.common.common_types import CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV, CACHE_SCALE_BATCH, CACHE_SCALE_SEQUENCE, CACHE_SCALE_HEADS, CACHE_SCALE_KV @@ -148,94 +145,6 @@ def einsum_fn_with_rhs_qtensor_and_dequant(self): ) -def kv_cache_as_linen( - *, - max_prefill_length: int, - max_target_length: int, - batch: int, - key_seq_len: int, - value_seq_len: int, - key_heads: int, - value_heads: int, - key_head_size: int, - value_head_size: int, - dtype: DType, - kv_quant: None | KVQuant = None, - prefill_cache_logical_axis_names: AxisNames = (CACHE_BATCH_PREFILL, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV), - cache_logical_axis_names: AxisNames = (CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV), - cache_scale_logical_axis_names: AxisNames = ( - CACHE_SCALE_BATCH, - CACHE_SCALE_SEQUENCE, - CACHE_SCALE_HEADS, - CACHE_SCALE_KV, - ), - prefill_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - ar_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - key_axis_order: AxisIdxes = (2, 0, 1, 3), - use_chunked_prefill: bool = False, - model_mode: str = MODEL_MODE_PREFILL, - is_gdn: bool = False, - conv_kernel_size: int = 0, - conv_dim: int = 0, - name: str | None = None, -): - """Initializes the KVCache module and returns it as a Linen module. - - Args: - max_prefill_length: The maximum prefill length. - max_target_length: The maximum target length. - batch: The batch size. - key_seq_len: The key sequence length. - value_seq_len: The value sequence length. - key_heads: The number of key heads. - value_heads: The number of value heads. - key_head_size: The key head size. - value_head_size: The value head size. - dtype: The data type. - kv_quant: The KVQuant configuration. - prefill_cache_logical_axis_names: The logical axis names for the prefill cache. - cache_logical_axis_names: The logical axis names for the cache. - cache_scale_logical_axis_names: The logical axis names for the cache scale. - prefill_cache_axis_order: The axis order for the prefill cache. - ar_cache_axis_order: The axis order for the autoregressive cache. - key_axis_order: The axis order for the key. - use_chunked_prefill: Whether to use chunked prefill. - model_mode: The model mode. - name: The name of the Linen module. - - Returns: - A Linen module that wraps the NNX `KVCache` module. - """ - return nnx_wrappers.to_linen( - KVCache, - max_prefill_length=max_prefill_length, - max_target_length=max_target_length, - batch=batch, - key_seq_len=key_seq_len, - value_seq_len=value_seq_len, - key_heads=key_heads, - value_heads=value_heads, - key_head_size=key_head_size, - value_head_size=value_head_size, - dtype=dtype, - kv_quant=kv_quant, - prefill_cache_logical_axis_names=prefill_cache_logical_axis_names, - cache_logical_axis_names=cache_logical_axis_names, - cache_scale_logical_axis_names=cache_scale_logical_axis_names, - prefill_cache_axis_order=prefill_cache_axis_order, - ar_cache_axis_order=ar_cache_axis_order, - key_axis_order=key_axis_order, - use_chunked_prefill=use_chunked_prefill, - model_mode=model_mode, - is_gdn=is_gdn, - conv_kernel_size=conv_kernel_size, - conv_dim=conv_dim, - metadata_fn=variable_to_logically_partitioned, - name=name, - abstract_init=False, - ) - - class BaseCache(nnx.Module): """Abstract base class for Caches.""" @@ -961,69 +870,49 @@ def __call__( raise ValueError(f"Model Mode isn't supported! {model_mode=}") -def mla_kv_cache_as_linen( - *, - max_prefill_length: int, - max_target_length: int, - batch: int, - key_seq_len: int, - value_seq_len: int, - key_head_size: int, - value_head_size: int, - dtype: DType, - key_heads: int = 1, - value_heads: int = 1, - kv_quant: None | KVQuant = None, - prefill_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - ar_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - use_chunked_prefill: bool = False, - model_mode: str = MODEL_MODE_PREFILL, - name: str | None = None, -): - """Initializes the MlaKVCache module and returns it as a Linen module. - - Args: - max_prefill_length: The maximum prefill length. - max_target_length: The maximum target length. - batch: The batch size. - key_seq_len: The key sequence length. - value_seq_len: The value sequence length. - key_head_size: The key head size. - value_head_size: The value head size. - dtype: The data type. - key_heads: The number of key heads. - value_heads: The number of value heads. - kv_quant: The KVQuant configuration. - prefill_cache_axis_order: The axis order for the prefill cache. - ar_cache_axis_order: The axis order for the autoregressive cache. - use_chunked_prefill: Whether to use chunked prefill. - model_mode: The model mode. - name: The name of the Linen module. - - Returns: - A Linen module that wraps the NNX `MlaKVCache` module. +class GatedDeltaNetCache(BaseCache): + """Cache for Linear Attention (Gated Delta Net). + + Stores the fixed-size recurrent state and the sliding window state for convolution. """ - return nnx_wrappers.to_linen( - MlaKVCache, - max_prefill_length=max_prefill_length, - max_target_length=max_target_length, - batch=batch, - key_seq_len=key_seq_len, - value_seq_len=value_seq_len, - key_head_size=key_head_size, - value_head_size=value_head_size, - dtype=dtype, - key_heads=key_heads, - value_heads=value_heads, - kv_quant=kv_quant, - prefill_cache_axis_order=prefill_cache_axis_order, - ar_cache_axis_order=ar_cache_axis_order, - use_chunked_prefill=use_chunked_prefill, - model_mode=model_mode, - metadata_fn=variable_to_logically_partitioned, - name=name, - abstract_init=False, - ) + + def __init__( + self, + batch: int, + num_heads: int, + k_head_dim: int, + v_head_dim: int, + conv_kernel_size: int, + conv_dim: int, + dtype: DType, + cache_batch_axis_name: str = CACHE_BATCH, + cache_heads_axis_name: str = CACHE_HEADS, + ): + super().__init__() + self.batch = batch + self.dtype = dtype + + # 1. Recurrent State (S) for the Delta Rule + # Shape: [Batch, Heads, K_Dim, V_Dim] + # We maintain the running state matrix. + self.recurrent_state = nnx.Cache( + jnp.zeros((int(batch), num_heads, k_head_dim, v_head_dim), dtype=dtype), + # Sharding: Batch, Heads, None (K), None (V) + out_sharding=(cache_batch_axis_name, cache_heads_axis_name, None, None), + ) + + # 2. Convolution State for the 1D Conv + # Shape: [Batch, Kernel_Size - 1, Conv_Dim] + # We store the last (K-1) inputs to perform the sliding window conv during decoding. + self.conv_state = nnx.Cache( + jnp.zeros((int(batch), conv_kernel_size - 1, conv_dim), dtype=dtype), + # Sharding: Batch, None (Time), None (Dim) + out_sharding=(cache_batch_axis_name, None, None), + ) + + def __call__(self): + """Returns the cache variables for the layer to use.""" + return self class MlaKVCache(KVCache): diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 5377e86ec3..92c3695aa5 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -31,18 +31,14 @@ else: from jax.experimental.layout import DeviceLocalLayout as DLL # type: ignore -from flax import linen as nn from flax import nnx from flax import struct from flax.linen import partitioning as nn_partitioning -import flax from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_PKG_DIR -from maxtext.models import models from maxtext.layers import quantizations from maxtext.inference import inference_utils -from maxtext.multimodal import processor as mm_processor from maxtext.utils import lora_utils from maxtext.utils import max_logging from maxtext.utils import max_utils @@ -117,40 +113,35 @@ def __init__(self, config: Any, devices: Any | None = None): # Model and Optimizer definition. quant = quantizations.configure_quantization(config) - if config.pure_nnx: - # `serve` only when the on-disk checkpoint already carries `qrhs.frozen` - # (no full-precision kernel). For `checkpoint_is_quantized=False` with - # quant enabled we stay in `train` mode and let AQT quantize per-forward - # against the full-precision kernel — same numerical result as `serve` - # for absmax calibration, just slower. - nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train" - # We need both PREFILL and AR abstract models because the cache vars inherit - # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and - # bulk_insert searches for the substring "cache_batch" in the AR-mode names. - # Calling nnx.eval_shape directly (instead of create_nnx_abstract_model) avoids - # the jax.set_mesh wrap that trips Flax 0.12.6 on logical-only axes like "norm". - _create_model = model_creation_utils.get_nnx_create_model_fn( - config, mesh=self._mesh, model_mode=MODEL_MODE_PREFILL, quant_mode_str=nnx_quant_mode_str - ) - _create_model_ar = model_creation_utils.get_nnx_create_model_fn( - config, mesh=self._mesh, model_mode=MODEL_MODE_AUTOREGRESSIVE, quant_mode_str=nnx_quant_mode_str - ) - self._nnx_quant_mode_str = nnx_quant_mode_str - with nn_partitioning.axis_rules(config.logical_axis_rules): - abstract_model = nnx.eval_shape(_create_model) - abstract_model_ar = nnx.eval_shape(_create_model_ar) - self.model = abstract_model - self.model_ar = abstract_model_ar - # 3-way split so JIT bodies can pass (params, cache, rest) separately to - # nnx.merge. `rest` (RNG state etc.) is materialized in load_params. - graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...) - self.graphdef = graphdef - self._create_model_fn = _create_model - self._nnx_rest_state = None - else: - self.model = models.transformer_as_linen(config, mesh=self._mesh, quant=quant, model_mode=MODEL_MODE_PREFILL) - self.graphdef = None - self._create_model_fn = None + # `serve` only when the on-disk checkpoint already carries `qrhs.frozen` + # (no full-precision kernel). For `checkpoint_is_quantized=False` with + # quant enabled we stay in `train` mode and let AQT quantize per-forward + # against the full-precision kernel — same numerical result as `serve` + # for absmax calibration, just slower. + nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train" + # We need both PREFILL and AR abstract models because the cache vars inherit + # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and + # bulk_insert searches for the substring "cache_batch" in the AR-mode names. + # Calling nnx.eval_shape directly (instead of create_nnx_abstract_model) avoids + # the jax.set_mesh wrap that trips Flax 0.12.6 on logical-only axes like "norm". + _create_model = model_creation_utils.get_nnx_create_model_fn( + config, mesh=self._mesh, model_mode=MODEL_MODE_PREFILL, quant_mode_str=nnx_quant_mode_str + ) + _create_model_ar = model_creation_utils.get_nnx_create_model_fn( + config, mesh=self._mesh, model_mode=MODEL_MODE_AUTOREGRESSIVE, quant_mode_str=nnx_quant_mode_str + ) + self._nnx_quant_mode_str = nnx_quant_mode_str + with nn_partitioning.axis_rules(config.logical_axis_rules): + abstract_model = nnx.eval_shape(_create_model) + abstract_model_ar = nnx.eval_shape(_create_model_ar) + self.model = abstract_model + self.model_ar = abstract_model_ar + # 3-way split so JIT bodies can pass (params, cache, rest) separately to + # nnx.merge. `rest` (RNG state etc.) is materialized in load_params. + graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...) + self.graphdef = graphdef + self._create_model_fn = _create_model + self._nnx_rest_state = None self.replicated_sharding = jax.sharding.NamedSharding(self._mesh, P(None)) self.abstract_params = None @@ -321,65 +312,7 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar if rng is None: rng = jax.random.PRNGKey(0) - if self.config.pure_nnx: - return self._load_params_nnx(params=params, rng=rng) - - if self.model.quant and self.config.checkpoint_is_quantized: - print("Loading from the quantized checkpoint...") - self.model.quant.quant_mode = quantizations.get_quant_mode("serve") - - rng1, rng2, rng3 = jax.random.split(rng, 3) - if params: - print("Resharding given params") - init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng) - _, self.state_mesh_annotations, state_mesh_shardings = maxtext_utils.get_abstract_state( - self.config, self._mesh, init_state_fn, False - ) - # reshard given params based on shardings from config in MaxEngine - params = jax.device_put(params, state_mesh_shardings.params) - state = maxtext_utils.init_decode_state(None, params) - state = max_utils.unbox_logicallypartioned(state) - else: - init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng1) - state, self.state_mesh_annotations = maxtext_utils.setup_decode_state(self.config, self._mesh, None, init_state_fn) - # pylint: disable=isinstance-second-argument-not-valid-type - self.abstract_params = jax.tree_util.tree_map( - lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding) - if isinstance(x, jax.Array) - else None, - state.params, - ) - - self.prefill_kv_cache_annotations = maxtext_utils.get_prefill_kv_cache_annotations( - self.model, self.config, rng2, self._mesh - ) - self.prefill_kv_cache_shardings = jax.tree_util.tree_map( - lambda x: jax.sharding.NamedSharding(self._mesh, x), - self.prefill_kv_cache_annotations, - ) - - if self.config.stack_prefill_result_cache: - # Add extra axis for the axis generated by the stack. - self.prefill_kv_cache_shardings = jax.tree_util.tree_map( - lambda x: jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec(None, *x.spec)), - self.prefill_kv_cache_shardings, - ) - self.prefill_kv_cache_shardings = self.prefill_kv_cache_shardings["decoder"]["layers_0"] - - self.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(self.model, self.config, rng2, self._mesh) - self.kv_cache_shardings = jax.tree_util.tree_map( - lambda x: jax.sharding.NamedSharding(self._mesh, x), - self.kv_cache_annotations, - ) - - if self.model.quant and not self.config.checkpoint_is_quantized: - params = self.quantize_params(state, rng3) - else: - params = state.params - - self.print_stats("After load_params") - - return params + return self._load_params_nnx(params=params, rng=rng) def _load_params_nnx(self, params, rng): """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings. @@ -530,122 +463,57 @@ def apply_adapter(self, base_params, adapter_config, adapter_params): lora_rank = int(adapter_config["r"]) lora_scale_factor = float(adapter_config["lora_alpha"]) / lora_rank - if self.config.pure_nnx: - lora_utils.apply_lora_on_base_params_nnx(base_params, adapter_params, lora_scale_factor) - else: - lora_utils.apply_lora_on_base_params(base_params, adapter_params, lora_scale_factor) + lora_utils.apply_lora_on_base_params_nnx(base_params, adapter_params, lora_scale_factor) def unapply_adapter(self, base_params, adapter_config, adapter_params): """Unapply the adapter params from the merged params to get back the base params.""" lora_rank = int(adapter_config["r"]) lora_scale_factor = float(adapter_config["lora_alpha"]) / lora_rank - if self.config.pure_nnx: - lora_utils.unapply_lora_from_base_params_nnx(base_params, adapter_params, lora_scale_factor) - else: - lora_utils.unapply_lora_from_base_params(base_params, adapter_params, lora_scale_factor) + lora_utils.unapply_lora_from_base_params_nnx(base_params, adapter_params, lora_scale_factor) def quantize_params(self, state, rng: PRNGKeyType | None = None): """Forward pass to quantize decode params.""" if rng is None: rng = jax.random.PRNGKey(0) - if self.config.pure_nnx: - # NNX takes a different code path: convert-on-load lives in `_load_params_nnx` - # via `_convert_and_quantize_nnx`, which runs the dummy forward against a - # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model. - # The standalone `quantize_params(state, rng)` API expects a Linen-shape - # `state.params` dict and isn't reachable on the NNX pathway in maxengine - # (load_params already dispatched to _load_params_nnx). - raise NotImplementedError( - "Use load_params() on NNX — the convert step runs inside _load_params_nnx via " - "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API." - ) - - self.model.quant.quant_mode = quantizations.get_quant_mode("convert") - - @jax.jit - def model_apply(_p, _rng): - image_shape = mm_processor.get_dummy_image_shape_for_init( - model_name=self.config.model_name, - batch_size=self.config.micro_batch_size_to_train_on, - ) - audio_shape = mm_processor.get_dummy_audio_shape_for_init(self.config) - return self.model.apply( - _p | {"aqt": {}}, - jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - encoder_images=jnp.ones(image_shape, dtype=jnp.float32) if self.config.use_multimodal else None, - # encoder_image_masks indicates valid tiles if image tiling + padding is used in vision encoder input. - encoder_image_masks=jnp.ones(image_shape[:2], dtype=jnp.int32) - if self.config.use_multimodal and "llama4" in self.config.model_name - else None, - encoder_audios=jnp.ones(audio_shape, dtype=jnp.float32) if self.config.use_audio else None, - decoder_segment_ids=jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - rngs={"params": _rng}, - mutable=True, - ) - - _, new_vars = model_apply(state.params, rng) - # Remove param values which have corresponding qtensors in aqt to save memory. - params = {} - params["aqt"] = new_vars["aqt"] - params["params"] = quantizations.remove_quantized_params(state.params["params"], new_vars["aqt"]) - self.abstract_params = jax.tree_util.tree_map( - lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding), - params, + # NNX takes a different code path: convert-on-load lives in `_load_params_nnx` + # via `_convert_and_quantize_nnx`, which runs the dummy forward against a + # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model. + # The standalone `quantize_params(state, rng)` API expects a Linen-shape + # `state.params` dict and isn't reachable on the NNX pathway in maxengine + # (load_params already dispatched to _load_params_nnx). + raise NotImplementedError( + "Use load_params() on NNX — the convert step runs inside _load_params_nnx via " + "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API." ) - maxtext_utils.save_quantized_checkpoint_if_configured(self.config, params) - self.model.quant.quant_mode = quantizations.get_quant_mode("serve") - return params def _maybe_stack_prefill_result_cache(self, cache): """Stack the caches across the layers.""" if not self.config.stack_prefill_result_cache: return cache - if self.config.pure_nnx: - if self.config.scan_layers: - # scan_layers already stacks the per-layer KV cache on axis 0; nothing to restack. - return cache - # scan_layers=False: stack the per-layer subtrees under decoder/layers into one - # subtree with a leading layer axis (matching the scan_layers=True shape). - layers = cache["decoder"]["layers"] - stacked = jax.tree.map(lambda *c: jnp.stack(c), *[layers[i] for i in range(self.config.num_decoder_layers)]) - return {"decoder": {"layers": stacked}} - - layer_keys = [] - for i in range(self.config.num_decoder_layers): - layer_keys.append(f"layers_{i}") - - layer_cache = [cache["decoder"][layer_key] for layer_key in layer_keys] - - return jax.tree.map(lambda *c: jnp.stack(c), *layer_cache) + if self.config.scan_layers: + # scan_layers already stacks the per-layer KV cache on axis 0; nothing to restack. + return cache + # scan_layers=False: stack the per-layer subtrees under decoder/layers into one + # subtree with a leading layer axis (matching the scan_layers=True shape). + layers = cache["decoder"]["layers"] + stacked = jax.tree.map(lambda *c: jnp.stack(c), *[layers[i] for i in range(self.config.num_decoder_layers)]) + return {"decoder": {"layers": stacked}} def _maybe_unstack_prefill_result_cache(self, cache): """Unstack the caches across the layers.""" if not self.config.stack_prefill_result_cache: return cache - if self.config.pure_nnx: - if self.config.scan_layers: - # Mirror _maybe_stack_prefill_result_cache: the cache already carries the - # layer axis, so there is nothing to unstack. - return cache - # scan_layers=False: split the leading layer axis back into per-layer subtrees. - stacked = cache["decoder"]["layers"] - layers = {i: jax.tree.map(lambda x, i=i: x[i], stacked) for i in range(self.config.num_decoder_layers)} - return {"decoder": {"layers": layers}} - - flat_cache, treedef = jax.tree.flatten(cache) - layer_cache = [jax.tree.unflatten(treedef, flat_cache_vars) for flat_cache_vars in zip(*flat_cache, strict=True)] - res_cache = {"decoder": {}} - - for i in range(self.config.num_decoder_layers): - res_cache["decoder"][f"layers_{i}"] = layer_cache[i] - - return res_cache + if self.config.scan_layers: + # Mirror _maybe_stack_prefill_result_cache: the cache already carries the + # layer axis, so there is nothing to unstack. + return cache + # scan_layers=False: split the leading layer axis back into per-layer subtrees. + stacked = cache["decoder"]["layers"] + layers = {i: jax.tree.map(lambda x, i=i: x[i], stacked) for i in range(self.config.num_decoder_layers)} + return {"decoder": {"layers": layers}} def prefill_aot( # pylint: disable=too-many-positional-arguments self, @@ -734,10 +602,7 @@ def _prefill_jit( if existing_prefix is not None: if not self.use_chunked_prefill: raise ValueError("Using chunked prefill is needed for existing_prefix.") - # NNX threads existing_prefix.cache via the nnx_cache local below; only - # the Linen path merges cache into input_params (params is a dict there). - if not self.config.pure_nnx: - input_params = params | {"cache": existing_prefix.cache} + # NNX threads existing_prefix.cache via the nnx_cache local below. start_position = existing_prefix.common_prefix_tokens.shape[0] # TODO(yuyanpeng): rename previous_chunk previous_chunk = jnp.expand_dims(existing_prefix.common_prefix_tokens, 0) @@ -769,50 +634,29 @@ def _prefill_jit( sequence_indicator = jnp.expand_dims(one_d_output, 0) rng, new_rng = jax.random.split(rng) - if self.config.pure_nnx: - # Prefill always operates on batch=1 (one padded prompt at a time). - nnx_cache = ( - existing_prefix.cache if existing_prefix is not None else self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL) + # Prefill always operates on batch=1 (one padded prompt at a time). + nnx_cache = ( + existing_prefix.cache if existing_prefix is not None else self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL) + ) + with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + flat_logits, new_cache_dict = self._nnx_run_model( + params=input_params, + cache_dict=nnx_cache, + decoder_input_tokens=input_tokens, + decoder_positions=positions, + decoder_segment_ids=sequence_indicator, + encoder_images=images, + encoder_image_masks=image_masks, + encoder_videos=videos, + encoder_video_masks=video_masks, + encoder_audios=audio_values, + enable_dropout=False, + model_mode=MODEL_MODE_PREFILL, + previous_chunk=previous_chunk, + true_length=true_length, + slot=slot, ) - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - flat_logits, new_cache_dict = self._nnx_run_model( - params=input_params, - cache_dict=nnx_cache, - decoder_input_tokens=input_tokens, - decoder_positions=positions, - decoder_segment_ids=sequence_indicator, - encoder_images=images, - encoder_image_masks=image_masks, - encoder_videos=videos, - encoder_video_masks=video_masks, - encoder_audios=audio_values, - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - previous_chunk=previous_chunk, - true_length=true_length, - slot=slot, - ) - new_vars = {"cache": new_cache_dict} - else: - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - flat_logits, new_vars = self.model.apply( - input_params, - input_tokens, - positions, - encoder_images=images, - encoder_image_masks=image_masks, - encoder_videos=videos, - encoder_video_masks=video_masks, - encoder_audios=audio_values, - decoder_segment_ids=sequence_indicator, - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - rngs={"params": new_rng}, - mutable=["cache"], - previous_chunk=previous_chunk, - true_length=true_length, - slot=slot, - ) + new_vars = {"cache": new_cache_dict} if return_prompt_logp: prompt_logp = inference_utils.prompt_logprobs_from_prefill(flat_logits, input_tokens, true_length) else: @@ -1027,32 +871,19 @@ def _prefill_multisampling_jit( sequence_indicator = jnp.expand_dims(one_d_output, 0) rng, new_rng = jax.random.split(rng) - if self.config.pure_nnx: - # Prefill is batch=1 (one prompt); multi-sampling only draws several first - # tokens from the shared logits below. Mirror the _prefill_jit NNX branch. - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - flat_logits, new_cache_dict = self._nnx_run_model( - params=params, - cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL), - decoder_input_tokens=input_tokens, - decoder_positions=positions, - decoder_segment_ids=sequence_indicator, - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - ) - new_vars = {"cache": new_cache_dict} - else: - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - flat_logits, new_vars = self.model.apply( - params, - input_tokens, - positions, - decoder_segment_ids=sequence_indicator, - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - rngs={"params": new_rng}, - mutable=["cache"], - ) + # Prefill is batch=1 (one prompt); multi-sampling only draws several first + # tokens from the shared logits below. + with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + flat_logits, new_cache_dict = self._nnx_run_model( + params=params, + cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL), + decoder_input_tokens=input_tokens, + decoder_positions=positions, + decoder_segment_ids=sequence_indicator, + enable_dropout=False, + model_mode=MODEL_MODE_PREFILL, + ) + new_vars = {"cache": new_cache_dict} next_pos = jnp.full((1, 1), true_length, dtype=jnp.int32) selected_logits = jax.lax.dynamic_slice( @@ -1163,33 +994,20 @@ def prefill_concat( input_tokens = jnp.expand_dims(padded_tokens, 0) # [BATCH, SEQUENCE] decoder_positions = jnp.expand_dims(decoder_positions, 0) decoder_segment_ids = jnp.expand_dims(decoder_segment_ids, 0) - rng, new_rng = jax.random.split(rng) - if self.config.pure_nnx: - # Packed prompts run as a single batch=1 prefill; the packed positions and - # segment ids keep the prompts separated. Mirror the _prefill_jit NNX branch. - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - flat_logits, new_cache_dict = self._nnx_run_model( - params=params, - cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL), - decoder_input_tokens=input_tokens, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - ) - new_vars = {"cache": new_cache_dict} - else: - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - flat_logits, new_vars = self.model.apply( - params, - input_tokens, - decoder_positions, - decoder_segment_ids=decoder_segment_ids, - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - rngs={"params": new_rng}, - mutable=["cache"], - ) + rng, _ = jax.random.split(rng) + # Packed prompts run as a single batch=1 prefill; the packed positions and + # segment ids keep the prompts separated. + with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + flat_logits, new_cache_dict = self._nnx_run_model( + params=params, + cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL), + decoder_input_tokens=input_tokens, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + enable_dropout=False, + model_mode=MODEL_MODE_PREFILL, + ) + new_vars = {"cache": new_cache_dict} cache = new_vars["cache"] cache = self._maybe_stack_prefill_result_cache(cache) if return_prompt_logp: @@ -1337,28 +1155,16 @@ def _generate_jit( previous_token = decode_state["tokens"] rng, new_rng = jax.random.split(rng) # run one step generation - if self.config.pure_nnx: - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - out_logits, new_cache_dict = self._nnx_run_model( - params=params, - cache_dict=decode_state["cache"], - decoder_input_tokens=previous_token, - decoder_positions=decode_state["next_pos"], - enable_dropout=False, - model_mode=MODEL_MODE_AUTOREGRESSIVE, - ) - new_vars = {"cache": new_cache_dict} - else: - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - out_logits, new_vars = self.model.apply( - params | {"cache": decode_state["cache"]}, - previous_token, - decode_state["next_pos"], - enable_dropout=False, - model_mode=MODEL_MODE_AUTOREGRESSIVE, - rngs={"params": new_rng}, - mutable=["cache"], - ) + with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + out_logits, new_cache_dict = self._nnx_run_model( + params=params, + cache_dict=decode_state["cache"], + decoder_input_tokens=previous_token, + decoder_positions=decode_state["next_pos"], + enable_dropout=False, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + ) + new_vars = {"cache": new_cache_dict} out_logits = jax.lax.with_sharding_constraint(out_logits, self.replicated_sharding) new_cache = jax.lax.with_sharding_constraint(new_vars["cache"], self.kv_cache_shardings) # sampling tokens @@ -1858,103 +1664,7 @@ def init_decode_state( **kwargs, # pylint: disable=unused-argument ) -> DecodeState: """Initialises any state which a generation step transforms.""" - if rng is None: - rng = jax.random.PRNGKey(0) - - if self.config.pure_nnx: - return self._init_decode_state_nnx() - - # pylint: disable=unused-argument - def init(abstract_params): - x = jnp.ones( - (int(self.config.per_device_batch_size * self.mesh.size), 1), - dtype=jnp.int32, - ) - dummy_image = jnp.ones( - mm_processor.get_dummy_image_shape_for_init( - model_name=self.config.model_name, batch_size=int(self.config.per_device_batch_size * self.mesh.size) - ), - dtype=jnp.int32, - ) - dummy_audio = jnp.ones( - mm_processor.get_dummy_audio_shape_for_init(self.config), - dtype=jnp.float32, - ) - _, cache = self.model.apply( - abstract_params, - x, - x, - encoder_images=dummy_image if self.config.use_multimodal else None, - encoder_audios=dummy_audio if self.config.use_audio else None, - enable_dropout=False, - model_mode=MODEL_MODE_AUTOREGRESSIVE, - rngs={"params": rng}, - mutable=["cache"], - slot=0, - ) - - next_pos = jnp.zeros( - (int(self.config.per_device_batch_size * self.mesh.size), 1), - dtype=jnp.int32, - ) - generated_tokens = jnp.zeros( - (int(self.config.per_device_batch_size * self.mesh.size), 1), - dtype=jnp.int32, - ) - tokens = jnp.zeros( - (int(self.config.per_device_batch_size * self.mesh.size), 1), - dtype=jnp.int32, - ) - token_logp = jnp.zeros( - (int(self.config.per_device_batch_size * self.mesh.size), 1), - dtype=jnp.float32, - ) - return { - "logits": jnp.zeros( - ( - int(self.config.per_device_batch_size * self.mesh.size), - 1, - self.config.vocab_size, - ) - ), - "cache": cache["cache"], - "next_pos": next_pos, - "generated_tokens": generated_tokens, - "tokens": tokens, - "token_logp": token_logp, - } - - with nn_partitioning.axis_rules(self.config.logical_axis_rules): - abstract_outputs = jax.eval_shape(init, self.abstract_params) - logical_annotations = nn.get_partition_spec(abstract_outputs) - - with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): - mesh_annotations = nn.logical_to_mesh(logical_annotations) - - shardings = jax.tree_util.tree_map( - lambda mesh_annotation: jax.sharding.NamedSharding(self._mesh, mesh_annotation), - mesh_annotations, - ) - - @functools.partial(jax.jit, out_shardings=shardings) - def initialize(): - return jax.tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), abstract_outputs) - - init_state = initialize() - cache = init_state["cache"] - - def is_lp(k): - return isinstance(k, flax.linen.spmd.LogicallyPartitioned) - - self.kv_cache_annotations_named = jax.tree_util.tree_map( - lambda x: tuple(x.logical_axes) - if hasattr(x, "logical_axes") - else (tuple(x.names) if hasattr(x, "names") else ()), - cache, - is_leaf=is_lp, - ) - zeroed = max_utils.unbox_logicallypartioned(init_state) - return zeroed + return self._init_decode_state_nnx() def _init_decode_state_nnx(self) -> DecodeState: """NNX equivalent of init_decode_state. Returns a decode_state dict with a pure-dict cache.""" @@ -2043,18 +1753,10 @@ def set_engine_vars_from_base_engine( """Set internal vars from base_engine, which has already loaded the checkpoint and has sharding, mesh, and kv cache related vars set. """ - if not engine.config.pure_nnx and base_engine.model.quant: - # NNX bakes the quant mode in at construction (via _nnx_quant_mode_str) rather - # than mutating model.quant.quant_mode, so there's nothing to copy on that path. - engine.model.quant.quant_mode = base_engine.model.quant.quant_mode engine.state_mesh_annotations = base_engine.state_mesh_annotations engine.abstract_params = base_engine.abstract_params - if engine.config.pure_nnx: - # Linen's get_kv_cache_annotations calls model.init(); NNX modules have no - # .init, so use the abstract-model variant (mirrors _load_params_nnx). - engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations_nnx(engine.model_ar, engine.config, engine.mesh) - else: - engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(engine.model, engine.config, rng, engine.mesh) # pylint: disable=protected-access + # NNX modules have no .init, so use the abstract-model variant (mirrors _load_params_nnx). + engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations_nnx(engine.model_ar, engine.config, engine.mesh) engine.kv_cache_shardings = jax.tree_util.tree_map( lambda x: jax.sharding.NamedSharding(engine.mesh, x), engine.kv_cache_annotations, # pylint: disable=protected-access diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py index 5908ab2e2a..9f9232c53e 100644 --- a/src/maxtext/inference/vllm_decode.py +++ b/src/maxtext/inference/vllm_decode.py @@ -98,8 +98,6 @@ def decode_with_vllm(config: Config) -> None: "debug_sharding": config.debug_sharding, "prefuse_moe_weights": config.prefuse_moe_weights, "scan_layers": config.scan_layers, - "enable_nnx": config.enable_nnx, - "pure_nnx_decoder": config.pure_nnx_decoder, }, "sharding": { "sharding_strategy": { @@ -201,9 +199,7 @@ def decode_with_tunix( overrides = config.vllm_hf_overrides if isinstance(overrides, str) and "MaxTextForCausalLM" in overrides: use_no_op_mappings = True - elif isinstance(overrides, dict) and "MaxTextForCausalLM" in overrides.get( - "architectures", [] - ): + elif isinstance(overrides, dict) and "MaxTextForCausalLM" in overrides.get("architectures", []): use_no_op_mappings = True tunix_model = TunixMaxTextAdapter( @@ -238,14 +234,8 @@ def decode_with_tunix( ) # MaxText uses -1 to mean "disabled"; vLLM requires top_p in (0, 1]. - top_p = ( - config.decode_sampling_nucleus_p - if config.decode_sampling_nucleus_p > 0 - else 1.0 - ) - top_k = ( - config.decode_sampling_top_k if config.decode_sampling_top_k > 0 else -1 - ) + top_p = config.decode_sampling_nucleus_p if config.decode_sampling_nucleus_p > 0 else 1.0 + top_k = config.decode_sampling_top_k if config.decode_sampling_top_k > 0 else -1 rollout_vllm_additional_config = { "maxtext_config": { @@ -255,8 +245,6 @@ def decode_with_tunix( "debug_sharding": config.debug_sharding, "prefuse_moe_weights": config.prefuse_moe_weights, "scan_layers": config.scan_layers, - "enable_nnx": config.enable_nnx, - "pure_nnx_decoder": config.pure_nnx_decoder, } } @@ -280,17 +268,9 @@ def decode_with_tunix( rollout_vllm_hbm_utilization=config.hbm_utilization_vllm, rollout_vllm_init_with_random_weights=True, rollout_vllm_tpu_backend_type="jax", - tensor_parallel_size=( - config.ici_tensor_parallelism - if config.ici_tensor_parallelism > 0 - else 1 - ), + tensor_parallel_size=(config.ici_tensor_parallelism if config.ici_tensor_parallelism > 0 else 1), data_parallel_size=jax.device_count() - // ( - config.ici_tensor_parallelism - if config.ici_tensor_parallelism > 0 - else 1 - ), + // (config.ici_tensor_parallelism if config.ici_tensor_parallelism > 0 else 1), rollout_vllm_additional_config=rollout_vllm_additional_config, rollout_vllm_kwargs={"hf_overrides": config.vllm_hf_overrides}, ) @@ -330,9 +310,7 @@ def main(argv: Sequence[str]) -> None: if FLAGS.use_tunix: maxtext_model, mesh = model_creation_utils.from_pretrained(config) if config.lora.enable_lora: - maxtext_model = lora_utils.apply_lora_to_model( - maxtext_model, mesh, config - ) + maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config) if config.lora.lora_restore_path: lora_utils.restore_lora_from_path(maxtext_model, config) decode_with_tunix(config, model=maxtext_model, mesh=mesh) diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 7174cf763a..b3254d0cb5 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -849,25 +849,22 @@ def maybe_quantize_model(model, config): if config.use_qwix_quantization and not config.use_batch_split_schedule: quantization_provider = get_qt_provider(config) if quantization_provider: - if config.pure_nnx: - input_shape = (config.micro_batch_size_to_train_on, config.max_target_length) - dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32) - dummy_positions = jnp.ones(input_shape, dtype=jnp.int32) - dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32) - model = qwix.quantize_model( - model, - quantization_provider, - dummy_tokens, - dummy_positions, - dummy_segment_ids, - enable_dropout=False, - ) - # Qwix quantization runs a forward pass during tracing, which sows transient nnx.Intermediate variables - # (e.g. max_logits from QK-Clip, MTP losses) into the model. Popping them here prevents structural mismatches - # between the initial setup GraphDef/state_mesh_shardings and the stripped states during train steps. - nnx.pop(model, nnx.Intermediate) - else: - model = qwix.quantize_model(model, quantization_provider) + input_shape = (config.micro_batch_size_to_train_on, config.max_target_length) + dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32) + dummy_positions = jnp.ones(input_shape, dtype=jnp.int32) + dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32) + model = qwix.quantize_model( + model, + quantization_provider, + dummy_tokens, + dummy_positions, + dummy_segment_ids, + enable_dropout=False, + ) + # Qwix quantization runs a forward pass during tracing, which sows transient nnx.Intermediate variables + # (e.g. max_logits from QK-Clip, MTP losses) into the model. Popping them here prevents structural mismatches + # between the initial setup GraphDef/state_mesh_shardings and the stripped states during train steps. + nnx.pop(model, nnx.Intermediate) return model diff --git a/src/maxtext/trainers/diloco/diloco.py b/src/maxtext/trainers/diloco/diloco.py index 1b7253c41d..9b0677e0e4 100644 --- a/src/maxtext/trainers/diloco/diloco.py +++ b/src/maxtext/trainers/diloco/diloco.py @@ -153,18 +153,11 @@ def add_diloco_dim(x): momentum=config.diloco_outer_momentum, nesterov=True, ) - # For NNX, model params (Param variables only) live under abstract_state.model; - # for Linen under abstract_state.params. - if config.pure_nnx: - _, model_params, _ = nnx.split(abstract_state.model, nnx.Param, ...) - model_params = model_params.to_pure_dict() - _, model_params_sharding, _ = nnx.split( - state_mesh_shardings.model, nnx.Param, ... - ) - model_params_sharding = model_params_sharding.to_pure_dict() - else: - model_params = abstract_state.params - model_params_sharding = state_mesh_shardings.params + # Model params (Param variables only) live under abstract_state.model. + _, model_params, _ = nnx.split(abstract_state.model, nnx.Param, ...) + model_params = model_params.to_pure_dict() + _, model_params_sharding, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) + model_params_sharding = model_params_sharding.to_pure_dict() outer_opt_state = jax.eval_shape(outer_optimizer.init, model_params) # Create abstract step @@ -217,17 +210,13 @@ def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]: # mesh automatically when jax.set_mesh is used. inner_state = drjax.broadcast(state, mesh=mesh) # Outer state retains a single copy of the model parameters and optimizer state. - # For NNX, model params (Param variables only) live under state.model; - # for Linen under state.params. - if config.pure_nnx: - _, outer_params, _ = nnx.split(state.model, nnx.Param, ...) - outer_params = outer_params.to_pure_dict() - else: - outer_params = state.params + # Model params (Param variables only) live under state.model. + _, outer_params, _ = nnx.split(state.model, nnx.Param, ...) + outer_params = outer_params.to_pure_dict() outer_opt_state = outer_optimizer.init(outer_params) outer_opt_state_sharding = jax.tree_util.tree_map(lambda x: x.sharding, outer_opt_state) - # For NNX, the step counter lives at state.optimizer.step; for Linen at state.step. - step = state.optimizer.step if config.pure_nnx else state.step + # The step counter lives at state.optimizer.step. + step = state.optimizer.step return ( DiLoCoTrainState(inner_state=inner_state, params=outer_params, outer_opt_state=outer_opt_state, step=step), outer_opt_state_sharding, @@ -264,68 +253,55 @@ def synchronize(state): # Calculate the delta between the current replica's state and the global # state (since last synchronization). broadcast_outer_params = drjax.broadcast(state.params, mesh=mesh) - # For NNX, model Param vars live under inner_state.model; for Linen under inner_state.params. - if config.pure_nnx: - _, inner_model_params, _ = nnx.split( - state.inner_state.model, nnx.Param, ... - ) - inner_model_params = inner_model_params.to_pure_dict() - else: - inner_model_params = state.inner_state.params + # Model Param vars live under inner_state.model. + _, inner_model_params, _ = nnx.split(state.inner_state.model, nnx.Param, ...) + inner_model_params = inner_model_params.to_pure_dict() model_delta = jax.tree.map(lambda x, y: y - x, inner_model_params, broadcast_outer_params) # Treat the average delta as the outer optimizer's gradient and apply to # the global (outer) model params. averaged_pseudo_grad = drjax.reduce_mean(model_delta) updates, new_opt_state = outer_optimizer.update(averaged_pseudo_grad, state.outer_opt_state, state.params) new_outer_params = optax.apply_updates(state.params, updates) + # Replace inner model params with the new global model params. # NOTE: inner optimizer state is retained despite the change in parameters, # see section 6.1 in https://arxiv.org/pdf/2311.08105. - if config.pure_nnx: - # For NNX: merge new Param vars back with the non-Param model vars (e.g. RNG state). - def replace_nnx_model_params(s, new_params): - s_model = s["model"] if hasattr(s, "keys") else s.model - s_opt = s["optimizer"] if hasattr(s, "keys") else s.optimizer - - graphdef, _, non_param_state = nnx.split(s_model, nnx.Param, ...) - new_model = nnx.merge(graphdef, new_params, non_param_state) - - if type(s_model).__name__ == "State": - new_model = nnx.state(new_model) - elif isinstance(s_model, dict): - new_model = nnx.to_pure_dict(new_model) - - if hasattr(s, "keys"): - # Replace "model" leaves by path, keeping s's treedef. Picking by position - # (leaves[N:]) breaks if a key sorts before "model"; reconstructing via - # type(s)({...}) breaks the lax.cond match — nnx.State recursive-wraps. - leaves_with_paths, treedef = jax.tree_util.tree_flatten_with_path(s) - new_model_iter = iter(jax.tree_util.tree_leaves(new_model)) - - def _is_model_leaf(path): - if not path: - return False - k = path[0] - return ( - getattr(k, "key", None) == "model" - or getattr(k, "name", None) == "model" - ) - - new_leaves = [ - next(new_model_iter) if _is_model_leaf(p) else leaf - for p, leaf in leaves_with_paths - ] - return jax.tree_util.tree_unflatten(treedef, new_leaves) - else: - return TrainStateNNX(new_model, s_opt) - - new_inner_state = drjax.map_fn( - lambda s: replace_nnx_model_params(s, new_outer_params), - state.inner_state, - mesh=mesh, - ) - else: - new_inner_state = drjax.map_fn(lambda s: s.replace(params=new_outer_params), state.inner_state, mesh=mesh) + # Merge new Param vars back with the non-Param model vars (e.g. RNG state). + def replace_nnx_model_params(s, new_params): + s_model = s["model"] if hasattr(s, "keys") else s.model + s_opt = s["optimizer"] if hasattr(s, "keys") else s.optimizer + + graphdef, _, non_param_state = nnx.split(s_model, nnx.Param, ...) + new_model = nnx.merge(graphdef, new_params, non_param_state) + + if type(s_model).__name__ == "State": + new_model = nnx.state(new_model) + elif isinstance(s_model, dict): + new_model = nnx.to_pure_dict(new_model) + + if hasattr(s, "keys"): + # Replace "model" leaves by path, keeping s's treedef. Picking by position + # (leaves[N:]) breaks if a key sorts before "model"; reconstructing via + # type(s)({...}) breaks the lax.cond match — nnx.State recursive-wraps. + leaves_with_paths, treedef = jax.tree_util.tree_flatten_with_path(s) + new_model_iter = iter(jax.tree_util.tree_leaves(new_model)) + + def _is_model_leaf(path): + if not path: + return False + k = path[0] + return getattr(k, "key", None) == "model" or getattr(k, "name", None) == "model" + + new_leaves = [next(new_model_iter) if _is_model_leaf(p) else leaf for p, leaf in leaves_with_paths] + return jax.tree_util.tree_unflatten(treedef, new_leaves) + else: + return TrainStateNNX(new_model, s_opt) + + new_inner_state = drjax.map_fn( + lambda s: replace_nnx_model_params(s, new_outer_params), + state.inner_state, + mesh=mesh, + ) return state.replace( params=new_outer_params, outer_opt_state=new_opt_state, @@ -343,8 +319,8 @@ def diloco_train_step(state, batch, prng): broadcast_rng = drjax.broadcast(prng, mesh=mesh) inner_state, metrics = drjax.map_fn(train_step, (state.inner_state, batch, broadcast_rng), mesh=mesh) avg_metrics = typed_reduce_mean(metrics) - # For NNX, the step counter lives at inner_state.optimizer.step; for Linen at inner_state.step. - new_step = inner_state.optimizer.step[0] if config.pure_nnx else inner_state.step[0] + # The step counter lives at inner_state.optimizer.step. + new_step = inner_state.optimizer.step[0] state = state.replace( inner_state=inner_state, step=new_step, diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index e7f5ea84b6..951ba2a90b 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -25,9 +25,6 @@ from absl import app - -import optax - import pathwaysutils # pylint: disable=unused-import import tensorflow as tf @@ -36,7 +33,7 @@ import jax.numpy as jnp from jax.sharding import NamedSharding -from flax import linen as nn, nnx +from flax import nnx from flax.linen import partitioning as nn_partitioning from flax.nnx import variablelib @@ -70,17 +67,14 @@ 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 VertexTensorboardManager, _vertex_tb_is_stub = vertex_tensorboard_modules() def get_first_step(model, state): - if isinstance(model, nn.Module): - return int(state.step) - if hasattr( - state, "inner_state" - ): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var + del model # NNX-only; kept for call-site signature parity + if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var return int(state.step.get_value()) return int(state.optimizer.step.get_value()) @@ -94,17 +88,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) # decimate proportion of data when per_device_batch_size<1 if is_train: for k, v in data.items(): @@ -112,140 +107,64 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr else: for k, v in data.items(): data[k] = v[: config.micro_batch_size_to_eval_on, :] - 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") - 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_images=data["images"] if config.use_multimodal else None, - encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, - enable_dropout=config.enable_dropout if is_train else False, - rngs={"dropout": rng1, "params": aqt_rng}, - mutable=mutable_collections, - decoder_target_tokens=data["targets"], - decoder_target_mask=data["targets_segmentation"], - ) - - 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: - 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, - ) - - # Mask out paddings at the end of each example. - 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_images=data["images"] if config.use_multimodal else None, + encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, + enable_dropout=config.enable_dropout if is_train else False, + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + ) + intermediates = nnx.pop(model, nnx.Intermediate) + intermediate_outputs = intermediates.to_pure_dict() + + # MTP sows mtp_losses/mtp_acceptance as custom Variable subclasses, not + # Intermediate, so the nnx.pop above misses them. Pop them here under their + # collection names so calculate_mtp_loss / calculate_mtp_acceptance_rate + # find them. Otherwise the MTP loss is silently zeroed. They are also + # excluded from the returned state below so they don't leak into + # out_shardings. + if config.mtp_num_layers > 0: + intermediate_outputs["mtp_losses"] = nnx.pop(model, mtp_losses).to_pure_dict() + intermediate_outputs["mtp_acceptance"] = nnx.pop(model, mtp_acceptance).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_images=data["images"] if config.use_multimodal else None, - encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, - enable_dropout=config.enable_dropout if is_train else False, - decoder_target_tokens=data["targets"], - decoder_target_mask=data["targets_segmentation"], + 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, ) - intermediates = nnx.pop(model, nnx.Intermediate) - intermediate_outputs = intermediates.to_pure_dict() - - # MTP sows mtp_losses/mtp_acceptance as custom Variable subclasses, not - # Intermediate, so the nnx.pop above misses them. Pop them here under their - # collection names so calculate_mtp_loss / calculate_mtp_acceptance_rate - # find them. Otherwise the MTP loss is silently zeroed. They are also - # excluded from the returned state below so they don't leak into - # out_shardings. - if config.mtp_num_layers > 0: - intermediate_outputs["mtp_losses"] = nnx.pop(model, mtp_losses).to_pure_dict() - intermediate_outputs["mtp_acceptance"] = nnx.pop(model, mtp_acceptance).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: - 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, - ) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + # Mask out paddings at the end of each example. + 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) total_weights = jnp.sum(data["targets_segmentation"] != 0) # If gradient accumulation is enabled, we don't need to divide xent_sum @@ -297,25 +216,21 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # get MoE routed bias term updates 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 update is a 2-D matrix that's transposed at the apply site below. - moe_bias_updates = next( - ( - val - for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs) - if tuple(k.key for k in path if hasattr(k, "key"))[-1:] == ("moe_bias_updates",) - ), - None, - ) - if moe_bias_updates is not None: - # The Linen path returns the sow tuple and indexes [0] downstream; tree_leaves - # already descended that tuple, so wrap it back so the apply site is uniform. - moe_bias_updates = (moe_bias_updates,) + # NNX intermediates are model-rooted (no "intermediates" prefix), so match by + # suffix. Unlike collect_intermediates_by_suffix we must not ravel: the update + # is a 2-D matrix that's transposed at the apply site below. + moe_bias_updates = next( + ( + val + for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs) + if tuple(k.key for k in path if hasattr(k, "key"))[-1:] == ("moe_bias_updates",) + ), + None, + ) + if moe_bias_updates is not None: + # The apply site indexes [0] downstream; tree_leaves already descended the sow + # tuple, so wrap it back so the apply site is uniform. + moe_bias_updates = (moe_bias_updates,) # Add the model's primary output to the intermediates dict so it can be used # by the acceptance rate calculation in eval_step. @@ -336,28 +251,24 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr 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. """ - # --- 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 + del dropout_rng # unused for NNX (kept for jit signature parity) + 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: @@ -371,61 +282,39 @@ 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, - params_shardings, - ) - sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - pure_params = params["params"] if sparsity_enabled else params - 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) + model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, nnx.Param, 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) - model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, nnx.Param, 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), - ) - 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: - curr_params = jax.tree.map( - functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode), - curr_params, - params_shardings, - ) - 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) - _, _, _, new_rest = nnx.split(local_model, nnx.Param, custom_param_filter, ...) - return loss, (aux, new_rest) - - grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True) - (loss, (aux, new_rest)), (raw_grads, custom_grads) = grad_func(curr_params, custom_params, rest, config, data) - nnx.update(state.model, nnx.State.merge(custom_grads, new_rest)) + 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: + curr_params = jax.tree.map( + functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode), + curr_params, + params_shardings, + ) + 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) + _, _, _, new_rest = nnx.split(local_model, nnx.Param, custom_param_filter, ...) + return loss, (aux, new_rest) + + grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True) + (loss, (aux, new_rest)), (raw_grads, custom_grads) = grad_func(curr_params, custom_params, rest, config, data) + nnx.update(state.model, nnx.State.merge(custom_grads, new_rest)) raw_grads = jax.tree_util.tree_map( lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x, @@ -448,97 +337,35 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): mtp_loss = aux.get("mtp_loss", 0.0) new_opt_state = None - 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) - else: - state.apply_gradients(grads) - new_state = state + 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 - if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: - target_bias = new_state.model.decoder.moe_layers.DeepSeekMoeBlock_0.MoeBlock_0.gate.bias - target_bias.value = target_bias.value + jnp.array(moe_bias_updates[0]).transpose() + # 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_bias = new_state.model.decoder.moe_layers.DeepSeekMoeBlock_0.MoeBlock_0.gate.bias + target_bias.value = target_bias.value + jnp.array(moe_bias_updates[0]).transpose() lm_loss = xent_sum / (total_weights + EPS) scalar_metrics = { @@ -552,10 +379,7 @@ def move(path, value): "learning/total_weights": total_weights, } 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: @@ -564,21 +388,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 = { @@ -588,8 +405,6 @@ def move(path, value): if config.record_internal_nn_metrics: 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. @@ -598,16 +413,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: @@ -646,10 +454,8 @@ 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_eval_step = jax_device_state["p_eval_step"] - model = jax_device_state["model"] # Unpack python_vars step = python_vars["step"] @@ -666,8 +472,6 @@ def training_loop_iteration( # Unpack immutable_data config = immutable_data["config"] # for helpers logical_axis_rules = immutable_data["logical_axis_rules"] - 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"] @@ -685,16 +489,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): - if shard_optimizer_over_data and isinstance(model, nn.Module): - state = sharding.maybe_shard_with_name(state, state_mesh_shardings, shard_mode) state, metrics = p_train_step(state, example_batch, *step_rng_args) step_time_delta = datetime.datetime.now() - last_step_completion @@ -771,24 +573,18 @@ def train_loop(config, recorder, state=None): start_step = get_first_step(model, state) # this is the start_step for training train_utils.validate_completed_steps(start_step, config.steps) - 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: jit_model, state = nnx.split(state) - 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 else: - params_shardings, state_mesh_shardings = ( - sharding.maybe_update_params_sharding_with_opt( - config, state_mesh_shardings - ) - ) + params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings) p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( config, @@ -805,13 +601,11 @@ def train_loop(config, recorder, state=None): with jax.set_mesh(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) @@ -825,12 +619,8 @@ 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: - setup_params = ( - state.params - ) # DiLoCoTrainState.params: the outer (global) params + if config.enable_diloco: + setup_params = state.params # DiLoCoTrainState.params: the outer (global) params else: _, setup_params, _ = nnx.split(state.model, nnx.Param, ...) metric_logger_instance.write_setup_info_to_tensorboard(setup_params) diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py index d472bb68d4..b72071db39 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,63 +129,44 @@ 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() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - def create_train_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - 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) return ( @@ -387,20 +367,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/src/maxtext/utils/generate_param_only_checkpoint.py b/src/maxtext/utils/generate_param_only_checkpoint.py index 8a2488abe1..2c932cb9ea 100644 --- a/src/maxtext/utils/generate_param_only_checkpoint.py +++ b/src/maxtext/utils/generate_param_only_checkpoint.py @@ -22,7 +22,6 @@ The output "parameter state" is output to the checkpoint directory. Additionally it is cast down to bf16. """ -import functools import os.path from typing import Sequence @@ -35,11 +34,8 @@ from jax.sharding import Mesh from maxtext.configs import pyconfig from maxtext.common import checkpointing -from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN -from maxtext.layers import quantizations +from maxtext.common.common_types import DecoderBlockType from maxtext.common import train_state_nnx -from maxtext.models import models -from maxtext.optimizers import optimizers from maxtext.utils import gcs_utils from maxtext.utils import lora_utils from maxtext.utils import max_logging @@ -54,47 +50,7 @@ def _possibly_unroll_params(config, training_state, training_state_annotations, """Unroll scanned input layers when force_unroll is set.""" if not config.scan_layers or not config.force_unroll: return - if config.pure_nnx: - _possibly_unroll_params_nnx(config, training_state, training_state_annotations, mesh) - return - - def unroll_layer_group(num_layers, layer_name="layers"): - """Helper function to unroll layers (e.g. dense or MoE) into individual layers.""" - layers = training_state.params["params"]["decoder"].get(layer_name, None) - layers_annotations = training_state_annotations.params["params"]["decoder"].get(layer_name, None) - - if layers is None or layers_annotations is None: - raise ValueError(f"Missing {layer_name} in training_state or training_state_annotations.") - - def new_pspec(x): - return jax.sharding.PartitionSpec(*(x[0 : config.param_scan_axis] + x[config.param_scan_axis + 1 :])) - - new_layer_annotation = jax.tree_util.tree_map(new_pspec, layers_annotations) - new_layer_sharding = jax.tree_util.tree_map(lambda x: jax.sharding.NamedSharding(mesh, x), new_layer_annotation) - - for i in range(num_layers): - - def slice_ith(input_layers): - return jax.tree_util.tree_map(lambda x: jax.numpy.take(x, i, axis=config.param_scan_axis), input_layers) - - # pylint: disable=not-callable - new_layer = jax.jit(slice_ith, out_shardings=new_layer_sharding)(layers) - - training_state.params["params"]["decoder"][f"{layer_name}_{i}"] = new_layer - training_state_annotations.params["params"]["decoder"][f"{layer_name}_{i}"] = new_layer_annotation - - # Remove the original layer collection - del training_state.params["params"]["decoder"][layer_name] - del training_state_annotations.params["params"]["decoder"][layer_name] - - jax.tree_util.tree_map(lambda x: x.delete(), layers) - - if config.decoder_block == DecoderBlockType.DEEPSEEK: - # Unroll dense and MoE layers separately - unroll_layer_group(config.first_num_dense_layers, layer_name="dense_layers") - unroll_layer_group(config.num_decoder_layers - config.first_num_dense_layers, layer_name="moe_layers") - else: - unroll_layer_group(config.num_decoder_layers, layer_name="layers") + _possibly_unroll_params_nnx(config, training_state, training_state_annotations, mesh) def _possibly_unroll_params_nnx(config, state, state_mesh_shardings, mesh): @@ -149,91 +105,34 @@ def slice_ith(input_layers): def _read_train_checkpoint(config, checkpoint_manager, mesh): """Read training checkpoint at path defined by load_full_state_path.""" rng = random.PRNGKey(0) - if config.pure_nnx: - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=rng) - model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) - _, tx = train_utils.create_training_optimizer(config, model) - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) - - def init_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=rng) + model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs) + _, tx = train_utils.create_training_optimizer(config, model) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) - else: - quant = quantizations.configure_quantization(config) - model = models.transformer_as_linen(config, mesh, quant, MODEL_MODE_TRAIN) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - tx = optimizers.get_optimizer(config, learning_rate_schedule) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng) + def init_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - state, state_mesh_notations, state_mesh_shardings, _ = maxtext_utils.setup_training_state( + state, _, state_mesh_shardings, _ = maxtext_utils.setup_training_state( None, config, mesh, checkpoint_manager, init_state_fn ) - if config.pure_nnx: - # On NNX, state is a flat nnx.State; params live under state.model and the - # legacy notations are unused (callers receive shardings directly). - num_params = max_utils.calculate_num_params_from_pytree(state.model) - max_logging.log(f"In input checkpoint Number of model params={num_params/1e9:.3f} billion") - return state, state_mesh_shardings - num_params = max_utils.calculate_num_params_from_pytree(state.params) + # On NNX, state is a flat nnx.State; params live under state.model and the + # legacy notations are unused (callers receive shardings directly). + num_params = max_utils.calculate_num_params_from_pytree(state.model) max_logging.log(f"In input checkpoint Number of model params={num_params/1e9:.3f} billion") - return state, state_mesh_notations + return state, state_mesh_shardings def _generate_lora_decode_checkpoints(config, mesh): """Read lora checkpoints checkpoint at path defined by load_full_state_path.""" - if config.pure_nnx: - _generate_lora_decode_checkpoints_nnx(config, mesh) - return - quant = quantizations.configure_quantization(config) - model = models.transformer_as_linen(config, mesh, quant, MODEL_MODE_TRAIN) - rng = random.PRNGKey(0) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - tx = optimizers.get_optimizer(config, learning_rate_schedule) - - lora_adapters = gcs_utils.gcs_list_directories(config.lora_input_adapters_path) - for lora_id in lora_adapters: - # Expected lora_checkpoint_dir = /loras/ - lora_checkpoint_dir = os.path.join(config.checkpoint_dir, "loras", lora_id, "") - - lora_adapter_path = os.path.join(config.lora_input_adapters_path, lora_id, "") - - # Create a checkpoint manager to save decode checkpoint at lora_checkpoint_dir - checkpoint_manager = checkpointing.create_orbax_checkpoint_manager( - lora_checkpoint_dir, - config.enable_checkpointing, - config.async_checkpointing, - config.checkpoint_period, - use_ocdbt=config.checkpoint_storage_use_ocdbt, - use_zarr3=config.checkpoint_storage_use_zarr3, - ) - - lora_config, lora_state, lora_state_annotations = lora_utils.setup_initial_lora_state( - model, None, tx, config, rng, mesh, checkpoint_manager, lora_adapter_path - ) - - _possibly_unroll_params(config, lora_state, lora_state_annotations, mesh) - - gcs_utils.write_dict_to_gcs_json(lora_config, os.path.join(lora_checkpoint_dir, "adapter_config.json")) - - # Save decode state to config's checkpoint directory at step 0 - _save_decode_checkpoint(config, lora_state, checkpoint_manager) - max_logging.log(f"Successfully saved LoRA checkpoint at: {os.path.join(lora_checkpoint_dir, '0', 'items')}") + _generate_lora_decode_checkpoints_nnx(config, mesh) def _save_decode_checkpoint(config, state, checkpoint_manager): """Generate checkpoint for decode from the training_state.""" - if config.pure_nnx: - _save_decode_checkpoint_nnx(config, state, checkpoint_manager) - return - decode_state = maxtext_utils.init_decode_state( - None, jax.tree_util.tree_map(lambda x: x.astype(jax.numpy.bfloat16), state.params) - ) - if checkpoint_manager is not None: - if checkpointing.save_checkpoint(checkpoint_manager, 0, decode_state): - max_logging.log(f"saved an decode checkpoint at {config.checkpoint_dir}") - checkpoint_manager.wait_until_finished() + _save_decode_checkpoint_nnx(config, state, checkpoint_manager) def _save_decode_checkpoint_nnx(config, state, checkpoint_manager): @@ -398,15 +297,8 @@ def generate_decode_checkpoint(config): # Read training state from config.load_paramaters_path max_logging.log(f"Read training checkpoint from: {config.load_full_state_path}") training_state, training_state_annotations = _read_train_checkpoint(config, checkpoint_manager, mesh) - if config.pure_nnx: - # NNX state is a flat nnx.State; opt_state lives under the optimizer sub-state. - assert ( - training_state.optimizer.opt_state - ), "missing opt_state in training checkpoint" - else: - assert ( - training_state.opt_state != {} - ), "missing opt_state in training checkpoint" + # NNX state is a flat nnx.State; opt_state lives under the optimizer sub-state. + assert training_state.optimizer.opt_state, "missing opt_state in training checkpoint" _possibly_unroll_params(config, training_state, training_state_annotations, mesh) diff --git a/src/maxtext/utils/layerwise_quantization.py b/src/maxtext/utils/layerwise_quantization.py index 3279b59a6f..e8d4f161b9 100644 --- a/src/maxtext/utils/layerwise_quantization.py +++ b/src/maxtext/utils/layerwise_quantization.py @@ -30,131 +30,25 @@ """ -import functools import os from typing import Any, Sequence from absl import app -from aqt.jax.v2 import aqt_tensor from flax import nnx from flax.linen import partitioning as nn_partitioning import jax import jax.numpy as jnp from maxtext.common import common_types -from maxtext.common import checkpointing from maxtext.layers import quantizations -from maxtext.models import deepseek, models from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx from maxtext.utils import model_creation_utils import orbax.checkpoint as ocp -from tqdm import tqdm from maxtext.configs import pyconfig -IGNORE = ocp.PLACEHOLDER PRNGKeyType = Any -DictKey = jax.tree_util.DictKey - - -def get_original_path_key(aqt_k_tuple: tuple[DictKey, ...]) -> tuple[DictKey, ...] | None: - """ - Maps an AQT PyTree path (tuple of keys) to its corresponding original parameter path. - Only returns a path if it corresponds to a parameter to be removed. - """ - aqt_k = list(aqt_k_tuple) - str_path = jax.tree_util.keystr(aqt_k_tuple) - if "AqtEinsum_" in str_path: - return None - if "AqtDotGeneral_" not in str_path: - return None - aqt_module_index = -1 - for i, key in enumerate(aqt_k): - if isinstance(key, DictKey) and key.key.startswith("AqtDotGeneral_"): - aqt_module_index = i - break - if aqt_module_index == -1: - return None - if ( - len(aqt_k) > aqt_module_index + 2 - and isinstance(aqt_k[aqt_module_index + 1], DictKey) - and aqt_k[aqt_module_index + 1].key == "qrhs" - and isinstance(aqt_k[aqt_module_index + 2], DictKey) - and aqt_k[aqt_module_index + 2].key == "frozen" - ): - parent_path = tuple(aqt_k[:aqt_module_index]) - return parent_path + (DictKey("kernel"),) - return None - - -def get_quantized_param_paths(aqt_params: Any, params: Any) -> set[tuple[DictKey, ...]]: - """ - Identifies the set of paths in the original params tree that have been quantized. - """ - - def is_qtensor(x): - return isinstance(x, aqt_tensor.QTensor) - - aqt_param_flat, _ = jax.tree_util.tree_flatten_with_path(aqt_params, is_leaf=is_qtensor) - if not aqt_param_flat: - return set() - param_tree_flat_with_path, _ = jax.tree_util.tree_flatten_with_path(params) - params_path_set: set[tuple[DictKey, ...]] = {tuple(k) for k, _ in param_tree_flat_with_path} - original_param_paths_to_remove: set[tuple[DictKey, ...]] = set() - for aqt_k_tuple, _ in aqt_param_flat: - original_k_tuple = get_original_path_key(aqt_k_tuple) - if original_k_tuple is None: - continue - if original_k_tuple in params_path_set: - original_param_paths_to_remove.add(original_k_tuple) - continue - params_keys_str = {jax.tree_util.keystr(k) for k in params_path_set} - raise ValueError( - f"Mapped AQT path {jax.tree_util.keystr(aqt_k_tuple)} to {jax.tree_util.keystr(original_k_tuple)}," - f" but not found in params. Available: {params_keys_str}" - ) - return original_param_paths_to_remove - - -def remove_quantized_params(params: Any, aqt_vars: Any) -> Any: - """Replaces the values in the original params tree that are now quantized with empty dicts.""" - quantized_param_path_set = get_quantized_param_paths(aqt_vars, params) - if not quantized_param_path_set: - return params - - def _map_fn(path, value): - return {} if tuple(path) in quantized_param_path_set else value - - return jax.tree_util.tree_map_with_path(_map_fn, params) - - -# --- Function to restructure NNX-Run AQT tree to match PURE LINEN saved format --- -def insert_deepseekmoeblock_scope(aqt_layer_tree: dict[str, Any]) -> dict[str, Any]: - """ - Moves top-level AqtEinsum_* entries into the existing 'DeepSeekMoeBlock_0' - dict to match the pure Linen AQT structure. - """ - if not isinstance(aqt_layer_tree, dict): - return aqt_layer_tree - - new_tree = dict(aqt_layer_tree) # Start with a copy - - einsum_items = {key: new_tree.pop(key) for key in list(new_tree.keys()) if key.startswith("AqtEinsum_")} - - if einsum_items: - if "DeepSeekMoeBlock_0" not in new_tree: - # This case indicates the MoE block itself was missing, which is unexpected - max_logging.log("Error: 'DeepSeekMoeBlock_0' not found in AQT vars for MoE layer.") - new_tree["DeepSeekMoeBlock_0"] = {} - elif not isinstance(new_tree["DeepSeekMoeBlock_0"], dict): - max_logging.log(f"Error: 'DeepSeekMoeBlock_0' is not a dict, type: {type(new_tree['DeepSeekMoeBlock_0'])}") - new_tree["DeepSeekMoeBlock_0"] = {} - - # Merge einsum_items into the DeepSeekMoeBlock_0 dict - new_tree["DeepSeekMoeBlock_0"].update(einsum_items) - - return new_tree class LayerwiseQuantization: @@ -166,123 +60,18 @@ def __init__(self, config: Any, rng: PRNGKeyType): self.config = config self.rng = rng - # The Linen path runs layer-by-layer (memory-efficient for big DeepSeek - # models) and is DeepSeek-specific because it relies on the per-layer - # `DeepSeek*ToLinen` wrappers. The NNX path runs whole-model convert - # forward and is model-agnostic — see `_load_and_quantize_nnx`. - if not config.pure_nnx: - assert config.decoder_block == common_types.DecoderBlockType.DEEPSEEK, ( - f"Linen layerwise quantization only supports {common_types.DecoderBlockType.DEEPSEEK}, " - f"got {config.decoder_block}." - ) # Mesh definition devices_array = maxtext_utils.create_device_mesh(config=config) self._mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) self.quant = quantizations.configure_quantization(config) - if config.pure_nnx: - # NNX takes a separate code path that builds the model via from_pretrained; - # no Linen abstract-state bookkeeping is needed here. - self.unboxed_abstract_state = None - return - model = models.transformer_as_linen( - config, mesh=self._mesh, quant=self.quant, model_mode=common_types.MODEL_MODE_TRAIN - ) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, self.config, False, self.rng) - - self.unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(self.config, self._mesh, init_state_fn, False) + self.unboxed_abstract_state = None def load_and_quantize(self) -> None: """ Load parameters layer by layer and quantize them. """ - if self.config.pure_nnx: - self._load_and_quantize_nnx() - return - quantized_params = {} - quantized_params["params"] = {"decoder": {}} - quantized_params["aqt"] = {"decoder": {}} - config = self.config - self.quant.quant_mode = quantizations.get_quant_mode("convert") - model_mode = common_types.MODEL_MODE_PREFILL - _, rng_quant_params = jax.random.split(self.rng) - - layers = [ - deepseek.DeepSeekDenseLayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - deepseek.DeepSeekMoELayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - ] - layer_prefixes = [ - "dense_layers", - "moe_layers", - ] - num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers - num_layers_list = [ - config.first_num_dense_layers, - num_moe_layers, - ] - - def model_apply(_p, _rng, layer): - return layer.apply( - _p | {"aqt": {}}, - jnp.ones((1, self.config.max_prefill_predict_length, self.config.base_emb_dim), dtype=jnp.int32), - None, - jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - True, - model_mode=model_mode, - rngs={"params": _rng}, - mutable=True, - ) - - for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes): - for index in tqdm(range(num_layers)): - layer_name = f"{layer_prefix}_{index}" - max_logging.log(f"Processing layer: {layer_name}") - - params = self._load_layer(layer_name) - params["params"] = params["params"]["decoder"][layer_name] - - _, new_vars = model_apply(params, rng_quant_params, layer) - - if "aqt" not in new_vars: - max_logging.log( - f"Warning: 'aqt' not found in new_vars for {layer_name}. Skipping AQT processing for this layer." - ) - quantized_params["params"]["decoder"][layer_name] = params["params"] # Keep original params - continue - - aqt_vars = new_vars["aqt"] - - try: - removed_params = remove_quantized_params(params["params"], aqt_vars) - quantized_params["params"]["decoder"][layer_name] = removed_params - except Exception as e: - max_logging.log(f"ERROR: Failed to remove quantized params for {layer_name}: {e}") - max_logging.log(f"Dumping params['params'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path( - lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), params["params"] - ) - max_logging.log(f"Dumping new_vars['aqt'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path(lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), aqt_vars) - raise - - # Restructure the aqt_vars for this layer to match pure Linen format for saving - if layer_prefix == "moe_layers": - structured_aqt = insert_deepseekmoeblock_scope(aqt_vars) - else: - structured_aqt = aqt_vars - quantized_params["aqt"]["decoder"][layer_name] = structured_aqt - - unquantized_layers = ["decoder_norm", "logits_dense"] - for unquantized_layer in unquantized_layers: - params = self._load_layer(unquantized_layer) - quantized_params["params"]["decoder"][unquantized_layer] = params["params"]["decoder"][unquantized_layer] - quantized_params["params"]["token_embedder"] = self._load_layer("token_embedder")["params"]["token_embedder"] - - maxtext_utils.save_quantized_checkpoint_if_configured(self.config, quantized_params) + self._load_and_quantize_nnx() def _load_and_quantize_nnx(self) -> None: """Whole-model NNX convert: load full-precision via TRAIN-mode `from_pretrained`, @@ -410,44 +199,6 @@ def _strip_kernels_at_quantized_paths(state_dict): out[k] = LayerwiseQuantization._strip_kernels_at_quantized_paths(v) if isinstance(v, dict) else v return out - def _load_layer(self, layer_name): - """Loads a specific layer's parameters from the checkpoint.""" - - config = self.config - with nn_partitioning.axis_rules(config.logical_axis_rules): - - params = checkpointing.load_params_from_path( - config.load_parameters_path, - self._create_partial_abstract_params(self.unboxed_abstract_state.params, layer_name), - config.checkpoint_storage_concurrent_gb, - config.checkpoint_storage_use_ocdbt, - config.checkpoint_storage_use_zarr3, - ) - return params - - def _create_partial_abstract_params(self, abstract_unboxed_params, layer): - """Creates a partial abstract params structure using ocp.PLACEHOLDER.""" - - def _should_keep(path, _): - # True if the layer name is part of the path - return any(isinstance(key, jax.tree_util.DictKey) and key.key == layer for key in path) - - def _map_fn(path, value): - if not _should_keep(path, value): - return IGNORE - if isinstance(value, jax.ShapeDtypeStruct): - zeros_array = jnp.zeros(value.shape, value.dtype) - if value.sharding is not None: - try: - return jax.device_put(zeros_array, value.sharding) - except Exception as e: # pylint: disable=broad-except - max_logging.log(f"Error applying sharding for path {path}: {e}") - return zeros_array - return zeros_array - return value - - return jax.tree_util.tree_map_with_path(_map_fn, abstract_unboxed_params) - def main(argv: Sequence[str]) -> None: jax.config.update("jax_default_prng_impl", "unsafe_rbg") diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 1aa443f093..384a0bd8ed 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/src/maxtext/utils/lora_utils.py @@ -15,7 +15,6 @@ """Common LoRA utils needed to support LoRA adapters.""" from collections.abc import Mapping -from functools import partial import json import os import re @@ -122,9 +121,8 @@ def unapply_lora_recursively(base_params, lora_params, module_name): def load_adapter(config, base_abstract_state_params, adapter_config_path, adapter_weights_path): """Load a LoRA adapter from disk and return its parameters. - When `config.pure_nnx` is True, `base_abstract_state_params` is the NNX - abstract param state (no outer `params` wrapper) and the returned - `lora_params` follows the same shape. Otherwise both use the Linen tree. + `base_abstract_state_params` is the NNX abstract param state (no outer + `params` wrapper) and the returned `lora_params` follows the same shape. Args: config: Top-level MaxText config. @@ -152,10 +150,7 @@ def load_adapter(config, base_abstract_state_params, adapter_config_path, adapte if not gcs_utils.gcs_path_exists(f"{adapter_weights_path}/commit_success.txt"): raise FileNotFoundError(f"Failed to read lora_weights from {adapter_weights_path}.") - if config.pure_nnx: - lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config) - else: - lora_state, _ = get_lora_abstract_state(base_abstract_state_params, lora_config) + lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config) with nn_partitioning.axis_rules(config.logical_axis_rules): lora_params = checkpointing.load_params_from_path( @@ -199,32 +194,26 @@ def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, checkp if lora_adapter_path: max_logging.log(f"Setting initial state of LoRA with lora_adapter_path = {lora_adapter_path}") - if config.pure_nnx: - # pylint: disable=import-outside-toplevel - from maxtext.common import train_state_nnx - from maxtext.utils import model_creation_utils + # pylint: disable=import-outside-toplevel + from maxtext.common import train_state_nnx + from maxtext.utils import model_creation_utils - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) - def create_train_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + def create_train_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - init_state_fn = create_train_state_fn - else: - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, rng) + init_state_fn = create_train_state_fn unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True) lora_config_path = lora_adapter_path + "adapter_config.json" lora_config = gcs_utils.read_json_from_gcs(lora_config_path) - if config.pure_nnx: - base_abstract_params = _nnx_param_subtree(unboxed_abstract_state) - lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config) - else: - lora_state, lora_state_annotations = get_lora_abstract_state(unboxed_abstract_state.params, lora_config) + base_abstract_params = _nnx_param_subtree(unboxed_abstract_state) + lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config) lora_weights_path = f"{lora_adapter_path}/0/items" diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 47c2c7e587..d306a3a5f3 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -96,10 +96,7 @@ def get_functional_train_with_signature( """Get the shardings (both state and data) for `train_step`.""" functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings) functional_train.__name__ = "train_step" - if config.pure_nnx: - in_shardings = (state_mesh_shardings, data_sharding) # State, batch - else: - in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng + in_shardings = (state_mesh_shardings, data_sharding) # State, batch out_shardings = (state_mesh_shardings, None) # State, metrics static_argnums = () # We partial out the static argnums of model and config donate_argnums = 0 # This is the index of the state - we allow the compiler to make use of this memory. @@ -110,10 +107,7 @@ def get_functional_eval_with_signature(eval_step, data_sharding, state_mesh_shar """Get the shardings (both state and data) for `eval_step`.""" functional_eval = functools.partial(eval_step, model, config) functional_eval.__name__ = "eval_step" - if config.pure_nnx: - in_shardings = (state_mesh_shardings, data_sharding) # State, batch (NNX: no rng) - else: - in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng + in_shardings = (state_mesh_shardings, data_sharding) # State, batch out_shardings = None # metrics static_argnums = () # We partial out the static argnums of model, config donate_argnums = () # state will be kept instead of being donated in eval_step @@ -213,11 +207,7 @@ def get_train_input_output_trees(func, input_args, input_kwargs): serialized_compiled = load_serialized_compiled(config.compiled_trainstep_file) shaped_batch = get_shaped_batch(config) - if config.pure_nnx: - shaped_input_args = (state, shaped_batch) - else: - example_rng = jax.random.PRNGKey(0) - shaped_input_args = (state, shaped_batch, example_rng) + shaped_input_args = (state, shaped_batch) shaped_input_kwargs = {} in_tree, out_tree = get_train_input_output_trees(partial_train, shaped_input_args, shaped_input_kwargs) p_train_step = deserialize_and_load(serialized_compiled, in_tree, out_tree, execution_devices=execution_devices) @@ -1507,47 +1497,20 @@ def setup_initial_state( # The update of data_iterator state happens in place, no need to assign explicitly state = restored["items"] - # For NNX, convert the pure dict to nnx.State using the abstract state as template - if config.pure_nnx: - nnx.replace_by_pure_dict(unboxed_abstract_state, state) - state = unboxed_abstract_state + # Convert the restored pure dict to nnx.State using the abstract state as template. + nnx.replace_by_pure_dict(unboxed_abstract_state, state) + state = unboxed_abstract_state else: init_state_partial = init_state_fn init_state_partial.__name__ = "initialize_state" - if config.pure_nnx: - state = jax.jit( - lambda: nnx.state(init_state_partial()), # Get state only, mapping to out_sharding structure - in_shardings=None, - out_shardings=state_mesh_shardings, - )() - else: - # pylint: disable=not-callable - state = jax.jit( - init_state_partial, - in_shardings=None, - out_shardings=state_mesh_shardings, - )() + state = jax.jit( + lambda: nnx.state(init_state_partial()), # Get state only, mapping to out_sharding structure + in_shardings=None, + out_shardings=state_mesh_shardings, + )() if raw_params: # If we loaded a partial state, we need to merge it. - if config.pure_nnx: - # raw_params should have the same sharding info as in the model - nnx.update(state.model, raw_params) - else: - sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - if sparsity_enabled: - # Sparsity-init keeps freshly initialized params for any leaf still - # represented as an abstract ShapeDtypeStruct in raw_params (i.e. not - # actually restored), and uses the restored value otherwise. - def _merge_params(p_raw, p_init): - if isinstance(p_raw, jax.ShapeDtypeStruct): - return p_init - return p_raw - - merged_params = jax.tree_util.tree_map(_merge_params, raw_params, state.params) - state = state.replace(params=merged_params) - else: - state = state.replace(params=raw_params) - if not config.pure_nnx: - state = max_utils.unbox_logicallypartioned(state) + # raw_params should have the same sharding info as in the model + nnx.update(state.model, raw_params) return state, state_mesh_annotations, state_mesh_shardings, data_iterator @@ -1561,51 +1524,8 @@ def get_logical_annotations(config, mesh, init_state_fn): def get_abstract_state(config, mesh, init_state_fn, is_training=True): - """Get a shaped abstraction of the state (including optimizer)""" - if config.pure_nnx: - return get_abstract_state_nnx(config, mesh, init_state_fn, is_training) - - init_state_partial = init_state_fn - - with nn_partitioning.axis_rules(config.logical_axis_rules): - abstract_state = jax.eval_shape(init_state_partial) - - state_logical_annotations = nn.get_partition_spec(abstract_state) - - state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, mesh, config.logical_axis_rules) - if is_training and config.shard_optimizer_over_data: - # Add data to sharding for optimizer state - state_mesh_shardings = state_mesh_shardings.replace( - opt_state=jax.tree.map_with_path( - functools.partial(sharding.add_data_to_sharding, mesh), - max_utils.unbox_logicallypartioned(abstract_state).opt_state, - state_mesh_shardings.opt_state, - ) - ) - if is_training and config.optimizer_memory_host_offload: - opt_state = jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="pinned_host"), state_mesh_shardings.opt_state) - state_mesh_shardings = state_mesh_shardings.replace(opt_state=opt_state) - if is_training and config.parameter_memory_host_offload: - assert config.param_scan_axis == 0, "You must set the scan axis 0 to enable parameter offloading." - - def move(path, x): - max_logging.log(f"max_utils.py: Moving {path} to host") - return x.with_memory_kind(kind="pinned_host") - - params = jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params) - state_mesh_shardings = state_mesh_shardings.replace(params=params) - - abstract_sharded_state = jax.jit(init_state_partial, in_shardings=None, out_shardings=state_mesh_shardings).eval_shape() - - unboxed_abstract_sharded_state = max_utils.unbox_logicallypartioned(abstract_sharded_state) - # Initialization - with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): - state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations) - return ( - unboxed_abstract_sharded_state, - state_mesh_annotations, - state_mesh_shardings, - ) + """Get a shaped abstraction of the state (including optimizer).""" + return get_abstract_state_nnx(config, mesh, init_state_fn, is_training) def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=True): diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index c0248783eb..d3d2b72c8b 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -865,14 +865,11 @@ def from_pretrained( ) config = pyconfig.HyperParameters(new_config) - if config.pure_nnx: - _create_model, abstract_model = create_nnx_abstract_model( - config, mesh, devices, model_mode, rng_key, quant_mode_str=quant_mode_str - ) - model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, _create_model, mesh=mesh) - # TODO: print debug_sharding info - else: - model = create_nnx_sharded_model_hybrid(config, mesh, devices, model_mode, rng_key) + _create_model, abstract_model = create_nnx_abstract_model( + config, mesh, devices, model_mode, rng_key, quant_mode_str=quant_mode_str + ) + model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, _create_model, mesh=mesh) + # TODO: print debug_sharding info # Compute logical-axis specs for downstream checkpoint alignment. # The model-creation helpers above resolve specs internally for sharding, but diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index 1c92d7815c..f456ac973b 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -33,8 +33,6 @@ import jax from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_PKG_DIR -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import maxtext_utils, model_creation_utils from optax.contrib._muon import MuonDimensionNumbers as mdn @@ -106,26 +104,17 @@ def get_transform_tree(tree, path=()): def get_muon_weight_dimension_numbers(model, config, verbose=False): """Extract muon dimension number from model structure.""" - if isinstance(model, nnx.Module): - _, abstract_param, _ = nnx.split(model, nnx.Param, ...) + _, abstract_param, _ = nnx.split(model, nnx.Param, ...) - def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf): - # Convert jax.tree_util.KeyEntry path to Tuple[str, ...] - path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey)) - return transform_logic(path_strings) + def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf): + # Convert jax.tree_util.KeyEntry path to Tuple[str, ...] + path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey)) + return transform_logic(path_strings) - # NNX abstract_param is an nnx.State (not Linen's dict of LogicallyPartitioned leaves); - # tree_map_with_path round-trips that structure so each Param.value holds the mdn result. - muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path( - apply_transform_nnx, nnx.to_pure_dict(abstract_param) - ) - muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers) - - else: # Linen - # quickly get param structure without materialization - abstract_param = maxtext_utils.get_abstract_param(model, config) - # get muon dimension number from param - muon_weight_dimension_numbers = get_transform_tree(abstract_param) + # tree_map_with_path handles NNX's PyTree structure; result is an nnx.State with the + # same structure, where each Param's value holds the mdn result. + muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path(apply_transform_nnx, nnx.to_pure_dict(abstract_param)) + muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers) if verbose: _print_structure_debug(abstract_param, muon_weight_dimension_numbers) @@ -157,7 +146,7 @@ def get_leaf_info(leaf): print("\nIs this reasonable?") -def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False): +def get_model_mdn(model_name, scan_layers=True, verbose=False): """Initializes a model and retrieves its Muon dimension numbers. This function sets up the configuration for a given model, initializes the @@ -181,22 +170,15 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False): f"model_name={model_name}", f"scan_layers={scan_layers}", "attention=dot_product", - f"pure_nnx={pure_nnx}", ] config = pyconfig.initialize(argv) # Setup model devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - quant = quantizations.configure_quantization(config) - if pure_nnx: - _, model = model_creation_utils.create_nnx_abstract_model(config, mesh) - else: - model = models.transformer_as_linen(config, mesh=mesh, quant=quant) + _, model = model_creation_utils.create_nnx_abstract_model(config, mesh) # Get dimension number muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config, verbose=verbose) - if pure_nnx: - muon_weight_dimension_numbers = {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)} - return muon_weight_dimension_numbers + return {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)} if __name__ == "__main__": @@ -205,4 +187,4 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False): sys.exit(1) model_name_arg = sys.argv[1] scan_layers_arg = sys.argv[2].lower() == "true" - get_model_mdn(model_name_arg, scan_layers_arg, verbose=True, pure_nnx=False) + get_model_mdn(model_name_arg, scan_layers_arg, verbose=True) diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index dc2f70ae14..5d8aa15a78 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -176,9 +176,7 @@ def remove_size_one_mesh_axis(spec, mesh): return P(*new_spec, unreduced=spec.unreduced, reduced=spec.reduced) -def get_nnx_var_named_sharding_with_scan_axis( - v: nnx.Variable, mesh -) -> nnx.Variable: +def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Variable: """Compute NamedSharding for an NNX variable, correctly handling the scan axis.""" val = v.get_value() if not hasattr(val, "shape"): @@ -190,11 +188,7 @@ def get_nnx_var_named_sharding_with_scan_axis( return v.replace(jax.tree.map(lambda _: replicated, val)) return v metadata = v.get_metadata() - out_sharding = ( - metadata.get("out_sharding") - or metadata.get("sharding_names") - or metadata.get("sharding") - ) + out_sharding = metadata.get("out_sharding") or metadata.get("sharding_names") or metadata.get("sharding") if not out_sharding: pspec = P() else: @@ -202,11 +196,7 @@ def get_nnx_var_named_sharding_with_scan_axis( if nnx.PARTITION_NAME in metadata: partition_name = metadata[nnx.PARTITION_NAME] scan_axis = metadata.get("param_scan_axis", 0) - out_sharding = ( - [out_sharding] - if isinstance(out_sharding, str) - else list(out_sharding) - ) + out_sharding = [out_sharding] if isinstance(out_sharding, str) else list(out_sharding) if partition_name not in out_sharding: out_sharding.insert(scan_axis, partition_name) out_sharding = tuple(out_sharding) @@ -215,9 +205,7 @@ def get_nnx_var_named_sharding_with_scan_axis( local_rules = metadata.get("sharding_rules", ()) if context_rules or local_rules: local_rules_list = list(local_rules) if local_rules is not None else [] - context_rules_list = ( - list(context_rules) if context_rules is not None else [] - ) + context_rules_list = list(context_rules) if context_rules is not None else [] rules = local_rules_list + context_rules_list pspec = logical_to_mesh_axes(out_sharding, mesh, rules=rules) else: @@ -565,24 +553,7 @@ def maybe_update_params_sharding_with_opt(config, state_mesh_shardings): - updated_state_mesh_shardings: State mesh shardings with updated params field (unchanged if shard_optimizer_over_data is False) """ - if config.pure_nnx: - return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings) - prev_params_shardings = state_mesh_shardings.params - if config.shard_optimizer_over_data: - if isinstance(state_mesh_shardings.opt_state, optax.ScaleByAdamState): - sharded_fp32_params = state_mesh_shardings.opt_state.mu - elif isinstance(state_mesh_shardings.opt_state, tuple) and isinstance( - state_mesh_shardings.opt_state[0], optax.ScaleByAdamState - ): - sharded_fp32_params = state_mesh_shardings.opt_state[0].mu - else: - raise NotImplementedError(f"Could not find optimizer state shardings from {type(state_mesh_shardings.opt_state)}") - if "params" not in sharded_fp32_params.keys(): - # When quantization=fp8 is enabled the sharded_fp32_params - # are not wrapped in `params`. Here we wrap them back. - sharded_fp32_params = {"params": sharded_fp32_params} - state_mesh_shardings = state_mesh_shardings.replace(params=dict(prev_params_shardings, **sharded_fp32_params)) - return prev_params_shardings, state_mesh_shardings + return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings) def maybe_update_params_sharding_with_opt_nnx( @@ -701,9 +672,7 @@ def _update_model_var(path, var): return prev_params_shardings, updated_state -def build_zero1_input_state_mesh_shardings( - config, state_mesh_shardings, params_shardings -): +def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_shardings): """Build the train-step input shardings under shard_optimizer_over_data (Zero-1). Model params on input use the original pre-Zero-1 sharding (params_shardings), @@ -716,8 +685,6 @@ def build_zero1_input_state_mesh_shardings( """ if not config.shard_optimizer_over_data: return state_mesh_shardings - if not config.pure_nnx: - return state_mesh_shardings.replace(params=params_shardings) # nnx.State has no .replace: shallow-copy via tree_map (preserves nested container # types) and overlay params_shardings under input_state.model. input_state = jax.tree_util.tree_map( diff --git a/src/maxtext/utils/standalone_checkpointer.py b/src/maxtext/utils/standalone_checkpointer.py index 6b1aa264c2..d6be1a533c 100644 --- a/src/maxtext/utils/standalone_checkpointer.py +++ b/src/maxtext/utils/standalone_checkpointer.py @@ -19,7 +19,6 @@ # See github.com/google/maxtext/issues/20 for more import datetime -from functools import partial import os from typing import Sequence @@ -52,23 +51,17 @@ def checkpoint_loop(config, state=None): on the configured cadence. Works on both Linen and NNX state shapes. """ init_rng = jax.random.PRNGKey(config.init_weights_seed) - if config.pure_nnx: - mesh = maxtext_utils.get_mesh_from_config(config) - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) - model = from_config(config, mesh=mesh, rngs=rngs) - _, tx = train_utils.create_training_optimizer(config, model) - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) - - def init_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - - else: - model = from_config(config) - mesh = model.mesh - _, tx = train_utils.create_training_optimizer(config, model) - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, init_rng) + mesh = maxtext_utils.get_mesh_from_config(config) + rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng) + model = from_config(config, mesh=mesh, rngs=rngs) + _, tx = train_utils.create_training_optimizer(config, model) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) + + def init_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn) unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, is_training=True) diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 73e5d06b05..aa9ff98ecd 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -19,7 +19,6 @@ import jax import functools import orbax.checkpoint.pathways as ocp_pathways -from functools import partial from flax import nnx from flax.linen import partitioning as nn_partitioning @@ -226,27 +225,20 @@ def setup_train_loop(config, recorder, devices=None): from maxtext.input_pipeline.input_pipeline_interface import create_data_iterator with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT): - is_training = True init_rng = jax.random.PRNGKey(config.init_weights_seed) mesh = maxtext_utils.get_mesh_from_config(config, devices) context_parallel_size = mesh.shape.get(config.context_sharding, 1) - if config.pure_nnx: - # Create abstract NNX model. - _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) - else: - model = model_creation_utils.from_config(config, devices) + # Create abstract NNX model. + _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) learning_rate_schedule, tx = create_training_optimizer(config, model) - if config.pure_nnx: - # For NNX, the train state is wrapped in the TrainStateNNX module. - def create_train_state_fn(): - model = _create_model_partial() - optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(model, optimizer) + # The train state is wrapped in the TrainStateNNX module. + def create_train_state_fn(): + model = _create_model_partial() + optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(model, optimizer) - init_state_fn = create_train_state_fn - else: - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, is_training, init_rng) + init_state_fn = create_train_state_fn checkpoint_manager = create_checkpoint_manager(config, mesh, init_state_fn) if checkpoint_manager is not None: checkpoint_step = checkpoint_manager.latest_step() @@ -302,42 +294,30 @@ def create_train_state_fn(): state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state( data_iterator, config, mesh, checkpoint_manager, init_state_fn ) - if config.pure_nnx: - with nn_partitioning.axis_rules(config.logical_axis_rules): - # We only need the graphdef here; it's merged with state below. Avoid - # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable - # under jax.set_mesh(mesh) and rejects any logical name missing from - # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes - # without a mesh skips sharding resolution, so it avoids the crash. - state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn)) - _, state_params, _ = nnx.split(state.model, nnx.Param, ...) - _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) - else: - state_params = state.params - state_mesh_shardings_params = state_mesh_shardings.params + with nn_partitioning.axis_rules(config.logical_axis_rules): + # We only need the graphdef here; it's merged with state below. Avoid + # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable + # under jax.set_mesh(mesh) and rejects any logical name missing from + # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes + # without a mesh skips sharding resolution, so it avoids the crash. + state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn)) + _, state_params, _ = nnx.split(state.model, nnx.Param, ...) + _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) if config.enable_diloco: with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh) # create state_mesh_shardings for the DilocoState - step_mesh = ( - state_mesh_shardings.optimizer.step.mesh - if config.pure_nnx - else state_mesh_shardings.step.mesh - ) + step_mesh = state_mesh_shardings.optimizer.step.mesh inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings) state_mesh_shardings = diloco.DiLoCoTrainState( inner_state_shardings, # Match the outer params' pure-dict structure (build_diloco_state stores # outer_params via to_pure_dict), so the sharding tree matches the state tree. - state_mesh_shardings_params.to_pure_dict() - if config.pure_nnx - else state_mesh_shardings_params, + state_mesh_shardings_params.to_pure_dict(), outer_opt_state_sharding, - jax.sharding.NamedSharding( - mesh=step_mesh, spec=jax.sharding.PartitionSpec() - ), + jax.sharding.NamedSharding(mesh=step_mesh, spec=jax.sharding.PartitionSpec()), ) # TODO(aireenmei, hengtaoguo): support sharding in vit for multimodal @@ -347,28 +327,21 @@ def create_train_state_fn(): # print weights sharding info under debug sharding mode if config.debug_sharding: - if config.pure_nnx: - # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we - # probably already lost the logical partition info at this moment. - logical_annotations_params = None - else: - logical_annotations = maxtext_utils.get_logical_annotations(config, mesh, init_state_fn) - logical_annotations_params = logical_annotations.params + # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we + # probably already lost the logical partition info at this moment. + logical_annotations_params = None max_utils.print_non_trivial_mesh_axis(model.mesh) maxtext_utils.print_shardings_params(state_params, state_mesh_shardings_params, mesh, logical_annotations_params) - if config.pure_nnx: - if config.enable_diloco: - # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner - # train step needs that graphdef as jit_model; the wrapper passes through as state. - train_state = state - model = state_graphdef - else: - train_state = nnx.merge(state_graphdef, state) - model = train_state.model - else: + if config.enable_diloco: + # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner + # train step needs that graphdef as jit_model; the wrapper passes through as state. train_state = state + model = state_graphdef + else: + train_state = nnx.merge(state_graphdef, state) + model = train_state.model return ( init_rng, From ab5140dbf397c3b1d2b1d8310c9f59164a1ac896 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Thu, 4 Jun 2026 18:43:53 +0000 Subject: [PATCH 2/4] [NNX] Delete Linen (2/4): remove the Linen decoder stack and dead *_as_linen wrappers Delete TransformerLinenPure, the Linen Decoder/DecoderLayer/SequentialBlockDecoderLayers stack (decoders.py), and the dead *_as_linen ToLinen wrappers across the layer/model files. The wrapped NNX classes are unchanged; transformer_as_linen (the NNX->Linen bridge) is kept for the checkpoint-conversion tools. --- src/maxtext/layers/attention_mla.py | 138 +- src/maxtext/layers/attention_op.py | 95 - src/maxtext/layers/attentions.py | 10 - src/maxtext/layers/decoders.py | 1748 ------------------ src/maxtext/layers/embeddings.py | 348 +--- src/maxtext/layers/encoders.py | 34 - src/maxtext/layers/multi_token_prediction.py | 37 - src/maxtext/layers/nnx_decoders.py | 24 +- src/maxtext/models/gemma3.py | 37 - src/maxtext/models/gemma4_vision.py | 15 - src/maxtext/models/llama4.py | 22 - src/maxtext/models/models.py | 359 +--- src/maxtext/models/qwen3.py | 49 +- 13 files changed, 45 insertions(+), 2871 deletions(-) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index d5f79a1a4b..92f6e90f70 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -63,9 +63,8 @@ DEFAULT_MASK_VALUE, ) -from maxtext.layers import nnx_wrappers from maxtext.layers.attentions import Attention -from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned +from maxtext.layers.initializers import nd_dense_init, NdInitializer from maxtext.layers.linears import DenseGeneral from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant @@ -388,141 +387,6 @@ def __call__( return indexer_mask, topk_indices, indexer_score -def mla_as_linen( - *, - config: Config, - num_query_heads: int, - num_kv_heads: int, - head_dim: int, - max_target_length: int, - mesh: Mesh, - attention_kernel: str, - inputs_q_shape: Tuple, - inputs_kv_shape: Tuple, - dtype: DType = jnp.float32, - weight_dtype: DType = jnp.float32, - max_prefill_predict_length: int = -1, - dropout_rate: float = 0.0, - kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal"), - float32_qk_product: bool = False, # computes logits in float32 for stability. - float32_logits: bool = False, # cast logits in float32 for stability. - quant: Optional[Quant] = None, - kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.MLA, # Default to MLA attention - attn_logits_soft_cap: float | None = None, - sliding_window_size: int | None = None, - use_ragged_attention: bool = False, - ragged_block_size: int = 256, - use_qk_norm: bool = False, - query_pre_attn_scalar: float | None = None, - use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections - # Temperature tuning parameters used for Llama4 - temperature_tuning: bool = False, - temperature_tuning_scale: float = 0.1, - temperature_tuning_floor_scale: float = 8192.0, - # Shard the query activation as the same as the key and value. - # TODO: Find a better sharding axis name. - # TODO: Further break down the Training and Inference axes for the q, k, v. - prefill_query_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), - prefill_key_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), - prefill_value_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), - query_axis_names: AxisNames = (KV_BATCH, LENGTH, KV_HEAD, KV_HEAD_DIM), - key_axis_names: AxisNames = (KV_BATCH, LENGTH, KV_HEAD, KV_HEAD_DIM), - value_axis_names: AxisNames = (KV_BATCH, LENGTH, KV_HEAD, KV_HEAD_DIM), - input_axis_names: AxisNames = (BATCH_ATTN, LENGTH, EMBED), - out_axis_names: AxisNames = (BATCH_ATTN, LENGTH, HEAD, D_KV), - prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, EMBED), - decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, EMBED), - prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV), - decode_out_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV), - prefill_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - ar_cache_axis_order: AxisIdxes = (1, 2, 0, 3), - compute_axis_order: AxisIdxes = (0, 1, 2, 3), - reshape_q: bool = False, - is_nope_layer: bool = False, - is_vision: bool = False, - model_mode: str = MODEL_MODE_TRAIN, - q_lora_rank: int = 0, - kv_lora_rank: int = 512, - qk_nope_head_dim: int = 128, - qk_rope_head_dim: int = 64, - v_head_dim: int = 128, - max_position_embeddings: int = 4096 * 4, - original_max_position_embeddings: int = 4096, - mscale: float = 1.0, # scaling factor for softmax - rope_factor: float = 40.0, # rotary embedding factor - name: str | None = None, -): - """A factory function to create an MLA as a Linen module. - - This function serves as a bridge to use the NNX-based `MLA` within a - Linen model. - """ - return nnx_wrappers.to_linen( - MLA, - config=config, - num_query_heads=num_query_heads, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - max_target_length=max_target_length, - mesh=mesh, - attention_kernel=attention_kernel, - inputs_q_shape=inputs_q_shape, - inputs_kv_shape=inputs_kv_shape, - dtype=dtype, - weight_dtype=weight_dtype, - max_prefill_predict_length=max_prefill_predict_length, - dropout_rate=dropout_rate, - kernel_init=kernel_init, - float32_qk_product=float32_qk_product, - float32_logits=float32_logits, - quant=quant, - kv_quant=kv_quant, - attention_type=attention_type, - attn_logits_soft_cap=attn_logits_soft_cap, - sliding_window_size=sliding_window_size, - use_ragged_attention=use_ragged_attention, - ragged_block_size=ragged_block_size, - use_qk_norm=use_qk_norm, - query_pre_attn_scalar=query_pre_attn_scalar, - use_bias_in_projections=use_bias_in_projections, - temperature_tuning=temperature_tuning, - temperature_tuning_scale=temperature_tuning_scale, - temperature_tuning_floor_scale=temperature_tuning_floor_scale, - prefill_query_axis_names=prefill_query_axis_names, - prefill_key_axis_names=prefill_key_axis_names, - prefill_value_axis_names=prefill_value_axis_names, - query_axis_names=query_axis_names, - key_axis_names=key_axis_names, - value_axis_names=value_axis_names, - input_axis_names=input_axis_names, - out_axis_names=out_axis_names, - prefill_input_axis_names=prefill_input_axis_names, - decode_input_axis_names=decode_input_axis_names, - prefill_out_axis_names=prefill_out_axis_names, - decode_out_axis_names=decode_out_axis_names, - prefill_cache_axis_order=prefill_cache_axis_order, - ar_cache_axis_order=ar_cache_axis_order, - compute_axis_order=compute_axis_order, - reshape_q=reshape_q, - is_nope_layer=is_nope_layer, - is_vision=is_vision, - model_mode=model_mode, - q_lora_rank=q_lora_rank, - kv_lora_rank=kv_lora_rank, - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=qk_rope_head_dim, - v_head_dim=v_head_dim, - max_position_embeddings=max_position_embeddings, - original_max_position_embeddings=original_max_position_embeddings, - mscale=mscale, - rope_factor=rope_factor, - name=name, - metadata_fn=variable_to_logically_partitioned, - abstract_init=False, - ) - - class MLA(Attention): """Multi-Head Latent Attention (MLA) layer.""" diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 810944e9eb..78816ea0bd 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -69,7 +69,6 @@ from maxtext.kernels.attention.ragged_attention import ragged_gqa from maxtext.kernels.attention.ragged_attention import ragged_mha from maxtext.layers import nnx_wrappers -from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec @@ -272,100 +271,6 @@ def _make_bidirectional_block_mask(bidirectional_mask): return bidirectional_block_mask -def attention_op_as_linen( - *, - config: Config, - mesh: Mesh, - attention_kernel: str, - max_target_length: int, - num_query_heads: int, - num_kv_heads: int, - float32_qk_product: bool = False, - max_prefill_predict_length: int = -1, - float32_logits: bool = False, - flash_axis_names_q: AxisNames = (BATCH_ATTN, HEAD, LENGTH, D_KV), - flash_axis_names_kv: AxisNames = (BATCH_ATTN, HEAD, KV_LENGTH, D_KV), - flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH), - prefill_cache_logical_axis_names: AxisNames = ( - CACHE_BATCH_PREFILL, - CACHE_SEQUENCE, - CACHE_HEADS, - CACHE_KV, - ), - cache_logical_axis_names: AxisNames = ( - CACHE_BATCH, - CACHE_SEQUENCE, - CACHE_HEADS, - CACHE_KV, - ), - cache_scale_logical_axis_names: AxisNames = ( - CACHE_SCALE_BATCH, - CACHE_SCALE_SEQUENCE, - CACHE_SCALE_HEADS, - CACHE_SCALE_KV, - ), - ragged_qkv_axis_names: AxisNames = ( - CACHE_BATCH, - CACHE_HEADS, - CACHE_SEQUENCE, - CACHE_KV, - ), - ragged_lengths_names: AxisNames = (CACHE_BATCH,), - compute_axis_order: AxisIdxes = (0, 1, 2, 3), - key_axis_order: AxisIdxes = (2, 0, 1, 3), - reshape_q: bool = False, - dropout_rate: float = 0.0, - dtype: DType = jnp.float32, - quant: Optional[Quant] = None, - kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention - attn_logits_soft_cap: float | None = None, - sliding_window_size: int | None = None, - chunk_attn_window_size: int | None = None, - use_ragged_attention: bool = False, - ragged_block_size: int = 256, -): - """A factory function to create an AttentionOp as a Linen module. - - This function serves as a bridge to use the NNX-based `AttentionOp` within a - Linen model. - """ - return nnx_wrappers.to_linen( - AttentionOp, - config=config, - mesh=mesh, - attention_kernel=attention_kernel, - max_target_length=max_target_length, - num_query_heads=num_query_heads, - num_kv_heads=num_kv_heads, - float32_qk_product=float32_qk_product, - max_prefill_predict_length=max_prefill_predict_length, - float32_logits=float32_logits, - flash_axis_names_q=flash_axis_names_q, - flash_axis_names_kv=flash_axis_names_kv, - flash_axis_names_splash_kernel=flash_axis_names_splash_kernel, - prefill_cache_logical_axis_names=prefill_cache_logical_axis_names, - cache_logical_axis_names=cache_logical_axis_names, - cache_scale_logical_axis_names=cache_scale_logical_axis_names, - ragged_qkv_axis_names=ragged_qkv_axis_names, - ragged_lengths_names=ragged_lengths_names, - compute_axis_order=compute_axis_order, - key_axis_order=key_axis_order, - reshape_q=reshape_q, - dropout_rate=dropout_rate, - dtype=dtype, - quant=quant, - kv_quant=kv_quant, - attention_type=attention_type, - attn_logits_soft_cap=attn_logits_soft_cap, - sliding_window_size=sliding_window_size, - chunk_attn_window_size=chunk_attn_window_size, - use_ragged_attention=use_ragged_attention, - ragged_block_size=ragged_block_size, - metadata_fn=variable_to_logically_partitioned, - ) - - class AttentionOp(nnx.Module): """Attention operation""" diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index b08a059938..b9b2af3b24 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -89,16 +89,6 @@ def __call__(self, x): return x * jax.lax.rsqrt(jnp.mean(x**2, axis=-1, keepdims=True) + self.eps) -def l2_norm_as_linen(self, eps: float = 1e-6): - """ - Initializes the L2Norm module and returns it as a Linen module. - - Args: - eps: float, epsilon used for numerical stability (default value should be ok for most cases). - """ - return nnx_wrappers.to_linen(L2Norm, eps=eps, metadata_fn=variable_to_logically_partitioned) - - def attention_as_linen( *, config: Config, diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index fc570e69bb..d8465b5e8f 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -13,262 +13,14 @@ # limitations under the License. """Module for decoder layers""" -# pylint: disable=arguments-differ -# pylint: disable=no-name-in-module -import functools -from typing import Any -import warnings - -from flax import linen as nn -from flax import nnx -from flax.linen.partitioning import ScanIn -import jax -from jax.ad_checkpoint import checkpoint_name import jax.numpy as jnp -from jax.sharding import Mesh -from maxtext.common.common_types import Config, DecoderBlockType, ShardMode -from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN -from maxtext.layers import linears -from maxtext.layers import mhc -from maxtext.layers import normalizations -from maxtext.layers import pipeline -from maxtext.layers.nnx_decoders import NNXDecoderLayer, NNXSequentialPipelineStage, NNXScannedPipelineStage -from maxtext.layers import quantizations -from maxtext.layers.attentions import attention_as_linen -from maxtext.layers.embeddings import attend_on_embedding, embed_as_linen, positional_embedding_as_linen -from maxtext.layers.normalizations import rms_norm -from maxtext.layers.quantizations import AqtQuantization as Quant -from maxtext.models import ( - deepseek, - deepseek4, - deepseek_batchsplit, - deepseek_batchsplit_fp8, - gemma, - gemma2, - gemma3, - gemma4, - gemma4_small, - gpt3, - gpt_oss, - llama2, - llama4, - mistral, - mixtral, - olmo3, - qwen2, - qwen3, - qwen3_custom, - qwen3_5, - simple_layer, -) -from maxtext.multimodal import utils as mm_utils -from maxtext.utils.sharding import create_sharding -from maxtext.utils import max_logging -from maxtext.utils import max_utils -from maxtext.utils import maxtext_utils -from maxtext.utils import sharding # ------------------------------------------------------------------------------ # The network: Decoder Definitions # ------------------------------------------------------------------------------ -class DecoderLayer(nn.Module): - """ - Transformer decoder layer that attends to the encoder. - This is the core, reusable building block for both the main model's - decoder stack and the auxiliary MTP layers. - """ - - config: Config - mesh: Mesh - model_mode: str - quant: None | Quant = None - - @nn.compact - def __call__( - self, - inputs, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk=None, - slot: None | int = None, - kv_cache: jax.Array | None = None, - attention_metadata: dict[str, Any] | None = None, - ): - cfg = self.config - mesh = self.mesh - _maybe_shard_with_logical = functools.partial( - sharding.maybe_shard_with_logical, - mesh=mesh, - shard_mode=cfg.shard_mode, - debug_sharding=cfg.debug_sharding, - ) - - if self.model_mode == MODEL_MODE_PREFILL: - logical_axis_names = ("activation_batch", "prefill_activation_length", "activation_embed") - else: - logical_axis_names = ("activation_batch", "activation_length", "activation_embed") - - if model_mode == MODEL_MODE_PREFILL: - inputs = _maybe_shard_with_logical(inputs, logical_axis_names) - else: - inputs = _maybe_shard_with_logical(inputs, logical_axis_names) - - inputs = checkpoint_name(inputs, "decoder_layer_input") - # inputs: embedded inputs to the decoder with shape [batch, length, emb_dim] - lnx = rms_norm( - num_features=inputs.shape[-1], - dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, - name="pre_self_attention_norm", - epsilon=cfg.normalization_layer_epsilon, - kernel_axes=("norm",), - )(inputs) - if model_mode == MODEL_MODE_PREFILL: - lnx = _maybe_shard_with_logical(lnx, logical_axis_names) - else: - lnx = _maybe_shard_with_logical(lnx, logical_axis_names) - - attention_layer = attention_as_linen( - config=self.config, - num_query_heads=cfg.num_query_heads, - num_kv_heads=cfg.num_kv_heads, - head_dim=cfg.head_dim, - max_target_length=cfg.max_target_length, - max_prefill_predict_length=cfg.max_prefill_predict_length, - attention_kernel=cfg.attention, - inputs_q_shape=lnx.shape, - inputs_kv_shape=lnx.shape, - mesh=mesh, - dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, - dropout_rate=cfg.dropout_rate, - name="self_attention", - float32_qk_product=cfg.float32_qk_product, - float32_logits=cfg.float32_logits, - quant=self.quant, - kv_quant=quantizations.configure_kv_quant(cfg), - prefill_cache_axis_order=tuple(map(int, cfg.prefill_cache_axis_order.split(","))), - ar_cache_axis_order=tuple(map(int, cfg.ar_cache_axis_order.split(","))), - compute_axis_order=tuple(map(int, cfg.compute_axis_order.split(","))), - reshape_q=cfg.reshape_q, - use_mrope=cfg.use_mrope, - mrope_section=cfg.mrope_section, - share_kv_projections=cfg.share_kv_projections, - model_mode=model_mode, - ) - - attention_lnx, kv_cache = attention_layer( - lnx, - lnx, - decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=deterministic, - model_mode=model_mode, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - ) - - if model_mode == MODEL_MODE_PREFILL: - attention_lnx = _maybe_shard_with_logical(attention_lnx, logical_axis_names) - else: - attention_lnx = _maybe_shard_with_logical(attention_lnx, logical_axis_names) - - # MLP block. - mlp_lnx = linears.mlp_block( - in_features=lnx.shape[-1], - intermediate_dim=cfg.mlp_dim, - activations=cfg.mlp_activations, - intermediate_dropout_rate=cfg.dropout_rate, - dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, - name="mlp", - model_mode=model_mode, - config=cfg, - quant=self.quant, - mesh=self.mesh, - )(lnx, deterministic=deterministic) - if model_mode == MODEL_MODE_PREFILL: - mlp_lnx = _maybe_shard_with_logical(mlp_lnx, logical_axis_names) - else: - mlp_lnx = _maybe_shard_with_logical(mlp_lnx, logical_axis_names) - - next_layer_addition = mlp_lnx + attention_lnx - - next_layer_addition_dropped_out = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))( - next_layer_addition, deterministic=deterministic - ) - - layer_output = next_layer_addition_dropped_out + inputs - if model_mode == MODEL_MODE_PREFILL: - layer_output = _maybe_shard_with_logical( - layer_output, - logical_axis_names, - ) - else: - layer_output = _maybe_shard_with_logical( - layer_output, - logical_axis_names, - ) - - if cfg.record_internal_nn_metrics: - self.sow("intermediates", "activation_mean", jnp.mean(layer_output)) - self.sow("intermediates", "activation_stdev", jnp.std(layer_output)) - self.sow( - "intermediates", - "activation_fraction_zero", - jnp.sum(layer_output == 0) / jnp.size(layer_output), - ) - - if cfg.scan_layers: - return layer_output, None - else: - return layer_output, kv_cache - - -class SequentialBlockDecoderLayers(nn.Module): - """Sequential unscanned series of decoder layers.""" - - decoder_layer: Any - num_decoder_layers: int - config: Config - mesh: Mesh - quant: Quant - model_mode: str - - @nn.compact - def __call__( - self, - inputs: jnp.ndarray, - decoder_segment_ids, - decoder_positions, - deterministic: bool, - model_mode, - slot: None | int = None, - ) -> jnp.ndarray: - for lyr in range(self.num_decoder_layers): - inputs = self.decoder_layer( - config=self.config, mesh=self.mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=model_mode - )( - inputs, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - slot=slot, - ) - if self.config.scan_layers: - inputs = inputs[0] # When scan_layers is True the decoder layers return (outputs, None). - if self.config.scan_layers: - return inputs, None # pytype: disable=bad-return-type - else: - return inputs - - def deepstack_process(hidden_states, bidirectional_mask, visual_embeds): """Process deepstack visual embeddings by adding them to hidden states at visual token positions. @@ -292,1503 +44,3 @@ def deepstack_process(hidden_states, bidirectional_mask, visual_embeds): # Only add where mask is True: hidden_states += visual_embeds * mask hidden_states = hidden_states + visual_embeds_scattered * mask_expanded return hidden_states - - -class Decoder(nn.Module): - """A stack of decoder layers as a part of an encoder-decoder architecture.""" - - config: Config - mesh: Mesh - quant: None | Quant = None - model_mode: str = MODEL_MODE_TRAIN - - def setup(self): - """Initialize decoder layer.""" - self.decoder_layer = self.get_decoder_layers() - self.norm_layer = self.get_norm_layer(num_features=self.config.emb_dim) - if self.config.using_pipeline_parallelism: - remat_policy = self.get_remat_policy() - nnx_blocks = self._get_nnx_decoder_block_classes() - - # Per-stage builder handed to create_pipeline: the pipeline transform invokes it - # once per pipeline stage, passing that stage's rngs, to construct the stage's - # decoder block(s). nnx_blocks (the selected decoder block classes) is captured from - # this setup scope; remat_policy is applied separately by create_pipeline. - def build_pipeline_stage_layers(rngs): - """Builds one pipeline stage module from the selected NNX decoder block classes.""" - return self._build_nnx_pipeline_stage(nnx_blocks, rngs) - - self.pipeline_module = pipeline.create_pipeline( - config=self.config, layers=build_pipeline_stage_layers, mesh=self.mesh, remat_policy=remat_policy - ) - - def minimal_policy(self, with_context=False, with_quantization=False): - """Helper for creating minimal checkpoint policies.""" - names = [ - "query_proj", - "value_proj", - "key_proj", - "qkv_proj", - "out_proj", - "mlpwi_0", - "mlpwi_1", - "mlpwi", - "mlpwo", - ] - if with_context: - names.append("context") - if with_quantization: - names.append("quantization") - return jax.checkpoint_policies.save_only_these_names(*names) - - def get_remat_policy(self): - """Get remat policy""" - policy = None - cfg = self.config - if cfg.remat_policy != "none": - if cfg.remat_policy in ("minimal_with_context", "minimal_flash"): - # save all - if cfg.remat_policy == "minimal_flash": - max_logging.log("WARNING: 'minimal_flash' will be deprecated soon, please use 'minimal_with_context' instead.") - max_logging.log("WARNING: 'minimal_flash' will be deprecated soon, please use 'minimal_with_context' instead.") - policy = self.minimal_policy(with_context=True) - elif cfg.remat_policy == "minimal": - # save all except context - policy = self.minimal_policy() - elif cfg.remat_policy == "minimal_with_quantization": - if cfg.scan_layers: - warnings.warn( - "Scan layers can introduce overhead to checkpointed values that in some configurations is slower" - "than not checkpointing at all. If you are using scan layers, benchmark with and without quantization " - "checkpointing in your workflow to see which is faster. Without scan layers, checkpointing quantizations is " - "beneficial for performance." - ) - policy = self.minimal_policy(with_context=False, with_quantization=True) - elif cfg.remat_policy == "minimal_with_context_and_quantization": - if cfg.scan_layers: - warnings.warn( - "Scan layers can introduce overhead to checkpointed values that in some configurations is slower" - "than not checkpointing at all. If you are using scan layers, benchmark with and without quantization " - "checkpointing in your workflow to see which is faster. Without scan layers, checkpointing quantizations is " - "beneficial for performance." - ) - policy = self.minimal_policy(with_context=True, with_quantization=True) - elif cfg.remat_policy == "save_dot_with_context_except_mlp": - policy = jax.checkpoint_policies.save_only_these_names( - "query_proj", - "value_proj", - "key_proj", - "qkv_proj", - "context", - "out_proj", - ) - elif cfg.remat_policy == "save_dot_except_mlpwi": - policy = jax.checkpoint_policies.save_only_these_names( - "query_proj", - "value_proj", - "key_proj", - "qkv_proj", - "out_proj", - "mlpwo", - ) - elif cfg.remat_policy == "save_dot_except_mlp": - policy = jax.checkpoint_policies.save_only_these_names( - "query_proj", - "value_proj", - "key_proj", - "qkv_proj", - "out_proj", - ) - elif cfg.remat_policy == "save_qkv_proj": - policy = jax.checkpoint_policies.save_only_these_names( - "query_proj", - "value_proj", - "key_proj", - "qkv_proj", - ) - elif cfg.remat_policy == "qkv_proj_offloaded": - policy = jax.checkpoint_policies.save_and_offload_only_these_names( - names_which_can_be_saved=[], - names_which_can_be_offloaded=["query_proj", "value_proj", "key_proj"], - offload_src="device", - offload_dst="pinned_host", - ) - elif cfg.remat_policy == "minimal_offloaded": - # offload all except context - policy = jax.checkpoint_policies.save_and_offload_only_these_names( - names_which_can_be_saved=[], - names_which_can_be_offloaded=[ - "query_proj", - "value_proj", - "key_proj", - "qkv_proj", - "out_proj", - "mlpwi_0", - "mlpwi_1", - "mlpwi", - "mlpwo", - ], - offload_src="device", - offload_dst="pinned_host", - ) - elif cfg.remat_policy == "custom": - policy = jax.checkpoint_policies.save_and_offload_only_these_names( - names_which_can_be_saved=cfg.tensors_on_device, - names_which_can_be_offloaded=cfg.tensors_to_offload, - offload_src="device", - offload_dst="pinned_host", - ) - elif cfg.remat_policy == "save_out_proj": - policy = jax.checkpoint_policies.save_only_these_names( - "out_proj", - ) - else: - assert cfg.remat_policy == "full", "Remat policy needs to be on list of remat policies" - policy = None - return policy - - def get_decoder_layers(self): - """Retrieves a list of decoder layer classes based on the `decoder_block` config. - - Returns: - A list containing one or more `nn.Module` classes for the decoder. - """ - match self.config.decoder_block: - case DecoderBlockType.DEFAULT: - return [DecoderLayer] - case DecoderBlockType.LLAMA2: - return [llama2.LlamaDecoderLayerToLinen] - case DecoderBlockType.MISTRAL: - # TODO(ranran): update to Mistral with sliding window attention - return [mistral.MistralDecoderLayerToLinen] - case DecoderBlockType.MIXTRAL: - return [mixtral.MixtralDecoderLayerToLinen] - case DecoderBlockType.DEEPSEEK: - return [ - deepseek.DeepSeekDenseLayerToLinen, - deepseek.DeepSeekMoELayerToLinen, - ] - case DecoderBlockType.DEEPSEEK4: - return ( - [deepseek4.DeepSeek4ScannableBlockToLinen] if self.config.scan_layers else [deepseek4.DeepSeek4LayerToLinen] - ) - case DecoderBlockType.GEMMA: - return [gemma.GemmaDecoderLayerToLinen] - case DecoderBlockType.GEMMA2: - return [gemma2.Gemma2DecoderLayerToLinen] - case DecoderBlockType.GEMMA3: - return [gemma3.Gemma3DecoderLayerToLinen] - case DecoderBlockType.GEMMA4: - return [gemma4.Gemma4ScannableBlockToLinen] if self.config.scan_layers else [gemma4.Gemma4DecoderLayerToLinen] - case DecoderBlockType.GEMMA4_SMALL: - # PLE input + KV-share donor threading requires per-layer-index state, - # which is not expressible inside ``nn.scan``. - return [gemma4_small.Gemma4SmallDecoderLayerToLinen] - case DecoderBlockType.GPT3: - return [gpt3.Gpt3DecoderLayerToLinen] - case DecoderBlockType.GPT_OSS: - return [gpt_oss.GptOssScannableBlockToLinen] if self.config.scan_layers else [gpt_oss.GptOssDecoderLayerToLinen] - case DecoderBlockType.QWEN2: - return [qwen2.Qwen2DecoderLayerToLinen] - case DecoderBlockType.QWEN3: - return [qwen3.Qwen3DecoderLayerToLinen] - case DecoderBlockType.QWEN3_MOE: - return [qwen3.Qwen3MoeDecoderLayerToLinen] - case DecoderBlockType.QWEN3_CUSTOM_MOE: - return [qwen3_custom.Qwen3CustomMoeDecoderLayerToLinen] - case DecoderBlockType.QWEN3_NEXT: - return [qwen3.Qwen3NextScannableBlockToLinen] if self.config.scan_layers else [qwen3.Qwen3NextDecoderLayerToLinen] - case DecoderBlockType.QWEN3_5: - return [qwen3_5.Qwen3_5ScannableBlockToLinen] if self.config.scan_layers else [qwen3_5.Qwen3_5DecoderLayerToLinen] - case DecoderBlockType.SIMPLE: - return [simple_layer.SimpleDecoderLayerToLinen] - case DecoderBlockType.SIMPLE_MLP: - return [simple_layer.SimpleMlpDecoderLayerToLinen] - case DecoderBlockType.LLAMA4: - return [llama4.Llama4ScannableBlockToLinen] if self.config.scan_layers else [llama4.Llama4DecoderLayerToLinen] - case DecoderBlockType.OLMO3: - return [olmo3.Olmo3ScannableBlockToLinen] if self.config.scan_layers else [olmo3.Olmo3DecoderLayerToLinen] - - case _: - # Default case to handle any unknown decoder block types. - raise ValueError(f"Incorrect decoder_block name {self.config.decoder_block.value=}") - - def _get_nnx_decoder_block_classes(self): - """Returns NNX decoder block classes for pipeline stage creation.""" - cfg = self.config - - def get_scannable(normal_cls, scannable_cls): - return [scannable_cls] if cfg.scan_layers else [normal_cls] - - layer_map = { - DecoderBlockType.DEFAULT: [NNXDecoderLayer], - DecoderBlockType.LLAMA2: [llama2.LlamaDecoderLayer], - DecoderBlockType.MISTRAL: [mistral.MistralDecoderLayer], - DecoderBlockType.MIXTRAL: [mixtral.MixtralDecoderLayer], - DecoderBlockType.GEMMA: [gemma.GemmaDecoderLayer], - DecoderBlockType.GEMMA2: [gemma2.Gemma2DecoderLayer], - DecoderBlockType.GEMMA3: [gemma3.Gemma3DecoderLayer], - DecoderBlockType.GEMMA4: get_scannable(gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock), - DecoderBlockType.GEMMA4_SMALL: [gemma4_small.Gemma4SmallDecoderLayer], - DecoderBlockType.GPT3: [gpt3.Gpt3DecoderLayer], - DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock), - DecoderBlockType.QWEN2: [qwen2.Qwen2DecoderLayer], - DecoderBlockType.QWEN3: [qwen3.Qwen3DecoderLayer], - DecoderBlockType.QWEN3_MOE: [qwen3.Qwen3MoeDecoderLayer], - DecoderBlockType.QWEN3_CUSTOM_MOE: [qwen3_custom.Qwen3CustomMoeDecoderLayer], - DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock), - DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock), - DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], - DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], - DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer], - DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), - DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), - } - - if cfg.decoder_block not in layer_map: - raise ValueError(f"Incorrect decoder_block name {cfg.decoder_block.value=}") - return layer_map[cfg.decoder_block] - - def set_remat_policy(self, block_layers, policy): - """Set remat policy""" - RemattedBlockLayers = [] - for block_layer in block_layers: - if self.config.parameter_memory_host_offload: - # Define parameter movement with mesh-based sharding - def move_to_device(variables): - """Move parameters to device with proper sharding.""" - - def map_fn(path, value): - max_logging.log(f"models.py: Moving parameter {path} to device") - return jax.device_put(value, max_utils.device_space()) - - return jax.tree_util.tree_map_with_path(map_fn, variables) - - # Transform layer class before remat - block_layer = nn.map_variables(block_layer, ["params"], move_to_device, mutable=True) - - # Apply remat policy to layer - layer = nn.remat( - block_layer, - prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config), - policy=policy, - static_argnums=(4, 5), # Deterministic and model mode are static arguments. - ) - RemattedBlockLayers.append(layer) - return RemattedBlockLayers - - def _build_nnx_pipeline_stage(self, decoder_blocks, rngs): - """Creates a single NNX pipeline stage module.""" - cfg = self.config - base_stage_cls = decoder_blocks[1] if cfg.decoder_block == DecoderBlockType.DEEPSEEK else decoder_blocks[0] - - if cfg.num_layers_per_pipeline_stage == 1: - return base_stage_cls(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) - elif cfg.scan_layers_per_stage: - return NNXScannedPipelineStage( - base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs - ) - return NNXSequentialPipelineStage( - base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs - ) - - def get_pipeline_stage_module(self, decoder_blocks): - """get pipeline stage module""" - - def get_layer_to_pipeline(blocks, cfg): - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - return blocks[1] # return the sparse block - else: - return blocks[0] - - cfg = self.config - base_stage = get_layer_to_pipeline(decoder_blocks, cfg) - if cfg.set_remat_policy_on_layers_per_stage: - policy = self.get_remat_policy() - base_stage = self.set_remat_policy([base_stage], policy)[0] - if cfg.num_layers_per_pipeline_stage == 1: - stage_module = base_stage(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode) - elif cfg.scan_layers_per_stage: - stage_module = self.scan_decoder_layers( - cfg, - base_stage, - cfg.num_layers_per_pipeline_stage, - "layers_per_stage", - self.mesh, - in_axes_tuple=(nn.broadcast,) * 4, - model_mode=self.model_mode, - ) - else: - stage_module = SequentialBlockDecoderLayers( - decoder_layer=base_stage, - num_decoder_layers=cfg.num_layers_per_pipeline_stage, - config=cfg, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, - ) - return stage_module - - def get_norm_layer(self, num_features: int): - """get normalization layer (return type inherits from nn.Module)""" - if self.config.decoder_block in ( - DecoderBlockType.DEFAULT, - DecoderBlockType.LLAMA2, - DecoderBlockType.MISTRAL, - DecoderBlockType.MIXTRAL, - DecoderBlockType.DEEPSEEK, - DecoderBlockType.DEEPSEEK4, - DecoderBlockType.GEMMA, - DecoderBlockType.GEMMA2, - DecoderBlockType.GEMMA3, - DecoderBlockType.GEMMA4, - DecoderBlockType.GEMMA4_SMALL, - DecoderBlockType.QWEN2, - DecoderBlockType.QWEN3, - DecoderBlockType.QWEN3_MOE, - DecoderBlockType.QWEN3_CUSTOM_MOE, - DecoderBlockType.GPT_OSS, - DecoderBlockType.SIMPLE, - DecoderBlockType.SIMPLE_MLP, - DecoderBlockType.LLAMA4, - DecoderBlockType.OLMO3, - ): - return functools.partial(rms_norm, num_features=num_features, shard_mode=self.config.shard_mode) - elif self.config.decoder_block == DecoderBlockType.GPT3: - return functools.partial(gpt3.gpt3_layer_norm, num_features=num_features, reductions_in_fp32=False, use_bias=True) - elif self.config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5): - return functools.partial( - normalizations.Qwen3NextRMSNormLinen, num_features=num_features, shard_mode=self.config.shard_mode - ) - else: - raise ValueError(f"Incorrect decoder_block name {self.config.decoder_block.value=}") - - def scan_decoder_layers(self, cfg, decoder_layer, length, metadata_axis_name, mesh, in_axes_tuple, **kwargs): - """scan decoder layers, calls `flax.linen.transforms.scan`""" - initializing = self.is_mutable_collection("params") - params_spec = cfg.param_scan_axis if initializing else ScanIn(cfg.param_scan_axis) - cache_spec = 0 - scan_fn = nn.scan( - decoder_layer, - variable_axes={ - "params": params_spec, - "cache": cache_spec, - "intermediates": 0, - "aqt": 0, - "batch_stats": 0, - "_overwrite_with_gradient": 0, - }, - split_rngs={ - "params": True, - "dropout": cfg.enable_dropout, - }, - in_axes=in_axes_tuple, - length=length, - metadata_params={nn.PARTITION_NAME: metadata_axis_name}, - ) - return scan_fn( - config=cfg, mesh=mesh, name=metadata_axis_name, quant=self.quant, **kwargs # pytype: disable=wrong-keyword-args - ) - - @nn.compact - def _apply_embedding( - self, - shared_embedding: nn.Module | nnx.Module, - decoder_input_tokens, - decoder_positions, - deterministic, - model_mode, - multimodal_input=None, - ): - """Applies token and positional embeddings to the input tokens.""" - cfg = self.config - - y = shared_embedding(decoder_input_tokens.astype("int32"), model_mode=model_mode) - - # Merge the image embeddings with the text embeddings for multimodal models - if multimodal_input is not None: - image_embeddings = multimodal_input.image_embeddings - bidirectional_mask = multimodal_input.bidirectional_mask - image_masks = multimodal_input.image_masks - video_embeddings = getattr(multimodal_input, "video_embeddings", None) - video_masks = getattr(multimodal_input, "video_masks", None) - bidirectional_mask_video = getattr(multimodal_input, "bidirectional_mask_video", None) - audio_embeddings = multimodal_input.audio_embeddings - audio_masks = multimodal_input.audio_masks - - if image_embeddings is not None and cfg.use_multimodal: - if cfg.model_name in [ - "gemma3-4b", - "gemma3-12b", - "gemma3-27b", - "gemma4-26b", - "gemma4-31b", - "gemma4-e2b", - "gemma4-e4b", - "llama4-17b-16e", - "llama4-17b-128e", - "qwen3-omni-30b-a3b", - "qwen3-vl-2b", - "qwen3-vl-4b", - "qwen3.5-35b-a3b", - "qwen3.5-397b-a17b", - ]: - y = mm_utils.merge_mm_embeddings( - text_embeddings=y, - multimodal_embeddings=image_embeddings, - mask=bidirectional_mask, - token_masks=image_masks, - ) - # TODO(hengtaoguo): Add support for other multimodal models such as Llama4, refactor if needed - else: - raise ValueError(f"Unsupported model_name for multimodal: {cfg.model_name}") - - if video_embeddings is not None and cfg.use_multimodal: - if cfg.model_name in ["qwen3-omni-30b-a3b", "qwen3-vl-2b", "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]: - y = mm_utils.merge_mm_embeddings( - text_embeddings=y, - multimodal_embeddings=video_embeddings, - mask=bidirectional_mask_video, - token_masks=video_masks, - ) - else: - raise ValueError(f"Unsupported model_name for video: {cfg.model_name}") - - if audio_embeddings is not None and cfg.use_audio: - if cfg.model_name in ["qwen3-omni-30b-a3b"]: - y = mm_utils.merge_mm_embeddings( - text_embeddings=y, - multimodal_embeddings=audio_embeddings, - mask=audio_masks, - token_masks=None, - ) - else: - raise ValueError(f"Unsupported model_name for audio: {cfg.model_name}") - - y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic) - y = y.astype(cfg.dtype) - - if cfg.use_untrainable_positional_embedding: - y += positional_embedding_as_linen(embedding_dims=cfg.base_emb_dim)(y.shape[1], decoder_positions) - - if cfg.trainable_position_size > 0: - y += embed_as_linen( - num_embeddings=cfg.trainable_position_size, - num_features=cfg.emb_dim, - dtype=cfg.dtype, - embedding_init=nn.initializers.normal(stddev=1.0), - name="position_embedder", - config=cfg, - mesh=self.mesh, - )(decoder_positions.astype("int32"), model_mode=model_mode) - return y - - @nn.compact - def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, deterministic, model_mode): - """Applies final normalization and projects hidden states to logits.""" - - cfg = self.config - if cfg.shard_mode == ShardMode.EXPLICIT: - norm_out_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_embed")) - else: - norm_out_sharding = None - - y = self.get_norm_layer(num_features=y.shape[-1])( - dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, - name="decoder_norm", - epsilon=cfg.normalization_layer_epsilon, - kernel_axes=("norm",), - parameter_memory_host_offload=cfg.parameter_memory_host_offload, - )(y, out_sharding=norm_out_sharding) - y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic) - - if model_mode in (MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE): - out_sharding = create_sharding(self.mesh, (None, None, "activation_vocab")) - else: - out_sharding = create_sharding( - self.mesh, ("activation_embed_and_logits_batch", "activation_length", "activation_vocab") - ) - - # [batch, length, emb_dim] -> [batch, length, vocab_size] - if cfg.logits_via_embedding: - # Use the transpose of embedding matrix for logit transform. - if isinstance(shared_embedding, nnx.Module): - embedding_table = shared_embedding.embedding.value - else: - embedding_table = shared_embedding.variables["params"]["embedding"] - if isinstance(embedding_table, nn.spmd.LogicallyPartitioned): - embedding_table = embedding_table.unbox() - attend_dtype = jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype - logits = attend_on_embedding(y, embedding_table, attend_dtype, self.config, out_sharding) - - if self.config.normalize_embedding_logits: - # Correctly normalize pre-softmax logits for this shared case. - logits = logits / jnp.sqrt(y.shape[-1]) - if cfg.final_logits_soft_cap: - logits = logits / cfg.final_logits_soft_cap - logits = jnp.tanh(logits) * cfg.final_logits_soft_cap - else: - logits = linears.dense_general( - inputs_shape=y.shape, - out_features_shape=cfg.vocab_size, - weight_dtype=cfg.weight_dtype, - dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype, # for logit training stability - kernel_axes=("embed_vocab", "vocab"), - shard_mode=cfg.shard_mode, - name="logits_dense", - matmul_precision=self.config.matmul_precision, - parameter_memory_host_offload=cfg.parameter_memory_host_offload, - )( - y, - out_sharding=out_sharding, - ) # We do not quantize the logits matmul. - - if self.config.cast_logits_to_fp32: - logits = logits.astype(jnp.float32) - - return logits - - @nn.compact - def __call__( - self, - shared_embedding: nn.Module | nnx.Module, - decoder_input_tokens, - decoder_positions, - decoder_segment_ids=None, - deterministic=False, - model_mode=MODEL_MODE_TRAIN, - previous_chunk=None, - slot: None | int = None, - multimodal_input=None, - kv_caches: list[jax.Array] | None = None, - attention_metadata=None, - deepstack_visual_embeds: None | list[jnp.ndarray] = None, - ): - cfg = self.config - mesh = self.mesh - assert decoder_input_tokens.ndim == 2 # [batch, len] - - # [batch, length] -> [batch, length, emb_dim] - y = self._apply_embedding( - shared_embedding, - decoder_input_tokens, - decoder_positions, - deterministic, - model_mode, - multimodal_input=multimodal_input, - ) - - mhc_expand, mhc_reduce = mhc.get_functions(cfg.mhc_expansion_rate) - if cfg.mhc_expansion_rate > 1: - # (batch, length, emb_dim) --> (batch, length, mhc_expansion_rate, emb_dim) - y = mhc_expand(y) - - policy = self.get_remat_policy() - RemattedBlockLayers = self.set_remat_policy(self.decoder_layer, policy) - # scan does not support kwargs in layer call, passing broadcast_args as positional arg - broadcast_args = ( - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - ) - if cfg.using_pipeline_parallelism: - logical_partition_spec = ( - self.pipeline_module.get_weight_sharding(y, decoder_segment_ids, decoder_positions, deterministic, model_mode) - if cfg.pipeline_fsdp_ag_once or cfg.pipeline_fsdp_ag_per_repeat - else None - ) - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek." - dense_layer = RemattedBlockLayers[0] - moe_layer = RemattedBlockLayers[1] - num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers - num_moe_layers_outside_pp = num_moe_layers - self.config.pipeline_parallel_layers - logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules) - # We chose not to pipeline the dense layers, only sparse for SPMD. - with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp): - y, _ = self.scan_decoder_layers( - cfg, - dense_layer, - cfg.first_num_dense_layers, - "dense_layers", - mesh, - in_axes_tuple=(nn.broadcast,) * len(broadcast_args), - model_mode=model_mode, - )(y, *broadcast_args) - if num_moe_layers_outside_pp > 0: - y, _ = self.scan_decoder_layers( - cfg, - moe_layer, - num_moe_layers_outside_pp, - "moe_layers", - mesh, - in_axes_tuple=(nn.broadcast,) * len(broadcast_args), - model_mode=model_mode, - )(y, *broadcast_args) - y = self.pipeline_module(y, *broadcast_args, logical_partition_spec=logical_partition_spec) - else: # Not DeepSeek - y = self.pipeline_module(y, *broadcast_args, logical_partition_spec=logical_partition_spec) - remaining_layers = self.config.num_decoder_layers - self.config.pipeline_parallel_layers - if remaining_layers > 0: - logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules) - with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp): - y, _ = self.scan_decoder_layers( - cfg, - RemattedBlockLayers[0], - remaining_layers, - "layers_outside_pipeline", - mesh, - in_axes_tuple=(nn.broadcast,) * len(broadcast_args), - model_mode=model_mode, - )(y, *broadcast_args) - else: - if cfg.scan_layers: - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek." - layer_call_kwargs = { - "previous_chunk": previous_chunk, - "slot": slot, - } - dense_layer = RemattedBlockLayers[0] - moe_layer = RemattedBlockLayers[1] - if cfg.engram_layers: - original_dense_call = dense_layer.__call__ - original_moe_call = moe_layer.__call__ - dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs) - moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs) - - common_kwargs = { - "dense_layer": dense_layer, - "moe_layer": moe_layer, - "original_dense_call": original_dense_call, - "original_moe_call": original_moe_call, - "layer_call_kwargs": layer_call_kwargs, - "decoder_segment_ids": decoder_segment_ids, - "decoder_positions": decoder_positions, - "deterministic": deterministic, - "model_mode": model_mode, - "decoder_input_tokens": decoder_input_tokens, - "broadcast_args": broadcast_args, - } - - # Apply Dense Layers - y = self._apply_interleaved_scanned_layers( - y, - layer_type="dense", - start_idx=0, - end_idx=cfg.first_num_dense_layers, - engram_indices=cfg.engram_layers, - **common_kwargs, - ) - - # Apply MoE Layers - y = self._apply_interleaved_scanned_layers( - y, - layer_type="moe", - start_idx=cfg.first_num_dense_layers, - end_idx=cfg.num_decoder_layers, - engram_indices=cfg.engram_layers, - **common_kwargs, - ) - else: - dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs) - y, _ = self.scan_decoder_layers( - cfg, - dense_layer, - cfg.first_num_dense_layers, - "dense_layers", - mesh, - in_axes_tuple=(nn.broadcast,) * len(broadcast_args), - model_mode=model_mode, - )(y, *broadcast_args) - moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs) - num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers - - # If batch-split schedule is used and initialization is complete, - # as detected by immutable params, use deepseek_batchsplit custom - # scan with initialized parameters. - if cfg.use_batch_split_schedule and not self.is_mutable_collection("params"): - # old version of batch-split that fully uses qwix quantization. - if cfg.use_qwix_quantization and not cfg.use_manual_quantization: - y = deepseek_batchsplit_fp8.scan_batch_split_layers( - y, - self.variables["params"]["moe_layers"], - decoder_positions, - decoder_segment_ids, - model_mode=model_mode, - mesh=mesh, - quant=self.quant, - cfg=cfg, - policy=policy, - ) - else: - # bf16 and fp8 code path for pure-JAX batch-split. - # fp8 code path supports both manual quantization and qwix - # quantization. - y = deepseek_batchsplit.scan_batch_split_layers( - y, - self.variables["params"]["moe_layers"], - decoder_positions, - mesh=mesh, - cfg=cfg, - num_layers=num_moe_layers, - ) - else: - y, _ = self.scan_decoder_layers( - cfg, - moe_layer, - num_moe_layers, - "moe_layers", - mesh, - in_axes_tuple=(nn.broadcast,) * len(broadcast_args), - model_mode=model_mode, - )(y, *broadcast_args) - elif cfg.decoder_block == DecoderBlockType.GEMMA3: - bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None - y = self._apply_gemma3_scanned_blocks( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - bidirectional_mask_value, - previous_chunk, - slot, - kv_caches=kv_caches, - attention_metadata=attention_metadata, - ) - elif cfg.decoder_block == DecoderBlockType.GEMMA4: - bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None - y = self._apply_gemma4_scanned_blocks( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - bidirectional_mask_value, - previous_chunk, - slot, - kv_caches=kv_caches, - attention_metadata=attention_metadata, - ) - elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4: - y = self._apply_deepseek4_scanned_blocks( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk, - slot, - decoder_input_tokens, - ) - else: - RemattedBlockLayer = RemattedBlockLayers[0] - scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval) - layer_kwargs = {} - if cfg.decoder_block == DecoderBlockType.LLAMA4: - layer_kwargs = { - "nope_layer_interval": self.config.nope_layer_interval, - "interleave_moe_layer_step": self.config.interleave_moe_layer_step, - } - # Update broadcast_args and in_axes_tuple for vLLM RPA - in_axes_tuple = (nn.broadcast,) * len(broadcast_args) - current_broadcast_args = list(broadcast_args) - current_in_axes_tuple = list(in_axes_tuple) - - if kv_caches is not None: - # Stack kv_caches for scan: [num_layers, ...] - stacked_kv_cache = jnp.stack(kv_caches, axis=0) - - # We pass (y, stacked_kv_cache, 0) as the carry - carry = (y, stacked_kv_cache, 0) - - # We don't pass kv_cache as a scanned argument anymore - - # Pass None for previous_chunk, slot, kv_cache to align with __call__ signature - current_broadcast_args.extend([None, None, None, attention_metadata]) - current_in_axes_tuple.extend([nn.broadcast] * 4) - - max_logging.info(f"DEBUG: len(current_broadcast_args)={len(current_broadcast_args)}") - max_logging.info(f"DEBUG: current_broadcast_args={[type(a) for a in current_broadcast_args]}") - - final_carry, _ = self.scan_decoder_layers( - cfg, - RemattedBlockLayer, - scan_length, - "layers", - mesh, - in_axes_tuple=tuple(current_in_axes_tuple), - model_mode=model_mode, - **layer_kwargs, - )(carry, *current_broadcast_args) - - y, returned_kv_cache, _ = final_carry - - # Update the list of KV caches from the scanned results - for i in range(cfg.num_decoder_layers): - kv_caches[i] = returned_kv_cache[i] - else: - # Fallback to old behavior if kv_caches is None (not vLLM RPA) - current_broadcast_args.append(None) - current_in_axes_tuple.append(nn.broadcast) - - y, _ = self.scan_decoder_layers( - cfg, - RemattedBlockLayer, - scan_length, - "layers", - mesh, - in_axes_tuple=tuple(current_in_axes_tuple), - model_mode=model_mode, - **layer_kwargs, - )(y, *current_broadcast_args) - else: - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - assert len(RemattedBlockLayers) == 2, "Unscanned layers must have a length of 2 using deepseek." - dense_layer = RemattedBlockLayers[0] - moe_layer = RemattedBlockLayers[1] - - layers = [dense_layer, moe_layer] - layer_prefixes = ["dense_layers", "moe_layers"] - num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers - num_layers_list = [cfg.first_num_dense_layers, num_moe_layers] - # Iterate over the two layer groups (dense and MoE) and apply layer transformation - global_layer_idx_offset = 0 - for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes): - for index in range(num_layers): - global_layer_idx = global_layer_idx_offset + index - kv_cache = kv_caches[index] if kv_caches is not None else None - input_tokens = decoder_input_tokens if cfg.engram_layers else None - y, kv_cache = layer( - config=cfg, - mesh=mesh, - name=f"{layer_prefix}_{index}", - quant=self.quant, - model_mode=self.model_mode, - layer_idx=global_layer_idx, - )( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk=previous_chunk, - slot=slot, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - decoder_input_tokens=input_tokens, - ) - if kv_caches is not None and kv_cache is not None: - kv_caches[index] = kv_cache - global_layer_idx_offset += num_layers - elif cfg.decoder_block == DecoderBlockType.GEMMA4_SMALL: - y, kv_caches = self._apply_gemma4_small_layers( - y, - decoder_input_tokens, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - multimodal_input=multimodal_input, - kv_caches=kv_caches, - attention_metadata=attention_metadata, - previous_chunk=previous_chunk, - slot=slot, - ) - else: - for lyr in range(cfg.num_decoder_layers): - RemattedBlockLayer = RemattedBlockLayers[0] - layer_kwargs = {} - layer_call_kwargs = {} - if cfg.decoder_block == DecoderBlockType.GEMMA3: - # Gemma3 uses both global and sliding window attention depending on the layer index. - bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None - layer_kwargs = {"attention_type": gemma3.get_attention_type(layer_id=lyr)} - layer_call_kwargs = {"bidirectional_mask": bidirectional_mask_value} - if cfg.decoder_block == DecoderBlockType.GEMMA4: - # Gemma4 uses both global and sliding window attention depending on the layer index. - bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None - layer_kwargs = {"attention_type": gemma4.get_attention_type(layer_id=lyr)} - layer_call_kwargs = {"bidirectional_mask": bidirectional_mask_value} - if cfg.decoder_block == DecoderBlockType.LLAMA4: - layer_kwargs = { - "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval), - "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step), - } - if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4): - layer_kwargs = {"layer_idx": lyr} - if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: - layer_call_kwargs["decoder_input_tokens"] = decoder_input_tokens - kv_cache = None - if kv_caches is not None: - # For all decoder blocks (including QWEN3_NEXT/QWEN3_5 with vLLM flat-list - # kv_caches), pass the per-layer cache directly. For hybrid attention+GDN - # models, kv_caches[lyr] is a regular attention cache for attention layers - # and a (conv_state, recurrent_state) paged-mamba tuple for GDN layers. - kv_cache = kv_caches[lyr] - - if cfg.decoder_block == DecoderBlockType.GPT_OSS: - layer_kwargs = {"attention_type": gpt_oss.get_attention_type(layer_id=lyr)} - if cfg.decoder_block == DecoderBlockType.OLMO3: - layer_kwargs = {"attention_type": olmo3.get_attention_type(layer_id=lyr)} - layer = RemattedBlockLayer( - config=cfg, mesh=mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=self.model_mode, **layer_kwargs - ) - y, returned_cache = layer( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk=previous_chunk, - slot=slot, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - **layer_call_kwargs, - ) - if kv_caches is not None and returned_cache is not None: - kv_caches[lyr] = returned_cache - - if deepstack_visual_embeds is not None and lyr < len(deepstack_visual_embeds): - visual_embeds = deepstack_visual_embeds[lyr] - # Use bidirectional_mask to identify visual token positions - bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None - if bidirectional_mask_value is not None and visual_embeds is not None: - y = deepstack_process(y, bidirectional_mask_value, visual_embeds) - - assert isinstance(y, jax.Array) - - # After the final transformer layer, `y` holds the raw, un-normalized hidden state. - if cfg.mhc_expansion_rate > 1: - # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) - hidden_state = mhc_reduce(y) - else: - hidden_state = y - - # When initializing with vLLM RPA attention, we need to run the output head to - # initialize any parameters associated with it. - # Same case applicable to vocab tiling - if self.is_initializing() and (cfg.num_vocab_tiling > 1 or cfg.attention == "vllm_rpa"): - _ = self.apply_output_head(shared_embedding, hidden_state, deterministic, model_mode) - - # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage. - if cfg.attention == "vllm_rpa": - logits = None - # When in the Indexer Dense Warm-up stage, skip the expensive output head projection - # for efficiency, as the main model is frozen and the LM loss is not needed. - # TODO(b/501446870): Investigate model_mode as train at beginning for decoding stage - elif ( - cfg.use_indexer and cfg.indexer_loss_scaling_factor > 0.0 and not cfg.indexer_sparse_training - ) and model_mode == MODEL_MODE_TRAIN: - logits = None - # When vocab tiling is enabled in training mode, full logits won't generate to reduce memory - # Instead, we keep track on the hidden states, which has smaller size compared to full logits - elif cfg.num_vocab_tiling > 1 and model_mode == MODEL_MODE_TRAIN: - logits = None - self.sow("intermediates", "hidden_states", hidden_state) - - else: - logits = self.apply_output_head(shared_embedding, hidden_state, deterministic, model_mode) - - # The API of the Decoder is now a tuple, providing both the main output - # and the raw hidden state needed for auxiliary tasks. - return logits, hidden_state, kv_caches - - def _apply_gemma3_scanned_blocks( - self, - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - bidirectional_mask, - previous_chunk, - slot, - kv_caches=None, - attention_metadata=None, - ): - """Applies Gemma3 scanned decoder blocks, handling main scan and remainders.""" - - cfg = self.config - mesh = self.mesh - - # Define the repeating pattern length and calculate how many full blocks to scan - attention_pattern_length = len(gemma3.GEMMA3_ATTENTION_PATTERN) - scan_length = cfg.num_decoder_layers // attention_pattern_length - - policy = self.get_remat_policy() - RemattedGemma3Block = self.set_remat_policy([gemma3.Gemma3ScannableBlockToLinen], policy)[0] - - layer_kwargs = {"num_of_layers": attention_pattern_length} - - # Apply the main scan over the full blocks - if scan_length > 0: - kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( - kv_caches, scan_length, attention_pattern_length, stack=True - ) - - broadcast_args_spec = [ - (decoder_segment_ids, nn.broadcast), - (decoder_positions, nn.broadcast), - (deterministic, nn.broadcast), - (model_mode, nn.broadcast), - (slot, nn.broadcast), - (None, nn.broadcast), # page_state - (previous_chunk, nn.broadcast), - (bidirectional_mask, nn.broadcast), - (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), - (attention_metadata, nn.broadcast), - ] - broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) - in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) - - y, returned_kv_cache = self.scan_decoder_layers( - cfg, - RemattedGemma3Block, - scan_length, - "layers", - mesh, - in_axes_tuple=in_axes_tuple, - model_mode=self.model_mode, - **layer_kwargs, - )(y, *broadcast_args) - - maxtext_utils.update_kv_caches_after_scan( - kv_caches, returned_kv_cache, scan_length, attention_pattern_length, stacked=True - ) - - # Apply any remaining layers that did not fit into a full scanned block - num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length - if num_remaining_layers > 0: - # We name the remainder block with a 'remainder' suffix to avoid parameter name collisions - rem_layer_kwargs = {"num_of_layers": num_remaining_layers} - layer = RemattedGemma3Block( - config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, name="layers_remainder", **rem_layer_kwargs - ) # pytype: disable=wrong-keyword-args - - remainder_kv = None - if kv_caches is not None: - start_idx = scan_length * attention_pattern_length - remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) - - y_and_kv = layer( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk=previous_chunk, - slot=slot, - bidirectional_mask=bidirectional_mask, - kv_cache=remainder_kv, - attention_metadata=attention_metadata, - ) - - if isinstance(y_and_kv, tuple): - y = y_and_kv[0] - updated_remainder_kv = y_and_kv[1] - else: - y = y_and_kv - updated_remainder_kv = None - - if kv_caches is not None and updated_remainder_kv is not None: - start_idx = scan_length * attention_pattern_length - for offset, updated_item in enumerate(updated_remainder_kv): - kv_caches[start_idx + offset] = updated_item - - return y - - def _apply_gemma4_scanned_blocks( - self, - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - bidirectional_mask, - previous_chunk, - slot, - kv_caches=None, - attention_metadata=None, - ): - """Applies Gemma4 scanned decoder blocks, handling main scan and remainders.""" - - cfg = self.config - mesh = self.mesh - - # Define the repeating pattern length and calculate how many full blocks to scan - block_pattern_len = len(gemma4.GEMMA4_ATTENTION_PATTERN) - num_full_blocks = cfg.num_decoder_layers // block_pattern_len - remainder_layers = cfg.num_decoder_layers % block_pattern_len - - if num_full_blocks > 0: - ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen - policy = self.get_remat_policy() - RemattedGemma4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0] - - kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( - kv_caches, num_full_blocks, block_pattern_len, stack=True - ) - - broadcast_args_spec = [ - (decoder_segment_ids, nn.broadcast), - (decoder_positions, nn.broadcast), - (deterministic, nn.broadcast), - (model_mode, nn.broadcast), - (slot, nn.broadcast), - (None, nn.broadcast), # page_state - (previous_chunk, nn.broadcast), - (bidirectional_mask, nn.broadcast), - (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), - (attention_metadata, nn.broadcast), - ] - broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) - in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) - - # For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks - y, returned_kv_cache = nn.scan( - RemattedGemma4Block, - variable_axes={ - "params": cfg.param_scan_axis, - "cache": 0, - "intermediates": 0, - "aqt": 0, - "_overwrite_with_gradient": 0, - }, - split_rngs={"params": True, "dropout": cfg.enable_dropout}, - in_axes=in_axes_tuple, - length=num_full_blocks, - metadata_params={ - nn.PARTITION_NAME: "layers", - "abstract_init": False, - }, - )( - config=cfg, - mesh=mesh, - quant=self.quant, - model_mode=model_mode, - num_of_layers=block_pattern_len, - name="scanned_blocks", - )( - y, *broadcast_args - ) - - maxtext_utils.update_kv_caches_after_scan( - kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True - ) - - # Process any remaining layers that don't fit into a full scanned block - for layer_id in range(cfg.num_decoder_layers - remainder_layers, cfg.num_decoder_layers): - attention_type = gemma4.get_attention_type(layer_id) - layer = gemma4.Gemma4DecoderLayerToLinen( - config=cfg, - mesh=mesh, - model_mode=model_mode, - quant=self.quant, - attention_type=attention_type, - layer_idx=layer_id, - ) - kv_cache = kv_caches[layer_id] if kv_caches is not None else None - - # inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, - # previous_chunk, page_state, slot, bidirectional_mask, kv_cache, attention_metadata - remainder_args = ( - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk, - None, # page_state - slot, - bidirectional_mask, - kv_cache, - attention_metadata, - ) - - y_and_kv = layer(y, *remainder_args) - if isinstance(y_and_kv, tuple): - y = y_and_kv[0] - new_kv = y_and_kv[1] - else: - y = y_and_kv - new_kv = None - - if kv_caches is not None and new_kv is not None: - kv_caches[layer_id] = new_kv - - return y - - def _apply_deepseek4_scanned_blocks( - self, - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk, - slot, - decoder_input_tokens, - ): - """Applies DeepSeek V4 scanned decoder blocks. - - DeepSeek V4 has some number of prefix layers (defined by `first_num_hash_layers`) - that use static Hash Routing. The remaining layers alternate `compress_ratio=128` (HCA) - and `compress_ratio=4` (CSA) and are evaluated in a single `nn.scan` block. - - For DeepSeek4-Flash (43 hidden layers total): - - 3 Prefix layers (Indices 0, 1, 2) - - 40 Scanned layers: 20 perfectly repeating chunks of [128, 4] - """ - - cfg = self.config - mesh = self.mesh - - broadcast_args = ( - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - slot, - previous_chunk, - ) - - layer_call_kwargs = { - "previous_chunk": previous_chunk, - "slot": slot, - "decoder_input_tokens": decoder_input_tokens, - } - - # 1. Prefix Unrolling - # Prefix layers are unrolled (unscanned) for two architectural reasons: - # 1. Heterogeneous Attention: JAX nn.scan requires identical computation graphs, but the first few layers - # use different attention configurations (e.g., DeepSeek-V4 uses compress_ratios [0, 0, 4] for layers 0, 1, 2). - # 2. Static Hash Routing: The first `first_num_hash_layers` (which is 3 for DeepSeek-V4) use deterministic - # token-to-expert Hash Routing instead of learned top-k routing. - # Therefore, these prefix layers are instantiated individually before we scan the remaining uniform blocks. - num_hash_layers = cfg.first_num_hash_layers - for layer_idx in range(num_hash_layers): - prefix_layer = deepseek4.DeepSeek4LayerToLinen( - config=cfg, - mesh=mesh, - name=f"layers_{layer_idx}", - quant=self.quant, - model_mode=self.model_mode, - layer_idx=layer_idx, - ) - y, _ = prefix_layer( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - **layer_call_kwargs, - ) - - # 2. Chunked Scanning - # The remaining layers perfectly alternate HCA (128) and CSA (4). - num_remaining_layers = cfg.num_decoder_layers - num_hash_layers - num_full_blocks = num_remaining_layers // 2 - - if num_full_blocks > 0: - ScannableBlockToLinen = deepseek4.DeepSeek4ScannableBlockToLinen - policy = self.get_remat_policy() - RemattedDeepSeek4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0] - - y, _ = nn.scan( - RemattedDeepSeek4Block, - variable_axes={ - "params": cfg.param_scan_axis, - "cache": 0, - "intermediates": 0, - "aqt": 0, - "_overwrite_with_gradient": 0, - }, - split_rngs={"params": True, "dropout": cfg.enable_dropout}, - in_axes=(nn.broadcast,) * len(broadcast_args), - length=num_full_blocks, - metadata_params={ - nn.PARTITION_NAME: "layers", - "abstract_init": False, - }, - )(config=cfg, mesh=mesh, quant=self.quant, model_mode=model_mode, name="scanned_blocks",)(y, *broadcast_args) - - return y - - def _apply_gemma4_small_layers( - self, - y, - decoder_input_tokens, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - multimodal_input=None, - kv_caches=None, - attention_metadata=None, - previous_chunk=None, - slot=None, - ): - """Apply Gemma 4 small (E2B / E4B) decoder layers. - - Threads per-call state through the layer loop: - * ``per_layer_inputs`` from PLE, sliced per layer. - * ``shared_kv_states``: donor-layer-index → (key, value) for - downstream KV-shared layers to consume. - * ``kv_caches``: when running via the vLLM RPA path, the per-layer - cache buffer threaded back from the kernel. KV-shared layers - redirect to the donor's cache slot via ``cache_index_of``. - - Returns ``(y, kv_caches)``. Scan-over-layers and pipeline - parallelism are not supported. - """ - cfg = self.config - mesh = self.mesh - bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None - - per_layer_inputs = None - if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0: - per_layer_inputs = gemma4_small.PLEToLinen( - config=cfg, - mesh=mesh, - name="per_layer_embedder", - )(decoder_input_tokens, y) - - layer_types = gemma4_small.build_layer_types(cfg.num_decoder_layers, cfg.model_name) - num_kv_shared = cfg.num_kv_shared_layers - shared_kv_states: dict[int, tuple[jax.Array, jax.Array]] = {} - - # tpu-inference allocates one `kv_caches` slot per non-shared layer; - # KV-shared layers reuse the donor's slot. - cache_index_of = gemma4_small.kv_cache_slot_map(layer_types, num_kv_shared) - - for lyr in range(cfg.num_decoder_layers): - attention_type = layer_types[lyr] - donor_idx = gemma4_small.kv_donor_layer_idx(lyr, layer_types, num_kv_shared) - is_donor = gemma4_small.is_kv_donor_layer(lyr, layer_types, num_kv_shared) - - shared_key = None - shared_value = None - if donor_idx is not None: - if donor_idx not in shared_kv_states: - raise RuntimeError( - f"KV-shared layer {lyr} references donor {donor_idx} but no donor K/V " - f"have been recorded yet. This indicates the layer iteration order is wrong." - ) - shared_key, shared_value = shared_kv_states[donor_idx] - - layer = gemma4_small.Gemma4SmallDecoderLayerToLinen( - config=cfg, - mesh=mesh, - name=f"layers_{lyr}", - quant=self.quant, - model_mode=self.model_mode, - attention_type=attention_type, - layer_idx=lyr, - ) - - # Donor layers expose their rotated, normed K / V to downstream - # shared layers via the decoder layer's compute_shared_kv method. - if is_donor: - donor_k, donor_v = layer(y, decoder_positions, nnx_method="compute_shared_kv") - shared_kv_states[lyr] = (donor_k, donor_v) - # Reuse the just-computed K / V in the layer's own forward pass to - # avoid double-computing the K / V projection / norm / RoPE. - shared_key, shared_value = donor_k, donor_v - - ple_slice = per_layer_inputs[..., lyr, :] if per_layer_inputs is not None else None - - cache_idx = cache_index_of[lyr] - kv_cache = kv_caches[cache_idx] if kv_caches is not None else None - y, kv_cache = layer( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk=previous_chunk, - slot=slot, - bidirectional_mask=bidirectional_mask_value, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - per_layer_input=ple_slice, - shared_key=shared_key, - shared_value=shared_value, - ) - if kv_caches is not None and kv_cache is not None: - kv_caches[cache_idx] = kv_cache - - return y, kv_caches - - # TODO(b/490118813): Relocate the following functions to their designated directories - # once the plug-in strategy is implemented: _find_next_boundary(), _apply_single_engram_layer() - # _apply_scanned_chunk() and _apply_interleaved_scanned_layers(). - def _find_next_boundary(self, current_idx, end_idx, engram_indices): - """Finds the next index boundary, either the next Engram layer index or the overall end index.""" - next_engrams = [l for l in engram_indices if l > current_idx] - if next_engrams: - return min(end_idx, *next_engrams) - return end_idx - - def _apply_single_engram_layer(self, y, current_idx, layer_type, **kwargs): - """Applies a single, unscanned Engram layer.""" - layer = kwargs["dense_layer"] if layer_type == "dense" else kwargs["moe_layer"] - layer_prefix = "dense_layers" if layer_type == "dense" else "moe_layers" - original_call = kwargs["original_dense_call"] if layer_type == "dense" else kwargs["original_moe_call"] - layer_call_kwargs = kwargs["layer_call_kwargs"] - - layer.__call__ = original_call - y, _ = layer( - config=self.config, - mesh=self.mesh, - name=f"{layer_prefix}_engram_{current_idx}", - quant=self.quant, - model_mode=self.model_mode, - layer_idx=current_idx, - )( - y, - kwargs["decoder_segment_ids"], - kwargs["decoder_positions"], - kwargs["deterministic"], - kwargs["model_mode"], - decoder_input_tokens=kwargs["decoder_input_tokens"], - **layer_call_kwargs, - ) - layer.__call__ = functools.partial(original_call, **layer_call_kwargs) - return y - - def _apply_scanned_chunk(self, y, current_idx, next_boundary, layer_type, **kwargs): - """Applies a contiguous chunk of layers using the scan operation.""" - layer = kwargs["dense_layer"] if layer_type == "dense" else kwargs["moe_layer"] - layer_prefix = "dense_layers" if layer_type == "dense" else "moe_layers" - broadcast_args = kwargs["broadcast_args"] - scan_length = next_boundary - current_idx - - if scan_length > 0: - y, _ = self.scan_decoder_layers( - self.config, - layer, - scan_length, - f"{layer_prefix}_{current_idx}_{next_boundary - 1}", - self.mesh, - in_axes_tuple=(nn.broadcast,) * len(broadcast_args), - model_mode=kwargs["model_mode"], - )(y, *broadcast_args) - return y - - def _apply_interleaved_scanned_layers(self, y, layer_type, start_idx, end_idx, engram_indices, **kwargs): - """Applies a mix of scanned standard layers and unscanned Engram layers.""" - current_idx = start_idx - while current_idx < end_idx: - if current_idx in engram_indices: - # Handle individual unscanned Engram layer - y = self._apply_single_engram_layer(y, current_idx, layer_type, **kwargs) - current_idx += 1 - else: - # Find next boundary and scan the chunk - next_boundary = self._find_next_boundary(current_idx, end_idx, engram_indices) - y = self._apply_scanned_chunk(y, current_idx, next_boundary, layer_type, **kwargs) - current_idx = next_boundary - return y diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index ad6b171f2f..45e9292afb 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -26,8 +26,7 @@ from flax import nnx from maxtext.common.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType -from maxtext.layers import nnx_wrappers -from maxtext.layers.initializers import Initializer, default_embed_init, variable_to_logically_partitioned +from maxtext.layers.initializers import Initializer, default_embed_init from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils.sharding import logical_to_mesh_axes, create_sharding @@ -43,52 +42,6 @@ def _maybe_move_embedding_to_device(embedding_table: Array, config: Config) -> A return embedding_table -def embed_as_linen( - *, - num_embeddings: int, - num_features: int, - config: Config, - mesh: Mesh, - cast_input_dtype: None | DType = None, - dtype: DType = jnp.float32, - attend_dtype: None | DType = None, - embedding_init: Initializer = default_embed_init, - name: str | None = None, -): - """Initializes the Embed NNX module and returns it as a Linen module. - - This function serves as a bridge to use the NNX-based `Embed` module within - a Linen model. It wraps the `Embed` module using `nnx.bridge.to_linen`, - making it compatible with the Linen API. - - Args: - num_embeddings: The number of embeddings. - num_features: The number of feature dimensions for each embedding. - config: The model configuration. - cast_input_dtype: The dtype to cast the input to, if any. - dtype: The dtype of the embedding vectors. - attend_dtype: The dtype for the `attend` method. - embedding_init: The initializer for the embedding matrix. - name: The name of the Linen module. - - Returns: - A Linen module that wraps the NNX `Embed` module. - """ - return nnx_wrappers.to_linen( - Embed, - num_embeddings=num_embeddings, - num_features=num_features, - config=config, - mesh=mesh, - cast_input_dtype=cast_input_dtype, - dtype=dtype, - attend_dtype=attend_dtype, - embedding_init=embedding_init, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) - - class Embed(nnx.Module): """A parameterized function from integers [0, n) to d-dimensional vectors.""" @@ -236,39 +189,6 @@ def attend_on_embedding( ) -def rotary_embedding_as_linen( - *, - min_timescale: int, - max_timescale: int, - embedding_dims: int = 0, - cast_as_fprop_dtype: bool = True, - fprop_dtype: DType = jnp.bfloat16, - name: str | None = None, -): - """Initializes the RotaryEmbedding module and returns it as a Linen module. - - Args: - min_timescale: Start of the geometric index. Determines the periodicity of - the added signal. - max_timescale: End of the geometric index. Determines the frequency of the - added signal. - embedding_dims: Dimension of the embedding to be generated. - cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. - fprop_dtype: The dtype of the output. - name: Name of the Linen module. - """ - return nnx_wrappers.to_linen( - RotaryEmbedding, - min_timescale=min_timescale, - max_timescale=max_timescale, - embedding_dims=embedding_dims, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) - - class RotaryEmbedding(nnx.Module): """Rotary Position Embedding.""" @@ -371,82 +291,6 @@ def __call__( return x_out -def llama_rotary_embedding_as_linen( - *, - min_timescale: int, - max_timescale: int, - embedding_dims: int = 0, - cast_as_fprop_dtype: bool = True, - fprop_dtype: DType = jnp.bfloat16, - use_scale: bool = True, - name: str | None = None, -): - """Initializes the LLaMARotaryEmbedding module and returns it as a Linen module. - - Args: - min_timescale: Start of the geometric index. Determines the periodicity of - the added signal. - max_timescale: End of the geometric index. Determines the frequency of the - added signal. - embedding_dims: Dimension of the embedding to be generated. - cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. - fprop_dtype: The dtype of the output. - use_scale: Whether to apply LLaMA3.1 scaling factor. - name: Name of the Linen module. - """ - return nnx_wrappers.to_linen( - LLaMARotaryEmbedding, - min_timescale=min_timescale, - max_timescale=max_timescale, - embedding_dims=embedding_dims, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - use_scale=use_scale, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) - - -def partial_rotary_embedding_as_linen( - *, - min_timescale: int, - max_timescale: int, - mesh: Mesh, - embedding_dims: int = 0, - partial_rotary_factor: float = 0.25, - cast_as_fprop_dtype: bool = True, - fprop_dtype: DType = jnp.bfloat16, - shard_mode: ShardMode = ShardMode.AUTO, - name: str | None = None, -): - """Initializes the PartialRotaryEmbedding module and returns it as a Linen module. - - Args: - min_timescale: Start of the geometric index. Determines the periodicity of - the added signal. - max_timescale: End of the geometric index. Determines the frequency of the - added signal. - embedding_dims: Dimension of the embedding to be generated. - partial_rotary_factor: Ratio of dimensions to apply ROPE to. - cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. - fprop_dtype: The dtype of the output. - name: Name of the Linen module. - """ - return nnx_wrappers.to_linen( - PartialRotaryEmbedding, - min_timescale=min_timescale, - max_timescale=max_timescale, - mesh=mesh, - embedding_dims=embedding_dims, - partial_rotary_factor=partial_rotary_factor, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - shard_mode=shard_mode, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) - - class PartialRotaryEmbedding(RotaryEmbedding): """Rotary Position Embedding applied to a partial fraction of dimensions.""" @@ -723,59 +567,6 @@ def __call__(self, inputs: jax.Array, position: None | jax.Array = None) -> jax. return outputs -def yarn_rotary_embedding_as_linen( - *, - embedding_dims: int, - mesh: Mesh, - max_position_embeddings: int = 4096 * 4, - original_max_position_embeddings: int = 4096, - beta_fast: float = 32, - beta_slow: float = 1, - rope_theta: float = 10000.0, - rope_factor: float = 40, - cast_as_fprop_dtype: bool = True, - fprop_dtype: DType = jnp.bfloat16, - name: str | None = None, - interleave: bool = True, - truncate: bool = True, - attention_scaling: bool = False, - shard_mode: ShardMode = ShardMode.AUTO, -): - """Initializes the YarnRotaryEmbedding module and returns it as a Linen module. - - Args: - embedding_dims: The dimension of the embeddings. - max_position_embeddings: The maximum number of positions. - original_max_position_embeddings: The original maximum number of positions. - beta_fast: The fast beta parameter for YaRN. - beta_slow: The slow beta parameter for YaRN. - rope_theta: The base for the rotary frequencies. - rope_factor: The scaling factor for RoPE. - cast_as_fprop_dtype: Whether to cast the output to `fprop_dtype`. - fprop_dtype: The forward pass dtype. - name: The name of the module. - """ - return nnx_wrappers.to_linen( - YarnRotaryEmbedding, - embedding_dims=embedding_dims, - max_position_embeddings=max_position_embeddings, - mesh=mesh, - original_max_position_embeddings=original_max_position_embeddings, - beta_fast=beta_fast, - beta_slow=beta_slow, - rope_theta=rope_theta, - rope_factor=rope_factor, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - metadata_fn=variable_to_logically_partitioned, - name=name, - interleave=interleave, - truncate=truncate, - attention_scaling=attention_scaling, - shard_mode=shard_mode, - ) - - class YarnRotaryEmbedding(nnx.Module): """Yarn rotary embedding. @@ -998,31 +789,6 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array: return output -def positional_embedding_as_linen( - *, - embedding_dims: int, - max_wavelength: int = _MAX_WAVELENGTH, - cast_as_fprop_dtype: bool = False, - fprop_dtype: DType = jnp.bfloat16, -): - """Initializes the PositionalEmbedding module and returns it as a Linen module. - - Args: - embedding_dims: The dimension of the embeddings. - max_wavelength: The maximum wavelength for the sinusoidal positional embeddings. - cast_as_fprop_dtype: Whether to cast output to fprop_dtype. - fprop_dtype: The dtype of the output when cast_as_fprop_dtype is True. - """ - return nnx_wrappers.to_linen( - PositionalEmbedding, - embedding_dims=embedding_dims, - max_wavelength=max_wavelength, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - metadata_fn=variable_to_logically_partitioned, - ) - - @dataclasses.dataclass(repr=False) class PositionalEmbedding(nnx.Module): """Sinusoidal positional embeddings supporting both uniform and per-batch positions. @@ -1108,43 +874,6 @@ def __call__( return self._compute_embeddings(position) -def llama_vision_rotary_embedding_as_linen( - *, - image_size: int, - patch_size: int, - hidden_size: int, - num_attention_heads: int, - rope_theta: float = 10000.0, - cast_as_fprop_dtype: bool = True, - fprop_dtype: DType = jnp.bfloat16, - name: str | None = None, -): - """Initializes the LlamaVisionRotaryEmbedding module and returns it as a Linen module. - - Args: - image_size: The size of the input image. - patch_size: The size of the image patches. - hidden_size: The size of the hidden dimension. - num_attention_heads: The number of attention heads. - rope_theta: The base theta value for the frequency computation. - cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. - fprop_dtype: The dtype of the output. - name: The name of the Linen module. - """ - return nnx_wrappers.to_linen( - LlamaVisionRotaryEmbedding, - image_size=image_size, - patch_size=patch_size, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - rope_theta=rope_theta, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) - - @dataclasses.dataclass(repr=False) class LlamaVisionRotaryEmbedding(nnx.Module): """Rotary position embedding for Llama4 vision encoder. @@ -1422,47 +1151,6 @@ def __call__(self, inputs: Array, num_frames: int, height: int, width: int) -> A return rotated -def qwen3omnimoe_vision_pos_embed_interpolate_as_linen( - *, - num_position_embeddings: int, - hidden_size: int, - spatial_merge_size: int, - dtype: DType = jnp.float32, - cast_as_fprop_dtype: bool = True, - fprop_dtype: DType = jnp.bfloat16, - name: str | None = None, -): - """Initializes Qwen3OmniMoe bilinear position embedding interpolation as Linen module. - - This implements fast bilinear interpolation of learned 2D positional embeddings - for dynamic input sizes. The embeddings are learned on a fixed grid and interpolated - to match the actual image/video dimensions. - - Args: - num_position_embeddings: Number of position embeddings in the fixed grid (e.g., 1024 for 32x32) - hidden_size: Hidden dimension size - spatial_merge_size: Size of spatial merging blocks - dtype: Data type for embeddings - cast_as_fprop_dtype: Whether to cast the output to the fprop dtype - fprop_dtype: The dtype of the output - name: Module name - - Returns: - A Linen module that wraps the NNX Qwen3OmniMoeVisionPosEmbedInterpolate module. - """ - return nnx_wrappers.to_linen( - Qwen3OmniMoeVisionPosEmbedInterpolate, - num_position_embeddings=num_position_embeddings, - hidden_size=hidden_size, - spatial_merge_size=spatial_merge_size, - dtype=dtype, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) - - class Qwen3OmniMoeVisionPosEmbedInterpolate(nnx.Module): """Bilinear interpolation of learned 2D positional embeddings for Qwen3OmniMoe vision. @@ -1769,40 +1457,6 @@ def __call__( return x_out -def qwen3_omni_mrope_embedding_as_linen( - *, - min_timescale: int, - max_timescale: int, - embedding_dims: int = 0, - cast_as_fprop_dtype: bool = True, - fprop_dtype: DType = jnp.bfloat16, - mrope_section: tuple[int, int, int] | None = None, - name: str | None = None, -): - """Initializes Qwen3OmniMoeThinkerTextRotaryEmbedding and returns it as a Linen module. - - Args: - min_timescale: Start of the geometric index. - max_timescale: End of the geometric index (rope_theta). - embedding_dims: Dimension of the embedding (head_dim). - cast_as_fprop_dtype: Whether to cast output to fprop dtype. - fprop_dtype: The dtype of the output. - mrope_section: Tuple of (temporal_dim, height_dim, width_dim) for MRoPE. - name: Name of the Linen module. - """ - return nnx_wrappers.to_linen( - Qwen3OmniMoeThinkerTextRotaryEmbedding, - min_timescale=min_timescale, - max_timescale=max_timescale, - embedding_dims=embedding_dims, - cast_as_fprop_dtype=cast_as_fprop_dtype, - fprop_dtype=fprop_dtype, - mrope_section=mrope_section, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) - - class DeepSeekV4RotaryEmbedding(RotaryEmbedding): """DeepSeek-V4 partial rotary embedding with interleaved frequencies. diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py index 19a7490016..9666c072fa 100644 --- a/src/maxtext/layers/encoders.py +++ b/src/maxtext/layers/encoders.py @@ -19,8 +19,6 @@ from jax.sharding import Mesh from maxtext.common.common_types import Config -from maxtext.layers import nnx_wrappers -from maxtext.layers import initializers class VisionEncoder(nnx.Module): @@ -151,35 +149,3 @@ def __call__(self, input_audio, deterministic=False): embeddings = projector(embeddings) return embeddings - - -def vision_encoder_as_linen( - config: Config, - mesh: Mesh, -): - """Creates a VisionEncoder module.""" - module = nnx_wrappers.to_linen( - VisionEncoder, - config=config, - mesh=mesh, - name="vision_encoder", - abstract_init=False, - metadata_fn=initializers.variable_to_logically_partitioned, - ) - return module - - -def audio_encoder_as_linen( - config: Config, - mesh: Mesh, -): - """Creates an AudioEncoder module.""" - module = nnx_wrappers.to_linen( - AudioEncoder, - config=config, - mesh=mesh, - name="audio_encoder", - abstract_init=False, - metadata_fn=initializers.variable_to_logically_partitioned, - ) - return module diff --git a/src/maxtext/layers/multi_token_prediction.py b/src/maxtext/layers/multi_token_prediction.py index 6f9af709fd..8d02a43112 100644 --- a/src/maxtext/layers/multi_token_prediction.py +++ b/src/maxtext/layers/multi_token_prediction.py @@ -16,14 +16,11 @@ from typing import Type -from flax import linen as nn from flax import nnx import jax import jax.numpy as jnp from jax.sharding import Mesh from maxtext.common.common_types import Config, DecoderBlockType, MODEL_MODE_TRAIN, ShardMode -from maxtext.layers.decoders import DecoderLayer -from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral from maxtext.layers.nnx_decoders import NNXDecoderLayer from maxtext.layers.normalizations import RMSNorm @@ -385,37 +382,3 @@ def calculate_mtp_acceptance_rate(intermediate_outputs, config): total_valid_tokens = jnp.sum(valid_mask) return (correct_predictions / (total_valid_tokens + EPS)) * 100 - - -def multi_token_prediction_block_as_linen( - *, - config: Config, - mesh: Mesh, - transformer_layer_module: Type[DecoderLayer], - decoder: nnx.Module, - rngs: nnx.Rngs, - name: str | None = None, -) -> nn.Module: - """Initializes MultiTokenPredictionBlock as a Linen module. - - Args: - config: Configuration object containing model hyperparameters. - mesh: JAX Mesh for model parallelism. - transformer_layer_module: The Transformer Decoder Layer class to use. - decoder: The decoder module that provides embedding and output head. - rngs: Random number generators for initialization. - name: Optional name for the module. - - Returns: - An instance of MultiTokenPredictionBlock wrapped as a Linen module. - """ - return nnx.bridge.to_linen( - MultiTokenPredictionBlock, - config=config, - mesh=mesh, - transformer_layer_module=transformer_layer_module, - decoder=decoder, - rngs=rngs, - metadata_fn=variable_to_logically_partitioned, - name=name, - ) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1b14803d94..5cb113dc1a 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -36,7 +36,7 @@ MultimodalInput, ShardMode, ) -from maxtext.layers import initializers, linears, mhc, normalizations, quantizations +from maxtext.layers import linears, mhc, normalizations, quantizations from maxtext.layers import nnx_wrappers from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding @@ -2150,25 +2150,3 @@ def _apply_gemma4_small_layers( kv_caches[cache_idx] = kv_cache return y, kv_caches - - -def decoder_as_linen( - config: Config, - mesh: Mesh, - rngs: nnx.Rngs, - model_mode: str, - quant: None | Quant = None, -): - """Creates a Decoder module""" - module = nnx_wrappers.to_linen( - NNXDecoder, - config=config, - mesh=mesh, - model_mode=model_mode, - rngs=rngs, - quant=quant, - name="decoder", - abstract_init=False, - metadata_fn=initializers.variable_to_logically_partitioned, - ) - return module diff --git a/src/maxtext/models/gemma3.py b/src/maxtext/models/gemma3.py index c33272d277..d7b75a45e4 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -30,7 +30,6 @@ from maxtext.layers.linears import DenseGeneral, MlpBlock, Dropout from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant -from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.utils import max_utils @@ -574,21 +573,6 @@ def __call__(self, x: jax.Array, eqn: str = "...tm,md->...td") -> jax.Array: return x -def visionembedder_as_linen( - config: Config, - mesh: Mesh, -): - """Creates a VisionEmbedder module.""" - return nnx_wrappers.to_linen( - VisionEmbedder, - config, - mesh=mesh, - name="VisionEmbedder_0", - abstract_init=False, - metadata_fn=variable_to_logically_partitioned, - ) - - class VisionExit(nnx.Module): """The vision exit layer. @@ -621,11 +605,6 @@ def __call__(self, x): return jnp.reshape(x, (batch_size, height * width, embed_dim)) -def vision_exit_as_linen(x: jax.Array, output_length: int) -> jax.Array: - """A wrapper to use VisionExit as a function.""" - return nnx.bridge.to_linen(VisionExit, output_length=output_length)(x) - - class Gemma3VisionEncoderLayer(nnx.Module): """gemma 3 vision encoder layer""" @@ -712,19 +691,3 @@ def __call__(self, inputs, deterministic, train=False): bn, l, c = x.shape x = jnp.reshape(x, [b, n, l, c]) return x - - -def gemma3visionencoder_as_linen( - config: Config, - mesh: Mesh, -): - """Creates a Gemma3VisionEncoder module.""" - module = nnx_wrappers.to_linen( - Gemma3VisionEncoderLayer, - config=config, - mesh=mesh, - name="Gemma3VisionEncoderLayer_0", - abstract_init=False, - metadata_fn=variable_to_logically_partitioned, - ) - return module diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py index 72cd841f81..f673488608 100644 --- a/src/maxtext/models/gemma4_vision.py +++ b/src/maxtext/models/gemma4_vision.py @@ -18,15 +18,12 @@ from typing import cast import jax import jax.numpy as jnp -from flax import linen as nn from flax import nnx from jax.sharding import Mesh from maxtext.common.common_types import Config, AttentionType from maxtext.layers import attentions -from maxtext.layers import initializers from maxtext.layers import linears -from maxtext.layers import nnx_wrappers from maxtext.layers import normalizations @@ -636,15 +633,3 @@ def __call__(self, x: jax.Array) -> jax.Array: x_normed = self.norm(x) x_projected = self.projection(x_normed) return x_projected - - -def gemma4_vision_encoder_as_linen(config: Config, mesh: Mesh) -> nn.Module: - """Wraps the Gemma 4 Vision Encoder as a Linen module.""" - return nnx_wrappers.to_linen( - Gemma4VisionEncoderLayer, - config=config, - mesh=mesh, - name="Gemma4VisionEncoderLayer", - abstract_init=False, - metadata_fn=initializers.variable_to_logically_partitioned, - ) diff --git a/src/maxtext/models/llama4.py b/src/maxtext/models/llama4.py index 26fd4d322d..da5d727866 100644 --- a/src/maxtext/models/llama4.py +++ b/src/maxtext/models/llama4.py @@ -253,17 +253,6 @@ def __call__(self, image_features: Array) -> Array: return hidden_states -def llama4multimodalprojector_as_linen(config: Config, mesh: Mesh): - return nnx_wrappers.to_linen( - Llama4MultiModalProjector, - config=config, - mesh=mesh, - name="Llama4MultiModalProjector_0", - abstract_init=False, - metadata_fn=initializers.variable_to_logically_partitioned, - ) - - def determine_is_nope_layer(layer_id: int, nope_layer_interval: int) -> bool: """ Determines whether the given layer at `layer_id` should use RoPE or not (NoPE). @@ -812,14 +801,3 @@ def __call__( hidden_states = jnp.reshape(hidden_states, [b, t, patch_num, patch_dim]) return hidden_states - - -def llama4visionmodel_as_linen(config: Config, mesh: Mesh) -> nn.Module: - return nnx_wrappers.to_linen( - Llama4VisionModel, - config=config, - mesh=mesh, - name="Llama4VisionModel_0", - abstract_init=False, - metadata_fn=initializers.variable_to_logically_partitioned, - ) diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index ac908c0f96..9c59c6b80a 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -29,230 +29,17 @@ from maxtext.layers.nnx_decoders import NNXDecoder from maxtext.layers import initializers from maxtext.layers import nnx_wrappers -from maxtext.layers.decoders import Decoder -from maxtext.layers.embeddings import Embed, embed_as_linen -from maxtext.layers.encoders import AudioEncoder, VisionEncoder, audio_encoder_as_linen, vision_encoder_as_linen -from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock, multi_token_prediction_block_as_linen +from maxtext.layers.embeddings import Embed +from maxtext.layers.encoders import AudioEncoder, VisionEncoder +from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.multimodal import processor as mm_processor -from maxtext.utils import max_utils # ------------------------------------------------------------------------------ # The network: Transformer Definitions # ------------------------------------------------------------------------------ -class TransformerLinenPure(nn.Module): - """An autoregressive transformer model.""" - - # Make new attributes required, so that all Transformer dependencies (train, decode, - # compile, etc) will error instead of silently use defaults. - # pylint: disable=attribute-defined-outside-init - config: Config - mesh: Mesh - quant: Quant - # Possible model_mode values can be found in maxtext.common.common_types. - # We generally use maxtext.common.common_types.MODEL_MODE_TRAIN or - # maxtext.common.common_types.MODEL_MODE_PREFILL for initializations here. - # TODO: Make model_mode required after confirming no users are affected. - model_mode: str = MODEL_MODE_TRAIN # May be different than the model_mode passed to __call__ - # pylint: enable=attribute-defined-outside-init - - def init(self, *args, model_mode: str = MODEL_MODE_TRAIN, **kwargs): - """Initializes the model.""" - module = self.clone(model_mode=model_mode) - kwargs["model_mode"] = model_mode - return nn.Module.init(module, *args, **kwargs) - - def apply(self, *args, model_mode: str = MODEL_MODE_TRAIN, **kwargs): - """Applies the model.""" - module = self.clone(model_mode=model_mode) - kwargs["model_mode"] = model_mode - return nn.Module.apply(module, *args, **kwargs) - - def setup(self): - """Initialize shared_embedding & decoder layers.""" - - cfg = self.config - mesh = self.mesh - self.shared_embedding = embed_as_linen( - num_embeddings=cfg.vocab_size, - num_features=cfg.emb_dim, - dtype=cfg.dtype, - attend_dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype, # for logit training stability - embedding_init=nn.initializers.normal(stddev=1.0), - name="token_embedder", - config=cfg, - mesh=self.mesh, - ) - self.vision_encoder = vision_encoder_as_linen(config=cfg, mesh=mesh) if cfg.use_multimodal else None - self.audio_encoder = audio_encoder_as_linen(config=cfg, mesh=mesh) if cfg.use_audio else None - self.decoder = Decoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode) - - # If MTP is enabled via config, set up the MTP block. - if self.config.mtp_num_layers > 0: - # Get the list of layer blueprints for the current model. - # For MTP, we use the DecoderLayer blueprint to ensure architectural consistency. - # By convention, this is the last layer in the list. - layer_types = self.decoder.get_decoder_layers() - mtp_layer_linen = layer_types[-1] - # UNWRAP: The MTP block is pure NNX. If the decoder returned a Linen wrapper, - # extract the native NNX class to preserve parameter tracing/scoping. - mtp_layer_nnx = getattr(mtp_layer_linen, "module_class", mtp_layer_linen) - self.mtp_block = multi_token_prediction_block_as_linen( - config=self.config, - mesh=self.mesh, - transformer_layer_module=mtp_layer_nnx, - decoder=self.decoder, - rngs=self.make_rng("mtp_block"), - ) - - def logits_from_hidden_states_for_vocab_tiling(self, hidden_states, deterministic, model_mode): - """ - Compute logits from hidden states (wrapping decoder.apply_output_head). - This function is only used for vocabulary tiling. - """ - logits = self.decoder.apply_output_head( - shared_embedding=self.shared_embedding, - y=hidden_states, - deterministic=deterministic, - model_mode=model_mode, - ) - return logits - - def __call__( - self, - decoder_input_tokens: jnp.ndarray, - decoder_positions: jnp.ndarray, - decoder_segment_ids=None, - encoder_images: None | jnp.ndarray = None, - encoder_image_masks: None | jnp.ndarray = None, - encoder_videos: None | jnp.ndarray = None, - encoder_video_masks: None | jnp.ndarray = None, - encoder_audios: None | jnp.ndarray = None, - enable_dropout=True, - model_mode=MODEL_MODE_TRAIN, - previous_chunk=None, - true_length: None | int = None, - slot: None | int = None, - decoder_target_tokens: None | jnp.ndarray = None, - decoder_target_mask: None | jnp.ndarray = None, - nnx_method=None, - kv_caches: list[jax.Array] | None = None, - attention_metadata: dict[str, Any] | None = None, - ): - """Applies Transformer decoder-branch on encoded-input and target. - - Args: - true_length: (Optional) Prompt length before padding - slot: (Optional) An integer representing the decode batch index selected - for this request. - """ - - if decoder_segment_ids is not None and model_mode == MODEL_MODE_AUTOREGRESSIVE: - raise ValueError( - f"During autoregressive decoding we assume the tokens are in the active sequence" - f" which is always {DECODING_ACTIVE_SEQUENCE_INDICATOR}." - ) - - bidirectional_mask_image = None - bidirectional_mask_video = None - image_embeddings = None - video_embeddings = None - audio_embeddings = None - deepstack_visual_embeds = None - - if self.config.use_multimodal and encoder_images is not None: - image_embeddings, deepstack_visual_embeds = self.vision_encoder( - input_images=encoder_images, deterministic=not enable_dropout - ) - bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision( - self.config, decoder_input_tokens, is_video=False - ) - - if self.config.use_multimodal and encoder_videos is not None: - video_embeddings, deepstack_visual_embeds = self.vision_encoder( - input_images=encoder_videos, deterministic=not enable_dropout - ) - bidirectional_mask_video = mm_processor.get_bidirectional_mask_vision( - self.config, decoder_input_tokens, is_video=True - ) - - if self.config.use_multimodal and encoder_audios is not None and self.audio_encoder is not None: - audio_embeddings = self.audio_encoder(input_audio=encoder_audios, deterministic=not enable_dropout) - - # Create audio mask for placeholder tokens (qwen3-omni models) - audio_masks = None - if audio_embeddings is not None: - audio_masks = mm_processor.get_bidirectional_mask_audio(self.config, decoder_input_tokens) - - multimodal_input = None - if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None: - multimodal_input = MultimodalInput( - image_embeddings=image_embeddings, - image_masks=encoder_image_masks, - video_embeddings=video_embeddings, - video_masks=encoder_video_masks, - audio_embeddings=audio_embeddings, - audio_masks=audio_masks, - bidirectional_mask=bidirectional_mask_image, - bidirectional_mask_video=bidirectional_mask_video, - ) - - logits, hidden_state, kv_caches = self.decoder( - shared_embedding=self.shared_embedding, - decoder_input_tokens=decoder_input_tokens, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=not enable_dropout, - model_mode=model_mode, - previous_chunk=previous_chunk, - slot=slot, - multimodal_input=multimodal_input, - kv_caches=kv_caches, - attention_metadata=attention_metadata, - deepstack_visual_embeds=deepstack_visual_embeds, - ) # pytype: disable=wrong-keyword-args - - # If we are initializing the model AND MTP is enabled, we must create - # dummy target tensors. This allows Flax to trace the MTPBlock and create - # all its necessary parameters, without requiring the main training pipeline - # to be aware of this initialization detail. - if self.is_initializing() and self.config.mtp_num_layers > 0: - if decoder_target_tokens is None: - dummy_shape = decoder_input_tokens.shape - decoder_target_tokens = jnp.ones(dummy_shape, dtype=jnp.int32) - decoder_target_mask = jnp.ones(dummy_shape, dtype=jnp.int32) - decoder_segment_ids = jnp.ones(dummy_shape, dtype=jnp.int32) - - # The Multi-Token Prediction (MTP) block functions as a "side-car" to the main - # model, active only during training. It computes an auxiliary loss based on - # predicting multiple future tokens, as described in the DeepSeek-V3 paper. - # To ensure architectural consistency, it uses two key components from the parent Transformer: - # 1. The same `DecoderLayer` blueprint for its internal transformer blocks. - # 2. The `shared_embedding` for both embedding future tokens and for its final - # logit projection. - # Its only effect is to "sow" these losses; it does not alter the primary logits output. - if self.config.mtp_num_layers > 0: - self.mtp_block( - shared_embedding=self.shared_embedding, - main_hidden_state=hidden_state, - input_ids=decoder_input_tokens, - target_ids=decoder_target_tokens, - target_mask=decoder_target_mask, - position_ids=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=not enable_dropout, - model_mode=model_mode, - ) - - if self.config.attention == "vllm_rpa": - # In vLLM, logits are computed separately after updating the KV cache. - return hidden_state, kv_caches - - return logits - - def transformer_as_linen( config: Config, mesh: Mesh, @@ -260,16 +47,13 @@ def transformer_as_linen( model_mode: str = MODEL_MODE_TRAIN, *, name: str | None = None, -) -> nnx_wrappers.ToLinen | TransformerLinenPure: - """Constructs a Transformer model as a Linen or NNX module. +) -> nnx_wrappers.ToLinen: + """Constructs an NNX Transformer wrapped as a Linen module. - This function returns an autoregressive Transformer model as either a Linen module - or an NNX-wrapped module, depending on the `config.enable_nnx` flag. The returned module - is suitable for training, evaluation, or decoding. - - If `config.enable_nnx` is True, returns a `TransformerLinen` that wraps the NNX-style - Transformer for integration with NNX-specific APIs and workflows. - Otherwise, returns a pure Flax Linen implementation (`TransformerLinenPure`). + Returns a `TransformerLinen` that wraps the NNX-style Transformer so it can be + driven through the Linen init/apply API (checkpoint conversion, AOT compile, + the inference engine). Pure-NNX call sites build `Transformer` directly via + `model_creation_utils.from_config`. Args: config (Config): The configuration object specifying model hyperparameters and options. @@ -277,29 +61,25 @@ def transformer_as_linen( quant (Quant): The quantization module or configuration to use. model_mode (str, optional): The operational mode for the model, e.g. training, prefill, or autoregressive. Defaults to `MODEL_MODE_TRAIN`. - name (str, optional): Optional module name for Linen/NNX construction. + name (str, optional): Optional module name for construction. Returns: - nnx_wrappers.ToLinen | TransformerLinenPure: - A constructed Transformer model compatible with the specified framework (Linen or NNX). + nnx_wrappers.ToLinen: An NNX Transformer wrapped as a Linen module. """ - if config.enable_nnx: - return TransformerLinen( - Transformer, - args=(), - kwargs=nn.FrozenDict( - { - "mesh": mesh, - "config": config, - "quant": quant, - "model_mode": model_mode, - } - ), - metadata_fn=initializers.variable_to_logically_partitioned, - name=name, - ) - else: - return TransformerLinenPure(config, mesh, quant, model_mode=model_mode, name=name) + return TransformerLinen( + Transformer, + args=(), + kwargs=nn.FrozenDict( + { + "mesh": mesh, + "config": config, + "quant": quant, + "model_mode": model_mode, + } + ), + metadata_fn=initializers.variable_to_logically_partitioned, + name=name, + ) class TransformerLinen(nnx_wrappers.ToLinen): @@ -355,41 +135,7 @@ def __init__( ) self.vision_encoder = VisionEncoder(config=cfg, mesh=mesh, rngs=rngs) if cfg.use_multimodal else None self.audio_encoder = AudioEncoder(config=cfg, mesh=mesh, rngs=rngs) if cfg.use_audio else None - if cfg.pure_nnx_decoder: - self.decoder = NNXDecoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) - else: - decoder_linen = Decoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode) - self.decoder = nnx_wrappers.ToNNX(decoder_linen, rngs=rngs) - - batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(config=cfg, model_mode=model_mode) - dummy_decoder_input_tokens = jnp.ones((batch_size, seq_len), dtype=jnp.int32) - dummy_decoder_positions = jnp.ones((batch_size, seq_len), dtype=jnp.int32) - - if self.config.attention == "vllm_rpa": - try: - # pylint: disable=import-outside-toplevel - from tpu_inference.layers.common.attention_metadata import AttentionMetadata # pytype: disable=import-error - except ImportError as e: - raise ImportError( - "vLLM RPA attention requires the vllm-tpu package. Please install it with `pip install vllm-tpu`." - ) from e - dummy_attention_metadata = AttentionMetadata( - input_positions=jnp.ones((batch_size * seq_len,), dtype=jnp.int32), - block_tables=jnp.ones((seq_len,), dtype=jnp.int32), - seq_lens=jnp.ones((1), dtype=jnp.int32), - query_start_loc=jnp.ones((2), dtype=jnp.int32), - request_distribution=jnp.ones((3), dtype=jnp.int32), - ) - else: - dummy_attention_metadata = None - - if not cfg.pure_nnx_decoder: - self.decoder.lazy_init( - shared_embedding=self.token_embedder, - decoder_input_tokens=dummy_decoder_input_tokens, - decoder_positions=dummy_decoder_positions, - attention_metadata=dummy_attention_metadata, - ) + self.decoder = NNXDecoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) # If MTP is enabled via config, set up the MTP block. if self.config.mtp_num_layers > 0: @@ -526,45 +272,20 @@ def __call__( bidirectional_mask_video=bidirectional_mask_video, ) - mutable_collections = [] - if self.config.record_internal_nn_metrics: - mutable_collections.append("intermediates") - if self.config.distill_beta > 0.0 and "intermediates" not in mutable_collections: - mutable_collections.append("intermediates") - if self.config.load_balance_loss_weight > 0.0 and "intermediates" not in mutable_collections: - mutable_collections.append("intermediates") - - if self.config.pure_nnx_decoder: - logits, hidden_state, kv_caches = self.decoder( - shared_embedding=self.token_embedder, - decoder_input_tokens=decoder_input_tokens, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=not enable_dropout, - model_mode=model_mode, - previous_chunk=previous_chunk, - slot=slot, - multimodal_input=multimodal_input, - kv_caches=kv_caches, - attention_metadata=attention_metadata, - deepstack_visual_embeds=deepstack_visual_embeds, - ) # pytype: disable=wrong-keyword-args - else: - logits, hidden_state, kv_caches = self.decoder( - shared_embedding=self.token_embedder, - decoder_input_tokens=decoder_input_tokens, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=not enable_dropout, - model_mode=model_mode, - previous_chunk=previous_chunk, - slot=slot, - multimodal_input=multimodal_input, - kv_caches=kv_caches, - attention_metadata=attention_metadata, - deepstack_visual_embeds=deepstack_visual_embeds, - mutable=mutable_collections, - ) # pytype: disable=wrong-keyword-args + logits, hidden_state, kv_caches = self.decoder( + shared_embedding=self.token_embedder, + decoder_input_tokens=decoder_input_tokens, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=not enable_dropout, + model_mode=model_mode, + previous_chunk=previous_chunk, + slot=slot, + multimodal_input=multimodal_input, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + deepstack_visual_embeds=deepstack_visual_embeds, + ) # pytype: disable=wrong-keyword-args # If we are initializing the model AND MTP is enabled, we must create # dummy target tensors. This allows Flax to trace the MTPBlock and create diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 3b5f08f0c1..e7b951a46e 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -45,7 +45,8 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.moe import RoutedMoE -from maxtext.layers.initializers import nd_dense_init, variable_to_logically_partitioned +from maxtext.layers.initializers import nd_dense_init + from maxtext.utils import max_utils from maxtext.inference import kvcache @@ -2167,29 +2168,6 @@ def __call__(self, hidden_states: Array) -> Array: return output -def qwen3omni_visionencoder_as_linen(config: Config, mesh: Mesh) -> nn.Module: - """Convert Qwen3OmniMoeVisionEncoder to Linen module.""" - return nnx_wrappers.to_linen( - Qwen3OmniMoeVisionEncoder, - config=config, - mesh=mesh, - name="Qwen3OmniMoeVisionEncoder_0", - abstract_init=False, - metadata_fn=max_initializers.variable_to_logically_partitioned, - ) - - -def qwen3omni_visionprojector_as_linen(config: Config, mesh: Mesh) -> nn.Module: - """Convert Qwen3OmniMoeVisionProjector to Linen module.""" - return nnx_wrappers.to_linen( - Qwen3OmniMoeVisionProjector, - config=config, - name="Qwen3OmniMoeVisionProjector_0", - abstract_init=False, - metadata_fn=max_initializers.variable_to_logically_partitioned, - ) - - class Qwen3OmniAudioEncoderLayer(nnx.Module): """Transformer encoder layer for audio model.""" @@ -2487,29 +2465,6 @@ def __call__(self, hidden_states: Array) -> Array: return hidden_states -def qwen3omni_audioencoder_as_linen(config: Config, mesh: Mesh): - """Convert AudioEncoder (convs + transformer layers, no projector) to Linen module.""" - return nnx_wrappers.to_linen( - Qwen3OmniAudioEncoder, - config=config, - mesh=mesh, - name="Qwen3OmniAudioEncoder_0", - abstract_init=False, - metadata_fn=variable_to_logically_partitioned, - ) - - -def qwen3omni_audioprojector_as_linen(config: Config, mesh: Mesh): - """Convert AudioProjector to Linen module.""" - return nnx_wrappers.to_linen( - Qwen3OmniAudioProjector, - config=config, - name="Qwen3OmniAudioProjector_0", - abstract_init=False, - metadata_fn=variable_to_logically_partitioned, - ) - - # Vision encoder Linen wrappers Qwen3OmniMoeVisionPatchMergerToLinen = nnx_wrappers.to_linen_class( Qwen3OmniMoeVisionPatchMerger, From f421c99fd6f39d821248a6b7a549eb7236df17ad Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Thu, 4 Jun 2026 18:43:53 +0000 Subject: [PATCH 3/4] [NNX] Delete Linen (3/4): drop obsolete Linen tests and flag references Remove obsolete Linen-only tests, drop redundant flag args from the rest, and compile the hlo_diff tests via base.yml + model_name so they exercise the real NNX path. --- docs/guides/distillation.md | 6 +- .../posttraining/knowledge_distillation.md | 4 +- docs/tutorials/posttraining/lora.md | 4 - .../posttraining/lora_on_multi_host.md | 4 - src/maxtext/common/checkpointing.py | 8 +- src/maxtext/common/train_state_nnx.py | 7 +- src/maxtext/examples/lora_llama3_demo.ipynb | 2 - .../distillation/scripts/run_distill_xpk.sh | 5 +- src/maxtext/utils/standalone_checkpointer.py | 4 +- .../golden_logits/golden_dpo_correctness.json | 18 +- .../generate_grpo_golden_logits.py | 93 +-- .../tpu/gemma3/4b/test_gemma3_lora.sh | 8 +- .../integration/deepseek_scan_engram_test.py | 94 +-- tests/integration/diloco_test.py | 193 ++----- tests/integration/maxengine_test.py | 23 +- .../integration/setup_train_loop_nnx_test.py | 7 +- .../integration/dpo_correctness_base.py | 3 - .../integration/grpo_correctness.py | 88 +-- .../grpo_trainer_correctness_test.py | 52 +- .../sft_trainer_correctness_test.py | 40 +- tests/post_training/unit/lora_utils_test.py | 2 - tests/unit/aqt_serve_roundtrip_nnx_test.py | 3 - .../correctness_tests_nnx_dispatch_test.py | 16 +- ...generate_param_only_checkpoint_nnx_test.py | 3 - tests/unit/grpo_nnx_test.py | 72 +-- tests/unit/maxengine_nnx_test.py | 3 - tests/unit/maxtext_utils_test.py | 152 ++--- tests/unit/muon_utils_test.py | 32 -- tests/unit/nnx_decoder_test.py | 3 - tests/unit/nnx_quant_guard_test.py | 31 +- tests/unit/optimizers_test.py | 2 +- tests/unit/quantizations_test.py | 183 ++---- tests/unit/sharding_compare_test.py | 328 +---------- tests/unit/sharding_nnx_test.py | 97 +--- tests/unit/state_dtypes_test.py | 50 +- tests/unit/tiling_test.py | 541 +----------------- tests/unit/train_compile_test.py | 55 +- tests/unit/train_state_nnx_checkpoint_test.py | 71 +-- tests/utils/forward_pass_logit_checker.py | 48 +- tests/utils/run_sharding_dump.py | 9 +- 40 files changed, 344 insertions(+), 2020 deletions(-) diff --git a/docs/guides/distillation.md b/docs/guides/distillation.md index 317521b0fe..d2c7eebfc2 100644 --- a/docs/guides/distillation.md +++ b/docs/guides/distillation.md @@ -63,9 +63,8 @@ Student and teacher must share the same vocabulary. The trainer asserts `student If `distill_beta > 0`, the model `sow`s the attention `out_projection` activations at every layer so the loss can read them. This requires: - `scan_layers: True` — activations are stacked along the leading scan axis; the loss does `jnp.take(features, layer_indices, axis=0)` over that axis. -- `enable_nnx: True` — `sow(nnx.Intermediate, ...)` is an NNX-specific call. -The trainer validates both at config initialization. Logit-only runs (`distill_beta = 0`) have no such constraint. +The trainer validates this at config initialization. Logit-only runs (`distill_beta = 0`) have no such constraint. ## Loss anatomy @@ -221,7 +220,6 @@ distill_feature_loss_type: cosine distill_layer_indices: [3, 7, 11, 15, 19, 23, 27, 31] # for 32-layer student scan_layers: True -enable_nnx: True ``` ### Logit-only baseline (cheapest; no feature extraction overhead) @@ -258,7 +256,7 @@ The trainer logs the following to TensorBoard (configured by `tensorboard_dir`, | Symptom | Likely cause | Fix | | ----------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------- | -| `ValueError: a value of self.distill_beta > 0.0 requires self.scan_layers = True` | Feature loss enabled without scanned layers. | Add `scan_layers=True enable_nnx=True` to your CLI / yml. | +| `ValueError: a value of self.distill_beta > 0.0 requires self.scan_layers = True` | Feature loss enabled without scanned layers. | Add `scan_layers=True` to your CLI / yml. | | `Vocab size mismatch! Student: X, Teacher: Y` | Different tokenizers. | Use teacher and student with the same vocab; the trainer cannot match logits across vocabularies. | | `Teacher model path is missing` | `teacher_overrides.load_parameters_path` not set in non-offline mode. | Set it in `teacher_overrides` in the yml or pass via CLI. | | `Features extracted from student or teacher model are None, but distill_beta > 0.0` | Model architecture doesn't sow `out_projection_activations` (e.g. uses an unsupported attention path). | Verify the attention layer in use sets `self.sow(nnx.Intermediate, "out_projection_activations", out)` (see `attentions.py`). | diff --git a/docs/tutorials/posttraining/knowledge_distillation.md b/docs/tutorials/posttraining/knowledge_distillation.md index 2f759d75e6..8bb2eda969 100644 --- a/docs/tutorials/posttraining/knowledge_distillation.md +++ b/docs/tutorials/posttraining/knowledge_distillation.md @@ -248,7 +248,7 @@ Key knobs (see the [Distillation guide](../../guides/distillation.md) for the fu ```yaml distill_alpha: 0.5 # weight on KL(teacher||student) distill_temperature: 1.0 -distill_beta: 0.0 # >0 enables feature distillation; requires scan_layers=True, enable_nnx=True +distill_beta: 0.0 # >0 enables feature distillation; requires scan_layers=True distill_layer_indices: None ``` @@ -312,7 +312,7 @@ python3 -m maxtext.trainers.post_train.distillation.train_distill \ distill_temperature=2.0 \ distill_beta=1.0 distill_beta_end=0.1 distill_beta_schedule=cosine \ distill_layer_indices=[2,5,8,11,14,17,20,23] \ - scan_layers=True enable_nnx=True \ + scan_layers=True \ profiler=xplane ``` diff --git a/docs/tutorials/posttraining/lora.md b/docs/tutorials/posttraining/lora.md index c309078593..bcd9ea9a51 100644 --- a/docs/tutorials/posttraining/lora.md +++ b/docs/tutorials/posttraining/lora.md @@ -117,8 +117,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \ max_target_length="${MAX_TARGET_LENGTH?}" \ learning_rate="${LEARNING_RATE?}" \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank="${LORA_RANK?}" \ lora.lora_alpha="${LORA_ALPHA?}" @@ -174,8 +172,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \ max_target_length="${MAX_TARGET_LENGTH?}" \ learning_rate="${LEARNING_RATE?}" \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank="${LORA_RANK?}" \ lora.lora_alpha="${LORA_ALPHA?}" diff --git a/docs/tutorials/posttraining/lora_on_multi_host.md b/docs/tutorials/posttraining/lora_on_multi_host.md index 4d0905ea13..8b3f140a80 100644 --- a/docs/tutorials/posttraining/lora_on_multi_host.md +++ b/docs/tutorials/posttraining/lora_on_multi_host.md @@ -186,8 +186,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ max_target_length=${MAX_TARGET_LENGTH?} \ learning_rate=${LEARNING_RATE?} \ chat_template_path=${CHAT_TEMPLATE_PATH?} \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank=${LORA_RANK?} \ lora.lora_alpha=${LORA_ALPHA?} \ @@ -267,8 +265,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ lora.lora_restore_path=${LORA_RESTORE_PATH?} \ learning_rate=${LEARNING_RATE?} \ chat_template_path=${CHAT_TEMPLATE_PATH?} \ - enable_nnx=True \ - pure_nnx_decoder=True \ lora.enable_lora=True \ lora.lora_rank=${LORA_RANK?} \ lora.lora_alpha=${LORA_ALPHA?} \ diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 5c515f567f..9feaded014 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -222,7 +222,7 @@ def _load_linen_checkpoint_into_nnx( use_ocdbt, use_zarr3, ): - """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume). + """Restores a Linen-layout checkpoint into an NNX state. Restores against a Linen-shape abstract, reshapes back via `from_linen_checkpoint_dict`, then fills NNX-only rngs/dropout with defaults. @@ -361,7 +361,7 @@ def combine_sharding(sds, shardings): else: raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}") else: - # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. + # Checkpoints use the Linen on-disk layout; reshape it back into the NNX state. if isinstance(abstract_unboxed_pre_state, nnx.State): return _load_linen_checkpoint_into_nnx( path, @@ -853,7 +853,7 @@ def map_to_pspec(data): ) ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True) - # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. + # Checkpoints use the Linen on-disk layout; reshape it back into the NNX state. # (Emergency managers use their own restore path below.) if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance( checkpoint_manager, @@ -1100,7 +1100,7 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.") return - # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. + # Save in the Linen on-disk layout so checkpoints stay interchangeable across formats. if config.enable_diloco: # DiLoCoTrainState: persist the synchronized global model (outer params). # The per-replica inner optimizer / outer-momentum state is not checkpointed. diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index a123956a73..aaaa851c37 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -58,10 +58,9 @@ def apply_gradients(self, grads: Any, **kwargs): # On-disk checkpoint format. # -# A pure_nnx run saves in the same on-disk layout as a Linen run, so the two are -# interchangeable. The NNX state pure dict differs from Linen's in three ways, -# all -# reshaped below at save time: +# NNX saves in the same on-disk layout as the old Linen format, so checkpoints +# stay interchangeable. The NNX state pure dict differs from that layout in three +# ways, all reshaped below at save time: # 1. top-level keys: {model, optimizer:{step, opt_state}} -> {params:{params:...}, step, opt_state} # 2. weights: model/... -> params/params/... (Linen `params` collection) # 3. opt_state: int-keyed dict (empty entries skipped) -> list with None for EmptyState, diff --git a/src/maxtext/examples/lora_llama3_demo.ipynb b/src/maxtext/examples/lora_llama3_demo.ipynb index d324419093..455df12a1e 100644 --- a/src/maxtext/examples/lora_llama3_demo.ipynb +++ b/src/maxtext/examples/lora_llama3_demo.ipynb @@ -307,8 +307,6 @@ " \"learning_rate=5e-5\", \n", " \"weight_dtype=bfloat16\",\n", " \"dtype=bfloat16\",\n", - " \"enable_nnx=true\",\n", - " \"pure_nnx_decoder=true\",\n", " \"scan_layers=true\",\n", " \"lora.enable_lora=true\",\n", " \"lora.lora_rank=16\",\n", diff --git a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh index 91078508d9..d0b0dda8f1 100644 --- a/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh +++ b/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh @@ -100,7 +100,7 @@ # DISTILL_ALPHA default: 0.5 # DISTILL_TEMPERATURE default: 1.0 # DISTILL_BETA default: 1.0 (>0 enables feature-map loss; -# requires scan_layers=True and enable_nnx=True) +# requires scan_layers=True) # DISTILL_LAYER_INDICES default: [0,1,2,3,4,5,6,7] (no spaces inside brackets) # # Image pinning (used by prep_image): @@ -185,8 +185,7 @@ LAST_WORKLOAD_FILE="${XPK_LAST_WORKLOAD_FILE:-${HOME}/.xpk_last_workload}" extra_cli="distill_alpha=${DISTILL_ALPHA} \ distill_temperature=${DISTILL_TEMPERATURE} \ distill_beta=${DISTILL_BETA} \ -distill_layer_indices=${DISTILL_LAYER_INDICES} \ -enable_nnx=True" +distill_layer_indices=${DISTILL_LAYER_INDICES}" if [ -n "${STEPS_OVERRIDE:-}" ]; then extra_cli="$extra_cli learning_rate_schedule_steps=${STEPS_OVERRIDE} steps=${STEPS_OVERRIDE}" fi diff --git a/src/maxtext/utils/standalone_checkpointer.py b/src/maxtext/utils/standalone_checkpointer.py index d6be1a533c..9ee0a44724 100644 --- a/src/maxtext/utils/standalone_checkpointer.py +++ b/src/maxtext/utils/standalone_checkpointer.py @@ -118,8 +118,8 @@ def add_entropy_to_checkpoint(state): * Linen `TrainState`: `state.params` + `state.opt_state` (tuple). * NNX `TrainStateNNX` (Module): `state.model` is an `nnx.Module`; the optimizer's `opt_state` is the optax tuple of NamedTuples. - * NNX `nnx.State` (post-split, what `setup_training_state` returns under - `pure_nnx`): `state.model` and `state.optimizer.opt_state` are sub-States; + * NNX `nnx.State` (post-split, what `setup_training_state` returns): + `state.model` and `state.optimizer.opt_state` are sub-States; `opt_state[0].mu`/`nu` are themselves States that can be reassigned. """ if hasattr(state, "model"): diff --git a/tests/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 9dc78f9df0..43389e30ee 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 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): @@ -113,19 +99,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", @@ -214,29 +193,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()) @@ -244,13 +210,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, @@ -271,7 +235,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, @@ -282,8 +245,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, ) @@ -298,13 +259,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/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh index 59408e8cdd..61972a236c 100755 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh @@ -49,8 +49,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ lora.enable_lora=True \ lora.lora_rank=16 \ lora.lora_alpha=32.0 \ - enable_nnx=True \ - pure_nnx_decoder=True \ enable_single_controller=True \ checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False @@ -68,8 +66,6 @@ python3 -m maxtext.inference.vllm_decode \ hbm_utilization_vllm=0.6 \ prompt="Suggest some famous landmarks in London." \ use_chat_template=True \ - enable_nnx=True \ - pure_nnx_decoder=True \ scan_layers=true # Step 5: Convert the checkpoint from MaxText format to Hugging Face format @@ -78,6 +74,4 @@ python3 -m maxtext.checkpoint_conversion.to_huggingface \ load_parameters_path=${SCANNED_CKPT_PATH} \ lora.lora_restore_path=${BASE_OUTPUT_DIRECTORY}/lora/${run_id}/checkpoints/5/model_params \ base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/unscanned/${run_id} \ - scan_layers=true \ - enable_nnx=True \ - pure_nnx_decoder=True + scan_layers=true diff --git a/tests/integration/deepseek_scan_engram_test.py b/tests/integration/deepseek_scan_engram_test.py index 6891fd83e2..d12b7d6321 100644 --- a/tests/integration/deepseek_scan_engram_test.py +++ b/tests/integration/deepseek_scan_engram_test.py @@ -14,22 +14,19 @@ """Unit tests for DeepSeek Engram across scanned decoder layers.""" -import gc -import os import unittest from unittest.mock import patch -import jax import jax.numpy as jnp -from jax.sharding import Mesh -from maxtext.configs import pyconfig -from maxtext.utils.globals import MAXTEXT_PKG_DIR -from maxtext.common.common_types import MODEL_MODE_TRAIN -from maxtext.layers.decoders import Decoder -from maxtext.utils import maxtext_utils import pytest +# The Linen Decoder this test exercised was removed in PR12 (Delete Linen). +# NNX decoder coverage is in tests/unit/nnx_decoders_test.py. +pytestmark = pytest.mark.skip( + reason="Linen Decoder removed in PR12 (Delete Linen); NNX decoder coverage is in tests/unit/nnx_decoders_test.py" +) + class DummyEmbedding: """Dummy embedding layer for testing.""" @@ -91,81 +88,10 @@ def _test_engram_pattern( base_num_decoder_layers=10, ): """Helper method to test different engram layer patterns.""" - - # Setup mock tokenizer - class MockTokenizer: - """Mock tokenizer for testing.""" - - pad_token_id = 0 - - def __len__(self): - return 128 - - def __call__(self, x): - return jnp.ones_like(x) - - def convert_ids_to_tokens(self, *args, **kwargs): - return "a" - - def decode(self, *args, **kwargs): - return "a" - - def batch_decode(self, token_ids, *args, **kwargs): - return ["a" for _ in token_ids] - - mock_from_pretrained.return_value = MockTokenizer() - - config_path = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml") - config = pyconfig.initialize( - [None, config_path] - + self._COMMON_CONFIG - + [ - f"engram_layers=[{engram_layers_str}]", - f"first_num_dense_layers={first_num_dense_layers}", - f"base_num_decoder_layers={base_num_decoder_layers}", - f"num_decoder_layers={base_num_decoder_layers}", - ] - ) - - devices_array = maxtext_utils.create_device_mesh(config) - mesh = Mesh(devices_array, config.mesh_axes) - - decoder = Decoder( - config=config, - mesh=mesh, - model_mode=MODEL_MODE_TRAIN, - ) - - batch_size = config.global_batch_size_to_load - seq_len = config.max_target_length - - decoder_input_tokens = jnp.ones((batch_size, seq_len), dtype=jnp.int32) - decoder_positions = jnp.ones((batch_size, seq_len), dtype=jnp.int32) - decoder_segment_ids = jnp.ones((batch_size, seq_len), dtype=jnp.int32) - - shared_embedding = DummyEmbedding(emb_dim=config.emb_dim) - - with jax.set_mesh(mesh), jax.disable_jit(): - variables = decoder.init( - {"params": jax.random.PRNGKey(0), "dropout": jax.random.PRNGKey(1), "aqt": jax.random.PRNGKey(2)}, - shared_embedding=shared_embedding, - decoder_input_tokens=decoder_input_tokens, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=True, - model_mode=MODEL_MODE_TRAIN, - ) - - self.assertIn("params", variables) - params = variables["params"] - for key in expected_keys: - self.assertIn(key, params) - - del variables - del params - del decoder - jax.clear_caches() - gc.collect() + # The Linen Decoder this exercised was removed in PR12 (Delete Linen); + # NNX decoder coverage lives in tests/unit/nnx_decoders_test.py. + del mock_from_pretrained, engram_layers_str, expected_keys, first_num_dense_layers, base_num_decoder_layers + raise unittest.SkipTest("Linen Decoder removed in PR12 (Delete Linen)") @pytest.mark.tpu_only @patch("transformers.AutoTokenizer.from_pretrained") diff --git a/tests/integration/diloco_test.py b/tests/integration/diloco_test.py index 77ed6dd923..052edb57de 100644 --- a/tests/integration/diloco_test.py +++ b/tests/integration/diloco_test.py @@ -21,7 +21,6 @@ import chex from flax.experimental import nnx -from flax.training import train_state import jax import jax.numpy as jnp import jax.sharding @@ -83,78 +82,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) - - # For NNX, train_step needs to take the TrainStateNNX and mutate it - - 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)) - - 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, - ) + + optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) + # diloco_test_state expects a TrainStateNNX instance. + initial_test_state = TrainStateNNX(model, optimizer) + + # train_step takes the TrainStateNNX and mutates it. + + 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)) + + 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( @@ -202,22 +159,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: @@ -247,22 +196,14 @@ def loss_fn(params, batch): chex.assert_equal(diloco_test_state.step, 2.0) chex.assert_trees_all_close(loss, 0.65) # 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: @@ -297,39 +238,21 @@ def loss_fn(params, batch): chex.assert_trees_all_close(loss, 0.4481) # 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/integration/maxengine_test.py b/tests/integration/maxengine_test.py index aaab11d924..17c8ae9491 100644 --- a/tests/integration/maxengine_test.py +++ b/tests/integration/maxengine_test.py @@ -71,20 +71,11 @@ def init_pyconfig(self, **kwargs): def test_stack_and_unstack_prefill_cache_nnx(self): """scan_layers=False: per-layer cache subtrees stack onto a leading layer axis and back.""" - cfg = self._init_nnx_pyconfig( - stack_prefill_result_cache=True, scan_layers=False - ) + cfg = self._init_nnx_pyconfig(stack_prefill_result_cache=True, scan_layers=False) engine = maxengine.MaxEngine(cfg, jax.devices()) num_layers = engine.config.num_decoder_layers # scan_layers=False keeps the per-layer subtrees under decoder/layers, keyed by layer index. - cache = { - "decoder": { - "layers": { - i: {"a": jnp.ones((1, 10)), "b": jnp.ones((1, 9))} - for i in range(num_layers) - } - } - } + cache = {"decoder": {"layers": {i: {"a": jnp.ones((1, 10)), "b": jnp.ones((1, 9))} for i in range(num_layers)}}} expected_stacked = { "decoder": { @@ -105,8 +96,8 @@ def test_stack_and_unstack_prefill_cache_nnx(self): # default. test_basic_prefill_nnx / test_basic_decode_nnx below cover the NNX path. def _init_nnx_pyconfig(self, **kwargs): - """Same as init_pyconfig but with the NNX flags turned on.""" - return self.init_pyconfig(pure_nnx=True, enable_nnx=True, pure_nnx_decoder=True, **kwargs) + """Same as init_pyconfig (NNX is the only path now).""" + return self.init_pyconfig(**kwargs) def _build_nnx_params(self, cfg, mesh): """Materialize an NNX Transformer and return its nnx.Param state.""" @@ -182,7 +173,7 @@ def test_basic_decode_nnx(self): ) def test_quantize_passes_gate_for_nnx(self): - """pure_nnx + quantization (convert-on-load) reaches the actual machinery in train mode.""" + """NNX + quantization (convert-on-load) reaches the actual machinery in train mode.""" # checkpoint_is_quantized defaults to False — full-precision on disk, AQT # quantizes per-forward against the loaded kernel (train mode). cfg = self._init_nnx_pyconfig(quantization="int8") @@ -196,7 +187,7 @@ def test_quantize_passes_gate_for_nnx(self): pass # any other failure (e.g. checkpoint not found) is fine for this test def test_load_pre_quantized_nnx_passes_quant_gate(self): - """pure_nnx + quantization + checkpoint_is_quantized=True clears the load gate.""" + """NNX + quantization + checkpoint_is_quantized=True clears the load gate.""" cfg = self._init_nnx_pyconfig(quantization="int8", checkpoint_is_quantized=True) engine = maxengine.MaxEngine(cfg, jax.devices()) self.assertEqual(engine._nnx_quant_mode_str, "serve") # pylint: disable=protected-access @@ -228,7 +219,7 @@ def test_quantized_prefill_nnx_train_mode(self): self.assertTrue(jnp.all(jnp.isfinite(prefill_result["logits"]))) def test_lora_load_single_adapter_reaches_loader_on_nnx(self): - """pure_nnx + LoRA: load_single_adapter dispatches to the NNX loader. + """NNX + LoRA: load_single_adapter dispatches to the NNX loader. A nonexistent adapter path should raise FileNotFoundError from the loader itself. A NotImplementedError here would mean the dispatch diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py index f446c8f9df..a5e7d1fca4 100644 --- a/tests/integration/setup_train_loop_nnx_test.py +++ b/tests/integration/setup_train_loop_nnx_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Integration test for setup_train_loop with pure_nnx=True. +"""Integration test for setup_train_loop on the NNX path. setup_train_loop wires together create_nnx_abstract_model, the training optimizer, @@ -43,7 +43,6 @@ def _tiny_nnx_pyconfig(**overrides): "enable_checkpointing": False, "dataset_type": "synthetic", "model_name": "default", - "pure_nnx": True, "per_device_batch_size": 1.0, "base_emb_dim": 8, "base_num_query_heads": 4, @@ -68,7 +67,7 @@ def _tiny_nnx_pyconfig(**overrides): class SetupTrainLoopNNXIntegrationTest(unittest.TestCase): """End-to-end check that setup_train_loop returns a usable TrainStateNNX.""" - def test_pure_nnx_setup_returns_train_state_nnx(self): + def test_setup_returns_train_state_nnx(self): config = _tiny_nnx_pyconfig() ( @@ -109,7 +108,7 @@ def test_pure_nnx_setup_returns_train_state_nnx(self): # flag them as unused — they're part of the public return contract. del checkpoint_manager, rampup_manager, eval_data_iterator - def test_pure_nnx_setup_param_only_split_matches_model(self): + def test_setup_param_only_split_matches_model(self): """nnx.split(state.model, nnx.Param, ...) must yield a non-empty Param tree whose structure matches state_mesh_shardings.model after the same split. 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 84ffdc85e0..581d3162d0 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 -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 @@ -63,17 +60,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) + # 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 self.tokenizer_model = transformers.AutoTokenizer.from_pretrained( "meta-llama/Llama-3.1-8B", add_bos_token=False, @@ -147,57 +138,24 @@ 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, + """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, using the cloned reference model.""" + return compute_log_probs_nnx( + self.reference_model, input_ids, input_position, input_segmentation, @@ -207,10 +165,8 @@ 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, using the cloned reference model.""" + return grpo_loss_fn_nnx(self.model, self.cfg, data, self.rng, None, self.reference_model) def test_logits(self): def _prepare_inputs(): @@ -306,12 +262,8 @@ 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. 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 c5a4392c98..1f11f48b2d 100644 --- a/tests/post_training/integration/grpo_trainer_correctness_test.py +++ b/tests/post_training/integration/grpo_trainer_correctness_test.py @@ -25,12 +25,10 @@ pytest tests/post_training/integration/grpo_trainer_correctness_test.py """ -import functools import os import subprocess import sys import unittest -from flax import linen as nn import jax import jax.numpy as jnp from jax.sharding import Mesh @@ -38,16 +36,13 @@ import maxtext as mt from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, 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 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,26 @@ 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. + 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.""" + 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, where `reference` is the frozen reference model.""" + 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 a94932a182..dc8673fbd7 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 subprocess import sys @@ -36,13 +35,10 @@ from jax.sharding import Mesh import jsonlines from maxtext.configs import pyconfig -from maxtext.common.common_types import MODEL_MODE_TRAIN from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from maxtext.utils.globals import MAXTEXT_PKG_DIR from maxtext.utils.globals import MAXTEXT_TEST_ASSETS_ROOT 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 @@ -120,40 +116,22 @@ 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 + # The model carries its own params, 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) + 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/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py index 0adb1a2768..63626b408b 100644 --- a/tests/post_training/unit/lora_utils_test.py +++ b/tests/post_training/unit/lora_utils_test.py @@ -55,8 +55,6 @@ "base_mlp_dim": 256, "max_prefill_predict_length": 4, "model_name": "llama2-7b", - "enable_nnx": True, - "pure_nnx_decoder": True, "override_model_config": True, "weight_dtype": "bfloat16", } diff --git a/tests/unit/aqt_serve_roundtrip_nnx_test.py b/tests/unit/aqt_serve_roundtrip_nnx_test.py index 00d6825de7..006a4580d2 100644 --- a/tests/unit/aqt_serve_roundtrip_nnx_test.py +++ b/tests/unit/aqt_serve_roundtrip_nnx_test.py @@ -74,9 +74,6 @@ def _init_cfg(self, ckpt_path, *, checkpoint_is_quantized): sys.argv[0], base_yml, "model_name=gpt3-52k", - "pure_nnx=true", - "enable_nnx=true", - "pure_nnx_decoder=true", "max_target_length=64", "max_prefill_predict_length=16", "per_device_batch_size=1", diff --git a/tests/unit/correctness_tests_nnx_dispatch_test.py b/tests/unit/correctness_tests_nnx_dispatch_test.py index c0177cfb14..ebfac32e2b 100644 --- a/tests/unit/correctness_tests_nnx_dispatch_test.py +++ b/tests/unit/correctness_tests_nnx_dispatch_test.py @@ -18,7 +18,7 @@ exercises the changed dispatch code that otherwise has no CPU coverage: - `mt.from_config` is exported (the GRPO trainer calls it) - the SFT correctness test's `setup_maxtext_model` / `get_maxtext_logits` run on - both paths (pure_nnx=True -> NNX, pure_nnx=False -> Linen) and stay finite + the NNX path and stay finite The GRPO NNX building blocks the other dispatch helpers call (`compute_log_probs_nnx`, `grpo_loss_fn_nnx`) are already covered by grpo_nnx_test. @@ -50,15 +50,12 @@ } -def _sft_config(pure_nnx): +def _sft_config(): return pyconfig.initialize( [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs/post_train", "sft.yml")], - run_name=f"unit-sft-{pure_nnx}", + run_name="unit-sft-nnx", model_name="default", enable_checkpointing=False, - pure_nnx=pure_nnx, - enable_nnx=pure_nnx, - pure_nnx_decoder=pure_nnx, **_SMALL, ) @@ -80,12 +77,7 @@ def test_from_config_is_exported(self): self.assertTrue(hasattr(mt, "from_config")) def test_sft_logits_nnx_path(self): - config = _sft_config(pure_nnx=True) - logits = sft.get_maxtext_logits(config, _fake_data(config)) - self.assertTrue(bool(jnp.isfinite(logits).all())) - - def test_sft_logits_linen_path(self): - config = _sft_config(pure_nnx=False) + config = _sft_config() logits = sft.get_maxtext_logits(config, _fake_data(config)) self.assertTrue(bool(jnp.isfinite(logits).all())) diff --git a/tests/unit/generate_param_only_checkpoint_nnx_test.py b/tests/unit/generate_param_only_checkpoint_nnx_test.py index 57995a6145..800610fb15 100644 --- a/tests/unit/generate_param_only_checkpoint_nnx_test.py +++ b/tests/unit/generate_param_only_checkpoint_nnx_test.py @@ -92,7 +92,6 @@ def test_unrolls_scanned_layers(self): config = SimpleNamespace( scan_layers=True, force_unroll=True, - pure_nnx=True, param_scan_axis=0, decoder_block=DecoderBlockType.LLAMA2, num_decoder_layers=num_layers, @@ -133,7 +132,6 @@ def __init__(self): config = SimpleNamespace( scan_layers=True, force_unroll=True, - pure_nnx=True, param_scan_axis=0, decoder_block=DecoderBlockType.DEEPSEEK, num_decoder_layers=5, @@ -184,7 +182,6 @@ def test_unrolls_scanned_lora_layers(self): config = SimpleNamespace( scan_layers=True, force_unroll=True, - pure_nnx=True, param_scan_axis=0, decoder_block=DecoderBlockType.LLAMA2, num_decoder_layers=num_layers, diff --git a/tests/unit/grpo_nnx_test.py b/tests/unit/grpo_nnx_test.py index 259ebae267..ea5de85246 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 @@ -91,7 +88,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 ) @@ -291,70 +288,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/maxengine_nnx_test.py b/tests/unit/maxengine_nnx_test.py index 08303f817a..074db10c22 100644 --- a/tests/unit/maxengine_nnx_test.py +++ b/tests/unit/maxengine_nnx_test.py @@ -41,9 +41,6 @@ def _nnx_config(self, **kwargs): "max_target_length": 8, "per_device_batch_size": 1, "enable_checkpointing": False, - "pure_nnx": True, - "enable_nnx": True, - "pure_nnx_decoder": True, } | kwargs return pyconfig.initialize([sys.argv[0], get_test_config_path()], **init_kwargs) diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 670fcce5e1..6f2e5c75ef 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -16,7 +16,6 @@ from collections.abc import Callable from dataclasses import dataclass, field -import functools from typing import Any, Sequence import unittest from unittest.mock import MagicMock, Mock, patch @@ -30,11 +29,9 @@ import jax.numpy as jnp from jax.sharding import AxisType, Mesh, NamedSharding, PartitionSpec from maxtext.common import train_state_nnx -from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN, ShardMode from maxtext.configs import pyconfig +from maxtext.common.common_types import DecoderBlockType, ShardMode from maxtext.inference import inference_utils -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx @@ -46,8 +43,6 @@ import optax import pytest -Transformer = models.transformer_as_linen - class TestGradientClipping(unittest.TestCase): """test class for gradient clipping""" @@ -350,49 +345,28 @@ def setUp(self): self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False) devices_array = maxtext_utils.create_device_mesh(self.config) self.mesh = Mesh(devices_array, self.config.mesh_axes) - quant = quantizations.configure_quantization(self.config) - if self.config.pure_nnx: - self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - else: - self.model = models.transformer_as_linen(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) def test_setup_decode_state(self): - rng = random.PRNGKey(0) - if self.config.pure_nnx: + def create_train_state_fn(): + nnx_model = self._create_model_partial() + return train_state_nnx.TrainStateNNX(nnx_model, None) - def create_train_state_fn(): - nnx_model = self._create_model_partial() - return train_state_nnx.TrainStateNNX(nnx_model, None) - - init_state_fn = create_train_state_fn - else: - init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng) + init_state_fn = create_train_state_fn state, _ = maxtext_utils.setup_decode_state(self.config, self.mesh, None, init_state_fn) - if self.config.pure_nnx: - self.assertNotIn("optimizer", state) - else: - self.assertEqual(state.tx, None) - self.assertEqual(state.opt_state, {}) + self.assertNotIn("optimizer", state) def test_setup_initial_state(self): - rng = random.PRNGKey(0) tx = optax.adam(learning_rate=0.001) - if self.config.pure_nnx: - def create_train_state_fn(): - nnx_model = self._create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + def create_train_state_fn(): + nnx_model = self._create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + 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, self.model, tx, self.config, True, rng) + init_state_fn = create_train_state_fn state, _, _, _ = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn) - if self.config.pure_nnx: - self.assertIsNotNone(state.optimizer) - else: - self.assertEqual(state.tx, tx) - self.assertNotEqual(state.opt_state, {}) + self.assertIsNotNone(state.optimizer) class MaxUtilsPpAsDp(unittest.TestCase): @@ -1020,9 +994,8 @@ def train_step(_model, _config, _state_shardings, _params_shardings, state, _bat return train_step - def _make_mock_config(self, pure_nnx=False): + def _make_mock_config(self): cfg = MagicMock() - cfg.pure_nnx = pure_nnx return cfg def test_returns_five_tuple(self): @@ -1039,20 +1012,11 @@ def test_functional_train_has_correct_name(self): ) self.assertEqual(fn.__name__, "train_step") - def test_linen_in_shardings_includes_rng(self): - """pure_nnx=False: in_shardings should be (state, batch, rng).""" - step = self._make_mock_step() - _, in_shardings, _, _, _ = maxtext_utils.get_functional_train_with_signature( - step, "data_sharding", "state_shardings", "model", self._make_mock_config(pure_nnx=False) - ) - self.assertEqual(len(in_shardings), 3) - self.assertIsNone(in_shardings[2]) # rng sharding is None - def test_nnx_in_shardings_excludes_rng(self): - """pure_nnx=True: in_shardings should be (state, batch) — no rng slot.""" + """in_shardings should be (state, batch) — no rng slot.""" step = self._make_mock_step() _, in_shardings, _, _, _ = maxtext_utils.get_functional_train_with_signature( - step, "data_sharding", "state_shardings", "model", self._make_mock_config(pure_nnx=True) + step, "data_sharding", "state_shardings", "model", self._make_mock_config() ) self.assertEqual(len(in_shardings), 2) @@ -1088,9 +1052,8 @@ def eval_step(_model, _config, _state, _batch, _rng=None): return eval_step - def _make_mock_config(self, pure_nnx=False): + def _make_mock_config(self): cfg = MagicMock() - cfg.pure_nnx = pure_nnx return cfg def test_returns_five_tuple(self): @@ -1118,21 +1081,13 @@ def test_donate_argnums_is_empty(self): self.assertEqual(donate_argnums, ()) def test_nnx_in_shardings_excludes_rng(self): - """pure_nnx=True: in_shardings should be (state, batch) — no rng slot.""" + """in_shardings should be (state, batch) — no rng slot.""" step = self._make_mock_eval_step() _, in_shardings, _, _, _ = maxtext_utils.get_functional_eval_with_signature( - step, "batch_sharding", "state_sharding", "model", self._make_mock_config(pure_nnx=True) + step, "batch_sharding", "state_sharding", "model", self._make_mock_config() ) self.assertEqual(len(in_shardings), 2) - def test_linen_in_shardings_includes_rng(self): - """pure_nnx=False: in_shardings should be (state, batch, rng).""" - step = self._make_mock_eval_step() - _, in_shardings, _, _, _ = maxtext_utils.get_functional_eval_with_signature( - step, "batch_sharding", "state_sharding", "model", self._make_mock_config(pure_nnx=False) - ) - self.assertEqual(len(in_shardings), 3) - @pytest.mark.cpu_only class TestGetShapedBatch(unittest.TestCase): @@ -1387,38 +1342,19 @@ def setUp(self): self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False) devices_array = maxtext_utils.create_device_mesh(self.config) self.mesh = Mesh(devices_array, self.config.mesh_axes) - quant = quantizations.configure_quantization(self.config) - if self.config.pure_nnx: - self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - else: - self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) def test_setup_training_state_returns_train_state(self): - rng = jax.random.PRNGKey(0) tx = optax.adam(learning_rate=0.001) - if self.config.pure_nnx: - def create_train_state_fn(): - nnx_model = self._create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + def create_train_state_fn(): + nnx_model = self._create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + 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, - self.model, - tx, - self.config, - True, - rng, - ) + init_state_fn = create_train_state_fn state, _, _, _ = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn) - if self.config.pure_nnx: - self.assertIsNotNone(state.optimizer) - else: - self.assertEqual(state.tx, tx) - self.assertNotEqual(state.opt_state, {}) + self.assertIsNotNone(state.optimizer) class TestGetLogicalAnnotations(unittest.TestCase): @@ -1428,36 +1364,20 @@ def setUp(self): self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False) devices_array = maxtext_utils.create_device_mesh(self.config) self.mesh = Mesh(devices_array, self.config.mesh_axes) - quant = quantizations.configure_quantization(self.config) - if self.config.pure_nnx: - self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - else: - self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) self.rng = jax.random.PRNGKey(0) self.tx = optax.adam(learning_rate=0.001) def test_returns_partition_spec_tree(self): - if self.config.pure_nnx: - - def create_train_state_fn(): - nnx_model = self._create_model_partial() - optimizer = nnx.Optimizer(nnx_model, self.tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - - init_state_fn = create_train_state_fn - annotations = maxtext_utils_nnx.get_partition_spec_nnx( - maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2] - ) - else: - init_state_fn = functools.partial( - maxtext_utils.init_initial_state, - self.model, - self.tx, - self.config, - True, - self.rng, - ) - annotations = maxtext_utils.get_logical_annotations(self.config, self.mesh, init_state_fn) + def create_train_state_fn(): + nnx_model = self._create_model_partial() + optimizer = nnx.Optimizer(nnx_model, self.tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + + init_state_fn = create_train_state_fn + annotations = maxtext_utils_nnx.get_partition_spec_nnx( + maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2] + ) # Result should be a pytree with PartitionSpec leaves leaves = jax.tree_util.tree_leaves(annotations) self.assertGreater(len(leaves), 0) diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py index deec290fdc..a04da3a39b 100644 --- a/tests/unit/muon_utils_test.py +++ b/tests/unit/muon_utils_test.py @@ -19,7 +19,6 @@ import io import contextlib import unittest -from unittest import mock import jax import jax.numpy as jnp @@ -163,37 +162,6 @@ def test_nnx_verbose_path_executes_print_debug(self): self.assertIn("Muon Dimension Numbers", buf.getvalue()) -class TestGetMuonWeightDimensionNumbersLinen(unittest.TestCase): - """Covers the Linen branch of get_muon_weight_dimension_numbers.""" - - def test_linen_branch_uses_get_abstract_param(self): - """Linen models dispatch to maxtext_utils.get_abstract_param + get_transform_tree.""" - # Build a Linen nn.Module so isinstance(model, nnx.Module) is False. - - class LinenStub(nn.Module): - - @nn.compact - def __call__(self, x): - return x - - model = LinenStub() - - # Mock the heavy get_abstract_param call with a pre-shaped dict that exercises - # both a standard weight path and an excluded path. - fake_abstract_param = { - "params": { - "self_attention": {"out": object()}, - "norm": {"scale": object()}, - }, - } - - with mock.patch.object(muon_utils.maxtext_utils, "get_abstract_param", return_value=fake_abstract_param): - result = muon_utils.get_muon_weight_dimension_numbers(model, config=mock.MagicMock()) - - self.assertEqual(result["params"]["self_attention"]["out"], mdn((0, -2), (-1,))) - self.assertIsNone(result["params"]["norm"]["scale"]) - - class TestPrintStructureDebug(unittest.TestCase): """Covers both branches of get_leaf_info inside _print_structure_debug.""" diff --git a/tests/unit/nnx_decoder_test.py b/tests/unit/nnx_decoder_test.py index 70d59eb6a9..daa6e4567d 100644 --- a/tests/unit/nnx_decoder_test.py +++ b/tests/unit/nnx_decoder_test.py @@ -57,9 +57,6 @@ def _model(self): enable_checkpointing=False, scan_layers=False, num_vocab_tiling=2, - pure_nnx=True, - enable_nnx=True, - pure_nnx_decoder=True, ) model = model_creation_utils.from_config(cfg, devices=jax.devices(), model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(0)) return cfg, model diff --git a/tests/unit/nnx_quant_guard_test.py b/tests/unit/nnx_quant_guard_test.py index 50cac5d349..8451dd1264 100644 --- a/tests/unit/nnx_quant_guard_test.py +++ b/tests/unit/nnx_quant_guard_test.py @@ -12,43 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""qwix + NNX coverage: the config guard and the ToNNX->Linen bridge. +"""qwix + NNX coverage for the ToNNX->Linen bridge. -- Config guard: qwix quantization under pure_nnx requires the pure NNX decoder. - The bridged Linen decoder (pure_nnx_decoder=False) is invisible to qwix, so - quantization/sparsity would silently no-op; validation must reject that combo. -- Bridge: nnx_attrs_to_linen_vars must skip qwix's non-Variable bookkeeping attrs - (qwix_path/qwix_rngs/disable_quant_stats_update) instead of raising. +nnx_attrs_to_linen_vars must skip qwix's non-Variable bookkeeping attrs +(qwix_path/qwix_rngs/disable_quant_stats_update) instead of raising. """ -import sys import unittest import jax.numpy as jnp from flax import nnx -from maxtext.configs import pyconfig from maxtext.layers import nnx_wrappers -from tests.utils.test_helpers import get_test_config_path - - -class QwixNnxQuantGuardTest(unittest.TestCase): - - def _init(self, **overrides): - overrides.setdefault("enable_checkpointing", False) - return pyconfig.initialize([sys.argv[0], get_test_config_path()], **overrides) - - def test_bridged_decoder_with_qwix_quant_raises(self): - with self.assertRaisesRegex(Exception, "pure_nnx_decoder"): - self._init(pure_nnx=True, pure_nnx_decoder=False, use_qwix_quantization=True, quantization="fp8_full") - - def test_pure_nnx_decoder_with_qwix_quant_ok(self): - cfg = self._init(pure_nnx=True, pure_nnx_decoder=True, use_qwix_quantization=True, quantization="fp8_full") - self.assertTrue(cfg.pure_nnx_decoder) - - def test_bridged_decoder_without_quant_ok(self): - cfg = self._init(pure_nnx=True, pure_nnx_decoder=False, quantization="") - self.assertEqual(cfg.quantization, "") class NnxAttrsToLinenVarsBridgeTest(unittest.TestCase): diff --git a/tests/unit/optimizers_test.py b/tests/unit/optimizers_test.py index cfcd8c45f7..b84df9d967 100644 --- a/tests/unit/optimizers_test.py +++ b/tests/unit/optimizers_test.py @@ -243,7 +243,7 @@ def test_model_integration(self, model_name, expected_output): Initializes the specified MaxText model and asserts that the generated Muon dimension numbers match the hardcoded reference. """ - actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True, pure_nnx=False) + actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True) self.assertEqual(actual_output, expected_output) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index c11f5ea3b7..497d47e44f 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -394,149 +394,71 @@ def quantization_config(self, quant, logits_tolerance=2e-1, grad_tolerance=5e-1, cfg = self.init_pyconfig(quantization=quant, **kwargs) ids, decoder_segment_ids, decoder_positions = self.get_data() - if cfg.pure_nnx: - qt_model = model_creation_utils.create_model(cfg, self.mesh, rngs=nnx.Rngs(0)) - if getattr(self.__class__, "_cached_base_results_nnx", None) is None: - base_cfg = self.init_pyconfig(quantization="", **kwargs) - base_model = model_creation_utils.create_model(base_cfg, self.mesh, rngs=nnx.Rngs(0)) - - def loss_base(model): - logits = model( - decoder_input_tokens=ids, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - enable_dropout=False, - ) - return jnp.mean((logits) ** 2) - - grads_base = nnx.grad(loss_base)(base_model) - logits_base = base_model( - decoder_input_tokens=ids, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - enable_dropout=False, - ) - self.__class__._cached_base_results_nnx = (grads_base, logits_base) - - grads_base, logits = self.__class__._cached_base_results_nnx + qt_model = model_creation_utils.create_model(cfg, self.mesh, rngs=nnx.Rngs(0)) + if getattr(self.__class__, "_cached_base_results_nnx", None) is None: + base_cfg = self.init_pyconfig(quantization="", **kwargs) + base_model = model_creation_utils.create_model(base_cfg, self.mesh, rngs=nnx.Rngs(0)) - def loss_quant(model): - logits_q = model( + def loss_base(model): + logits = model( decoder_input_tokens=ids, decoder_positions=decoder_positions, decoder_segment_ids=decoder_segment_ids, enable_dropout=False, ) - return jnp.mean((logits_q) ** 2) + return jnp.mean((logits) ** 2) - grads_quant = nnx.grad(loss_quant)(qt_model) - quant_logits = qt_model( + grads_base = nnx.grad(loss_base)(base_model) + logits_base = base_model( decoder_input_tokens=ids, decoder_positions=decoder_positions, decoder_segment_ids=decoder_segment_ids, enable_dropout=False, ) + self.__class__._cached_base_results_nnx = (grads_base, logits_base) - print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}") - assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance - - # nnx.grad returns a State object which is a mapping of paths to gradients. - # Flatten them to check for tolerance. - grads_base_flat = traversals.flatten_mapping(grads_base) - grads_quant_flat = traversals.flatten_mapping(grads_quant) - - # Filter for param collections to compare only parameters and not stats/buffers if any - # Note: NNX grads structure might contain variables like 'kernel', 'bias'. - # For simplicity we compare all matching keys. - def flatten_and_filter(grads_flat): - return {k: v for k, v in grads_flat.items() if hasattr(v, "shape") and "quant_stats" not in str(k)} - - gb_f = flatten_and_filter(grads_base_flat) - gq_f = flatten_and_filter(grads_quant_flat) - - for k in gb_f: - if k in gq_f: - diff = jnp.abs(gb_f[k] - gq_f[k]).mean() / (jnp.abs(gb_f[k]).mean() + 1e-8) - if diff > grad_tolerance: - print(f"Gradient mismatch for {k}: rel_error = {diff}") - assert diff <= grad_tolerance - else: - qt_model = model_creation_utils.create_model(cfg, self.mesh) - if not hasattr(self.__class__, "_cached_base_results"): - model = model_creation_utils.create_model(self.cfg, self.mesh) - var = model.init( - {"params": self.rng, "aqt": self.rng, "dropout": self.rng}, - ids, - decoder_positions, - decoder_segment_ids, - enable_dropout=False, - mutable=True, - ) - - def loss_base_linen(all_vars, inputs): - logits_b, _ = model.apply( - all_vars, - *inputs, - enable_dropout=False, - rngs={"params": self.rng}, - mutable=True, - ) - return jnp.mean((logits_b) ** 2) - - grads_base_linen = jax.grad(loss_base_linen)(var, (ids, decoder_positions, decoder_segment_ids)) - logits_b, _ = model.apply( - var, - ids, - decoder_positions, - decoder_segment_ids, - enable_dropout=False, - rngs={"params": self.rng}, - mutable=True, - ) - self.__class__._cached_base_results = (grads_base_linen, logits_b) + grads_base, logits = self.__class__._cached_base_results_nnx - grads_base_linen, logits = self.__class__._cached_base_results - - quantized_vars = qt_model.init( - {"params": self.rng, "aqt": self.rng, "dropout": self.rng}, - ids, - decoder_positions, - decoder_segment_ids, + def loss_quant(model): + logits_q = model( + decoder_input_tokens=ids, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, enable_dropout=False, - mutable=True, ) + return jnp.mean((logits_q) ** 2) + + grads_quant = nnx.grad(loss_quant)(qt_model) + quant_logits = qt_model( + decoder_input_tokens=ids, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + enable_dropout=False, + ) - def loss_quant_linen(all_vars, inputs): - logits_q, _ = qt_model.apply( - all_vars, - *inputs, - enable_dropout=False, - rngs={"params": self.rng}, - mutable=True, - ) - return jnp.mean((logits_q) ** 2) + print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}") + assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance - grads_quant_linen = jax.grad(loss_quant_linen)(quantized_vars, (ids, decoder_positions, decoder_segment_ids)) + # nnx.grad returns a State object which is a mapping of paths to gradients. + # Flatten them to check for tolerance. + grads_base_flat = traversals.flatten_mapping(grads_base) + grads_quant_flat = traversals.flatten_mapping(grads_quant) - quant_logits, _ = qt_model.apply( - quantized_vars, - ids, - decoder_positions, - decoder_segment_ids, - enable_dropout=False, - rngs={"params": self.rng}, - mutable=True, - ) - print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}") - assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance - self.print_grad_diff(grads_base_linen["params"], grads_quant_linen["params"]) - self.assertTrue( - self.pytree_allclose( - grads_base_linen["params"], - grads_quant_linen["params"], - tolerance=grad_tolerance, - ) - ) + # Filter for param collections to compare only parameters and not stats/buffers if any + # Note: NNX grads structure might contain variables like 'kernel', 'bias'. + # For simplicity we compare all matching keys. + def flatten_and_filter(grads_flat): + return {k: v for k, v in grads_flat.items() if hasattr(v, "shape") and "quant_stats" not in str(k)} + + gb_f = flatten_and_filter(grads_base_flat) + gq_f = flatten_and_filter(grads_quant_flat) + + for k in gb_f: + if k in gq_f: + diff = jnp.abs(gb_f[k] - gq_f[k]).mean() / (jnp.abs(gb_f[k]).mean() + 1e-8) + if diff > grad_tolerance: + print(f"Gradient mismatch for {k}: rel_error = {diff}") + assert diff <= grad_tolerance @pytest.mark.tpu_only def test_int8_quantization(self): @@ -544,7 +466,7 @@ def test_int8_quantization(self): @pytest.mark.tpu_only def test_int8_quantization_nnx(self): - self.quantization_config("int8", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True) + self.quantization_config("int8") @pytest.mark.tpu_only def test_fp8_quantization(self): @@ -552,7 +474,7 @@ def test_fp8_quantization(self): @pytest.mark.tpu_only def test_fp8_quantization_nnx(self): - self.quantization_config("fp8", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True) + self.quantization_config("fp8") @pytest.mark.tpu_only def test_fp8_full_quantization(self): @@ -560,7 +482,7 @@ def test_fp8_full_quantization(self): @pytest.mark.tpu_only def test_fp8_full_quantization_nnx(self): - self.quantization_config("fp8_full", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True) + self.quantization_config("fp8_full") @pytest.mark.gpu_only @pytest.mark.external_serving @@ -570,7 +492,7 @@ def test_fp8_gpu_quantization(self): @pytest.mark.gpu_only @pytest.mark.external_serving def test_fp8_gpu_quantization_nnx(self): - self.quantization_config("fp8_gpu", grad_tolerance=1.5, enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True) + self.quantization_config("fp8_gpu", grad_tolerance=1.5) @pytest.mark.gpu_only @pytest.mark.external_serving @@ -580,7 +502,7 @@ def test_fp8_nanoo_quantization(self): @pytest.mark.gpu_only @pytest.mark.external_serving def test_fp8_nanoo_quantization_nnx(self): - self.quantization_config("fp8_nanoo", grad_tolerance=1.5, enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True) + self.quantization_config("fp8_nanoo", grad_tolerance=1.5) @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available") @pytest.mark.gpu_only @@ -661,8 +583,6 @@ def test_maybe_quantize_model_pops_intermediates(self): quantization="int8", use_qwix_quantization=True, use_batch_split_schedule=False, - pure_nnx=True, - pure_nnx_decoder=True, micro_batch_size_to_train_on=1, max_target_length=2, ) @@ -699,7 +619,6 @@ def test_nnx_abstract_state_has_no_intermediates(self): enable_checkpointing=False, model_name="deepseek3-tiny", attention="dot_product", - pure_nnx=True, use_qwix_quantization=True, use_qk_clip=True, # This sows QK clip intermediates during the forward pass ) diff --git a/tests/unit/sharding_compare_test.py b/tests/unit/sharding_compare_test.py index 90641550c9..2331d54cf3 100644 --- a/tests/unit/sharding_compare_test.py +++ b/tests/unit/sharding_compare_test.py @@ -12,327 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Compare expected sharding of models with actual sharding of models.""" +"""Compare expected sharding of models with actual sharding of models. -import functools -import hashlib -import json -import os -import jax -import jax.numpy as jnp -from maxtext.configs import pyconfig -from maxtext.utils import maxtext_utils -from maxtext.utils.sharding import clear_input_shardings_dump -# import optax - -from maxtext.layers import quantizations -from maxtext.models import models -from maxtext.optimizers import optimizers -from maxtext.trainers.pre_train.train_compile import get_shaped_inputs, get_topology_mesh, validate_config -from tests.utils.sharding_dump import TEST_CASES, load_json, input_sharding_to_json, named_shardings_to_json, partition_specs_to_json -from tests.utils.test_helpers import get_test_config_path -import pytest - -Transformer = models.transformer_as_linen - - -def compute_checksum(d: dict) -> str: - """Compute a checksum (SHA256) of a dictionary.""" - # Serialize the dictionary into a JSON string (ensuring consistent ordering of keys) - json_str = json.dumps(d, sort_keys=True) - - # Compute the SHA256 checksum of the serialized string - checksum = hashlib.sha256(json_str.encode("utf-8")).hexdigest() - - return checksum - - -def compare_sharding_jsons(json1: dict, model1_name: str, json2: dict, model2_name: str) -> bool: - """Compare two json files and print the differences if any.""" - keys1 = set(json1.keys()) - keys2 = set(json2.keys()) - - only_in_1 = keys1 - keys2 - only_in_2 = keys2 - keys1 - shared_keys = keys1 & keys2 - - has_diff = False - - if only_in_1: - print(f"Keys only in {model1_name}:") - for k in sorted(only_in_1): - print(f" {k}") - has_diff = True - - if only_in_2: - print(f"Keys only in {model2_name}:") - for k in sorted(only_in_2): - print(f" {k}") - has_diff = True - - for key in sorted(shared_keys): - entry1 = json1[key] - entry2 = json2[key] - - if isinstance(entry1, dict) and isinstance(entry2, dict): - mesh1 = entry1.get("mesh", {}) - mesh2 = entry2.get("mesh", {}) - - spec1 = entry1.get("partition_spec", []) - spec2 = entry2.get("partition_spec", []) - - shape1 = entry1.get("shape") - shape2 = entry2.get("shape") - - if mesh1 != mesh2: - print(f"\nMesh mismatch at '{key}':") - print(f" {model1_name}: {mesh1}") - print(f" {model2_name}: {mesh2}") - has_diff = True - - if spec1 != spec2: - print(f"\nPartitionSpec mismatch at '{key}':") - print(f" {model1_name}: {spec1}") - print(f" {model2_name}: {spec2}") - has_diff = True - - if shape1 != shape2: - print(f"\nShape mismatch at '{key}':") - print(f" {model1_name}: {shape1}") - print(f" {model2_name}: {shape2}") - has_diff = True - - else: - print(f"\nFormat mismatch at '{key}':") - print(f" {model1_name} type: {type(entry1)}") - print(f" {model2_name} type: {type(entry2)}") - has_diff = True - - return has_diff - - -# Requires JAX TPU support to generate the simulated TPU topology. -@pytest.mark.cpu_only -@pytest.mark.tpu_backend -@pytest.mark.parametrize("model_name, topology, num_slice, custom_mesh_and_rule, overrides", TEST_CASES) -def test_sharding_dump_for_model( - model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple -) -> None: - """ - Test sharding configurations from train_compile.get_shaped_inputs. - This test verifies that the sharding configurations for various models and topologies remain consistent with golden files. - """ - params = [ - "/deps/MaxText/tests/unit/sharding_compare_test", - get_test_config_path(), - f"compile_topology={topology}", - f"compile_topology_num_slices={num_slice}", - f"model_name={model_name}", - "log_config=false", - "debug_sharding=true", # for input sharding dump - "pure_nnx=False", - "enable_nnx=False", - "pure_nnx_decoder=False", - ] - if custom_mesh_and_rule: - params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}") - if overrides: - params.extend(overrides) - - root_dir = "tests/utils/sharding_info" - rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default" - if overrides: - rule_name += "_" + "_".join(overrides) - base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name) - - named_json_path = os.path.join(base_path, "named_shardings.json") - logical_json_path = os.path.join(base_path, "logical_shardings.json") - input_json_path = os.path.join(base_path, "input_shardings.json") - - if not os.path.exists(named_json_path): - pytest.skip(f"Missing named_shardings.json for {model_name} {topology} slice {num_slice}") - return - if not os.path.exists(logical_json_path): - pytest.skip(f"Missing logical_shardings.json for {model_name} {topology} slice {num_slice}") - return - if not os.path.exists(input_json_path): - pytest.skip(f"Missing input_shardings.json for {model_name} {topology} slice {num_slice}") - return - - config = pyconfig.initialize(params) - validate_config(config) - - clear_input_shardings_dump() - topology_mesh = get_topology_mesh(config) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - optimizers.get_optimizer(config, learning_rate_schedule) - shaped_train_args, _, state_mesh_shardings, logical_shardings, _ = get_shaped_inputs(topology_mesh, config) - - error_messages = [] - - # 1. Compare Named Shardings - actual_named = named_shardings_to_json(state_mesh_shardings, shaped_train_args[0]) - expected_named = load_json(named_json_path) - # calculate checksum - actual_named_sum = compute_checksum(actual_named) - expected_named_sum = compute_checksum(expected_named) - named_match = actual_named_sum == expected_named_sum - - if not named_match: - print(f"\n[FAIL] Physical Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True) - compare_sharding_jsons(expected_named, "Expected (Physical)", actual_named, "Actual (Physical)") - error_messages.append(f" Physical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - # 2. Compare Logical Shardings - actual_logical = partition_specs_to_json(logical_shardings, shaped_train_args[0]) - expected_logical = load_json(logical_json_path) - # calculate checksum - actual_logical_sum = compute_checksum(actual_logical) - expected_logical_sum = compute_checksum(expected_logical) - logical_match = actual_logical_sum == expected_logical_sum - - if not logical_match: - print(f"\n[FAIL] Logical Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True) - compare_sharding_jsons(expected_logical, "Expected (Logical)", actual_logical, "Actual (Logical)") - error_messages.append(f"Logical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - # 3. Compare Input Shardings - actual_input = input_sharding_to_json() - expected_input = load_json(input_json_path) - # calculate checksum - actual_input_sum = compute_checksum(actual_input) - expected_input_sum = compute_checksum(expected_input) - - input_match = actual_input_sum == expected_input_sum - - if not input_match: - print(f"\n[FAIL] Input Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True) - # compare_sharding_jsons(expected_input, "Expected (Input)", actual_input, "Actual (Input)") - error_messages.append(f"Input sharding mismatch for {model_name} on {topology} slice {num_slice}") - - assert not error_messages, "\n".join(error_messages) - - -@pytest.fixture( - scope="module", - params=[pytest.param(case, id=f"{case[0]}-{case[1]}-{case[2]}-{case[3]}-{''.join(case[4])}") for case in TEST_CASES], -) -def abstract_state_and_shardings(request): - """Pytest fixture to set up model, config, and generate abstract state once per test case.""" - model_name, topology, num_slice, custom_mesh_and_rule, overrides = request.param - print( - f"Testing model: {model_name}, topology: {topology}, num_slices: {num_slice}, " - "rule: {custom_mesh_and_rule}, overrides: {overrides}", - flush=True, - ) - params = [ - "/deps/MaxText/tests/unit/sharding_compare_test", - get_test_config_path(), - f"compile_topology={topology}", - f"compile_topology_num_slices={num_slice}", - f"model_name={model_name}", - "weight_dtype=float32", - "pure_nnx=False", - "enable_nnx=False", - "pure_nnx_decoder=False", - ] - if custom_mesh_and_rule: - params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}") - if overrides: - params.extend(overrides) - config = pyconfig.initialize(params) - validate_config(config) - - topology_mesh = get_topology_mesh(config) - quant = quantizations.configure_quantization(config) - model = Transformer(config, mesh=topology_mesh, quant=quant) - - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - # tx = optax.adam(learning_rate=learning_rate_schedule) - tx = optimizers.get_optimizer(config, learning_rate_schedule) - rng = jax.random.PRNGKey(0) - - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng) - - # Get abstract state and physical shardings from maxtext_utils - abstract_state, _, state_mesh_shardings = maxtext_utils.get_abstract_state( - config, topology_mesh, init_state_fn, is_training=True - ) - - # Get logical shardings from maxtext_utils - logical_shardings = maxtext_utils.get_logical_annotations(config, topology_mesh, init_state_fn) - - return ( - model_name, - topology, - num_slice, - custom_mesh_and_rule, - overrides, - abstract_state, - state_mesh_shardings, - logical_shardings, - ) - - -@pytest.mark.cpu_only -@pytest.mark.tpu_backend -class TestGetAbstractState: - """Test class for get_abstract_state function and sharding comparison.""" - - # Requires JAX TPU support to generate the simulated TPU topology. - def test_get_abstract_state_sharding(self, abstract_state_and_shardings): # pylint: disable=redefined-outer-name - """Tests that get_abstract_state returns a state with the correct abstract structure and compares sharding.""" - - ( - model_name, - topology, - num_slice, - custom_mesh_and_rule, - overrides, - abstract_state, - state_mesh_shardings, - logical_shardings, - ) = abstract_state_and_shardings - - assert hasattr(abstract_state, "params") - assert hasattr(abstract_state, "opt_state") - param_leaf = jax.tree_util.tree_leaves(abstract_state.params)[0] - assert isinstance(param_leaf, jax.ShapeDtypeStruct) - assert param_leaf.dtype == jnp.float32 - - root_dir = "tests/utils/sharding_info" # Or your target directory - rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default" - if overrides: - rule_name += "_" + "_".join(overrides) - base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name) - os.makedirs(base_path, exist_ok=True) # Ensure directory exists for saving actual - - error_messages = [] - - # 1. Compare Physical/Named Shardings - named_json_path = os.path.join(base_path, "named_shardings.json") - if not os.path.exists(named_json_path): - pytest.skip(f"Missing named_shardings.json for {model_name} {topology} slice {num_slice}") - return - - # Use state_mesh_shardings from the fixture - actual_named = named_shardings_to_json(state_mesh_shardings, abstract_state) - expected_named = load_json(named_json_path) - - if compare_sharding_jsons(expected_named, "Expected (Physical)", actual_named, "Actual (Physical)"): - error_messages.append(f"Physical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - # 2. Compare Logical Shardings - logical_json_path = os.path.join(base_path, "logical_shardings.json") - if not os.path.exists(logical_json_path): - pytest.skip(f"Missing logical_shardings.json for {model_name} {topology} slice {num_slice}") - return - - # Use logical_shardings from the fixture - actual_logical = partition_specs_to_json(logical_shardings, abstract_state) - expected_logical = load_json(logical_json_path) - - if compare_sharding_jsons(expected_logical, "Expected (Logical)", actual_logical, "Actual (Logical)"): - error_messages.append(f"Logical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - assert not error_messages, "\n".join(error_messages) +The sharding-comparison tests in this file relied on Linen golden files and +Linen TrainState structure, which no longer exist after the NNX-only migration. +They were removed; this module is intentionally left without tests. +""" diff --git a/tests/unit/sharding_nnx_test.py b/tests/unit/sharding_nnx_test.py index fef2659f92..e9950b4e70 100644 --- a/tests/unit/sharding_nnx_test.py +++ b/tests/unit/sharding_nnx_test.py @@ -30,7 +30,6 @@ @dataclass class _Cfg: - pure_nnx: bool = True shard_optimizer_over_data: bool = False @@ -74,9 +73,7 @@ def _to_sharding(var): pspec = PartitionSpec("data", "model") return var.replace(NamedSharding(mesh, pspec)) - return jax.tree.map( - _to_sharding, state, is_leaf=lambda x: isinstance(x, nnx.Variable) - ) + return jax.tree.map(_to_sharding, state, is_leaf=lambda x: isinstance(x, nnx.Variable)) class TestMaybeUpdateParamsShardingWithOptNNX(unittest.TestCase): @@ -85,15 +82,11 @@ class TestMaybeUpdateParamsShardingWithOptNNX(unittest.TestCase): def setUp(self): self.model = _LinearNNX(rngs=nnx.Rngs(0)) - def test_dispatch_from_main_helper_when_pure_nnx(self): + def test_dispatch_from_main_helper(self): """maybe_update_params_sharding_with_opt should dispatch to the NNX variant.""" - cfg = _Cfg(pure_nnx=True, shard_optimizer_over_data=False) - state_mesh_shardings = _build_state_mesh_shardings( - self.model, optax.adam(1e-3) - ) - prev, updated = sharding.maybe_update_params_sharding_with_opt( - cfg, state_mesh_shardings - ) + cfg = _Cfg(shard_optimizer_over_data=False) + state_mesh_shardings = _build_state_mesh_shardings(self.model, optax.adam(1e-3)) + prev, updated = sharding.maybe_update_params_sharding_with_opt(cfg, state_mesh_shardings) # prev is the param-only view (no rngs / non-Param nodes) self.assertIsInstance(prev, nnx.State) self.assertIn("linear", prev) @@ -103,15 +96,9 @@ def test_dispatch_from_main_helper_when_pure_nnx(self): def test_extract_param_only_skips_non_param_variables(self): """prev_params_shardings must contain Params only — RngKey/RngCount/OptVariable filtered out.""" cfg = _Cfg(shard_optimizer_over_data=False) - state_mesh_shardings = _build_state_mesh_shardings( - self.model, optax.adam(1e-3) - ) - prev, _ = sharding.maybe_update_params_sharding_with_opt_nnx( - cfg, state_mesh_shardings - ) - leaves = jax.tree.leaves( - prev, is_leaf=lambda x: isinstance(x, nnx.Variable) - ) + state_mesh_shardings = _build_state_mesh_shardings(self.model, optax.adam(1e-3)) + prev, _ = sharding.maybe_update_params_sharding_with_opt_nnx(cfg, state_mesh_shardings) + leaves = jax.tree.leaves(prev, is_leaf=lambda x: isinstance(x, nnx.Variable)) # Every surviving leaf is wrapped as an nnx.Param. self.assertTrue(all(isinstance(leaf, nnx.Param) for leaf in leaves)) # The model has linear.kernel and linear.bias — exactly two Param leaves. @@ -120,20 +107,14 @@ def test_extract_param_only_skips_non_param_variables(self): def test_returns_unchanged_when_shard_optimizer_over_data_false(self): """When shard_optimizer_over_data=False, the second return value must be the input object.""" cfg = _Cfg(shard_optimizer_over_data=False) - state_mesh_shardings = _build_state_mesh_shardings( - self.model, optax.adam(1e-3) - ) - _, updated = sharding.maybe_update_params_sharding_with_opt_nnx( - cfg, state_mesh_shardings - ) + state_mesh_shardings = _build_state_mesh_shardings(self.model, optax.adam(1e-3)) + _, updated = sharding.maybe_update_params_sharding_with_opt_nnx(cfg, state_mesh_shardings) self.assertIs(updated, state_mesh_shardings) def test_zero1_propagates_mu_sharding_to_model_params(self): """Zero-1: model param shardings must be replaced with the optimizer mu shardings.""" cfg = _Cfg(shard_optimizer_over_data=True) - state_mesh_shardings = _build_state_mesh_shardings( - self.model, optax.adam(1e-3) - ) + state_mesh_shardings = _build_state_mesh_shardings(self.model, optax.adam(1e-3)) # Mutate the optimizer mu leaves in place so the function picks up a distinct PartitionSpec. mesh = _create_2d_test_mesh() @@ -145,20 +126,14 @@ def test_zero1_propagates_mu_sharding_to_model_params(self): # After _build_state_mesh_shardings, every leaf's value is a NamedSharding (no .shape), # so we just override every Variable leaf in mu in place via set_value (modern API). mu_state = state_mesh_shardings.optimizer.opt_state[0]["mu"] - for var in jax.tree.leaves( - mu_state, is_leaf=lambda x: isinstance(x, nnx.Variable) - ): + for var in jax.tree.leaves(mu_state, is_leaf=lambda x: isinstance(x, nnx.Variable)): if isinstance(var, nnx.Variable): var.set_value(new_mu_sharding) - _, updated = sharding.maybe_update_params_sharding_with_opt_nnx( - cfg, state_mesh_shardings - ) + _, updated = sharding.maybe_update_params_sharding_with_opt_nnx(cfg, state_mesh_shardings) # All Param leaves under updated.model must now share the new mu sharding. - param_leaves = jax.tree.leaves( - updated.model, is_leaf=lambda x: isinstance(x, nnx.Variable) - ) + param_leaves = jax.tree.leaves(updated.model, is_leaf=lambda x: isinstance(x, nnx.Variable)) param_leaves = [v for v in param_leaves if isinstance(v, nnx.Param)] self.assertGreater(len(param_leaves), 0) for leaf in param_leaves: @@ -167,13 +142,9 @@ def test_zero1_propagates_mu_sharding_to_model_params(self): def test_raises_when_no_adam_state_present(self): """Stateless optimizers (e.g., SGD) have no mu — function must raise NotImplementedError.""" cfg = _Cfg(shard_optimizer_over_data=True) - state_mesh_shardings = _build_state_mesh_shardings( - self.model, optax.sgd(1e-3) - ) + state_mesh_shardings = _build_state_mesh_shardings(self.model, optax.sgd(1e-3)) with self.assertRaises(NotImplementedError): - sharding.maybe_update_params_sharding_with_opt_nnx( - cfg, state_mesh_shardings - ) + sharding.maybe_update_params_sharding_with_opt_nnx(cfg, state_mesh_shardings) def test_chained_optimizer_recursion_finds_adam_mu(self): """A nested optax.chain(clip, adam) wraps mu under multiple containers — recursion must find it.""" @@ -182,22 +153,18 @@ def test_chained_optimizer_recursion_finds_adam_mu(self): state_mesh_shardings = _build_state_mesh_shardings(self.model, chained) # Should not raise; verify update happens (params replaced with mu shardings). - prev, updated = sharding.maybe_update_params_sharding_with_opt_nnx( - cfg, state_mesh_shardings - ) + prev, updated = sharding.maybe_update_params_sharding_with_opt_nnx(cfg, state_mesh_shardings) self.assertIsInstance(prev, nnx.State) self.assertIsInstance(updated, nnx.State) # Same number of Param leaves before and after. - n_prev = len( - jax.tree.leaves(prev, is_leaf=lambda x: isinstance(x, nnx.Variable)) + n_prev = len(jax.tree.leaves(prev, is_leaf=lambda x: isinstance(x, nnx.Variable))) + n_after = len( + [ + v + for v in jax.tree.leaves(updated.model, is_leaf=lambda x: isinstance(x, nnx.Variable)) + if isinstance(v, nnx.Param) + ] ) - n_after = len([ - v - for v in jax.tree.leaves( - updated.model, is_leaf=lambda x: isinstance(x, nnx.Variable) - ) - if isinstance(v, nnx.Param) - ]) self.assertEqual(n_prev, n_after) @@ -244,9 +211,7 @@ def test_scan_axis_inserted_at_param_scan_axis(self): result_sharding = out["w"].get_value() self.assertIsInstance(result_sharding, NamedSharding) # 'layers' resolves to physical axis 'stage' and is inserted at position 1 (param_scan_axis=1). - self.assertEqual( - result_sharding.spec, PartitionSpec(None, "stage", "fsdp") - ) + self.assertEqual(result_sharding.spec, PartitionSpec(None, "stage", "fsdp")) def test_scan_axis_not_inserted_when_already_present(self): """Guard against double-insertion when partition_name is already in out_sharding.""" @@ -260,9 +225,7 @@ def test_scan_axis_not_inserted_when_already_present(self): out = self._run(self._build_state(w=v)) result_sharding = out["w"].get_value() # 'stage' must appear exactly once — the same PartitionSpec we started with. - self.assertEqual( - result_sharding.spec, PartitionSpec("stage", None, "fsdp") - ) + self.assertEqual(result_sharding.spec, PartitionSpec("stage", None, "fsdp")) def test_masked_node_preserved_as_is(self): """Values without a .shape attribute (e.g., optax.MaskedNode) are returned unchanged.""" @@ -406,9 +369,7 @@ def test_removes_size_one_mesh_axes_no_rules(self): sharding.remove_size_one_mesh_axis = self._old_remove_size_one_mesh_axis try: # Use an explicit 1x1 mesh so physical axis 'fsdp' has size 1 deterministically. - mesh_1x1 = Mesh( - np.array(jax.local_devices()[:1]).reshape(1, 1), ("fsdp", "stage") - ) + mesh_1x1 = Mesh(np.array(jax.local_devices()[:1]).reshape(1, 1), ("fsdp", "stage")) with jax.set_mesh(mesh_1x1): v = nnx.Param( jnp.zeros((4,)), @@ -427,9 +388,7 @@ def test_removes_size_one_mesh_axes(self): sharding.remove_size_one_mesh_axis = self._old_remove_size_one_mesh_axis try: # Use an explicit 1x1 mesh so physical axes 'fsdp' and 'stage' have size 1 deterministically. - mesh_1x1 = Mesh( - np.array(jax.local_devices()[:1]).reshape(1, 1), ("fsdp", "stage") - ) + mesh_1x1 = Mesh(np.array(jax.local_devices()[:1]).reshape(1, 1), ("fsdp", "stage")) rules = (("embed", "fsdp"), ("layers", "stage")) with jax.set_mesh(mesh_1x1), nn_partitioning.axis_rules(rules): v = nnx.Param( diff --git a/tests/unit/state_dtypes_test.py b/tests/unit/state_dtypes_test.py index 21c87c8cae..a92394d99c 100644 --- a/tests/unit/state_dtypes_test.py +++ b/tests/unit/state_dtypes_test.py @@ -14,25 +14,20 @@ """Test that all weights are expected dtype (default float32)""" -from functools import partial import unittest from flax import nnx import jax import jax.numpy as jnp from jax.sharding import Mesh + from maxtext.common import train_state_nnx -from maxtext.common.common_types import MODEL_MODE_TRAIN from maxtext.configs import pyconfig -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.optimizers import optimizers from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils from tests.utils.test_helpers import get_test_config_path -Transformer = models.transformer_as_linen - class StateDtypes(unittest.TestCase): """Tests that state has expected dtypes, e.g. weights default to float32""" @@ -41,47 +36,30 @@ def get_state(self, argv): """Gets model state including weights and optimizer state""" # Setup necessary inputs to build a model state config = pyconfig.initialize(argv) - quant = quantizations.configure_quantization(config) devices_array = maxtext_utils.create_device_mesh(config) mesh = Mesh(devices_array, config.mesh_axes) - if config.pure_nnx: - _create_model_partial, model = ( - model_creation_utils.create_nnx_abstract_model(config, mesh) - ) - else: - model = Transformer( - config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN - ) + _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh) learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) tx = optimizers.get_optimizer(config, learning_rate_schedule, model) - _, example_rng = jax.random.split(jax.random.PRNGKey(0), 2) - - if config.pure_nnx: - def create_train_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + def create_train_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - init_state_fn = create_train_state_fn - else: - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, example_rng) + init_state_fn = create_train_state_fn abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True) - return abstract_state, config.pure_nnx + return abstract_state def get_weights(self, argv): - state, is_nnx = self.get_state(argv) - if is_nnx: - return state.model - return state.params + state = self.get_state(argv) + return state.model def get_mu(self, argv): - state, is_nnx = self.get_state(argv) - if is_nnx: - return state.optimizer.opt_state[0]["mu"] - return state.opt_state[0].mu + state = self.get_state(argv) + return state.optimizer.opt_state[0]["mu"] def assert_pytree_is_dtype(self, weights, expected_dtype): """Asserts that all valid parameter arrays within the PyTree match the expected dtype.""" @@ -96,9 +74,7 @@ def check_dtype(path, leaf): return # Skip PRNG keys - if type(leaf_dtype).__name__ == "KeyTy" or str(leaf_dtype).startswith( - "key<" - ): + if type(leaf_dtype).__name__ == "KeyTy" or str(leaf_dtype).startswith("key<"): return if jnp.issubdtype(leaf_dtype, jnp.integer): diff --git a/tests/unit/tiling_test.py b/tests/unit/tiling_test.py index b1d257c9c5..cdda374708 100644 --- a/tests/unit/tiling_test.py +++ b/tests/unit/tiling_test.py @@ -21,45 +21,22 @@ import unittest import pytest -from flax import linen as nn from flax import nnx from flax.linen import partitioning as nn_partitioning import jax import jax.numpy as jnp -from jax.sharding import Mesh from maxtext.configs import pyconfig -from maxtext.common.common_types import Config from maxtext.common.common_types import MODEL_MODE_TRAIN -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx from maxtext.utils import model_creation_utils -from maxtext.utils.vocabulary_tiling import vocab_tiling_linen_loss, vocab_tiling_nnx_loss +from maxtext.utils.vocabulary_tiling import vocab_tiling_nnx_loss from tests.utils.test_helpers import get_test_config_path -def compute_loss_linen(intermediate_outputs, logits, data, config, model, params, is_train): - """ - A loss function wrapper that deals with both vocab tiling or non-vocab tiling cases - """ - if 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] - total_loss, _ = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train) - else: - one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) - xent, _ = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) - xent = nn.with_logical_constraint(xent, ("activation_embed_and_logits_batch", "activation_length")) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - total_loss = jnp.sum(xent) - return total_loss - - class LossAndGradientCorrectnessTest(unittest.TestCase): """ Unit tests for verifying loss and gradient correctness of: @@ -71,207 +48,18 @@ def setUp(self): """ Set up common configurations and dummy data for the tests. """ - # vocab_tiling on the Linen path uses transformer_as_linen + model.apply, - # so this class must stay on Linen even when NNX defaults are flipped to - # True. The NNX-side equivalents live in VocabTilingNNXTest below. self.base_config = [ None, get_test_config_path(), "base_emb_dim=32", "vocab_size=128", - "enable_nnx=False", - "pure_nnx=False", - "pure_nnx_decoder=False", ] self.rng = jax.random.PRNGKey(1234) self.batch_size = 1 self.seq_len = 64 - self.dummy_inputs = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) self.rtol = 1e-2 self.atol = 1e-2 - def get_grads(self, cfg: Config, params, data): - """ - Computes and returns the gradients for a given configuration and set of parameters. - """ - quant = quantizations.configure_quantization(cfg) - devices_array = maxtext_utils.create_device_mesh(cfg) - mesh = Mesh(devices_array, cfg.mesh_axes) - model = models.transformer_as_linen(cfg, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - - @jax.jit - def grad_fn(p, d): - def loss_fn(train_params): - logits, intermediate_outputs = model.apply( - train_params, - decoder_input_tokens=self.dummy_inputs, - decoder_positions=self.dummy_inputs, - mutable=["intermediates"], - ) - return compute_loss_linen(intermediate_outputs, logits, d, cfg, model, train_params, is_train=True) - - return jax.value_and_grad(loss_fn)(p) - - return grad_fn(params, data) - - def assert_pytrees_all_close(self, tree1, tree2, error_message=None): - """Manually assert that two pytrees are all-close.""" - - # Map jnp.allclose to every leaf - # This creates a new pytree of the same structure, but with boolean leaves - leaves_are_close = jax.tree_util.tree_map( - lambda x, y: jnp.allclose(x, y, rtol=self.rtol, atol=self.atol), - tree1, - tree2, - ) - - # Get a flat list of all boolean leaves - all_leaves = jax.tree_util.tree_leaves(leaves_are_close) - - # Check if every single leaf was True - all_are_close = all(all_leaves) - - # Assert the final result - if not all_are_close: - # For a better error, you could find the first False leaf - print("Pytrees are not all-close. Mismatched leaves:") - # This part is more complex, which is why chex is preferred - mismatches = jax.tree_util.tree_map( - lambda x, y, z: "MISMATCH" if not z else "MATCH", tree1, tree2, leaves_are_close - ) - print(mismatches) - raise AssertionError(error_message) - - @pytest.mark.tpu_only - def test_gradient_accumulation(self): - """ - Tests GA loss and gradient correctness. - """ - cfg_non_ga = pyconfig.initialize( - self.base_config, - run_name="non_GA_grad_test", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=4, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - gradient_accumulation_steps=1, - ) - quant_non_ga = quantizations.configure_quantization(cfg_non_ga) - devices_array_non_ga = maxtext_utils.create_device_mesh(cfg_non_ga) - mesh_non_ga = Mesh(devices_array_non_ga, cfg_non_ga.mesh_axes) - model_non_ga = models.transformer_as_linen( - cfg_non_ga, mesh=mesh_non_ga, quant=quant_non_ga, model_mode=MODEL_MODE_TRAIN - ) - - rng_model, rng_targets = jax.random.split(self.rng) - - params = model_non_ga.init( - {"params": rng_model, "dropout": rng_model}, - self.dummy_inputs, - self.dummy_inputs, - ) - - data = { - "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_ga.vocab_size), - "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), - } - - loss_non_ga, grads_non_ga = self.get_grads(cfg_non_ga, params, data) - - cfg_ga = pyconfig.initialize( - self.base_config, - run_name="GA_grad_test", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=1, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - gradient_accumulation_steps=4, - ) - loss_ga, grads_ga = self.get_grads(cfg_ga, params, data) - # Loss correctness test - assert jnp.allclose(loss_non_ga, loss_ga, rtol=self.rtol), "Losses do not match for gradient accumulation test." - - # Gradient correctness test - self.assert_pytrees_all_close( - grads_non_ga, - grads_ga, - "Gradients of embedding table do not match for GA.", - ) - - @pytest.mark.tpu_only - def test_vocab_tiling_gradient_with_z_loss(self): - """ - Tests loss and gradient correctness when z-loss is enabled, comparing - standard computation vs. vocabulary tiling computation. - """ - cfg_non_tiling = pyconfig.initialize( - self.base_config, - run_name="grad_test_z_loss_no_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=False, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=1, - z_loss_multiplier=1e-4, # Enable z-loss - ) - quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) - devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) - mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) - model_non_tiling = models.transformer_as_linen( - cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN - ) - - rng_model, rng_targets = jax.random.split(self.rng) - - params = model_non_tiling.init( - {"params": rng_model, "dropout": rng_model}, - self.dummy_inputs, - self.dummy_inputs, - ) - - data = { - "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), - "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), - } - - loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) - - cfg_tiling = pyconfig.initialize( - self.base_config, - run_name="grad_test_z_loss_with_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=False, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=4, - z_loss_multiplier=1e-4, # Enable z-loss - ) - loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) - - # Loss correctness test - assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match when z-loss is enabled." - - # Gradient correctness test - self.assert_pytrees_all_close( - grads_non_tiling, - grads_tiling, - "Gradients do not match for vocab tiling when z-loss is enabled.", - ) - @pytest.mark.tpu_only def test_vocab_tiling_nnx_loss(self): """ @@ -291,8 +79,6 @@ def test_vocab_tiling_nnx_loss(self): matmul_precision="high", num_vocab_tiling=4, z_loss_multiplier=1e-4, - enable_nnx=True, - pure_nnx=True, ) rng_model, rng_hidden, rng_targets = jax.random.split(self.rng, 3) rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=rng_model) @@ -317,328 +103,6 @@ def test_vocab_tiling_nnx_loss(self): xent_sum_tiled, xent_sum_ref, rtol=self.rtol, atol=self.atol ), f"NNX vocab tiling loss {xent_sum_tiled} does not match non-tiled reference {xent_sum_ref}." - @pytest.mark.tpu_only - def test_vocab_tiling_gradient_non_tied_embedding(self): - """ - Tests loss and gradient correctness for a model with non-tied embeddings (FSDP). - """ - cfg_non_tiling = pyconfig.initialize( - self.base_config, - run_name="grad_test_non_tied_no_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=False, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=1, - ) - quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) - devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) - mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) - model_non_tiling = models.transformer_as_linen( - cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN - ) - - rng_model, rng_targets = jax.random.split(self.rng) - - params = model_non_tiling.init( - {"params": rng_model, "dropout": rng_model}, - self.dummy_inputs, - self.dummy_inputs, - ) - - data = { - "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), - "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), - } - - loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) - - cfg_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_non_tied_with_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=False, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=4, - ) - loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) - # Loss correctness test - assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for non-tied embeddings." - - # Gradient correctness test - self.assert_pytrees_all_close( - grads_non_tiling, - grads_tiling, - "Gradients of embedding table do not match for non-tied embeddings.", - ) - - @pytest.mark.tpu_only - @pytest.mark.scheduled_only - def test_vocab_tiling_gradient_tied_embedding(self): - """ - Tests loss and gradient correctness for a model with tied embeddings (FSDP). - """ - cfg_non_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_tied_no_tiling", - enable_checkpointing=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=True, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=1, - ) - - quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) - devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) - mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) - model_non_tiling = models.transformer_as_linen( - cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN - ) - - rng_model, rng_targets = jax.random.split(self.rng) - - params = model_non_tiling.init( - {"params": rng_model, "dropout": rng_model}, - self.dummy_inputs, - self.dummy_inputs, - ) - - data = { - "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), - "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), - } - - loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) - - cfg_tiling = pyconfig.initialize( - self.base_config, - run_name="grad_test_tied_with_tiling", - enable_checkpointing=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=True, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=4, - ) - loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) - - assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for tied embeddings." - - self.assert_pytrees_all_close( - grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for tied embeddings." - ) - - @pytest.mark.tpu_only - @pytest.mark.scheduled_only - def test_vocab_tiling_gradient_data_parallelism(self): - """ - Tests loss and gradient correctness for data parallelism sharding. - """ - cfg_non_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_dp_non_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - ici_data_parallelism=4, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=1, - ) - quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) - devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) - mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) - model_non_tiling = models.transformer_as_linen( - cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN - ) - - rng_model, rng_targets = jax.random.split(self.rng) - - params = model_non_tiling.init( - {"params": rng_model, "dropout": rng_model}, - self.dummy_inputs, - self.dummy_inputs, - ) - - data = { - "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), - "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), - } - - loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) - - cfg_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_dp_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=False, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - ici_data_parallelism=4, - num_vocab_tiling=4, - ) - loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) - # Loss correctness test - assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for data parallelism." - - # Gradient correctness test - self.assert_pytrees_all_close( - grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for data parallelism." - ) - - @pytest.mark.tpu_only - @pytest.mark.scheduled_only - def test_vocab_tiling_gradient_tensor_parallelism(self): - """ - Tests loss and gradient correctness for tensor parallelism sharding. - """ - cfg_non_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_tp_non_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - ici_tensor_parallelism=4, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=1, - ) - quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) - devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) - mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) - model_non_tiling = models.transformer_as_linen( - cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN - ) - - rng_model, rng_targets = jax.random.split(self.rng) - - params = model_non_tiling.init( - {"params": rng_model, "dropout": rng_model}, - self.dummy_inputs, - self.dummy_inputs, - ) - - data = { - "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), - "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), - } - - loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) - - cfg_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_tp_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=False, - base_num_decoder_layers=0, - dtype="float32", - matmul_precision="high", - ici_tensor_parallelism=4, - num_vocab_tiling=4, - ) - loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) - # Loss correctness test - assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for tensor parallelism." - - # Gradient correctness test - self.assert_pytrees_all_close( - grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for tensor parallelism." - ) - - @pytest.mark.tpu_only - @pytest.mark.scheduled_only - def test_vocab_tiling_gradient_context_parallelism(self): - """ - Tests loss and gradient correctness for context parallelism sharding. - """ - cfg_non_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_cp_non_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - ici_context_parallelism=4, - base_num_decoder_layers=0, - dataset_type="synthetic", - packing=False, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=1, - ) - quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) - devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) - mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) - model_non_tiling = models.transformer_as_linen( - cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN - ) - - rng_model, rng_targets = jax.random.split(self.rng) - - params = model_non_tiling.init( - {"params": rng_model, "dropout": rng_model}, - self.dummy_inputs, - self.dummy_inputs, - ) - - data = { - "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), - "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), - } - - loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) - - cfg_tiling = pyconfig.initialize( - self.base_config, - run_name="value_and_grad_test_cp_tiling", - enable_checkpointing=False, - enable_dropout=False, - max_target_length=self.seq_len, - per_device_batch_size=self.batch_size, - logits_via_embedding=False, - base_num_decoder_layers=0, - ici_context_parallelism=4, - dataset_type="synthetic", - packing=False, - dtype="float32", - matmul_precision="high", - num_vocab_tiling=4, - ) - loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) - - # Loss correctness test - assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for context parallelism." - - # Gradient correctness test - self.assert_pytrees_all_close( - grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for context parallelism." - ) - class VocabTilingNNXTest(unittest.TestCase): """Loss + gradient parity for the NNX vocab-tiling `custom_vjp` path. @@ -683,9 +147,6 @@ def _build_cfg_and_model( matmul_precision="high", num_vocab_tiling=num_vocab_tiling, z_loss_multiplier=z_loss_multiplier, - pure_nnx=True, - enable_nnx=True, - pure_nnx_decoder=True, ) mesh = maxtext_utils.get_mesh_from_config(cfg) rngs = maxtext_utils_nnx.create_nnx_rngs(cfg) diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index 2866a7b6f4..1957695c70 100644 --- a/tests/unit/train_compile_test.py +++ b/tests/unit/train_compile_test.py @@ -30,7 +30,6 @@ from tempfile import gettempdir, NamedTemporaryFile -from maxtext.configs import pyconfig from maxtext.trainers.pre_train.train_compile import main as train_compile_main from tests.utils.test_helpers import get_test_config_path @@ -515,10 +514,6 @@ def test_moe_dense_int8(self): @pytest.mark.cpu_only def test_moe_pp_bf16(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Pipeline parallelism not supported for pure_nnx_decoder=True") - temp_dir = gettempdir() compiled_trainstep_file = os.path.join(temp_dir, "test_moe_pp_bf16.pickle") train_compile_main( @@ -615,10 +610,6 @@ def test_moe_deepseek_with_device_limit(self): @pytest.mark.cpu_only def test_moe_deepseek_pipeline_subset(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Pipeline parallelism not supported for pure_nnx_decoder=True") - compiled_trainstep_file = "/tmp/test_moe_deepseek_pipeline_subset.pickle" train_compile_main( ( @@ -642,10 +633,7 @@ def test_moe_deepseek_pipeline_subset(self): @pytest.mark.cpu_only def test_pipeline_subset(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Test not supported for pure_nnx_decoder=True") - + pytest.skip("Pipeline parallelism not yet supported on NNX (pending NNX pipeline parallelism, PR11.5).") compiled_trainstep_file = "/tmp/test_pipeline_subset.pickle" train_compile_main( ( @@ -802,32 +790,6 @@ def test_deepseek32(self): ) ) - @parameterized.named_parameters( - {"testcase_name": "scanned", "scan_layers": "true"}, - ) - @pytest.mark.cpu_only - def test_deepseek4(self, scan_layers): - # test deepseek4 compile (Linen-only: DeepSeek NNX decoder rewrite is a follow-up PR). - compiled_trainstep_file = f"/tmp/test_deepseek4_{scan_layers}.pickle" - train_compile_main(( - "", - get_test_config_path(), - f"compiled_trainstep_file={compiled_trainstep_file}", - "compile_topology=v5p-256", - "use_iota_embed=true", - "compile_topology_num_slices=1", - "model_name=deepseek4-284b", - "per_device_batch_size=1", - "max_target_length=1024", - f"scan_layers={scan_layers}", - "attention=dot_product", - "dtype=bfloat16", - "weight_dtype=bfloat16", - "enable_nnx=False", - "pure_nnx=False", - "pure_nnx_decoder=False", - )) - @pytest.mark.cpu_only def test_indexer_dense_warmup(self): # test deepseek3.2 with sparse attention @@ -948,10 +910,6 @@ def test_engram_integration(self): @pytest.mark.cpu_only def test_circular_pipeline_ag_per_repeat_ep_ds(self): - cfg = pyconfig.initialize([None, get_test_config_path()]) - if getattr(cfg, "pure_nnx_decoder", False): - pytest.skip("Pipeline parallelism not supported for pure_nnx_decoder=True") - temp_dir = gettempdir() compiled_trainstep_file = os.path.join(temp_dir, "test_circular_pipeline_ag_per_repeat_ep_ds.pickle") train_compile_main( @@ -1145,13 +1103,7 @@ def test_zero1_optimizer_sharding(self): @pytest.mark.cpu_only def test_vocab_tiling_bf16_nnx(self): - """AOT compile vocab tiling on the NNX path (vocab_tiling_nnx_loss + custom_vjp). - - Sets `pure_nnx`/`enable_nnx`/`pure_nnx_decoder` explicitly so the NNX AOT - path is covered regardless of the default values. Once those defaults flip - to True, `test_vocab_tiling_bf16` above will already exercise this same - path via defaults. - """ + """AOT compile vocab tiling on the NNX path (vocab_tiling_nnx_loss + custom_vjp).""" compiled_trainstep_file = "/tmp/test_vocab_tiling_bf16_nnx.pickle" train_compile_main( ( @@ -1165,8 +1117,5 @@ def test_vocab_tiling_bf16_nnx(self): "max_target_length=1024", "num_vocab_tiling=4", "weight_dtype=bfloat16", - "pure_nnx=true", - "enable_nnx=true", - "pure_nnx_decoder=true", ) ) diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py index fdb94046c7..bdaaef2119 100644 --- a/tests/unit/train_state_nnx_checkpoint_test.py +++ b/tests/unit/train_state_nnx_checkpoint_test.py @@ -338,12 +338,10 @@ class TestMaybeSaveCheckpointStepAlignment(unittest.TestCase): """Verify maybe_save_checkpoint's fallback step matches the last completed step. When the training loop's final save calls maybe_save_checkpoint without an - explicit `step`, it derives `actual_step` from the state: - - NNX: int(state.optimizer.step) - 1 - - Linen: int(state.step) - 1 - Both TrainStateNNX.apply_gradients (via nnx.Optimizer.update) and Linen - TrainState.apply_gradients increment the counter by 1 per call, so after N - gradient applications the counter is N and the "last completed step" is N-1. + explicit `step`, it derives `actual_step` from int(state.optimizer.step) - 1. + TrainStateNNX.apply_gradients (via nnx.Optimizer.update) increments the + counter by 1 per call, so after N gradient applications the counter is N and + the "last completed step" is N-1. """ N_STEPS = 5 @@ -367,25 +365,10 @@ def loss_fn(m): # (train_step returns nnx.state(new_state)). return nnx.state(state) - def _build_linen_state(self, num_steps): - """Build a Linen TrainState after num_steps gradient applications.""" - model = LinenMockModel() - variables = model.init(jax.random.key(0), jnp.ones((1, 2))) - state = train_state.TrainState.create(apply_fn=model.apply, params=variables["params"], tx=self.tx) - grads = jax.tree.map(jnp.ones_like, state.params) - for _ in range(num_steps): - state = state.apply_gradients(grads=grads) - return state - - def _invoke_maybe_save(self, state, pure_nnx): + def _invoke_maybe_save(self, state): """Call maybe_save_checkpoint with save_checkpoint patched, return {step, state} captured.""" # checkpoint_period=1 keeps force_ckpt_save False regardless of actual_step. - config = SimpleNamespace( - pure_nnx=pure_nnx, - checkpoint_period=1, - async_checkpointing=False, - enable_diloco=False, - ) + config = SimpleNamespace(checkpoint_period=1, async_checkpointing=False, enable_diloco=False) mgr = mock.MagicMock() mgr.reached_preemption.return_value = False @@ -403,30 +386,15 @@ def fake_save_checkpoint(_mgr, step, state_arg, *_args, **_kwargs): def test_nnx_final_save_step_is_n_minus_1(self): state = self._build_nnx_state(self.N_STEPS) self.assertEqual(int(state.optimizer.step.value), self.N_STEPS) - captured = self._invoke_maybe_save(state, pure_nnx=True) - self.assertEqual(captured["step"], self.N_STEPS - 1) - - def test_linen_final_save_step_is_n_minus_1(self): - state = self._build_linen_state(self.N_STEPS) - self.assertEqual(int(state.step), self.N_STEPS) - captured = self._invoke_maybe_save(state, pure_nnx=False) + captured = self._invoke_maybe_save(state) self.assertEqual(captured["step"], self.N_STEPS - 1) - def test_nnx_and_linen_agree_on_actual_step(self): - """TrainStateNNX and Linen TrainState must yield the same fallback actual_step.""" - nnx_state = self._build_nnx_state(self.N_STEPS) - linen_state = self._build_linen_state(self.N_STEPS) - self.assertEqual( - self._invoke_maybe_save(nnx_state, pure_nnx=True)["step"], - self._invoke_maybe_save(linen_state, pure_nnx=False)["step"], - ) - def test_nnx_state_is_saved_in_linen_layout(self): - """For pure_nnx=True, maybe_save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" + """maybe_save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" state = self._build_nnx_state(self.N_STEPS) self.assertIsInstance(state, nnx.State) # precondition: NNX train_step returns an nnx.State - captured = self._invoke_maybe_save(state, pure_nnx=True) + captured = self._invoke_maybe_save(state) # save_checkpoint should receive a plain dict in Linen layout, not the nnx.State. self.assertIsInstance(captured["state"], dict) @@ -439,24 +407,12 @@ def test_nnx_state_is_saved_in_linen_layout(self): self.assertNotIn("optimizer", captured["state"]) self.assertIn("params", captured["state"]["params"]) - def test_linen_state_is_passed_through_unchanged(self): - """For pure_nnx=False, maybe_save_checkpoint must pass the original TrainState object through.""" - state = self._build_linen_state(self.N_STEPS) - captured = self._invoke_maybe_save(state, pure_nnx=False) - # Linen path must not invoke to_pure_dict(); state is forwarded as-is. - self.assertIs(captured["state"], state) - def test_maybe_save_checkpoint_skips_if_already_saved(self): """Verify maybe_save_checkpoint skips saving if latest_step matches actual_step.""" state = self._build_nnx_state(self.N_STEPS) actual_step = self.N_STEPS - 1 - config = SimpleNamespace( - pure_nnx=True, - checkpoint_period=1, - async_checkpointing=False, - enable_diloco=False, - ) + config = SimpleNamespace(checkpoint_period=1, async_checkpointing=False, enable_diloco=False) mgr = mock.MagicMock() mgr.reached_preemption.return_value = False # Mock latest_step to return the same actual_step @@ -475,12 +431,7 @@ def test_maybe_save_checkpoint_saves_if_not_already_saved(self): state = self._build_nnx_state(self.N_STEPS) actual_step = self.N_STEPS - 1 - config = SimpleNamespace( - pure_nnx=True, - checkpoint_period=1, - async_checkpointing=False, - enable_diloco=False, - ) + config = SimpleNamespace(checkpoint_period=1, async_checkpointing=False, enable_diloco=False) mgr = mock.MagicMock() mgr.reached_preemption.return_value = False # Mock latest_step to return a different step (or None) diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index 3236486f74..7fc92e429d 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -37,7 +37,6 @@ """Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs""" import argparse -import functools import os from pathlib import Path import sys @@ -49,8 +48,6 @@ from maxtext.utils.globals import MAXTEXT_TEST_ASSETS_ROOT from maxtext.checkpoint_conversion.utils.hf_utils import convert_jax_weight_to_torch from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_TRAIN -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import max_logging from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils @@ -247,26 +244,21 @@ def get_data(golden_data_point, config): def main(config, test_args): # pylint: disable=W0621 """Test the Whole Model of model_name""" init_rng = jax.random.PRNGKey(config.init_weights_seed) - init_rng, rng1 = jax.random.split(init_rng) + init_rng, _ = jax.random.split(init_rng) devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) if not test_args.run_hf_model: """Comparing maxtext/huggingface model with pre-loaded golden logitis""" max_logging.log("Initializing MaxText model") - quant = quantizations.configure_quantization(config) - if config.pure_nnx_decoder and config.enable_nnx: - model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) - - if config.lora.enable_lora: - model = lora_utils.apply_lora_to_model(model, mesh, config) - if config.lora.lora_restore_path: - lora_utils.restore_lora_from_path(model, config) - state = None - else: - model = models.transformer_as_linen(config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, config, False, rng1) - state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn) + model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + + if config.lora.enable_lora: + model = lora_utils.apply_lora_to_model(model, mesh, config) + if config.lora.lora_restore_path: + mock_trainer = type("MockTrainer", (), {"model": model, "train_steps": 0}) + lora_utils.restore_lora_from_path(mock_trainer, config) + state = None if test_args.golden_logits_path == "": input_golden_data_path = os.path.join( @@ -461,22 +453,14 @@ def main(config, test_args): # pylint: disable=W0621 if any(config.model_name.startswith(prefix) for prefix in pad_token_prefixes): tokenizer.pad_token = tokenizer.eos_token - quant = quantizations.configure_quantization(config) - if config.pure_nnx_decoder and config.enable_nnx: - maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) - if config.lora.enable_lora: - maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config) - if config.lora.lora_restore_path: - lora_utils.restore_lora_from_path(maxtext_model, config) - maxtext_state = None - else: - maxtext_model = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, rng1) - if test_args.ckpt_type == "linen": - maxtext_state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn) - else: - maxtext_state, _ = model_creation_utils.setup_decode_state_from_nnx(maxtext_model, config, rng1, mesh) + if config.lora.enable_lora: + maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config) + if config.lora.lora_restore_path: + mock_trainer = type("MockTrainer", (), {"model": maxtext_model, "train_steps": 0}) + lora_utils.restore_lora_from_path(mock_trainer, config) + maxtext_state = None prompts = ["I love to", "Today is a", "What is the"] all_data_to_save = [] diff --git a/tests/utils/run_sharding_dump.py b/tests/utils/run_sharding_dump.py index 62c71a9b5b..7d3156fe00 100644 --- a/tests/utils/run_sharding_dump.py +++ b/tests/utils/run_sharding_dump.py @@ -59,12 +59,9 @@ flags.DEFINE_string("topology", None, "Specific topology to dump.") flags.DEFINE_string("num_slice", None, "Specific number of slices to dump.") flags.DEFINE_string("custom_mesh_and_rule", None, "Specific custom_mesh_and_rule to dump.") -flags.DEFINE_bool("pure_nnx", False, "Use pure NNX model.") -def run_single_dump( - model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple, pure_nnx: bool = False -) -> None: +def run_single_dump(model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple) -> None: """Generate sharding json file for one specific model, topology, slice and rule.""" args = [ "python3", @@ -82,8 +79,6 @@ def run_single_dump( args.append(f"custom_mesh_and_rule={custom_mesh_and_rule}") if overrides: args.extend(overrides) - if pure_nnx: - args.append("pure_nnx=true") subprocess.run(args, check=True) @@ -122,7 +117,7 @@ def main(argv: Sequence[str]) -> None: print(" -> Sharding files already exist. Regenerating to overwrite.") try: - run_single_dump(model_name, topology, str(num_slice), custom_mesh_and_rule, overrides, pure_nnx=FLAGS.pure_nnx) + run_single_dump(model_name, topology, str(num_slice), custom_mesh_and_rule, overrides) except subprocess.CalledProcessError: print(f"!!! FAILED: {model_name} {topology} {num_slice} {custom_mesh_and_rule} overrides={overrides}") From 91dfd98eca4d7e08c008cc891e6203f1adbf628d Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Thu, 4 Jun 2026 18:43:53 +0000 Subject: [PATCH 4/4] [NNX] Delete Linen (4/4): remove the pure_nnx/enable_nnx/pure_nnx_decoder config flags Remove the three flags from types.py, base.yml, inference/vllm.yml, pyconfig, and the post-train distillation configs. NNX is the only path; the flags no longer exist. --- src/maxtext/configs/base.yml | 5 ----- src/maxtext/configs/inference/vllm.yml | 2 -- .../configs/post_train/distillation_gpt_oss_20b.yml | 1 - .../post_train/distillation_qwen3_30b_base.yml | 1 - .../distillation_qwen3_30b_base_pdbs8.yml | 1 - src/maxtext/configs/pyconfig_deprecated.py | 7 ++----- src/maxtext/configs/types.py | 13 ------------- 7 files changed, 2 insertions(+), 28 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82e448c530..84dc073e92 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1235,11 +1235,6 @@ position_id_per_seconds: 25 # Example: "8,8" to use a 8x8 subgrid (64 chips) of a full pod (16x16) of trillium. subslice_shape: "" -# NNX -enable_nnx: true -pure_nnx_decoder: true -pure_nnx: true - ################################## Qwen3-Next Specific Configs ################################## # Kernel size for the 1D convolution in the Gated Delta Net gdn_conv_kernel_dim: 4 diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml index 44efdcb161..cefba1d82f 100644 --- a/src/maxtext/configs/inference/vllm.yml +++ b/src/maxtext/configs/inference/vllm.yml @@ -16,8 +16,6 @@ base_config: "base.yml" attention: "vllm_rpa" model_call_mode: "inference" -# NNX required for vLLM integration -enable_nnx: true # Avoid re-initializing JAX distributed system when using vLLM skip_jax_distributed_system: true # Scanned layers are not supported with vLLM integration diff --git a/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml b/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml index b5cf9ff0f3..4ebed81f1e 100644 --- a/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml +++ b/src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml @@ -19,7 +19,6 @@ distill_alpha: 0.5 distill_temperature: 1.0 distill_beta: 0 distill_layer_indices: [] -enable_nnx: True load_balance_loss_weight: 0.001 # Megablox grouped-matmul m-tile (batch_seq). The k/n dims already default to diff --git a/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml b/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml index e3c12ff08c..3bd8aa826f 100644 --- a/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml +++ b/src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml @@ -21,7 +21,6 @@ distill_alpha: 0.6 distill_temperature: 1.0 distill_beta: 1.0 distill_layer_indices: [0,1,2,3,4,5,6,7] -enable_nnx: True load_balance_loss_weight: 0.001 ici_fsdp_parallelism: -1 diff --git a/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml b/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml index 36d48da2ed..0e4ee36a0d 100644 --- a/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml +++ b/src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml @@ -21,7 +21,6 @@ distill_alpha: 0.6 distill_temperature: 1.0 distill_beta: 1.0 distill_layer_indices: [0,1,2,3,4,5,6,7] -enable_nnx: True load_balance_loss_weight: 0.001 ici_fsdp_parallelism: -1 diff --git a/src/maxtext/configs/pyconfig_deprecated.py b/src/maxtext/configs/pyconfig_deprecated.py index 3e6329b7aa..1395d6ad1a 100644 --- a/src/maxtext/configs/pyconfig_deprecated.py +++ b/src/maxtext/configs/pyconfig_deprecated.py @@ -193,8 +193,7 @@ def validate_expert_shard_attention_option(expert_shard_attention_option: str) - ) -def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int, enable_nnx: bool): - del enable_nnx # NNX vocab tiling supported via vocab_tiling_nnx_loss in vocabulary_tiling.py +def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int): if (per_device_batch_size * max_target_length) % num_vocab_tiling != 0: raise ValueError("Per device batch size times sequence length should be divisible by the number of vocab tiles.") @@ -238,9 +237,7 @@ def validate_keys(keys): validate_model_call_mode(keys["model_call_mode"]) validate_prefill_and_target_lengths(keys["max_prefill_predict_length"], keys["max_target_length"]) validate_rope_type(keys["rope_type"]) - validate_vocab_tiling( - keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"], keys["enable_nnx"] - ) + validate_vocab_tiling(keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"]) if keys["enable_rampup_batch_size"]: validate_rampup_batch_size( keys["per_device_batch_size_start"], diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a8c1745..d29bed2572 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -988,11 +988,8 @@ class HardwareAndMesh(BaseModel): CustomRule.DEFAULT, description="Customized mesh and logical rules for granularity." ) allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.") - enable_nnx: bool = Field(True, description="Whether to use NNX for model definition.") optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.") shardy: bool = Field(True, description="Whether to use shardy XLA backend.") - pure_nnx_decoder: bool = Field(True, description="Whether to enable pure NNX decoder.") - pure_nnx: bool = Field(True, description="Whether to enable pure NNX mode.") remove_size_one_mesh_axis_from_type: bool = Field( True, description="Whether to remove size one mesh axis from type through jax.config." ) @@ -2708,16 +2705,6 @@ def validate_and_set_hlo_dump_defaults(): if self.distill_beta > 0.0: if not self.scan_layers: raise ValueError("a value of self.distill_beta > 0.0 requires self.scan_layers = True") - if not self.enable_nnx: - raise ValueError("a value of self.distill_beta > 0.0 requires self.enable_nnx = True") - - if self.pure_nnx and not self.pure_nnx_decoder and self.use_qwix_quantization and not self.use_batch_split_schedule: - if self.quantization: - raise ValueError( - f"quantization='{self.quantization}' with use_qwix_quantization=True under pure_nnx=True requires " - "pure_nnx_decoder=True. The bridged Linen decoder (pure_nnx_decoder=False) is invisible to Qwix, " - "so quantization (and weight sparsity) would silently have no effect. Set pure_nnx_decoder=True." - ) # Validate distillation schedule parameters if self.distill_alpha_end is not None and not 0.0 <= self.distill_alpha_end <= 1.0: