Add Indexer Loss [Deepseek v4]

[octatrifan]

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):

• #3364
• #3415

In DeepSeek-V3.2 (MLA), calculate_indexer_loss distills dense per-token attention into token-level indexer scores. DeepSeek-V4 uses block-compressed key-values where $m=4$ tokens share a single key vector. This change implements calculate_csa_indexer_loss to compute KL divergence between the student's block-level predictions and the teacher's compressed attention distribution.

Key Changes

• attention_compressed.py:
  • Implemented calculate_csa_indexer_loss with block causal masking and document boundary support.
  • Added jax.lax.stop_gradient to teacher projections and indexer inputs to isolate indexer gradients.
  • Added head-chunking scan support (mla_qk_head_chunk_size) to prevent HBM OOM during teacher attention.
  • Gated dense warm-up on indexer_loss_scaling_factor > 0.0 and not indexer_sparse_training; defaults to sparse CSA otherwise.
• train.py:
  • Gated LM loss bypass on indexer_loss_scaling_factor > 0.0 and not indexer_sparse_training, aligning with existing decoders.py logic across both 3.2 and 4.
• types.py:
  • Set CSA block size to 4 and updated validation to allow indexer_topk <= max_target_length // 4.
• Model Configs (deepseek4-tiny.yml, deepseek4-284b.yml):
  • Enabled use_indexer: true.
• Unit Tests (tests/unit/deepseek_v4_indexer_loss_test.py):
  • Added tests covering indexer KL divergence, gradient isolation, head chunking parity, mask routing across modes, and default scale=0 LM loss.

BUGS: b/548036332

Tests

1. Unit Tests

Ran on TPU VM (v5p8):

python3 tests/unit/deepseek_v4_indexer_loss_test.py

10/10 passed in 39.89s.

2. Dense Indexer Warm-up 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_warmup_100step steps=100 \
  use_indexer=True indexer_loss_scaling_factor=0.01 indexer_sparse_training=False \
  trainable_parameters_mask=['.*indexer.*'] \
  dataset_type=hf hf_path='HuggingFaceFW/fineweb-edu' tokenizer_path=deepseek-ai/DeepSeek-V3.2 \
  base_output_directory=/tmp/maxtext_output

Verified lm_loss = 0.000 across all 100 steps; non-indexer weights frozen via trainable_parameters_mask=['.*indexer.*'] (matching DeepSeek-V3.2); indexer KL loss stable at ~0.00295. Step time: 2.08s (~11,800 tokens/s/device).

Step	total_loss	lm_loss	indexer_loss	moe_lb_loss	Step Time
0	0.003	0.000	0.002969	0.000405	28.97s (compile)
1	0.003	0.000	0.002970	0.000405	8.08s
25	0.003	0.000	0.002971	0.000403	2.08s
50	0.003	0.000	0.002958	0.000402	2.08s
99	0.003	0.000	0.002950	0.000400	2.08s

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_output

Verified 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.

Step	total_loss	lm_loss	indexer_loss	moe_lb_loss	Step Time
0	12.297	12.279	0.01678	0.000405	29.66s (compile)
1	12.298	12.281	0.01673	0.000404	7.82s
25	12.155	12.138	0.01653	0.000404	2.27s
50	11.989	11.973	0.01577	0.000405	2.27s
99	11.869	11.854	0.01493	0.000405	2.27s

Full log
Full log for attention=flash

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

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.

[codecov]

Codecov Report

❌ Patch coverage is 83.48624% with 18 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_compressed.py	83.48%	13 Missing and 5 partials ⚠️

📢 Thoughts on this report? Let us know!