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 c4a270a567..8750327b91 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1015,6 +1015,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, " @@ -3531,8 +3532,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 6ae9abf2cc..50486714b1 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 @@ -2573,13 +2584,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 138bfe1ec7..37f841266d 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -35,7 +35,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 @@ -467,6 +467,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: @@ -556,9 +557,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 = { @@ -571,6 +600,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", ) )