From 250da311abfb806828ba718cb6aaf423ed608533 Mon Sep 17 00:00:00 2001 From: Octa Trifan Date: Wed, 9 Sep 2026 00:23:46 +0000 Subject: [PATCH] Support DeepSeek-V4 Compressed Sparse Attention (CSA) Indexer Loss - Implement calculate_csa_indexer_loss for DeepSeek-V4 CSA Indexer KL divergence distillation. - Add stop_gradient to teacher projections and indexer inputs to isolate indexer gradients. - Support dense warm-up and sparse pre-training modes with proper block causal masking and document boundaries. - Gate CSA indexer execution behind use_indexer in csa_compressor to support Stage 1 dense pre-training. - Optimize score tensor materialization in CompressedAttention via return_indexer_scores=should_compute_loss. - Enforce indexer_topk >= max_blocks boundary validation in types.py. - Set use_indexer: true in deepseek4-284b.yml and deepseek4-tiny.yml to preserve out-of-the-box sparse execution. - Ensure static return signatures in DeepseekV4Indexer (2-tuple) and DeepseekV4CSACompressor (3-tuple). - Document DeepSeek-V4 3-stage pre-training in Run_DeepSeek.md with override_model_config=true for Stage 1. - Add comprehensive unit test suite in deepseek_v4_indexer_loss_test.py and update reference parity tests. --- src/maxtext/configs/models/deepseek4-284b.yml | 1 + src/maxtext/configs/models/deepseek4-tiny.yml | 1 + src/maxtext/configs/types.py | 56 +- src/maxtext/layers/attention_compressed.py | 284 ++++++++++- tests/end_to_end/tpu/deepseek/Run_DeepSeek.md | 85 ++++ tests/unit/deepseek_v4_indexer_loss_test.py | 479 ++++++++++++++++++ tests/unit/deepseek_v4_vs_reference_test.py | 10 +- 7 files changed, 873 insertions(+), 43 deletions(-) create mode 100644 tests/unit/deepseek_v4_indexer_loss_test.py diff --git a/src/maxtext/configs/models/deepseek4-284b.yml b/src/maxtext/configs/models/deepseek4-284b.yml index 5689114145..064fd23dc1 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/src/maxtext/configs/models/deepseek4-284b.yml @@ -60,6 +60,7 @@ routed_scaling_factor: 1.5 # --- Attention configuration --- attention_type: 'compressed' +use_indexer: true q_lora_rank: 1024 o_groups: 8 o_lora_rank: 1024 diff --git a/src/maxtext/configs/models/deepseek4-tiny.yml b/src/maxtext/configs/models/deepseek4-tiny.yml index c406595ad9..da81e937a8 100644 --- a/src/maxtext/configs/models/deepseek4-tiny.yml +++ b/src/maxtext/configs/models/deepseek4-tiny.yml @@ -57,6 +57,7 @@ log_moe_bias_norms: false # --- Attention configuration --- attention_type: 'compressed' +use_indexer: true q_lora_rank: 16 o_groups: 4 o_lora_rank: 16 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 2a3b067498..e278e3686e 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -764,9 +764,9 @@ class CompressedAttention(BaseModel): class AttentionIndexer(BaseModel): - """Configuration for DeepSeek Sparse Attention (DSA): DeepSeek3.2-style MLA with indexer.""" + """Configuration for DeepSeek Sparse Attention (DSA): MLA or Compressed Attention with indexer.""" - use_indexer: bool = Field(False, description="Whether to use sparse indexer for MLA.") + use_indexer: bool = Field(False, description="Whether to use sparse indexer for MLA or Compressed Attention.") indexer_head_dim: NonNegativeInt = Field(128, description="Head dim for indexer query and key.") indexer_n_heads: NonNegativeInt = Field(64, description="Number of query heads in indexer.") indexer_topk: NonNegativeInt = Field(2048, description="Number of tokens selected by the query token in indexer.") @@ -4232,10 +4232,11 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ) if self.use_indexer: - if self.attention_type != AttentionType.MLA.value: + if self.attention_type not in (AttentionType.MLA.value, AttentionType.COMPRESSED.value): raise ValueError( - f"`use_indexer=True` requires `attention_type='{AttentionType.MLA.value}'`, since only the " - "MLA indexer produces this mask." + f"`use_indexer=True` requires `attention_type='{AttentionType.MLA.value}'` or " + f"`attention_type='{AttentionType.COMPRESSED.value}'`, since only MLA and " + "Compressed Attention indexers produce this mask." ) if self.q_lora_rank == 0: raise NotImplementedError("Sparse indexer has not implemented for q_lora_rank = 0.") @@ -4243,23 +4244,36 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de supports_flash_splash = self.attention == "flash" and self.use_tokamax_splash if not (supports_dot_product or supports_flash_splash): raise ValueError( - "Sparse indexer is only supported with dot_product attention or flash attention with tokamax splash." - ) - if ( - self.attention == "flash" - and self.context_parallel_strategy == "all_gather" - and self.ici_context_parallelism * self.dcn_context_parallelism > 1 - and self.attention_sink - ): - raise ValueError( - "Sparse indexer with all-gather context parallelism for flash attention does not support attention sinks." - ) - if self.indexer_loss_scaling_factor > 0.0 and self.indexer_topk >= self.max_target_length: - raise ValueError( - f"`indexer_topk` ({self.indexer_topk}) must be < `max_target_length` ({self.max_target_length}) " - "when indexer loss is enabled (`indexer_loss_scaling_factor > 0.0`); otherwise the indexer " - "short-circuits to select all tokens and no indexer loss is produced." + f"Sparse indexer with {self.attention_type} is only supported with dot_product attention or flash " + "attention with tokamax splash." ) + if self.attention_type == AttentionType.MLA.value: + if ( + self.attention == "flash" + and self.context_parallel_strategy == "all_gather" + and self.ici_context_parallelism * self.dcn_context_parallelism > 1 + and self.attention_sink + ): + raise ValueError( + "Sparse indexer with all-gather context parallelism for flash attention does not support attention sinks." + ) + if self.indexer_loss_scaling_factor > 0.0 and self.indexer_topk >= self.max_target_length: + raise ValueError( + f"`indexer_topk` ({self.indexer_topk}) must be < `max_target_length` ({self.max_target_length}) " + "when indexer loss is enabled (`indexer_loss_scaling_factor > 0.0`); otherwise the indexer " + "short-circuits to select all tokens and no indexer loss is produced." + ) + elif self.attention_type == AttentionType.COMPRESSED.value: + # DeepSeek-V4 CSA natively uses a compression rate of 4 for the indexer blocks. + compress_rate = 4 + max_blocks = self.max_target_length // compress_rate + if self.indexer_loss_scaling_factor > 0.0 and self.indexer_topk >= max_blocks: + raise ValueError( + f"`indexer_topk` ({self.indexer_topk}) must be < total compressed blocks ({max_blocks}) " + f"(max_target_length={self.max_target_length} // compress_rate={compress_rate}) " + "when indexer loss is enabled (`indexer_loss_scaling_factor > 0.0`); otherwise the indexer " + "short-circuits to select all compressed blocks and no indexer loss is produced." + ) if not self.use_indexer and self.indexer_cutoff_threshold != RematLocation.REMAT: raise ValueError( f"Setting `indexer_cutoff_threshold='{self.indexer_cutoff_threshold}'` is only valid when " diff --git a/src/maxtext/layers/attention_compressed.py b/src/maxtext/layers/attention_compressed.py index 2466e3bdd1..9fba97b79d 100644 --- a/src/maxtext/layers/attention_compressed.py +++ b/src/maxtext/layers/attention_compressed.py @@ -37,6 +37,7 @@ ) from maxtext.layers import nnx_wrappers +from maxtext.layers.attention_mla import indexer_losses from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import DeepSeekV4RotaryEmbedding from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned @@ -45,6 +46,7 @@ from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.inference.kvcache import KVQuant from maxtext.inference import kvcache +from maxtext.utils.globals import EPS class CSAPoolingConfig(enum.IntEnum): @@ -783,7 +785,8 @@ def __call__( attention_mask: Optional[Array] = None, model_mode: str = MODEL_MODE_TRAIN, cache: Optional[Any] = None, - ) -> Array: + return_scores: bool = False, + ) -> Tuple[Array, Optional[Array]]: """Forward pass for the DeepSeek-V4 Indexer. Args: @@ -793,11 +796,16 @@ def __call__( attention_mask: Optional attention mask. model_mode: Execution mode (train, prefill, or autoregressive). cache: Optional Indexer KV cache instance for inference. + return_scores: Whether to return (final_indices, index_scores). Returns: - Top-K selected indices for each query position. + A tuple (top_k_indices, indexer_scores) where indexer_scores is None when return_scores=False. """ batch_size, seq_len, _ = hidden_states.shape + # Stop gradient on indexer inputs so indexer loss does not backprop into main model projections + hidden_states = jax.lax.stop_gradient(hidden_states) + q_latent = jax.lax.stop_gradient(q_latent) + future_mask = None kv = self.kv_proj(hidden_states) gate = self.gate_proj(hidden_states) @@ -871,7 +879,8 @@ def indexer_compressor_fn(buf_kv, buf_gate): ) if compressed_len == 0: - return jnp.zeros((batch_size, seq_len, min(self.index_topk, compressed_len)), dtype=jnp.int32) + empty_indices = jnp.zeros((batch_size, seq_len, min(self.index_topk, compressed_len)), dtype=jnp.int32) + return empty_indices, (jnp.zeros((batch_size, seq_len, 0), dtype=jnp.float32) if return_scores else None) # --- TOP-K ROUTING MATH (Executes in both Prefill and AR) --- compressed_kv = jnp.expand_dims(compressed, axis=1) @@ -904,14 +913,14 @@ def indexer_compressor_fn(buf_kv, buf_gate): if attention_mask is not None: att_m = attention_mask[:, :, :compressed_len] index_scores += att_m - combined_invalid = combined_invalid | (att_m < -100.0) + combined_invalid = combined_invalid | (att_m < (DEFAULT_MASK_VALUE / 2)) top_k_indices = jax.lax.top_k(index_scores, k)[1] invalid = jnp.take_along_axis(combined_invalid, top_k_indices, axis=-1) final_indices = jnp.where(invalid, jnp.full_like(top_k_indices, -1), top_k_indices) - return final_indices + return final_indices, (index_scores if return_scores else None) class DeepseekV4CSACompressor(BaseDeepseekCompressor): @@ -979,7 +988,9 @@ def __call__( model_mode: str = MODEL_MODE_TRAIN, cache: Optional[Any] = None, indexer_cache: Optional[Any] = None, - ) -> Tuple[Array, Array]: + use_indexer: bool = True, + return_indexer_scores: bool = False, + ) -> Tuple[Array, Array, Optional[Array]]: """Forward pass for the CSA compressor. Args: @@ -990,15 +1001,30 @@ def __call__( model_mode: Execution mode (train, prefill, or autoregressive). cache: Optional CSA compressor KV cache instance for inference. indexer_cache: Optional Indexer KV cache instance for inference. + use_indexer: Whether to run the indexer to compute sparse attention masks. + return_indexer_scores: Whether to return indexer scores along with compressed KV and mask. Returns: compressed_kv: The pooled KV tensors. compressed_mask: The sparse attention mask computed by the indexer. + index_scores (optional): The raw indexer scores if return_indexer_scores is True. """ batch_size, seq_len, _ = hidden_states.shape - # 1. ALWAYS Run Indexer (It fetches its own history inside AR) - top_k_indices = self.indexer(hidden_states, q_latent, position_ids, attention_mask, model_mode, indexer_cache) + # 1. Run Indexer if use_indexer is True + if use_indexer: + top_k_indices, indexer_scores = self.indexer( + hidden_states, + q_latent, + position_ids, + attention_mask, + model_mode, + indexer_cache, + return_scores=return_indexer_scores, + ) + else: + top_k_indices = None + indexer_scores = None kv = self.kv_proj(hidden_states) gate = self.gate_proj(hidden_states) @@ -1072,7 +1098,12 @@ def csa_compressor_fn(buf_kv, buf_gate): ) if compressed_len == 0: - return compressed_kv, jnp.zeros((batch_size, 1, seq_len, 0), dtype=self.dtype) + empty_mask = jnp.zeros((batch_size, 1, seq_len, 0), dtype=self.dtype) + return compressed_kv, empty_mask, (indexer_scores if return_indexer_scores else None) + + if not use_indexer or top_k_indices is None: + compressed_mask = jnp.zeros((batch_size, 1, seq_len, compressed_len), dtype=self.dtype) + return compressed_kv, compressed_mask, None # 3. Apply Dynamic Masking Logic k = top_k_indices.shape[-1] @@ -1093,7 +1124,7 @@ def csa_compressor_fn(buf_kv, buf_gate): dtype=self.dtype, ) - return compressed_kv, compressed_mask + return compressed_kv, compressed_mask, (indexer_scores if return_indexer_scores else None) class CompressedAttention(Attention): @@ -1555,15 +1586,64 @@ def __call__( inputs_kv, q_normed, inputs_positions, model_mode, self.compressor_cache ) elif self.compress_ratio == 4: - compressed_kv, compressed_mask = self.csa_compressor( - inputs_kv, - q_normed, - inputs_positions, - compressed_segment_mask, - model_mode, - self.compressor_cache, - self.indexer_cache, - ) + use_indexer = self.config.use_indexer + scaling_factor = self.config.indexer_loss_scaling_factor + should_compute_loss = use_indexer and scaling_factor > 0.0 and model_mode == MODEL_MODE_TRAIN + + if use_indexer: + compressed_kv, sparse_compressed_mask, indexer_scores = self.csa_compressor( + inputs_kv, + q_normed, + inputs_positions, + compressed_segment_mask, + model_mode, + self.compressor_cache, + self.indexer_cache, + use_indexer=True, + return_indexer_scores=should_compute_loss, + ) + is_sparse_training = self.config.indexer_sparse_training + is_dense_warmup = (model_mode == MODEL_MODE_TRAIN) and (scaling_factor > 0.0) and (not is_sparse_training) + use_sparse_mask = not is_dense_warmup + compressed_mask = self.get_compressed_mask( + inputs_positions, + compressed_kv.shape[1], + sparse_compressed_mask=sparse_compressed_mask if use_sparse_mask else None, + ) + + if ( + should_compute_loss + and indexer_scores is not None + and compressed_kv is not None + and compressed_kv.shape[1] > 0 + ): + indexer_loss = self.calculate_csa_indexer_loss( + indexer_score=indexer_scores, + query=q, + compressed_kv=compressed_kv, + compressed_mask=sparse_compressed_mask, + segment_mask=compressed_segment_mask, + position_ids=inputs_positions, + sparse_loss=is_sparse_training, + scaling_factor=scaling_factor, + ) + self.indexer_loss = indexer_losses(indexer_loss) + else: + compressed_kv, _, _ = self.csa_compressor( + inputs_kv, + q_normed, + inputs_positions, + compressed_segment_mask, + model_mode, + self.compressor_cache, + self.indexer_cache, + use_indexer=False, + return_indexer_scores=False, + ) + compressed_mask = self.get_compressed_mask( + inputs_positions, + compressed_kv.shape[1], + ) # Apply segment masking to the compressed blocks if compressed_segment_mask is not None and compressed_mask is not None: @@ -1662,6 +1742,172 @@ def __call__( # Return the Tuple expected by the transformer block return final_out, current_kv_cache + def get_compressed_mask( + self, + inputs_positions: Array, + compressed_len: int, + sparse_compressed_mask: Optional[Array] = None, + ) -> Array: + """Builds compressed attention mask. Returns sparse_compressed_mask if provided, else dense causal mask.""" + if sparse_compressed_mask is not None: + return sparse_compressed_mask + usable_len = compressed_len * self.compress_ratio + block_positions = inputs_positions[:, : usable_len : self.compress_ratio] + is_future = (block_positions[:, None, :] + self.compress_ratio) > (inputs_positions[:, :, None] + 1) + dense_causal_mask = jnp.where(is_future, DEFAULT_MASK_VALUE, 0.0).astype(self.dtype) + return dense_causal_mask[:, None, :, :] + + def calculate_csa_indexer_loss( + self, + indexer_score: Array, + query: Array, + compressed_kv: Array, + compressed_mask: Array, + segment_mask: Optional[Array] = None, + position_ids: Optional[Array] = None, + sparse_loss: bool = False, + scaling_factor: float = 1.0, + ) -> Array: + """Calculates the indexer KL divergence loss for Compressed Attention (DeepSeek-V4). + + This loss trains the indexer to predict which compressed blocks are important by matching + the distribution of true attention scores from the main model over compressed KV blocks. + + The target distribution is derived through the following steps: + 1. Compute raw attention scores via Q @ K_comp^T (Q is already pre-scaled by 1/sqrt(head_dim)). + 2. Apply causal block masking and segment masking over compressed blocks. + 3. Softmax across compressed blocks dimension for each head. + 4. Aggregate probabilities by summing across all attention heads. + 5. Apply L1-normalization across the compressed block sequence dimension. + + target_distribution = L1_Normalize(Sum_h(Softmax_w(Q @ K_comp^T + teacher_mask))) + + Reference: + DeepSeek-V4 (https://arxiv.org/abs/2606.19348): + - Section 2.3.1 describes the Lightning Indexer and CSA forward architecture. + - Section 4.2.2 ("Training Setups") describes the 3-stage sparse attention pre-training pipeline: + Stage 1 (Dense Pre-training, first 1T tokens): + use_indexer=False, indexer_loss_scaling_factor=0.0 + Stage 2 (Lightning Indexer Warm-up): + use_indexer=True, indexer_sparse_training=False, indexer_loss_scaling_factor=1.0, + trainable_parameters_mask=['.*indexer.*'] + Stage 3 (Sparse Pre-training): + use_indexer=True, indexer_sparse_training=True, indexer_loss_scaling_factor=1.0 + DeepSeek-V3.2 Section 2.1, Eqs. 3–4 (https://arxiv.org/abs/2512.02556) - DeepSeek Sparse Attention (DSA) + KL divergence distillation loss, adapted here from uncompressed tokens to CSA compressed KV blocks. + + Note: + Teacher distribution normalizes over compressed blocks only (W), rather than + the full (S + W) joint keys used during forward attention. + + Args: + indexer_score: Scores predicted by indexer [batch, q_len, compressed_len]. + query: Query tensor from main model [batch, q_len, heads, dim]. + compressed_kv: Compressed KV tensor from main model [batch, compressed_len, 1, dim]. + compressed_mask: Indexer compressed mask [batch, 1, q_len, compressed_len] or [batch, q_len, compressed_len]. + segment_mask: Segment mask [batch, q_len, compressed_len] or [batch, 1, q_len, compressed_len] or None. + position_ids: Token position IDs [batch, q_len] or None. + sparse_loss: Whether to use sparse loss. + scaling_factor: The scaling factor for the loss. + + Returns: + The computed scalar KL divergence loss. + """ + if compressed_kv is None or indexer_score is None: + return jnp.array(0.0, dtype=jnp.float32) + + batch, q_len, _, _ = query.shape + compressed_len = compressed_kv.shape[1] + if compressed_len == 0: + return jnp.array(0.0, dtype=jnp.float32) + + # Detach main model components from the computational graph. + query = jax.lax.stop_gradient(query) + compressed_kv = jax.lax.stop_gradient(compressed_kv) + + # Construct complete teacher mask: causal block mask + segment packing mask + # 1. Causal block mask: query at position t can only attend to block w if (w+1)*r <= t+1 + if position_ids is not None: + usable_len = compressed_len * self.compress_ratio + block_positions = position_ids[:, : usable_len : self.compress_ratio] + future_mask = (block_positions[:, None, :] + self.compress_ratio) > (position_ids[:, :, None] + 1) + else: + q_pos = jnp.arange(q_len)[:, None] + block_end_pos = (jnp.arange(compressed_len)[None, :] + 1) * self.compress_ratio + future_mask = block_end_pos > (q_pos + 1) + future_mask = jnp.broadcast_to(future_mask[None, :, :], (batch, q_len, compressed_len)) + + # Ensure indexer_mask is 2D/3D [batch, q_len, compressed_len] + if compressed_mask.ndim == 4: + indexer_mask = compressed_mask[:, 0, :, :] + else: + indexer_mask = compressed_mask + + # Combine causal, segment, and sparse masks via boolean disjunction to prevent float32 additive overflow + is_invalid_block = future_mask + if segment_mask is not None: + if segment_mask.ndim == 4: + c_seg = segment_mask[:, 0, :, :compressed_len] + elif segment_mask.ndim == 3: + c_seg = segment_mask[:, :, :compressed_len] + else: + c_seg = segment_mask + is_invalid_block = is_invalid_block | (c_seg < (DEFAULT_MASK_VALUE / 2)) + + if sparse_loss: + is_invalid_block = is_invalid_block | (indexer_mask < (DEFAULT_MASK_VALUE / 2)) + + teacher_mask = jnp.where(is_invalid_block, DEFAULT_MASK_VALUE, 0.0) + c_mask = teacher_mask[:, None, :, :] # [batch, 1, q_len, compressed_len] + + # Valid tokens mask checks causal, segment, and sparse boundaries after all masks are combined + valid_tokens_mask = jnp.any(~is_invalid_block, axis=-1) # [batch, q_len] + + # In CSA, compressed KV is pooled into a single representation per block (num_kv_heads = 1), + # which is broadcast across all query heads. + k_vec = compressed_kv[:, :, 0, :] if compressed_kv.ndim == 4 else compressed_kv + + # Student scores: index_scores already contains causal future_mask and segment masking. + # In sparse training mode, mask non-selected blocks with DEFAULT_MASK_VALUE. + if sparse_loss: + student_scores = jnp.where(indexer_mask < (DEFAULT_MASK_VALUE / 2), DEFAULT_MASK_VALUE, indexer_score) + else: + student_scores = indexer_score + + safe_student_scores = jnp.where(valid_tokens_mask[:, :, None], student_scores, 0.0) + log_indexer_probs = jax.nn.log_softmax(safe_student_scores.astype(jnp.float32), axis=-1) + log_indexer_probs = jnp.where(valid_tokens_mask[:, :, None], log_indexer_probs, 0.0) + + # Query is already scaled by softmax_scale in compressed_query_projection; do not scale again + attention_scores = jnp.einsum("bthd, bwd -> bhtw", query, k_vec, precision=self.config.matmul_precision) + attention_scores = attention_scores + c_mask + + # Apply NaN shielding for pre-block tokens + safe_scores = jnp.where(valid_tokens_mask[:, None, :, None], attention_scores, 0.0) + # Softmax covers compressed blocks only (W), ignoring the sliding window keys (S). + raw_probs = jax.nn.softmax(safe_scores.astype(jnp.float32), axis=-1) + raw_probs = jnp.where(valid_tokens_mask[:, None, :, None], raw_probs, 0.0) + target_probs = jnp.sum(raw_probs, axis=1) + target_probs = jax.lax.optimization_barrier(target_probs) + + # L1 normalize aggregated target distribution across compressed blocks + target_probs = jnp.where(valid_tokens_mask[:, :, None], target_probs, 0.0) + target_probs = target_probs / (jnp.sum(target_probs, axis=-1, keepdims=True) + EPS) + + # KL Divergence: KL(attention || indexer) + log_target_probs = jnp.log(target_probs + EPS) + kl_element = jnp.where(target_probs > 0.0, target_probs * (log_target_probs - log_indexer_probs), 0.0) + kl_per_token = jnp.sum( + jnp.where(valid_tokens_mask[:, :, None], kl_element, 0.0), + axis=-1, + ) + + # Average loss across valid tokens only (ignoring pre-block tokens t < compress_rate) + num_valid_tokens = jnp.maximum(jnp.sum(valid_tokens_mask.astype(jnp.float32)), 1.0) + indexer_loss = (jnp.sum(kl_per_token) / num_valid_tokens) * scaling_factor + + return indexer_loss + def compressed_attention( *, diff --git a/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md b/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md index 661fb62bee..485aee96d6 100644 --- a/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md +++ b/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md @@ -259,6 +259,91 @@ python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ indexer_sparse_training=True ``` +## Pre-training for DeepSeek-V4 + +DeepSeek-V4 employs a hybrid attention architecture across its layers, interleaving Sliding Window Attention (SWA), Highly Compressed Attention (HCA), and **Compressed Sparse Attention (CSA)**. The CSA layers incorporate a **Lightning Indexer** that selects top-k compressed blocks. Note that the indexer is activated only if `max_target_length // 4` > `indexer_topk`. + +As described in the DeepSeek-V4 technical report (Section 4.2.2), sparse attention pre-training follows a three-stage strategy: **Dense Pre-training**, **Lightning Indexer Warm-up**, and **Sparse Pre-training**. + +1. **Dense Pre-training Stage** +The model is pre-trained with standard dense attention across all tokens (first 1T tokens) before attention sparsity is introduced. +```sh +python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + run_name=dsv4_dense_pretraining \ + model_name=deepseek4-284b \ + tokenizer_type=huggingface \ + tokenizer_path=deepseek-ai/DeepSeek-V4-Flash \ + per_device_batch_size=1 \ + enable_checkpointing=false \ + async_checkpointing=false \ + ici_fsdp_parallelism=-1 \ + opt_type=sgd \ + steps=20 \ + max_target_length=4096 \ + attention=dot_product \ + dtype=bfloat16 \ + weight_dtype=bfloat16 \ + dataset_type=synthetic \ + # Standard dense pre-training flags + override_model_config=true \ + use_indexer=false \ + indexer_loss_scaling_factor=0.0 +``` + +2. **Lightning Indexer Warmup Stage** +When attention sparsity is introduced, the lightning indexer undergoes a short warmup stage via KL divergence distillation while the main model parameters remain frozen and language modeling loss is skipped. +```sh +python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + run_name=dsv4_indexer_warmup \ + model_name=deepseek4-284b \ + tokenizer_type=huggingface \ + tokenizer_path=deepseek-ai/DeepSeek-V4-Flash \ + load_parameters_path=${SCANNED_CKPT_PATH?} \ + per_device_batch_size=1 \ + enable_checkpointing=false \ + async_checkpointing=false \ + ici_fsdp_parallelism=-1 \ + opt_type=sgd \ + steps=20 \ + max_target_length=4096 \ + attention=dot_product \ + dtype=bfloat16 \ + weight_dtype=bfloat16 \ + dataset_type=synthetic \ + # Indexer training specific flags (inherits use_indexer=true from deepseek4-284b.yml) + indexer_sparse_training=false \ + indexer_loss_scaling_factor=1.0 \ + trainable_parameters_mask=['.*indexer.*'] +``` + +3. **Sparse Pre-training Stage** +The model trains with sparse attention for the remainder of pre-training, where core attention attends only to the top-k selected compressed blocks and the indexer continues to train jointly. +```sh +python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + run_name=dsv4_sparse_pretraining \ + model_name=deepseek4-284b \ + tokenizer_type=huggingface \ + tokenizer_path=deepseek-ai/DeepSeek-V4-Flash \ + load_parameters_path=${SCANNED_CKPT_PATH?} \ + per_device_batch_size=1 \ + enable_checkpointing=false \ + async_checkpointing=false \ + ici_fsdp_parallelism=-1 \ + opt_type=sgd \ + steps=20 \ + max_target_length=4096 \ + attention=dot_product \ + dtype=bfloat16 \ + weight_dtype=bfloat16 \ + dataset_type=synthetic \ + # Indexer training specific flags (inherits use_indexer=true from deepseek4-284b.yml) + indexer_sparse_training=true \ + indexer_loss_scaling_factor=1.0 +``` + ## Decoding One example command to run decoding with V3 on v5p-256 with unscanned checkpoint for fast decoding. diff --git a/tests/unit/deepseek_v4_indexer_loss_test.py b/tests/unit/deepseek_v4_indexer_loss_test.py new file mode 100644 index 0000000000..9f83d2c122 --- /dev/null +++ b/tests/unit/deepseek_v4_indexer_loss_test.py @@ -0,0 +1,479 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Unit tests for DeepSeek-V4 Compressed Sparse Attention (CSA) Indexer loss and training.""" + +import unittest +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, DEFAULT_MASK_VALUE +from maxtext.configs import pyconfig +from maxtext.layers import attention_compressed +from maxtext.layers.attention_mla import indexer_losses +from maxtext.trainers.pre_train import train as pre_train + + +class _MockNnxDecoder(nnx.Module): + """Minimal mock NNX decoder for pre_train.loss_fn tests.""" + + def __init__(self, vocab_size: int): + self.vocab_size = vocab_size + self.mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) + + def __call__( + self, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + encoder_images=None, + encoder_image_masks=None, + enable_dropout=False, + decoder_target_tokens=None, + decoder_target_mask=None, + ): + del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks + del enable_dropout, decoder_target_tokens, decoder_target_mask + return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) + + +class DeepSeekV4IndexerLossTest(unittest.TestCase): + """Tests for DeepSeek-V4 CSA Indexer KL Divergence loss calculation and gradients.""" + + def setUp(self): + super().setUp() + self.batch_size = 2 + self.seq_len = 16 + self.base_emb_dim = 64 + self.base_num_query_heads = 4 + self.base_num_kv_heads = 1 + self.head_dim = 32 + self.compress_ratio = 4 + self.indexer_n_heads = 4 + self.indexer_head_dim = 32 + self.indexer_topk = 2 + self.q_lora_rank = 32 + self._get_config() + + def _get_config( + self, + use_indexer=True, + indexer_loss_scaling_factor=0.5, + indexer_sparse_training=False, + indexer_topk=None, + ): + """Constructs a test MaxTextConfig with CSA indexer configuration.""" + topk = indexer_topk if indexer_topk is not None else self.indexer_topk + argv = [ + "", + "src/maxtext/configs/base.yml", + "run_name=test_dsv4_indexer", + "decoder_block=deepseek4", + "attention_type=compressed", + "attention=dot_product", + f"use_indexer={use_indexer}", + f"indexer_loss_scaling_factor={indexer_loss_scaling_factor}", + f"indexer_sparse_training={indexer_sparse_training}", + f"max_target_length={self.seq_len}", + f"indexer_topk={topk}", + f"indexer_n_heads={self.indexer_n_heads}", + f"indexer_head_dim={self.indexer_head_dim}", + f"base_emb_dim={self.base_emb_dim}", + f"base_num_query_heads={self.base_num_query_heads}", + f"base_num_kv_heads={self.base_num_kv_heads}", + f"head_dim={self.head_dim}", + f"qk_rope_head_dim={self.head_dim}", + f"q_lora_rank={self.q_lora_rank}", + "o_groups=2", + "o_lora_rank=16", + "enable_checkpointing=False", + ] + return pyconfig.initialize(argv) + + def _init_csa_attention(self, config): + """Initializes a CompressedAttention module for testing.""" + rngs = nnx.Rngs(0) + mesh = jax.sharding.Mesh(jax.devices(), ("data",)) + attn = attention_compressed.CompressedAttention( + config=config, + num_query_heads=config.num_query_heads, + num_kv_heads=config.num_kv_heads, + head_dim=config.head_dim, + max_target_length=config.max_target_length, + mesh=mesh, + attention_kernel="dot_product", + inputs_q_shape=(self.batch_size, self.seq_len, config.emb_dim), + inputs_kv_shape=(self.batch_size, self.seq_len, config.emb_dim), + compress_ratio=self.compress_ratio, + q_lora_rank=config.q_lora_rank, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs, + ) + return attn + + def test_csa_indexer_loss_computation(self): + """Test that CSA forward pass computes and stores indexer_loss variable.""" + config = self._get_config(indexer_loss_scaling_factor=0.5, indexer_sparse_training=False) + attn = self._init_csa_attention(config) + + inputs_q = jax.random.normal(jax.random.PRNGKey(1), (self.batch_size, self.seq_len, config.emb_dim)) + inputs_kv = jax.random.normal(jax.random.PRNGKey(2), (self.batch_size, self.seq_len, config.emb_dim)) + positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len)) + segment_ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + + out, _ = attn( + inputs_q=inputs_q, + inputs_kv=inputs_kv, + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + self.assertIsNotNone(out) + self.assertEqual(out.shape, (self.batch_size, self.seq_len, config.emb_dim)) + self.assertTrue(hasattr(attn, "indexer_loss")) + self.assertIsInstance(attn.indexer_loss, indexer_losses) + + loss_val = attn.indexer_loss.get_value() + self.assertGreater(float(loss_val), 0.0) + + def test_csa_indexer_loss_sparse_training_mode(self): + """Test CSA forward pass and indexer loss in sparse pre-training mode.""" + config = self._get_config(indexer_loss_scaling_factor=0.5, indexer_sparse_training=True) + attn = self._init_csa_attention(config) + + inputs_q = jax.random.normal(jax.random.PRNGKey(3), (self.batch_size, self.seq_len, config.emb_dim)) + inputs_kv = jax.random.normal(jax.random.PRNGKey(4), (self.batch_size, self.seq_len, config.emb_dim)) + positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len)) + segment_ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + + out, _ = attn( + inputs_q=inputs_q, + inputs_kv=inputs_kv, + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + self.assertIsNotNone(out) + self.assertTrue(hasattr(attn, "indexer_loss")) + self.assertIsInstance(attn.indexer_loss, indexer_losses) + self.assertGreater(float(attn.indexer_loss.get_value()), 0.0) + + def test_csa_indexer_loss_kl_divergence_zero(self): + """Test KL divergence is 0 when predicted and target distributions match.""" + config = self._get_config(indexer_loss_scaling_factor=1.0) + attn = self._init_csa_attention(config) + + n_windows = self.seq_len // self.compress_ratio + query = jnp.zeros((self.batch_size, self.seq_len, config.num_query_heads, config.head_dim)) + compressed_kv = jnp.zeros((self.batch_size, n_windows, config.num_kv_heads, config.head_dim)) + compressed_mask = jnp.zeros((self.batch_size, 1, self.seq_len, n_windows)) + + # Causal block mask (matches what DeepseekV4Indexer produces on uniform logits) + q_pos = jnp.arange(self.seq_len)[:, None] + block_end_pos = (jnp.arange(n_windows)[None, :] + 1) * self.compress_ratio + future_mask = block_end_pos > (q_pos + 1) + indexer_score = jnp.where(future_mask[None, :, :], DEFAULT_MASK_VALUE, 0.0) + + loss = attn.calculate_csa_indexer_loss( + indexer_score=indexer_score, + query=query, + compressed_kv=compressed_kv, + compressed_mask=compressed_mask, + segment_mask=None, + position_ids=None, + sparse_loss=False, + scaling_factor=1.0, + ) + np.testing.assert_allclose(float(loss), 0.0, atol=1e-5) + + def test_csa_indexer_gradients_flow(self): + """Test that gradients flow to indexer parameters and do not leak into main projections or inputs.""" + for is_sparse in (False, True): + with self.subTest(indexer_sparse_training=is_sparse): + config = self._get_config(indexer_loss_scaling_factor=1.0, indexer_sparse_training=is_sparse) + attn = self._init_csa_attention(config) + + inputs_q = jax.random.normal(jax.random.PRNGKey(1), (self.batch_size, self.seq_len, config.emb_dim)) + inputs_kv = jax.random.normal(jax.random.PRNGKey(2), (self.batch_size, self.seq_len, config.emb_dim)) + positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len)) + segment_ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + + def loss_fn(attn_model, q, kv, seg=segment_ids, pos=positions): + attn_model( + inputs_q=q, + inputs_kv=kv, + decoder_segment_ids=seg, + inputs_positions=pos, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return attn_model.indexer_loss.get_value() + + # 1. Gradients with respect to model parameters (argnums=0) + grad_model_fn = nnx.grad(loss_fn, argnums=0) + grads = grad_model_fn(attn, inputs_q, inputs_kv) + + # Gradients must flow to indexer projection kernels + self.assertIsNotNone(grads.csa_compressor.indexer.q_proj.kernel) + self.assertIsNotNone(grads.csa_compressor.indexer.kv_proj.kernel) + self.assertIsNotNone(grads.csa_compressor.indexer.gate_proj.kernel) + self.assertIsNotNone(grads.csa_compressor.indexer.weights_proj.kernel) + + q_grad_norm = jnp.linalg.norm(grads.csa_compressor.indexer.q_proj.kernel.get_value()) + self.assertGreater(float(q_grad_norm), 0.0) + + # Gradients must not leak into main model projections + self.assertAlmostEqual(float(jnp.linalg.norm(grads.wq_a.kernel.get_value())), 0.0) + self.assertAlmostEqual(float(jnp.linalg.norm(grads.wq_b.kernel.get_value())), 0.0) + self.assertAlmostEqual(float(jnp.linalg.norm(grads.wkv.kernel.get_value())), 0.0) + + # 2. Gradients with respect to inputs (argnums=(1, 2)) must be zero (detached) + grad_inputs_fn = nnx.grad(loss_fn, argnums=(1, 2)) + grad_q, grad_kv = grad_inputs_fn(attn, inputs_q, inputs_kv) + self.assertAlmostEqual(float(jnp.linalg.norm(grad_q)), 0.0) + self.assertAlmostEqual(float(jnp.linalg.norm(grad_kv)), 0.0) + + def test_dense_warmup_forward_mask_is_causal_dense(self): + """Test that dense warm-up forward pass executes the dense causal path. + + Asserts mask values and compares against top-k=1 sparse mode. + """ + # Case A: Verify default pre-training (scale=0) executes cleanly without registering indexer loss + config_unscaled = self._get_config(indexer_loss_scaling_factor=0.0, indexer_sparse_training=False) + attn_unscaled = self._init_csa_attention(config_unscaled) + + inputs_q = jax.random.normal(jax.random.PRNGKey(1), (self.batch_size, self.seq_len, config_unscaled.emb_dim)) + inputs_kv = jax.random.normal(jax.random.PRNGKey(2), (self.batch_size, self.seq_len, config_unscaled.emb_dim)) + positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len)) + segment_ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + n_windows = self.seq_len // self.compress_ratio + + attn_unscaled( + inputs_q=inputs_q, + inputs_kv=inputs_kv, + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + # No indexer loss registered when scaling_factor == 0.0 + self.assertFalse(hasattr(attn_unscaled, "indexer_loss")) + self.assertIsNone(getattr(attn_unscaled, "indexer_loss", None)) + + # Directly assert the dense causal mask values across block boundaries + dense_mask = attn_unscaled.get_compressed_mask(positions, n_windows) + self.assertEqual(dense_mask.shape, (self.batch_size, 1, self.seq_len, n_windows)) + # Token t=0: All 4 blocks are future -> all masked + np.testing.assert_allclose(np.array(dense_mask[:, 0, 0, :]), DEFAULT_MASK_VALUE, atol=1e-5) + # Token t=3: Block 0 complete (valid 0.0), Blocks 1..3 future (DEFAULT_MASK_VALUE) + np.testing.assert_allclose(np.array(dense_mask[:, 0, 3, 0]), 0.0, atol=1e-5) + np.testing.assert_allclose(np.array(dense_mask[:, 0, 3, 1:]), DEFAULT_MASK_VALUE, atol=1e-5) + # Token t=7: Blocks 0..1 complete (valid 0.0), Blocks 2..3 future (DEFAULT_MASK_VALUE) + np.testing.assert_allclose(np.array(dense_mask[:, 0, 7, :2]), 0.0, atol=1e-5) + np.testing.assert_allclose(np.array(dense_mask[:, 0, 7, 2:]), DEFAULT_MASK_VALUE, atol=1e-5) + # Token t=15: All blocks 0..3 complete -> all 0.0 + np.testing.assert_allclose(np.array(dense_mask[:, 0, 15, :]), 0.0, atol=1e-5) + + # Case B: Output divergence between dense warm-up and top-1 sparse mode + config_dense = self._get_config(indexer_loss_scaling_factor=1.0, indexer_sparse_training=False, indexer_topk=1) + config_sparse = self._get_config(indexer_loss_scaling_factor=1.0, indexer_sparse_training=True, indexer_topk=1) + + attn_dense = self._init_csa_attention(config_dense) + attn_sparse = self._init_csa_attention(config_sparse) + + state_dense = nnx.state(attn_dense) + nnx.update(attn_sparse, state_dense) + + out_dense, _ = attn_dense( + inputs_q=inputs_q, + inputs_kv=inputs_kv, + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + out_sparse, _ = attn_sparse( + inputs_q=inputs_q, + inputs_kv=inputs_kv, + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + diff_t15 = jnp.linalg.norm(out_dense[:, 15, :] - out_sparse[:, 15, :]) + self.assertGreater(float(diff_t15), 0.05) + + # Loss must be populated in dense warm-up mode when scaling_factor > 0.0 + self.assertIsNotNone(attn_dense.indexer_loss) + self.assertGreater(float(attn_dense.indexer_loss.get_value()), 0.0) + + def test_teacher_causality_and_packing_on_loss_function(self): + """Test calculate_csa_indexer_loss directly on a 2-segment packed sequence with causal boundaries.""" + config = self._get_config(indexer_loss_scaling_factor=1.0) + attn = self._init_csa_attention(config) + + # 2 segments: Doc 1 = tokens 0..7 (blocks 0, 1), Doc 2 = tokens 8..15 (blocks 2, 3) + n_windows = self.seq_len // self.compress_ratio # 4 blocks + positions = jnp.array([[0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7]] * self.batch_size) + segment_ids = jnp.array([[1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2]] * self.batch_size) + + # Build compressed_segment_mask for the 2 documents + comp_seg_ids = jnp.array([[1, 1, 2, 2]] * self.batch_size) + valid_comp_seg = segment_ids[:, :, None] == comp_seg_ids[:, None, :] + compressed_segment_mask = jnp.where(valid_comp_seg, 0.0, DEFAULT_MASK_VALUE) + + query = jnp.zeros((self.batch_size, self.seq_len, config.num_query_heads, config.head_dim)) + compressed_kv = jnp.zeros((self.batch_size, n_windows, config.num_kv_heads, config.head_dim)) + compressed_mask = jnp.zeros((self.batch_size, 1, self.seq_len, n_windows)) + + # Ground truth student prediction matching causal + packed teacher distribution + usable_len = n_windows * attn.compress_ratio + block_positions = positions[:, : usable_len : attn.compress_ratio] + is_future = (block_positions[:, None, :] + attn.compress_ratio) > (positions[:, :, None] + 1) + causal_mask = jnp.where(is_future, DEFAULT_MASK_VALUE, 0.0) + ground_truth_student_scores = causal_mask + compressed_segment_mask + + loss_perfect = attn.calculate_csa_indexer_loss( + indexer_score=ground_truth_student_scores, + query=query, + compressed_kv=compressed_kv, + compressed_mask=compressed_mask, + segment_mask=compressed_segment_mask, + position_ids=positions, + sparse_loss=False, + scaling_factor=1.0, + ) + np.testing.assert_allclose(float(loss_perfect), 0.0, atol=1e-5) + + # Case B: Student predicts mass on a future block in Doc 1 (t=4 predicting block 1) + leaky_student_scores = ground_truth_student_scores.at[:, 4, 1].set(100.0) + loss_future_leak = attn.calculate_csa_indexer_loss( + indexer_score=leaky_student_scores, + query=query, + compressed_kv=compressed_kv, + compressed_mask=compressed_mask, + segment_mask=compressed_segment_mask, + position_ids=positions, + sparse_loss=False, + scaling_factor=1.0, + ) + self.assertGreater(float(loss_future_leak), 0.1) + + # Case C: Student in Doc 2 predicts mass on a block from Doc 1 (t=12 predicting block 0) + cross_doc_student_scores = ground_truth_student_scores.at[:, 12, 0].set(100.0) + loss_cross_doc = attn.calculate_csa_indexer_loss( + indexer_score=cross_doc_student_scores, + query=query, + compressed_kv=compressed_kv, + compressed_mask=compressed_mask, + segment_mask=compressed_segment_mask, + position_ids=positions, + sparse_loss=False, + scaling_factor=1.0, + ) + self.assertGreater(float(loss_cross_doc), 0.1) + + def test_csa_indexer_loss_jit_compile(self): + """Compile smoke test: verifies that jitting forward pass with CSA indexer loss executes cleanly.""" + config = self._get_config(indexer_loss_scaling_factor=0.5, indexer_sparse_training=False) + attn = self._init_csa_attention(config) + + inputs_q = jax.random.normal(jax.random.PRNGKey(1), (self.batch_size, self.seq_len, config.emb_dim)) + inputs_kv = jax.random.normal(jax.random.PRNGKey(2), (self.batch_size, self.seq_len, config.emb_dim)) + positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len)) + segment_ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + + @nnx.jit + def jitted_forward(attn_model, q, kv, seg, pos): + out, _ = attn_model( + inputs_q=q, + inputs_kv=kv, + decoder_segment_ids=seg, + inputs_positions=pos, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return out, attn_model.indexer_loss.get_value() + + out, loss_val = jitted_forward(attn, inputs_q, inputs_kv, segment_ids, positions) + self.assertEqual(out.shape, (self.batch_size, self.seq_len, config.emb_dim)) + self.assertGreater(float(loss_val), 0.0) + + def test_mask_routing_matrix(self): + """Verify mask routing selects sparse mask by default, and dense only during active dense warmup.""" + positions = jnp.broadcast_to(jnp.arange(self.seq_len)[None, :], (self.batch_size, self.seq_len)) + sparse_mask = jnp.full((self.batch_size, 1, self.seq_len, 4), DEFAULT_MASK_VALUE) + + # (mode, scaling_factor, sparse_training, expected_is_sparse) + test_matrix = [ + (MODEL_MODE_TRAIN, 0.0, False, True), # Default pre-training: sparse mask + (MODEL_MODE_TRAIN, 1.0, False, False), # Active dense warm-up: dense causal mask + (MODEL_MODE_TRAIN, 1.0, True, True), # Sparse training: sparse mask + (MODEL_MODE_PREFILL, 0.0, False, True), # Prefill inference: sparse mask + (MODEL_MODE_AUTOREGRESSIVE, 0.0, False, True), # AR decode inference: sparse mask + ] + + for mode, scale, sparse_training, expected_sparse in test_matrix: + config = self._get_config(indexer_loss_scaling_factor=scale, indexer_sparse_training=sparse_training) + attn = self._init_csa_attention(config) + is_dense_warmup = (mode == MODEL_MODE_TRAIN) and (scale > 0.0) and (not sparse_training) + use_sparse_mask = not is_dense_warmup + routed = attn.get_compressed_mask(positions, 4, sparse_compressed_mask=sparse_mask if use_sparse_mask else None) + + if expected_sparse: + np.testing.assert_allclose(np.array(routed), np.array(sparse_mask), atol=1e-5) + else: + np.testing.assert_allclose(np.array(routed[:, 0, 15, :]), 0.0, atol=1e-5) + + def test_pre_train_loss_fn_stages(self): + """Verify pre_train.loss_fn behavior across dense, warm-up, and sparse training stages.""" + data = { + "inputs": jnp.zeros((self.batch_size, self.seq_len), dtype=jnp.int32), + "inputs_position": jnp.broadcast_to(jnp.arange(self.seq_len), (self.batch_size, self.seq_len)), + "inputs_segmentation": jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32), + "targets": jnp.zeros((self.batch_size, self.seq_len), dtype=jnp.int32), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32), + } + # Stage 1: Standard dense pre-training (use_indexer=False) + # Must compute normal LM loss (xent_sum > 0) + cfg_dense = self._get_config(use_indexer=False, indexer_loss_scaling_factor=0.0, indexer_sparse_training=False) + mock_model = _MockNnxDecoder(vocab_size=cfg_dense.vocab_size) + loss_dense, aux_dense = pre_train.loss_fn(mock_model, cfg_dense, data, None, None, is_train=True) + self.assertGreater(float(aux_dense["xent_sum"]), 0.0) + self.assertGreater(float(loss_dense), 0.0) + + # Stage 2: Dense warm-up configuration (use_indexer=True, scaling_factor=1.0, sparse_training=False) + # Must zero out main model LM loss (xent_sum == 0.0) + cfg_warmup = self._get_config(use_indexer=True, indexer_loss_scaling_factor=1.0, indexer_sparse_training=False) + _, aux_warmup = pre_train.loss_fn(mock_model, cfg_warmup, data, None, None, is_train=True) + self.assertEqual(float(aux_warmup["xent_sum"]), 0.0) + self.assertEqual(float(aux_warmup["z_loss"]), 0.0) + + # Stage 3: Sparse training configuration (use_indexer=True, scaling_factor=1.0, sparse_training=True) + # Must compute normal LM loss (xent_sum > 0) + cfg_sparse = self._get_config(use_indexer=True, indexer_loss_scaling_factor=1.0, indexer_sparse_training=True) + loss_sparse, aux_sparse = pre_train.loss_fn(mock_model, cfg_sparse, data, None, None, is_train=True) + self.assertGreater(float(aux_sparse["xent_sum"]), 0.0) + self.assertGreater(float(loss_sparse), 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/deepseek_v4_vs_reference_test.py b/tests/unit/deepseek_v4_vs_reference_test.py index 15b44b91e7..15db139f5e 100644 --- a/tests/unit/deepseek_v4_vs_reference_test.py +++ b/tests/unit/deepseek_v4_vs_reference_test.py @@ -470,7 +470,7 @@ def setUp(self): attention_dropout=0.0, ) - def _build_maxtext_config(self, layer_type): + def _build_maxtext_config(self, layer_type, attention_kernel="dot_product"): """Builds a MaxText pyconfig for a specific layer_type.""" config_arguments = { @@ -496,6 +496,9 @@ def _build_maxtext_config(self, layer_type): "indexer_topk": self.pt_config.index_topk, "normalization_layer_epsilon": self.pt_config.rms_norm_eps, "use_tokamax_splash": True, + "attention_type": "compressed", + "attention": attention_kernel, + "use_indexer": True, } argv = [sys.argv[0], "src/maxtext/configs/base.yml"] @@ -549,7 +552,7 @@ def _run_e2e_test(self, layer_type, is_packed=False, attention_kernel="dot_produ rope_main = PTRope(self.pt_config) rope_compress = PTRope(self.pt_config) - mt_config = self._build_maxtext_config(layer_type) + mt_config = self._build_maxtext_config(layer_type, attention_kernel=attention_kernel) mesh = Mesh(mesh_utils.create_device_mesh((1,), devices=jax.devices()[:1]), axis_names=("fsdp",)) @@ -678,7 +681,7 @@ def _run_e2e_test(self, layer_type, is_packed=False, attention_kernel="dot_produ mt_q_latent = mt_attn.wq_a(x_mt) mt_q_residual = mt_attn.q_norm(mt_q_latent) - mt_top_k_indices = mt_attn.csa_compressor.indexer(x_mt, mt_q_residual, pos_mt) + mt_top_k_indices, _ = mt_attn.csa_compressor.indexer(x_mt, mt_q_residual, pos_mt) print(f"MaxText top_k_indices:\n{mt_top_k_indices[0]}") num_mismatches = np.sum(pt_top_k_indices.detach().numpy() != np.array(mt_top_k_indices)) @@ -1378,6 +1381,7 @@ def setUp(self): "dtype": "float32", "weight_dtype": "float32", "skip_jax_distributed_system": True, + "attention": "dot_product", "use_tokamax_splash": True, } argv = [sys.argv[0], "src/maxtext/configs/base.yml"]