diff --git a/docs/reference/core_concepts/moe_configuration.md b/docs/reference/core_concepts/moe_configuration.md index 8f999751ca..f1a4756007 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 f475afa115..09cc6eb989 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 self.config.mtp_num_layers > 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..5c9f1fd1b3 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/src/maxtext/configs/models/deepseek4-284b.yml @@ -51,6 +51,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 +adamw_mask: [".*gate.*bias.*"] routed_scaling_factor: 1.5 @@ -67,3 +70,14 @@ rope_max_timescale: 10000 # Main RoPE theta compressed_rope_max_timescale: 160000 # Compressed RoPE theta max_position_embeddings: 1048576 original_max_position_embeddings: 65536 + +# --- MLA Sharding Fix --- +# DeepSeek-V4 uses Multi-head Latent Attention with only 1 KV head. +# We override the default sharding rules to prevent the partitioner from attempting +# to slice the single KV head across tensor parallel chips. +logical_axis_rules: [ + ['activation_kv_heads', []], + ['kv_heads', []], + ['activation_kv', []], + ['kv', []] +] diff --git a/src/maxtext/configs/models/deepseek4-tiny.yml b/src/maxtext/configs/models/deepseek4-tiny.yml index 762df3b46e..dc7dc54dee 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 +adamw_mask: [".*gate.*bias.*"] # --- 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 22a2a5ff0b..bc840194c5 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3246,8 +3246,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("`local_checkpoint_period` must be > 0 for emergency checkpointing.") if self.moba and self.attention not in ("dot_product"): raise ValueError("MoBA is only supported with dot_product attention.") - if self.decoder_block == DecoderBlockType.DEEPSEEK4 and self.attention != "dot_product": - raise ValueError("DeepSeek4 decoder block currently only supports dot_product attention.") if self.use_indexer: if self.q_lora_rank == 0: raise NotImplementedError("Sparse indexer has not implemented for q_lora_rank = 0.") @@ -3309,8 +3307,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/attention_compressed.py b/src/maxtext/layers/attention_compressed.py index 72957dc177..0af5ea78c5 100644 --- a/src/maxtext/layers/attention_compressed.py +++ b/src/maxtext/layers/attention_compressed.py @@ -96,36 +96,36 @@ def csa_overlap_pooling( n_windows = chunk_kv.shape[1] // compress_rate - # Reshape flat sequence into discrete compression windows - # -> [batch, n_windows, compress_rate, 2 * head_dim] - chunk_kv = chunk_kv.reshape((batch_size, n_windows, compress_rate, 2 * head_dim)) - chunk_gate = chunk_gate.reshape((batch_size, n_windows, compress_rate, 2 * head_dim)) + position_bias + # DIPAK - OPTIMIZATION + # Split features immediately on the 3D sequence tensor to keep memory contiguous. + # Pad the sequence at the beginning by exactly one compress_rate window, then truncate the end. + # This replaces the 4D block concatenation with a highly contiguous 1D sequence slice, improving layout. - # Split the projections into Ca and Cb components for overlapping - # 2x [batch, n_windows, compress_rate, head_dim] + # Split features immediately on the 3D sequence tensor to keep memory contiguous a_kv, b_kv = jnp.split(chunk_kv, 2, axis=-1) a_gate, b_gate = jnp.split(chunk_gate, 2, axis=-1) - # Shift Ca forward by one window to align with the next Cb + # Pad the sequence at the beginning by exactly one compress_rate window, then truncate the end + # This replaces the 4D block concatenation with a highly contiguous 1D sequence slice a_kv_shifted = jnp.concatenate( - [ - jnp.zeros((batch_size, 1, compress_rate, head_dim), dtype=a_kv.dtype), - a_kv[:, :-1], - ], - axis=1, + [jnp.zeros((batch_size, compress_rate, head_dim), dtype=a_kv.dtype), a_kv[:, :-compress_rate]], axis=1 ) a_gate_shifted = jnp.concatenate( - [ - jnp.full((batch_size, 1, compress_rate, head_dim), -jnp.inf, dtype=a_gate.dtype), - a_gate[:, :-1], - ], - axis=1, + [jnp.full((batch_size, compress_rate, head_dim), -jnp.inf, dtype=a_gate.dtype), a_gate[:, :-compress_rate]], axis=1 ) - # Concatenate shifted Ca and unshifted Cb to form the final overlapping window + # Reshape into windows + a_kv_windows = a_kv_shifted.reshape((batch_size, n_windows, compress_rate, head_dim)) + b_kv_windows = b_kv.reshape((batch_size, n_windows, compress_rate, head_dim)) + + # Add position bias during reshape to fuse the operations + a_gate_windows = a_gate_shifted.reshape((batch_size, n_windows, compress_rate, head_dim)) + position_bias[:, :head_dim] + b_gate_windows = b_gate.reshape((batch_size, n_windows, compress_rate, head_dim)) + position_bias[:, head_dim:] + + # Concatenate within the window # -> [batch, n_windows, 2 * compress_rate, head_dim] - new_kv = jnp.concatenate([a_kv_shifted, b_kv], axis=2) - new_gate = jnp.concatenate([a_gate_shifted, b_gate], axis=2) + new_kv = jnp.concatenate([a_kv_windows, b_kv_windows], axis=2) + new_gate = jnp.concatenate([a_gate_windows, b_gate_windows], axis=2) # Apply softmax gating and sum across the overlapping window dimension gate_weights = jax.nn.softmax(new_gate, axis=2).astype(new_kv.dtype) @@ -482,30 +482,24 @@ def __call__( dtype=jnp.int32, ) - # Broadcast the compressed KV representations across all indexer heads - # -> [batch, 1, n_windows, index_head_dim] - compressed_kv = jnp.expand_dims(compressed, axis=1) - # -> [batch, index_n_heads, n_windows, index_head_dim] - compressed_kv = jnp.broadcast_to( - compressed_kv, - (batch_size, self.index_n_heads, compressed_len, self.index_head_dim), - ) + # OPTIMIZATION DIPAK: Removed jnp.expand_dims and jnp.broadcast_to on compressed_kv to save memory. # Project the latent query to match the Indexer's dimensions # [batch, seq_len, index_n_heads * index_head_dim] -> [batch, seq_len, index_n_heads, index_head_dim] q = self.q_proj(q_latent).reshape((batch_size, seq_len, self.index_n_heads, self.index_head_dim)) - # -> [batch, index_n_heads, seq_len, index_head_dim] - q = jnp.transpose(q, (0, 2, 1, 3)) + + # OPTIMIZATION DIPAK: Removed jnp.transpose on q to maintain bshd layout for fused einsum. # Apply standard Rotary Positional Embeddings to queries - q = self.rotary_emb(q, position_ids, unsqueeze_dim=1) + # OPTIMIZATION DIPAK: Adjusted unsqueeze_dim for bshd layout + q = self.rotary_emb(q, position_ids, unsqueeze_dim=2) q = q.astype(jnp.float32) - compressed_kv = compressed_kv.astype(jnp.float32) + compressed_kv = compressed.astype(jnp.float32) # Compute dot product between Queries and Compressed KV Blocks - # -> [batch, index_n_heads, seq_len, n_windows] - scores = jnp.einsum("bhsd,bhwd->bhsw", q, compressed_kv) + # OPTIMIZATION DIPAK: Native aligned einsum. q is [b, s, h, d], compressed_kv is [b, w, d] + scores = jnp.einsum("bshd,bwd->bshw", q, compressed_kv) scores = jax.nn.relu(scores) * self.softmax_scale # Compute routing weights to combine scores across indexer heads @@ -514,7 +508,7 @@ def __call__( # Combine individual head scores according to routing weights # -> [batch, seq_len, n_windows] - index_scores = jnp.einsum("bhsw,bsh->bsw", scores, weights) + index_scores = jnp.einsum("bshw,bsh->bsw", scores, weights) k = min(self.index_topk, compressed_len) @@ -523,18 +517,22 @@ def __call__( entry_indices = jnp.arange(compressed_len) future_mask = entry_indices[None, None, :] >= jnp.expand_dims(causal_threshold, axis=-1) - index_scores = jnp.where(future_mask, jnp.full_like(index_scores, -jnp.inf), index_scores) + # OPTIMIZATION DIPAK: Additive causal mask to avoid full_like tensor allocation + causal_mask = jnp.where(future_mask, -jnp.inf, 0.0) + index_scores += causal_mask # Apply standard segment attention mask (additive 0 and -inf) if attention_mask is not None: index_scores += attention_mask[:, :, :compressed_len] # Retrieve the top-k highest scoring block indices for each token - top_k_indices = jax.lax.top_k(index_scores, k)[1] + # OPTIMIZATION DIPAK: Replaced slow top_k with hardware-friendly approx_max_k + top_k_indices = jax.lax.approx_max_k(index_scores, k)[1] # Invalidate any top-k selections that point to future blocks (edge case safety) invalid = top_k_indices >= jnp.expand_dims(causal_threshold, axis=-1) - top_k_indices = jnp.where(invalid, jnp.full_like(top_k_indices, -1), top_k_indices) + # OPTIMIZATION DIPAK: Removed full_like in top_k_indices invalidation + top_k_indices = jnp.where(invalid, -1, top_k_indices) return top_k_indices @@ -654,12 +652,10 @@ def __call__( # Only compute and apply the complex block mask if top-k selections exist if k > 0: - valid = top_k_indices >= 0 - entry_indices = jnp.arange(compressed_len)[None, None, :] - is_in_topk = jnp.expand_dims(top_k_indices, axis=-1) == entry_indices[None, ...] - is_valid_and_in_topk = is_in_topk & jnp.expand_dims(valid, axis=-1) - - is_selected = jnp.any(is_valid_and_in_topk, axis=2) + # OPTIMIZATION DIPAK: Replaced memory-heavy 4D broadcast with boolean one_hot scatter. + # JAX natively maps -1 (padding) to all False in boolean one_hot, avoiding explicit validity checks. + one_hot = jax.nn.one_hot(top_k_indices, compressed_len, dtype=jnp.bool_) + is_selected = jnp.any(one_hot, axis=2) is_selected = jnp.expand_dims(is_selected, axis=1) compressed_mask = jnp.where(is_selected, 0.0, DEFAULT_MASK_VALUE).astype(self.dtype) @@ -1053,8 +1049,20 @@ def __call__( ) # Extend local KV tensors with the compressed blocks + decoder_segment_ids_q = decoder_segment_ids + decoder_segment_ids_kv = decoder_segment_ids + if compressed_kv is not None: kv = jnp.concatenate([kv, compressed_kv], axis=1) + if decoder_segment_ids is not None: + # Pad segment IDs to match the new KV sequence length. + # We pad with 0 because the actual masking for compressed blocks is + # explicitly handled by the compressed_mask boolean logic. + segment_padding = jnp.zeros( + (decoder_segment_ids.shape[0], compressed_kv.shape[1]), + dtype=decoder_segment_ids.dtype, + ) + decoder_segment_ids_kv = jnp.concatenate([decoder_segment_ids, segment_padding], axis=1) kv = checkpoint_name(kv, "kv_proj") @@ -1072,7 +1080,8 @@ def __call__( q, kv, kv, - decoder_segment_ids, + decoder_segment_ids_q, + decoder_segment_ids_kv, inputs_positions, model_mode, sinks=self.sinks.value, diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index b3f22fc336..a3b1df8c07 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -56,6 +56,7 @@ DType, D_KV, HEAD, + KV_HEAD, KV_LENGTH, LENGTH, MODEL_MODE_AUTOREGRESSIVE, @@ -382,7 +383,7 @@ def __init__( 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_kv: AxisNames = (BATCH_ATTN, KV_HEAD, KV_LENGTH, D_KV), flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH), prefill_cache_logical_axis_names: AxisNames = ( CACHE_BATCH_PREFILL, @@ -1020,7 +1021,8 @@ def apply_attention( query: Array, key: Array | KVTensor, value: Array | KVTensor, - decoder_segment_ids: Array | None, + decoder_segment_ids_q: Array | None, + decoder_segment_ids_kv: Array | None, segment_positions: Array | None, lengths: Array | None, model_mode: str, @@ -1048,7 +1050,7 @@ def apply_attention( if use_ragged_attention and model_mode == MODEL_MODE_AUTOREGRESSIVE: if lengths is None: - lengths = jnp.sum(decoder_segment_ids, axis=-1) + lengths = jnp.sum(decoder_segment_ids_q, axis=-1) if target_hardware == "tpu": impl = self.tpu_ragged_attention @@ -1075,7 +1077,7 @@ def apply_attention( query, key, value, - decoder_segment_ids, + decoder_segment_ids_q, model_mode, previous_chunk, segment_positions=segment_positions, @@ -1104,7 +1106,8 @@ def apply_attention( query, key, value, - decoder_segment_ids, + decoder_segment_ids_q, + decoder_segment_ids_kv, self.attn_logits_soft_cap, sinks, indexer_mask=indexer_mask, @@ -1112,6 +1115,7 @@ def apply_attention( previous_chunk=previous_chunk, bidirectional_mask=bidirectional_mask, use_ragged_attention=use_ragged_attention, + compressed_mask=compressed_mask, record_max_logits=record_max_logits, ) if max_logits is not None: @@ -1289,7 +1293,8 @@ def tpu_flash_attention( query: Array, key: Array, value: Array, - decoder_segment_ids: Array | None, + decoder_segment_ids_q: Array | None, + decoder_segment_ids_kv: Array | None, attn_logits_soft_cap: float | None = None, sinks: Array | None = None, indexer_mask: Array | None = None, @@ -1297,11 +1302,70 @@ def tpu_flash_attention( previous_chunk: Any = None, bidirectional_mask: Any = None, use_ragged_attention: bool = False, + compressed_mask: Array | None = None, record_max_logits: bool = False, ) -> tuple[Array, Array]: """TPU Flash Attention.""" use_tokamax_ring = tokamax_ring_attention.is_context_parallel_ring_requested(self.config) + + def create_csa_mask_info_jax(top_k_indices, s_len, c_len, sliding_window_size, block_shape=(128, 128)): + from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_mask_info + + block_q, block_kv = block_shape + q_len = top_k_indices.shape[0] + k = top_k_indices.shape[1] + kv_len = s_len + c_len + + q_blocks_count = q_len // block_q + kv_blocks_count = kv_len // block_kv + + q_b = jnp.arange(q_blocks_count)[:, None] + kv_b = jnp.arange(kv_blocks_count)[None, :] + min_q, max_q = q_b * block_q, q_b * block_q + block_q - 1 + min_kv, max_kv = kv_b * block_kv, kv_b * block_kv + block_kv - 1 + + sliding_window_size = sliding_window_size if sliding_window_size is not None else s_len + all_local = (max_kv <= min_q) & (min_kv > max_q - sliding_window_size) & (max_kv < s_len) + any_local = (min_kv <= max_q) & (max_kv > min_q - sliding_window_size) & (min_kv < s_len) + + block_mask = jnp.zeros((q_blocks_count, kv_blocks_count + 1), dtype=jnp.int32) + block_mask = block_mask.at[:, :-1].set(jnp.where(all_local, 2, jnp.where(any_local, 1, 0))) + + top_k = top_k_indices + q_indices = jnp.arange(q_len) // block_q + valid = top_k >= 0 + kv_indices = (s_len + top_k) // block_kv + + flat_q = jnp.broadcast_to(q_indices[:, None], (q_len, k)).flatten() + flat_kv = jnp.where(valid, kv_indices, kv_blocks_count).flatten() + + block_mask = block_mask.at[flat_q, flat_kv].set(1) + block_mask = block_mask[:, :-1] + + block_ids = jnp.arange(block_mask.size, dtype=jnp.int32).reshape(block_mask.shape) + active_mask = block_mask > 0 + num_active_blocks = active_mask.flatten().sum(keepdims=True) + active_indices = jnp.argwhere(active_mask, size=active_mask.size, fill_value=-1) + active_rows = active_indices[:, 0].astype(np.int32) + active_cols = active_indices[:, 1].astype(np.int32) + + block_mask_flat = block_mask[active_rows, active_cols] + mask_next = block_ids.at[active_rows, active_cols].get(wrap_negative_indices=False) + mask_next = jnp.where(block_mask_flat == 1, mask_next, 0) + + mask = (jnp.arange(block_mask.size) < num_active_blocks).astype(np.int32) + block_mask_flat = (block_mask_flat * mask).astype(np.int8) + + return splash_attention_mask_info.MaskInfo( + mask_next=mask_next, + active_rows=active_rows, + active_cols=active_cols, + num_active_blocks=num_active_blocks, + block_mask=block_mask_flat, + block_shape=block_shape, + ) + cp_size = self.mesh.shape.get(self.config.context_sharding, 1) load_balanced_context_parallel = self.config.context_parallel_load_balance if use_tokamax_ring: @@ -1322,7 +1386,7 @@ def tpu_flash_attention( segment_axis_names_q = None segment_axis_names_kv = None sink_axis_names = self._logical_to_mesh_axes((HEAD,)) - if decoder_segment_ids is not None: + if decoder_segment_ids_q is not None: segment_axis_names_q = self._logical_to_mesh_axes((BATCH_ATTN, Q_LENGTH)) segment_axis_names_kv = self._logical_to_mesh_axes((BATCH_ATTN, KV_LENGTH)) @@ -1690,8 +1754,8 @@ def kernel_fn(q, k, v, d, s): query = self._maybe_shard_with_pspec(query, axis_names_q) key = self._maybe_shard_with_pspec(key, axis_names_kv) value = self._maybe_shard_with_pspec(value, axis_names_kv) - decoder_segment_ids_q = self._maybe_shard_with_pspec(decoder_segment_ids, segment_axis_names_q) - decoder_segment_ids_kv = self._maybe_shard_with_pspec(decoder_segment_ids, segment_axis_names_kv) + decoder_segment_ids_q = self._maybe_shard_with_pspec(decoder_segment_ids_q, segment_axis_names_q) + decoder_segment_ids_kv = self._maybe_shard_with_pspec(decoder_segment_ids_kv, segment_axis_names_kv) sinks = self._maybe_shard_with_pspec(sinks, sink_axis_names) indexer_mask = self._maybe_shard_with_pspec(indexer_mask, indexer_mask_axis_names) @@ -2252,7 +2316,8 @@ def __call__( query, key, value, - decoder_segment_ids, + decoder_segment_ids_q, + decoder_segment_ids_kv, inputs_positions, model_mode, cached_values=None, @@ -2270,7 +2335,7 @@ def __call__( prefill_kv_cache, ar_kv_cache = cached_values[0], cached_values[1] if model_mode != MODEL_MODE_TRAIN: assert prefill_kv_cache - key, value, decoder_segment_ids = prefill_kv_cache + key, value, decoder_segment_ids_kv = prefill_kv_cache indexer_mask_prefill = None indexer_mask_ar = None @@ -2284,7 +2349,8 @@ def __call__( query=query, key=key, value=value, - decoder_segment_ids=decoder_segment_ids, + decoder_segment_ids_q=decoder_segment_ids_q, + decoder_segment_ids_kv=decoder_segment_ids_kv, segment_positions=inputs_positions, lengths=None, model_mode=model_mode, diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 7737734b09..009d8d2d93 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 @@ -2571,13 +2582,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 69d8800080..ebf9dc1b1e 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -428,10 +428,18 @@ def __init__( self.scanned_layers = None self.is_deepseek = self.config.decoder_block == DecoderBlockType.DEEPSEEK + self.is_deepseek4 = self.config.decoder_block == DecoderBlockType.DEEPSEEK4 self.is_gemma3 = self.config.decoder_block == DecoderBlockType.GEMMA3 self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4 self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL + if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: + self.hc_head = mhc.DeepSeek4HyperHead( + config=config, + mesh=self.mesh, + rngs=self.rngs, + ) + self._init_decoder_layers(decoder_block_classes, rngs, mesh) def _init_decoder_layers(self, decoder_block_classes, rngs, mesh): @@ -528,6 +536,8 @@ def _init_scanned_layers(self, decoder_block_classes, rngs, mesh): """Initializes decoder layers with scanning (non-pipeline).""" if self.is_deepseek: self._init_scanned_deepseek(decoder_block_classes, rngs) + elif self.is_deepseek4: + self._init_scanned_deepseek4(rngs) elif self.is_gemma3: self._init_scanned_gemma3(decoder_block_classes, rngs, mesh) elif self.is_gemma4: @@ -535,6 +545,30 @@ def _init_scanned_layers(self, decoder_block_classes, rngs, mesh): else: self._init_scanned_generic(decoder_block_classes, rngs) + def _init_scanned_deepseek4(self, rngs): + """Initializes DeepSeek V4 scanned layers: unrolls first_num_hash_layers prefix layers and scans remaining full blocks.""" + config = self.config + self.layers = nnx.List([]) + num_hash_layers = config.first_num_hash_layers + for layer_idx in range(num_hash_layers): + self._create_and_register_layer( + deepseek4.DeepSeek4DecoderLayer, + rngs, + "layers", + layer_idx, + layer_idx=layer_idx, + ) + + num_remaining_layers = config.num_decoder_layers - num_hash_layers + num_full_blocks = num_remaining_layers // 2 + if num_full_blocks > 0: + self.scanned_blocks = self._create_scanned_layers( + deepseek4.DeepSeek4ScannableBlock, + length=num_full_blocks, + metadata_axis_name="scanned_blocks", + rngs=rngs, + ) + def _init_scanned_deepseek(self, decoder_block_classes, rngs): """Initializes scanned DeepSeek layers with optional Engram support.""" config = self.config @@ -729,6 +763,7 @@ def _init_sequential_generic(self, decoder_block_classes, rngs): elif config.decoder_block in { DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, + DecoderBlockType.DEEPSEEK4, }: layer_kwargs = {"layer_idx": lyr} elif config.decoder_block == DecoderBlockType.GPT_OSS: @@ -870,7 +905,9 @@ def pure_layer_fn(state_in, y_in): return out - def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches_stacked=None, **kwargs): + def _apply_layers_sequentially( + self, layers, x_in, *args, length: int, kv_caches_stacked=None, metadata_axis_name: str = "layers", **kwargs + ): """Runs the layer stack using nnx.scan. Args: @@ -881,6 +918,10 @@ def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches kv_caches_stacked: Optional pytree whose leaves have shape [num_layers, ...]. When provided, the i-th slice is passed as `kv_cache=` to layer i and the updated caches are returned as a third element of the tuple. + 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: @@ -998,8 +1039,12 @@ def layer_fn(carry, scanned_vars): final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state)) returned_kv_stacked = None - # Ensure metadata rank matches the stacked values - scanned_state = maxtext_utils_nnx.nnx_add_scan_axis(scanned_state, "layers", 0) + # Ensure metadata rank matches the stacked values. + # We dynamically pass metadata_axis_name instead of hardcoding "layers" + # so that custom NNX Variables (like MoEBiasVar) which were initialized under + # a different axis name (e.g. "scanned_blocks") do not fail JAX's strict + # pjit PyTree metadata matching during lowering. + scanned_state = maxtext_utils_nnx.nnx_add_scan_axis(scanned_state, metadata_axis_name, 0) if scan_axis != 0: new_params, new_rest = scanned_state.split(nnx.Param, ...) @@ -1198,6 +1243,7 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs): DecoderBlockType.MISTRAL, DecoderBlockType.MIXTRAL, DecoderBlockType.DEEPSEEK, + DecoderBlockType.DEEPSEEK4, DecoderBlockType.GEMMA, DecoderBlockType.GEMMA2, DecoderBlockType.GEMMA3, @@ -1730,6 +1776,18 @@ def __call__( length=num_moe, **layer_kwargs, ) + + elif self.is_deepseek4: + y = self._apply_deepseek4_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot, + previous_chunk, + decoder_input_tokens, + ) elif self.is_gemma3: y = self._apply_gemma3_scanned_blocks( y, @@ -1809,7 +1867,9 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in): else: kv_cache = None - input_tokens = decoder_input_tokens if cfg.engram_layers else None + input_tokens = ( + decoder_input_tokens if (cfg.engram_layers or cfg.decoder_block == DecoderBlockType.DEEPSEEK4) else None + ) if input_tokens is not None: layer_kwargs["decoder_input_tokens"] = input_tokens @@ -1833,8 +1893,11 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in): # 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) + if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: + hidden_state = self.hc_head(y) + else: + # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) + hidden_state = mhc_reduce(y) else: hidden_state = y @@ -1860,6 +1923,56 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in): return logits, hidden_state, kv_caches + def _apply_deepseek4_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot=None, + previous_chunk=None, + decoder_input_tokens=None, + ): + cfg = self.config + num_hash_layers = cfg.first_num_hash_layers + + layer_call_kwargs = { + "previous_chunk": previous_chunk, + "slot": slot, + "decoder_input_tokens": decoder_input_tokens, + } + + # 1. Unrolled prefix layers (0, 1, 2) + for layer_idx in range(num_hash_layers): + layer = getattr(self, f"layers_{layer_idx}") + y, _ = layer( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + **layer_call_kwargs, + ) + + # 2. Scanned blocks + num_remaining_layers = cfg.num_decoder_layers - num_hash_layers + num_full_blocks = num_remaining_layers // 2 + if num_full_blocks > 0 and hasattr(self, "scanned_blocks"): + y, self.scanned_blocks, _ = self._apply_layers_sequentially( + self.scanned_blocks, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + length=num_full_blocks, + metadata_axis_name="scanned_blocks", + **layer_call_kwargs, + ) + + return y + def _apply_gemma3_scanned_blocks( self, y, diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 39d3fec6b5..3e5aca923f 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -36,7 +36,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 @@ -533,10 +533,37 @@ def move(path, value): state.apply_gradients(grads) new_state = state + bias_metrics = {} # 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)) + 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) + 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 = { @@ -549,6 +576,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/deepseek_v4_vs_reference_test.py b/tests/unit/deepseek_v4_vs_reference_test.py index 50e7ab2bf3..3e5b25e7c8 100644 --- a/tests/unit/deepseek_v4_vs_reference_test.py +++ b/tests/unit/deepseek_v4_vs_reference_test.py @@ -1512,6 +1512,18 @@ def test_hyper_head_parity(self): print(f"HYPER HEAD PARITY - MAX ABS DIFF: {max_diff:.6e}, MEAN ABS DIFF: {mean_diff:.6e}") np.testing.assert_allclose(mt_out, pt_out, rtol=5e-5, atol=5e-5) + def test_nnx_decoder_hyperhead_integration(self): + from maxtext.layers.nnx_decoders import NNXDecoder + + # Specifically instantiate full NNXDecoder to ensure it didn't bypass the HyperHead! + mt_decoder = NNXDecoder( + config=self.mx_config, + mesh=self.mesh, + rngs=self.rngs, + ) + self.assertTrue(hasattr(mt_decoder, "hc_head"), "NNXDecoder completely missed hc_head setup! Bug in nnx_decoders.py!") + self.assertIsInstance(mt_decoder.hc_head, DeepSeek4HyperHead, "NNXDecoder instantiated the wrong HyperHead class!") + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index 4585c9c3ee..bbfc08fa96 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -951,3 +951,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()