From 8fcb928c0db1f11407a65050634829bc2b0e1292 Mon Sep 17 00:00:00 2001 From: Dipak Gaikwad Date: Fri, 17 Jul 2026 22:10:40 +0000 Subject: [PATCH] Implement DSV4 auxillary loss free and sequence-wise load balancing via pure NNX MoEBiasVar This commit migrates the DeepSeek V4 auxiliary-loss-free routing bias to a pure nnx.Variable (MoEBiasVar), automatically isolating it from the optimizer and standard sequence-wise gradients without requiring the jax.lax.stop_gradient hacks or Optax global masking. DeepSeek V3 backward compatibility remains completely untouched and functional via the legacy nnx.Param paths. To prove the variance reduction and model convergence, the DeepSeek-V4-Flash (284B) model was trained for 300 steps on a Ironwood cluster. The following analysis was performed: * Run A (Baseline): Variance at Step 300 with NO load balancing active (routed_bias_update_rate=0.0, load_balance_loss_weight=0.0). Shows natural degradation and expert collapse. * Run B (With Load Balancing): Variance at Step 300 with FULL load balancing active (routed_bias_update_rate=0.001, load_balance_loss_weight=0.0001). Shows healthy token distribution and stable loss curve comparable or better than the PR-4497 base architecture. === DeepSeek V4 Load Balancing Variance Analysis (Step 300) === | Layer Index | Routing Type | Baseline (No LB) | With Load Balancing | Improvement | |-------------|--------------|------------------|---------------------|-------------| | 0 | Hash Routed | 587520.00 | 587520.00 | 0.00% | | 1 | Hash Routed | 587520.00 | 587520.00 | 0.00% | | 2 | Hash Routed | 587520.00 | 587520.00 | 0.00% | | 3 | Top-K Routed | 14623.82 | 2447.45 | 83.26% | | 4 | Top-K Routed | 15871.75 | 3436.97 | 78.35% | | 5 | Top-K Routed | 15261.11 | 3110.88 | 79.62% | | 6 | Top-K Routed | 15717.74 | 1812.94 | 88.47% | | 7 | Top-K Routed | 15728.96 | 2615.92 | 83.37% | | 8 | Top-K Routed | 15999.51 | 4967.23 | 68.95% | | 9 | Top-K Routed | 15240.05 | 3919.24 | 74.28% | | 10 | Top-K Routed | 16086.67 | 2348.59 | 85.40% | | 11 | Top-K Routed | 16092.98 | 2607.32 | 83.80% | | 12 | Top-K Routed | 14255.88 | 2896.68 | 79.68% | | 13 | Top-K Routed | 14632.31 | 4420.16 | 69.79% | | 14 | Top-K Routed | 15010.89 | 4668.00 | 68.90% | | 15 | Top-K Routed | 16004.66 | 2190.59 | 86.31% | | 16 | Top-K Routed | 14625.14 | 2506.02 | 82.86% | | 17 | Top-K Routed | 13549.16 | 2880.44 | 78.74% | | 18 | Top-K Routed | 15119.30 | 3142.03 | 79.22% | | 19 | Top-K Routed | 14413.22 | 4134.32 | 71.32% | | 20 | Top-K Routed | 14239.44 | 2910.45 | 79.56% | | 21 | Top-K Routed | 13616.19 | 2687.55 | 80.26% | | 22 | Top-K Routed | 14360.10 | 1947.84 | 86.44% | | 23 | Top-K Routed | 14887.30 | 3880.00 | 73.94% | | 24 | Top-K Routed | 14978.84 | 4074.12 | 72.80% | | 25 | Top-K Routed | 14996.25 | 1401.96 | 90.65% | | 26 | Top-K Routed | 14232.96 | 3889.90 | 72.67% | | 27 | Top-K Routed | 14977.17 | 1858.18 | 87.59% | | 28 | Top-K Routed | 14378.70 | 2400.12 | 83.31% | | 29 | Top-K Routed | 13691.07 | 2181.08 | 84.07% | | 30 | Top-K Routed | 15055.45 | 3521.05 | 76.61% | | 31 | Top-K Routed | 14677.20 | 5095.23 | 65.28% | | 32 | Top-K Routed | 16188.89 | 3978.03 | 75.43% | | 33 | Top-K Routed | 14369.13 | 2921.98 | 79.66% | | 34 | Top-K Routed | 14483.14 | 5714.06 | 60.55% | | 35 | Top-K Routed | 15321.41 | 2911.53 | 81.00% | | 36 | Top-K Routed | 13667.09 | 2495.93 | 81.74% | | 37 | Top-K Routed | 14704.00 | 3494.09 | 76.24% | | 38 | Top-K Routed | 14339.20 | 2385.22 | 83.37% | | 39 | Top-K Routed | 14203.30 | 2454.55 | 82.72% | | 40 | Top-K Routed | 13582.94 | 2371.51 | 82.54% | | 41 | Top-K Routed | 16522.23 | 3767.62 | 77.20% | | 42 | Top-K Routed | 16277.68 | 1494.26 | 90.82% | |-------------|--------------|------------------|---------------------|-------------| | TOTAL/AVG | Top-K Only | 595982.83 | 123941.04 | 79.20% | Raw data logs and loss curve trajectories collected for this analysis: [Run : With Load Balancing Logs]: https://paste.googleplex.com/5739467624284160 [Run : PR-4497 Legacy Baseline Logs]: https://paste.googleplex.com/6184775940440064 --- .../core_concepts/moe_configuration.md | 7 ++- src/maxtext/common/metric_logger.py | 2 +- src/maxtext/configs/models/deepseek4-284b.yml | 4 ++ src/maxtext/configs/models/deepseek4-tiny.yml | 52 ++++++++-------- src/maxtext/configs/types.py | 12 +++- src/maxtext/layers/moe.py | 29 ++++++++- src/maxtext/layers/nnx_decoders.py | 8 ++- src/maxtext/trainers/pre_train/train.py | 38 ++++++++++-- tests/unit/attention_compressed_test.py | 1 + tests/unit/nnx_decoders_test.py | 62 +++++++++++++++++++ tests/unit/train_compile_test.py | 2 + 11 files changed, 182 insertions(+), 35 deletions(-) diff --git a/docs/reference/core_concepts/moe_configuration.md b/docs/reference/core_concepts/moe_configuration.md index 171fd1680d..fc883d11a6 100644 --- a/docs/reference/core_concepts/moe_configuration.md +++ b/docs/reference/core_concepts/moe_configuration.md @@ -64,7 +64,12 @@ Dropping: `routed_bias`: If enabled, adds a learnable bias term to the gate logits to facilitate load balancing. -`routed_bias_update_rate`: Defines the update rate to routed bias term above. Applicable only to the DeepSeek decoder block. +`routed_bias_update_rate`: Defines the update rate to the routed bias term above. Applicable only to the DeepSeek decoder block. For DeepSeek V4, this enables a specialized, auxiliary-loss-free routing bias mechanism. This implementation utilizes a pure `nnx.Variable` (`MoEBiasVar`) instead of a standard `nnx.Param`, which completely isolates the bias update step from the global model optimizer state. The bias is updated directly at the end of the routing step to balance the token distribution mathematically across experts without compromising language modeling convergence. + +#### DeepSeek V4 Auxiliary-Loss-Free & Sequence-Wise Load Balancing +MaxText implements an exact, paper-aligned version of DeepSeek V4's load balancing strategies (as specified in [the DeepSeek-V4 technical report](https://arxiv.org/html/2606.19348v1)). The architecture employs two distinct mechanisms: +1. **Auxiliary-Loss-Free Strategy**: Handled via `MoEBiasVar`, this implementation utilizes a pure `nnx.Variable` instead of a standard `nnx.Param`. This strictly isolates the bias update step from the global model optimizer state. Unlike the Hugging Face reference implementation—which deviates from the paper by keeping the bias parameters coupled to the global optimizer—our implementation ensures the routing bias balances token distribution mathematically across experts without polluting the main model gradients. +2. **Sequence-Wise Balance Loss**: An augmenting auxiliary loss (`load_balance_loss_weight`) applied to prevent extreme routing imbalance within individual sequences. `routed_score_func`: Defines the scoring function for the router. diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py index d157fcaba5..a976af1698 100644 --- a/src/maxtext/common/metric_logger.py +++ b/src/maxtext/common/metric_logger.py @@ -216,7 +216,7 @@ def _log_training_metrics(self, metrics, step): if self.config.num_experts > 1: moe_lb_loss = scalars.get("learning/moe_lb_loss", 0.0) - log_parts.append(f"moe_lb_loss: {moe_lb_loss:.3f}") + log_parts.append(f"moe_lb_loss: {moe_lb_loss:.6f}") if getattr(self.config, "mtp_num_layers", 0) > 0: mtp_loss = scalars.get("learning/mtp_loss", 0.0) diff --git a/src/maxtext/configs/models/deepseek4-284b.yml b/src/maxtext/configs/models/deepseek4-284b.yml index 708a36e522..5689114145 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/src/maxtext/configs/models/deepseek4-284b.yml @@ -22,6 +22,7 @@ base_mlp_dim: 2048 base_moe_mlp_dim: 2048 vocab_size: 129280 head_dim: 512 +qk_rope_head_dim: 64 # --- Standard Defaults --- enable_dropout: false @@ -51,6 +52,9 @@ shared_experts: 1 routed_score_func: "sqrtsoftplus" norm_topk_prob: true routed_bias: true +routed_bias_update_rate: 0.001 +load_balance_loss_weight: 0.0001 +log_moe_bias_norms: false routed_scaling_factor: 1.5 diff --git a/src/maxtext/configs/models/deepseek4-tiny.yml b/src/maxtext/configs/models/deepseek4-tiny.yml index 762df3b46e..c406595ad9 100644 --- a/src/maxtext/configs/models/deepseek4-tiny.yml +++ b/src/maxtext/configs/models/deepseek4-tiny.yml @@ -11,16 +11,18 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# Test Model config for DeepSeek-V4-Flash 284B (https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash) -base_emb_dim: 4096 -base_num_query_heads: 64 +# Tiny model config for DeepSeek V4 for CPU execution and testing + +base_emb_dim: 64 +base_num_query_heads: 4 base_num_kv_heads: 1 -base_num_decoder_layers: 7 -base_mlp_dim: 2048 -base_moe_mlp_dim: 2048 +base_num_decoder_layers: 43 +base_mlp_dim: 64 +base_moe_mlp_dim: 64 vocab_size: 129280 -head_dim: 512 +head_dim: 32 +qk_rope_head_dim: 32 # --- Standard Defaults --- enable_dropout: false @@ -31,40 +33,38 @@ normalization_layer_epsilon: 1.0e-6 decoder_block: "deepseek4" mhc_expansion_rate: 4 first_num_hash_layers: 3 -indexer_head_dim: 128 -indexer_n_heads: 64 -indexer_topk: 512 +indexer_head_dim: 32 +indexer_n_heads: 4 +indexer_topk: 16 # Note: Layers (0, 1, 2) are prefix layers as `first_num_hash_layers=3`. -# The 6th layer (MTP module with compress_ratio=0) has been explicitly dropped for now. -# This leaves exactly 7 layers: 3 prefix [0,0,4] + 4 scanned. +# The 44th layer (MTP module with compress_ratio=0) has been explicitly dropped for now. +# This leaves exactly 43 layers: 3 prefix [0,0,4] + 40 scanned. # `compress_ratio=0` uses sliding window attention. In this case, layer (0, 1). -# This is a tiny version of deepseek4 with fewer layers and less experts for debugging. -compress_ratios: [0, 0, 4, 128, 4, 128, 4] +compress_ratios: [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4] # --- MoE configuration --- mlp_activations: ["silu", "linear"] -num_experts: 8 -num_experts_per_tok: 3 +num_experts: 16 +num_experts_per_tok: 4 mlp_activations_limit: 10 shared_experts: 1 routed_score_func: "sqrtsoftplus" routed_bias: true -routed_scaling_factor: 1.5 - +routed_bias_update_rate: 0.001 +load_balance_loss_weight: 0.0001 +log_moe_bias_norms: false # --- Attention configuration --- attention_type: 'compressed' -attention: 'dot_product' -q_lora_rank: 1024 -o_groups: 8 -o_lora_rank: 1024 -sliding_window_size: 128 +q_lora_rank: 16 +o_groups: 4 +o_lora_rank: 16 +sliding_window_size: 32 # --- RoPE --- + rope_type: "default" rope_max_timescale: 10000 # Main RoPE theta compressed_rope_max_timescale: 160000 # Compressed RoPE theta -max_position_embeddings: 1048576 -original_max_position_embeddings: 65536 - +max_position_embeddings: 4096 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a548422..dc536f9799 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1030,6 +1030,7 @@ class DeepSeekMoE(BaseModel): routed_score_func: str = Field("", description="Scoring function for routing (e.g., 'softmax', 'sigmoid').") routed_bias: bool = Field(False, description="Whether to add a bias term for routing.") routed_bias_update_rate: float = Field(0.0, description="Update rate applied to the router bias term.") + log_moe_bias_norms: bool = Field(False, description="Whether to log the norms of MoE router biases.") mlp_bias: bool = Field( False, description="Whether to add a learnable bias for MLP matmul, " @@ -3559,8 +3560,17 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ) if self.decoder_block == DecoderBlockType.GPT_OSS and not self.sparse_matmul and self.capacity_factor != -1: raise ValueError("GPT-OSS MoE only supports dropless (capacity_factor=-1) with dense matmul.") - if self.routed_bias and self.routed_bias_update_rate > 0.0 and self.decoder_block != DecoderBlockType.DEEPSEEK: + if ( + self.routed_bias + and self.routed_bias_update_rate > 0.0 + and self.decoder_block not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4) + ): raise ValueError("Loss-free load balancing is only supported for the DeepSeek decoder block.") + if not self.pure_nnx and self.routed_bias and self.decoder_block == DecoderBlockType.DEEPSEEK4: + raise ValueError( + "Auxiliary-loss-free routed bias for DeepSeek V4 is only supported in pure NNX mode. " + "Please set pure_nnx=True or disable routed_bias." + ) if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts: raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.") self.validate_ragged_buffer_factor() diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index b310d97fb5..f6e89f204d 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -232,6 +232,10 @@ class Tid2EidVar(nnx.Variable): """Custom variable to hold tid2eid without trainable param overhead.""" +class MoEBiasVar(nnx.Variable): + """Custom NNX Variable for Auxiliary-Loss-Free MoE Routing Bias (DSV4).""" + + class GateLogit(nnx.Module): """A layer used to compute gate logits, allowing to return the pre bias values for DeepSeek routing.""" @@ -307,10 +311,17 @@ def __init__( if self.use_bias: bias_axes = self.kernel_axes[-len(self.out_features_shape) :] bias_shape = kernel_shape[-len(self.out_features_shape) :] + # DSV3 was using nnx.Param and that code we are keeping the same self.bias = nnx.Param( default_bias_init(rngs.params(), bias_shape, self.weight_dtype), out_sharding=bias_axes, ) + if self.model_name.startswith("deepseek4"): + # DSV4 uses MoEBiasVar to naturally isolate from sequence-wise updates + self.bias = MoEBiasVar( + default_bias_init(rngs.params(), bias_shape, self.weight_dtype), + out_sharding=bias_axes, + ) else: self.bias = None @@ -2574,13 +2585,29 @@ def generate_masks(self, top_k_indices, softmax_probs): # See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. def load_balance_loss(self, top_k_indices, logits) -> jax.Array: - """Compute the load balance loss.""" + """Compute the sequence-wise load balance loss. + + For DeepSeek V4 like models, standard load balancing across an entire batch can + be inadequate due to heterogeneous prompt lengths and varying sequence + characteristics. This method implements sequence-wise load balancing by + computing the token density and routing probabilities on a per-sequence basis. + + The resulting loss is scaled by `self.config.load_balance_loss_weight`. + When this configuration value is set > 0, the computed loss is aggregated + into the total training loss. By minimizing this scaled auxiliary loss, + the optimizer updates the routing parameters to actively enforce an even + distribution of tokens to experts within each individual sequence. + """ expert_mask = jax.nn.one_hot(top_k_indices, num_classes=self.num_experts, dtype=jnp.int32) summed_expert_mask = jnp.sum(expert_mask, axis=2) # Get fraction of tokens dispatched to each expert + # jnp.mean over axis=1 (sequence length) isolates the token density per sequence. density = jnp.mean(summed_expert_mask, axis=1) # get fraction of probability allocated to each expert + # jnp.mean over axis=1 isolates the routing probability per sequence. density_prob = jnp.mean(logits, axis=1) + # The sequence-wise densities and probabilities are multiplied and then averaged + # over the batch dimension, scaled by the required constant. loss = jnp.mean(density * density_prob) * (self.num_experts**2) * self.config.load_balance_loss_weight return loss diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 895ea27c14..ee02407502 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -925,6 +925,7 @@ def _apply_layers_sequentially( kv_caches_stacked=None, skip_block_remat: bool = False, unroll: int = 1, + metadata_axis_name: str = "layers", **kwargs, ): """Runs the layer stack using nnx.scan. @@ -947,6 +948,10 @@ def _apply_layers_sequentially( e.g. per-layer) remat internally, to avoid double rematerialization. unroll: Number of scan iterations to unroll into straight-line code (forwarded to jax.lax.scan). unroll >= length fully unrolls the loop. + metadata_axis_name: The name of the scan axis used during layer initialization. + This must perfectly match the string passed to `_create_scanned_layers` + (e.g., "layers", "scanned_blocks") to prevent strict JAX `pjit` PyTree + metadata mismatch errors when using custom `nnx.Variable` types (like `MoEBiasVar`). **kwargs: Keyword args forwarded to the layer (filtered by the layer signature). Returns: @@ -1074,7 +1079,7 @@ def layer_fn(carry, scanned_vars): # Move the scan axis to each variable's param_scan_axis and restore its name # in the sharding metadata. jax.lax.scan emits it at position 0. - scanned_state = maxtext_utils_nnx.nnx_add_and_sync_scan_axis(scanned_state, "layers") + scanned_state = maxtext_utils_nnx.nnx_add_and_sync_scan_axis(scanned_state, metadata_axis_name) returned_kv_stacked = None @@ -2053,6 +2058,7 @@ def _apply_deepseek4_scanned_blocks( deterministic, model_mode, length=num_full_blocks, + metadata_axis_name="scanned_blocks", **layer_call_kwargs, ) diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 0906c3e786..a992ca6fab 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -40,7 +40,7 @@ import jax.numpy as jnp from jax.sharding import NamedSharding -from flax import linen as nn, nnx +from flax import linen as nn, nnx, traverse_util from flax.linen import partitioning as nn_partitioning from flax.nnx import variablelib @@ -472,6 +472,7 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): moe_bias_updates = aux.get("moe_bias_updates") mtp_loss = aux.get("mtp_loss", 0.0) new_opt_state = None + bias_metrics = {} if isinstance(model, nn.Module): if config.gradient_clipping_threshold > 0: @@ -561,9 +562,37 @@ def move(path, value): 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() + if config.routed_bias and config.routed_bias_update_rate > 0.0: + if config.model_name.startswith("deepseek4"): + max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.") + flat_intermediates = traverse_util.flatten_dict(aux.get("intermediate_outputs", {})) + for path, update in flat_intermediates.items(): + if path[-1] != "moe_bias_updates": + continue + target = new_state.model + prefix = path[1:-1] if path[0] == "intermediates" else path[:-1] + for key in prefix: + if hasattr(target, key): + target = getattr(target, key) + elif isinstance(target, dict) and key in target: + target = target[key] + else: + target = None + break + if target is None: + continue + for _, node in nnx.iter_graph(target): + if type(node).__name__ == "GateLogit" and hasattr(node, "bias") and node.bias is not None: + update_val = update[0] if isinstance(update, (tuple, list)) else update + name_prefix = "-".join(map(str, prefix)) + if getattr(config, "log_moe_bias_norms", False): + bias_metrics[f"learning/moe_bias_before_norm_{name_prefix}"] = jnp.linalg.norm(node.bias.value) + node.bias.value = node.bias.value + jnp.array(update_val) + if getattr(config, "log_moe_bias_norms", False): + bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val)) + elif 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 = { @@ -576,6 +605,7 @@ def move(path, value): "learning/mtp_loss": mtp_loss, "learning/total_weights": total_weights, } + scalar_metrics.update(bias_metrics) if config.use_qk_clip: if isinstance(model, nn.Module): new_state = qk_clip_utils.apply_qk_clip(new_state, intermediate_outputs, config) diff --git a/tests/unit/attention_compressed_test.py b/tests/unit/attention_compressed_test.py index 47dda737dc..36dd3c4d08 100644 --- a/tests/unit/attention_compressed_test.py +++ b/tests/unit/attention_compressed_test.py @@ -37,6 +37,7 @@ def setUp(self): "qk_rope_head_dim=16", "v_head_dim=16", "qk_nope_head_dim=16", + "override_model_config=True", ] ) self.mesh = Mesh(jax.devices(), ("data",)) diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index fcd5acb5cc..2f65225731 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -1259,3 +1259,65 @@ def mock_donor_idx(lyr, layer_types, num_kv_shared): model_mode=MODEL_MODE_TRAIN, kv_caches=kv_caches, ) + + +class TestApplyLayersSequentiallyMetadataAxisName(unittest.TestCase): + + def test_metadata_axis_name_parameterization(self): + from maxtext.layers.nnx_decoders import NNXDecoder + from maxtext.utils import maxtext_utils_nnx + import jax + from flax import nnx + from unittest.mock import MagicMock + + cfg = _make_config() + cfg.param_scan_axis = 0 + mesh = _make_mesh(cfg) + rngs = nnx.Rngs(params=0) + + decoder = NNXDecoder( + config=cfg, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs, + ) + + class DummyLayer(nnx.Module): + + def __init__(self, rngs): + self.p = nnx.Param(jax.numpy.zeros((2, 2))) + + def __call__(self, x, **kwargs): + return x + self.p.value, None + + # Manually create a stacked layer using NNX scan + stacked_layers = nnx.vmap(lambda: DummyLayer(rngs=rngs), in_axes=(), out_axes=0, axis_size=2)() + + x_in = jax.numpy.zeros((2,)) + + # We mock maxtext_utils_nnx.nnx_add_scan_axis to ensure the custom name is passed + original_add_scan_axis = maxtext_utils_nnx.nnx_add_scan_axis + mock_add_scan_axis = MagicMock(side_effect=original_add_scan_axis) + maxtext_utils_nnx.nnx_add_scan_axis = mock_add_scan_axis + + try: + # Use a custom metadata_axis_name + custom_axis_name = "custom_scanned_blocks" + out, layers, _ = decoder._apply_layers_sequentially( + layers=stacked_layers, x_in=x_in, length=2, metadata_axis_name=custom_axis_name + ) + + # Verify that the custom axis name was indeed passed down + found_custom_name = False + for call_args in mock_add_scan_axis.call_args_list: + if call_args[0][1] == custom_axis_name: + found_custom_name = True + break + + self.assertTrue(found_custom_name, "The custom metadata_axis_name was not passed to nnx_add_scan_axis!") + finally: + maxtext_utils_nnx.nnx_add_scan_axis = original_add_scan_axis + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index ff5c702023..53bc0dc702 100644 --- a/tests/unit/train_compile_test.py +++ b/tests/unit/train_compile_test.py @@ -852,6 +852,8 @@ def test_deepseek4(self, scan_layers, enable_nnx): f"enable_nnx={enable_nnx}", f"pure_nnx={enable_nnx}", f"pure_nnx_decoder={enable_nnx}", + "routed_bias=False", + "override_model_config=True", ) )