Add Indexer Loss [Deepseek v4] - #5033
Conversation
- Implement DeepSeek-V4 CSA Indexer KL divergence loss distillation against dense attention representations - Support memory-efficient head-chunking scan via mla_qk_head_chunk_size - Add NaN shielding (valid_tokens_mask) for pre-block queries prior to block 0 completion - Initialize dedicated DeepSeekV4RotaryEmbedding in DeepseekV4Indexer for index_head_dim - Update MaxTextConfig validators to allow use_indexer with AttentionType.COMPRESSED - Add comprehensive unit tests for CSA indexer loss computation, KL boundary condition, gradient isolation, head-chunking parity, and sparse mode
…double scaling, enforce causal block masking on teacher, revert RoPE rewrite, and isolate loss gating
…ent packing for CSA indexer - Ensure attention forward pass is dense over all blocks during dense warm-up stage (indexer_sparse_training=False) - Incorporate segment/packing mask with causal block mask in teacher_mask and valid_tokens_mask - Prevent redundant duplicate causal mask application on student indexer_score - Enforce strict teacher causality ensuring 0.0 probability mass on all future blocks when position_ids is set - Add unit tests for teacher causality/packing and dense warm-up forward pass
…gorous 2-segment packed test cases - Replace all-zero forward mask with dense causal block mask during dense warm-up stage in CompressedAttention - Verify that for query token t=4, block 0 is unmasked (0.0) and block 1 is masked (-1e9) - Add test_teacher_causality_and_packing_on_loss_function calling calculate_csa_indexer_loss on a 2-segment packed document sequence - Verify loss is 0.0 on matching ground truth and penalizes future/cross-document score leakage (> 0.1)
…inst top-k=1 sparse mode
…mpile smoke test - Update types.py: derive CSA rate from compress_ratios, clarify block bounds validation - Set use_indexer: true on deepseek4-tiny.yml and deepseek4-284b.yml - Refactor calculate_csa_indexer_loss: rename causal_mask to segment_mask, replace literal -1e9/-100.0 with DEFAULT_MASK_VALUE, line-for-line structure alignment with MLA - Add test_csa_indexer_loss_jit_compile smoke test and verify 8/8 unit tests pass on TPU VM
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements the indexer KL divergence loss for Compressed Sparse Attention (CSA) in DeepSeek-V4. It updates model configurations and validation logic to support the indexer with compressed attention, modifies the indexer and compressor layers to return scores, and introduces the calculate_csa_indexer_loss method to compute the KL divergence loss. Additionally, it adds a comprehensive suite of unit tests to verify the loss computation, gradient flow, and mask behavior. There are no review comments to address, so I have no feedback to provide.
ef954b3 to
4123d22
Compare
…tests - Decouple mask selection from indexer loss calculation in CompressedAttention; scale=0 and sparse=False stays dense causal without indexer loss - Add get_compressed_mask helper to CompressedAttention and assert mask tensor values across block boundaries in unit tests - Extend gradient detachment test to assert input grads are 0.0 under both sparse and dense modes - Update types.py docstrings to reference MLA or Compressed Attention for use_indexer
741f351 to
8209361
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
The implementation of the DeepSeek-V4 CSA Indexer KL divergence loss is well-structured, mathematically sound, and rigorously tested. The gradient isolation, causal/segment mask routing, and numerically stable loss calculations are implemented correctly.
🔍 General Feedback
- Testing: The unit tests provide excellent coverage for gradients, chunking parity, loss values, and masking logic.
- NaN Shielding: The use of
valid_tokens_maskand proper conditional replacements effectively preventNaN/Infduring softmax and KL calculations when dealing with fully masked tokens. - Robustness: Mask broadcasting and the manual head-chunking scan are implemented cleanly.
- I've left a minor, low-severity note regarding the hardcoded compression rate for documentation purposes, but everything else looks solid and ready to merge!
| ): | ||
| raise ValueError( | ||
| "Sparse indexer with all-gather context parallelism for flash attention does not support attention sinks." | ||
| ) |
There was a problem hiding this comment.
The compression rate for DeepSeek-V4 CSA is currently hardcoded to 4 here. While this is mathematically correct for the current architecture (CSA uses compress_ratio=4, HCA uses compress_ratio=128), adding a brief comment or dynamically resolving it from the config (if applicable) would improve long-term maintainability.
| ) | |
| elif self.attention_type == AttentionType.COMPRESSED.value: | |
| # DeepSeek-V4 CSA uses a hardcoded compression rate of 4 for the indexer blocks. | |
| compress_rate = 4 | |
| max_blocks = self.max_target_length // compress_rate |
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Thanks for the change!
| "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: |
There was a problem hiding this comment.
Will compressed attention have similar constraints?
supports_dot_product = self.attention == "dot_product"
supports_flash_splash = self.attention == "flash" and self.use_tokamax_splash
| final_indices = jnp.where(invalid, jnp.full_like(top_k_indices, -1), top_k_indices) | ||
|
|
||
| if return_scores: | ||
| return final_indices, index_scores |
There was a problem hiding this comment.
Should we always return a 2-element tuple regardless of return_scores (e.g., returning None for scores when false)? This keeps the return signature static and ensures graph compilation doesn't depend on the return_scores value
In this case you could:
top_k_indices, indexer_scores = self.indexer(
hidden_states,
q_latent,
position_ids,
attention_mask,
model_mode,
indexer_cache,
return_scores=return_indexer_scores,
)
| target_distribution = L1_Normalize(Sum_h(Softmax_w(Q @ K_comp^T + teacher_mask))) | ||
|
|
||
| Reference: | ||
| DeepSeek-V4 (CSA / Lightning Indexer) - Paper §2.3.1, Eqs. 13–17 |
There was a problem hiding this comment.
It seems Eqs. 13–17 is implementation for attention, may I know the source of this loss implementation? Or any pointer will be great!
There was a problem hiding this comment.
It seems Megatron LM previously also submit a PR, could you cross check if matches at high level?
A few comments from Gemini
1. Teacher Distribution Semantic Mismatch (Megatron-LM Issue #5776 / PR #5960):
DeepSeek-V4 CSA uses a shared softmax denominator across uncompressed sliding-window tokens, compressed KV blocks, and learnable attention sinks.
Restricting teacher softmax strictly to compressed blocks forces each head to allocate unit compressed mass ($1.0$), artificially amplifying local-attention heads and distorting the indexer distillation target.
2. Numerical & JAX Autodiff Stability Fixes:
Mask Additive Overflow: Adding multiple instances of DEFAULT_MASK_VALUE ($-2.38 \times 10^{38}$) overflows IEEE-754 float32 to -inf and crashes float16. Fixed via boolean conjunctions and a single mask application.
Autodiff Gradient Vanishing & NaN Leakage: Replacing P * (log(P + EPS) - log(Q + EPS)) with jax.nn.log_softmax combined with block-level finite shielding yields the exact analytical gradient $Q - P$, eliminating epsilon-distortion and reverse-mode VJP NaN contamination.
Sparse Top-$k$ Mask Ordering: Fixed the bug where valid_tokens_mask was computed before incorporating indexer_mask, which caused unselected/invalid tokens to evaluate rows of all-$-\infty$ into NaNs.
2 participants
Description
Adds Indexer KL divergence distillation loss for DeepSeek-V4 Compressed Sparse Attention (CSA), enabling two-stage training (Dense Warm-up and Sparse Pre-training) on TPUs.
Reference PRs (for DSV3.2):
In DeepSeek-V3.2 (MLA),$m=4$ tokens share a single key vector. This change implements
calculate_indexer_lossdistills dense per-token attention into token-level indexer scores. DeepSeek-V4 uses block-compressed key-values wherecalculate_csa_indexer_lossto compute KL divergence between the student's block-level predictions and the teacher's compressed attention distribution.Key Changes
attention_compressed.py:calculate_csa_indexer_losswith block causal masking and document boundary support.jax.lax.stop_gradientto teacher projections and indexer inputs to isolate indexer gradients.mla_qk_head_chunk_size) to prevent HBM OOM during teacher attention.indexer_loss_scaling_factor > 0.0 and not indexer_sparse_training; defaults to sparse CSA otherwise.train.py:indexer_loss_scaling_factor > 0.0 and not indexer_sparse_training, aligning with existingdecoders.pylogic across both 3.2 and 4.types.py:indexer_topk <= max_target_length // 4.deepseek4-tiny.yml,deepseek4-284b.yml):use_indexer: true.tests/unit/deepseek_v4_indexer_loss_test.py):BUGS: b/548036332
Tests
1. Unit Tests
Ran on TPU VM (v5p8):
10/10 passed in 39.89s.
2. Dense Indexer Warm-up Mode (100 steps,
HuggingFaceFW/fineweb-edu)Repro command:
Verified
lm_loss = 0.000across all 100 steps; non-indexer weights frozen viatrainable_parameters_mask=['.*indexer.*'](matching DeepSeek-V3.2); indexer KL loss stable at ~0.00295. Step time: 2.08s (~11,800 tokens/s/device).Full log
Full log for attention=flash
3. Sparse Pre-training Mode (100 steps,
HuggingFaceFW/fineweb-edu)Repro command:
python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml \ model_name=deepseek4-tiny run_name=dsv4_sparse_100step steps=100 \ use_indexer=True indexer_loss_scaling_factor=0.1 indexer_sparse_training=True \ dataset_type=hf hf_path='HuggingFaceFW/fineweb-edu' tokenizer_path=deepseek-ai/DeepSeek-V3.2 \ base_output_directory=/tmp/maxtext_outputVerified monotonic loss descent: Total loss dropped 12.297 -> 11.869 (
lm_loss: 12.279 -> 11.854,indexer_loss: 0.01678 -> 0.01493). Step time: 2.27s (~10,825 tokens/s/device). Checkpoint step 99 serialized and finalized.Full log
Full log for attention=flash
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.