fix(deepseek-v4): correct V4-Flash attention math, wire the MoE balance loss

[lhzhang333]

Summary

Correctness fixes to the DeepSeek-V4-Flash attention and MoE paths, plus two
config keys that had been declared for a long time without anything reading them,
plus a recipe rebalance for the 4-node run. Every fix is backed by the released
reference (deepseek-ai/DeepSeek-V4-Flash, inference/model.py), the DeepSeek-V4
paper, or the upstream open-source training reference.

Thirteen commits plus one merge, merged up to the tip of
dev/tas/deepseek-v4-phase2 (9eea27db). Net diff: 29 files, +4400 / −90 —
4 new production modules, 9 new test files. (The thirteenth is an isort / black
pass over the twelve below, with no behaviour change.)

Several of these change model math on almost every layer, so loss curves will
shift and checkpoints trained with the old code are no longer equivalent. That
is the intended effect: the old behaviour did not match the reference.

#	Commit	Change	Severity	Layers	Throughput
1	10524f85	Drop YaRN m_scale from the attention softmax scale	High	41 compressed	0%
2	70a0fc06	Rotate compressed KV at original-sequence positions; fix Flash compress_rope_theta	High	41 compressed	~0%
3	2255b56f	De-rotate the attention output before the O projection	Blocking	all 43	see 6g
4	4dbfb578	Keep the indexer QK in high precision by default	High	21 CSA	slightly faster
5	1ae566c5	Add opt-in indexer distillation loss	Blocking (gap)	21 CSA	0% (default off)
6	f133beb3	Align the V4 attention and CSA indexer with the reference (7 items)	Blocking	all 43	+0.26% net
7	c1431ac5	Wire the MoE balance loss and shared-expert overlap	Medium-high	all MoE	0% (overlap dormant)
8	5b833bf1	Pair the compressed RoPE base with its YaRN factor	High	41 compressed	0%
9	28410e30	Ship the keep-in-FP32 pinning disabled by default (fixes 6f)	Blocking	all 43	0%
10	44c613be	Correct the yaml comment on when the distillation loss should be on	Docs	—	0%
11	7b010cef	Build the indexer KL target without the dense per-head tensors	Perf	21 CSA	−2.7% when the loss is on
12	1458f9e4	Rebalance the 4N PP layout, default recompute off, make both overridable	Recipe	—	4N only

Three things ship disabled, dormant, or not working. Keep-in-FP32 (6f)
aborts every GPU at 4 nodes and is now behind PRIMUS_V4_KEEP_FP32=1 — see 9.
Shared-expert overlap (7b) cannot be reached by the production recipe. Both new
auxiliary losses train correctly but do not reach the log — read "Known
gap" before reading the balance-loss numbers.

Layer counts: compress_ratios has 44 entries; the trailing 0 is the MTP
slot, truncated by _normalize_compress_ratios. The 43 decoder layers are
2 SWA + 21 CSA + 20 HCA, so 41 are compressed and the last decoder layer is
CSA, not SWA.

---

Part 1 — Attention correctness

1. Drop YaRN m_scale from the softmax scale. _attention_scale()
multiplied 1/sqrt(head_dim) by the layer's YaRN magnitude factor, inflating
every logit on the 41 compressed layers by ~1.277x at
rotary_scaling_factor = 16. The reference uses a plain head_dim ** -0.5 for
both the core attention and the indexer. Paper §2.4 gives the reason: V4 RMSNorms
the queries and KV entries directly, so no temperature compensation is needed
(and for the same reason QK-Clip is not used). The YaRN frequency interpolation
on inv_freq is untouched.

2. Rotate compressed KV at original-sequence positions. Two defects in the
compressed (CSA/HCA) RoPE branch. Position basis: _build_compressed_pool
rotated entry s at position s, while queries are rotated at their original
token positions. Entry s covers the window starting at s*ratio, so the two
sides lived in different coordinate systems and the relative phase was off by a
factor of compress_ratio (4x CSA, 128x HCA). The reference samples
freqs_cis[:cutoff:ratio]; RoPECache.forward_arange grows a stride argument
folded into the memo key, so the table stays cached. Base value: Flash
inherited compress_rope_theta = 160000, which is the V4-Pro value; the
released Flash ModelArgs ships 40000.0.

Cross-checking against HuggingFace transformers: its DeepseekV4Config still
defaults compress_rope_theta to 160000, conflicting with the released
inference/model.py (huggingface/transformers issue 45910). This PR follows
the reference.

3. De-rotate the attention output before the O projection (blocking). V4
shares one latent between K and V, so the values entering the softmax are already
rotated and the core output o_t = sum_s P[t,s] * R(s) * v_s carries absolute
position information. Paper §2.3 applies RoPE with position −i on the last 64
dimensions of each output, right before wo_a. Primus went straight into the
grouped-O projection, so all 43 layers fed wo_a absolute positions — not
absorbable by training, since wo_a is a learned linear and what is missing is a
per-token rotation. All ten attention backends converge to [B,S,H,head_dim]
before the projection, so one insertion point covers them all with no kernel
changes; both exits now go through a single _project_output() helper so no
branch can reach wo_a un-derotated.

4. Keep the indexer QK in high precision by default.
run_deepseek_v4_flash.sh defaulted USE_V4_FP8_INDEXER=True, overriding both
run_deepseek_v4.sh and the yaml — that divergence is why nobody noticed. It
should not be a silent default: the indexer decides which entries each query
attends to, so quantization error there changes the selection rather than
perturbing a value; and it is a fake quant around a BF16 GEMM, so outside QAT
it is pure overhead. Item 6 closes the other half — the indexer's weight
projections were gated on PRIMUS_V4_FP8_ATTN_PROJ, so enabling FP8 for the
attention projections silently quantized the selector too. Now its own knob,
PRIMUS_V4_FP8_INDEXER_PROJ, default off.

5. Add opt-in indexer distillation loss. topk is not differentiable and the
forward discarded the indexer scores, so the indexer had no gradient path and was
frozen — a from-scratch pretrain selects at random across all 21 CSA layers. (The
freeze was deliberate: dead params in the grad buckets waste optimizer state and
break the grad-ready-hook invariant, so unfreezing and adding the loss must
happen together.) Adds KL(attention || indexer) over the selected entries per
V3.2 §2.1, sparse (top-k) variant, so the computation stays in [B, S, K]
and never materialises the dense [B, H, S, P] tensor.
v4_indexer_distill_loss_coeff (default 0.0, surfaced as
PRIMUS_V4_INDEXER_DISTILL_LOSS_COEFF) gates the loss and whether the indexer
is trainable, so the two can never disagree; 0.0 makes this a complete no-op
for existing runs.

When to turn it on: whenever training with sparse attention. V3.2 §2.1.1
keeps the KL on through the entire sparse stage — the main attention keeps
moving, so the indexer has to keep tracking it. Off is only right when nothing
is being trained. 1e-2 is a starting point. Known deviation: V3.2 gives
the indexer its own optimizer and LR (~137x the main model's); here the KL
rides the main loss, so coeff is standing in for that ratio.

6. Align the V4 attention and CSA indexer with the reference. Item 5 shipped
the loss but left it incomplete in ways that only bite once the coefficient is
positive. Re-reading the CSA side of the reference (item 5 had only consulted the
DSA side) turned up seven defects:

• 6a — scores were missing their temperature. I_{t,s} was the bare head-sum
  of w_h * ReLU(q_h · k_s); the reference scales per-head weights by
  index_n_heads ** -0.5 and folds in index_head_dim ** -0.5 — at Flash widths
  a factor of ~90.5. topk is invariant under a positive constant, which is
  why it went unnoticed; it only bites once the scores enter a softmax, where a
  distribution ~90x too sharp makes the KL gradient unusable. Folded into w_i
  so all scoring branches inherit it. A parity test cannot see this, so the tests
  assert the constant directly.
• 6b — the loss was not one-directional. Neither the KL target nor the
  indexer's input was detached, so the loss also rewarded the main Q projection
  and the compressor for becoming easier to predict, and leaked into every layer
  below. The reference detaches both. The assertion that matters is negative:
  coeff 0 vs 1e3 must give bit-identical gradients outside indexer.*.
• 6c — the aux loss scale had no driver. set_loss_scale was never called,
  so the seeded gradient stayed at 1.0 — under gradient accumulation the
  effective coefficient is num_microbatches times too large. Now follows the
  per-microbatch scale the schedule already installs for the MoE aux loss.
• 6d — the loss was invisible. Now routed through the MoE aux-loss tracker,
  every layer reporting (indexer-less layers report an explicit zero, since a key
  present only on CSA-owning ranks would diverge the cross-PP collective).
  Does not actually reach the log yet — see "Known gap".
• 6e — the indexer scored raw projections. No positional information and no
  Hadamard rotation. The reference applies partial RoPE at the compressed base
  (queries at their own positions, keys at s * compress_ratio) then a normalised
  Hadamard — indexer-only; the main compressor stays unrotated. The rotation is
  orthogonal so the inner product is unchanged in exact arithmetic; what it buys
  is spreading channel energy so nothing dominates the low-precision QK product.
  Prefers the reference's Hadamard extension, falls back to a cached Sylvester
  matrix. Applied before the scoring dispatch, so eager / Triton-tail /
  full-Triton all inherit it. Changes which entries the selector picks, and
  therefore the training trajectory; required for loading the released
  checkpoint.
• 6f — ape / attn_sink were not FP32. Both are FP32 in the released
  checkpoint and both feed a softmax directly, but module.bfloat16() took them
  down. Implemented module-side (keep_in_fp32.py) since the pinned Megatron has
  no equivalent, restoring from a saved FP32 copy rather than casting downgraded
  values back. Read item 9 first — it crashes at 4 nodes and now ships behind
  a flag.
• 6g — the de-rotation used the slower kernel path. Item 3 built cos/sin
  eagerly while Q/KV generate them in-kernel: three extra launches, a materialised
  tensor and a contiguous copy per layer, across 43 layers. An INVERSE flag on
  the in-kernel variant recovers it.

Part 2 — MoE

7a. Sequence-wise balance loss. moe_router_load_balancing_type: seq_aux_loss + moe_aux_loss_coeff: 0.001 had nothing reading them. The paper
pairs the aux-loss-free expert bias with a slight sequence-wise loss precisely
because the bias reacts over many steps and the whole batch, so it cannot stop a
single sequence collapsing onto a few experts. Not inherited from the framework's
built-in seq_aux_loss: that path only knows softmax / sigmoid and raises on
V4's sqrtsoftplus. Implemented from the DeepSeek-V3 formula:

s'_{i,t} = s_{i,t} / sum_j s_{j,t}
f_i      = (E / (K*T)) * sum_t 1(token t -> expert i)
P_i      = (1/T) * sum_t s'_{i,t}
L        = alpha * sum_i f_i * P_i

f sums to E and P to 1, so a perfectly balanced sequence scores exactly
alpha — the logged value is directly readable and the magnitude is cheap to
assert. f counts the actual routing (the tokens really dispatched); P uses
the unbiased normalised affinities, the only part carrying a gradient.

Two recipe caveats before reading this as a working balancing mechanism,
both pre-existing. (1) run_deepseek_v4.sh hard-codes
--moe_router_force_load_balancing True, and _apply_force_load_balancing
replaces the routing actually dispatched — so in the benchmark recipe the
objective optimises a discarded decision. (2)
PRIMUS_MOE_ENABLE_EXPERT_BIAS defaults to False and the CLI wins over the
yaml, so the aux-loss-free half of the pairing is not running at all. Turning
it on is not safe yet: finalize_model_grads._update_router_expert_bias reads
module.local_tokens_per_expert on every module with an expert_bias
attribute, and the standalone V4 router has the latter but not the former.

7b. Shared-expert / all-to-all overlap. moe_shared_expert_overlap: true was
declared while the shared expert ran serially. The dispatcher hooks sit in exactly
the decomposed calls the V4 MoE forward already makes. The load-bearing detail:
combine_postprocess adds the shared output into the combine result, so
handing the shared expert over also means the layer must stop adding it or it is
double-counted — that is what the new helper's return value gates, with its own
test.

Dormant in the V4-Flash recipe and not exercised by the numbers below.
Three pre-existing gates: the recipe yaml already overrides
moe_shared_expert_overlap: false; run_deepseek_v4.sh pins
--moe_shared_expert_overlap False under USE_TURBO_DEEPEP=True; and the
pinned Megatron's MoEFlexTokenDispatcher.set_shared_experts raises
NotImplementedError. The third is version-specific — newer upstream drops
that raise, so bumping the submodule unblocks this path, at which point the
double-count guard is what keeps it correct. Until then treat the wiring as
untested end-to-end: only the unit test covers it.

Part 3 — Config and recipe

8. Pair the compressed RoPE base with its YaRN factor. Item 2 pinned
compress_rope_theta to 40000 while rotary_scaling_factor still inherited the
Pro value of 16. Those two parameterise the same inv_freq table, so aligning
only the base leaves the low-frequency end off by 4x. This records a known gap
rather than papering over it: the released inference/model.py sets
original_seq_len = 0 and only interpolates when that is positive, so the
released model effectively runs with YaRN frequency interpolation off. The
pair set here follows the upstream training recipe, not the released
checkpoint's effective configuration; resolving that needs the official
config.json, so the yaml carries the caveat inline.

9. Ship the keep-in-FP32 pinning disabled by default. Item 6f is correct in
isolation and aborts every GPU at 4 nodes: on 4N / PP4 / EP8 with
use_precision_aware_optimizer + store_param_remainders, every GPU died with
Memory access fault ... Reason: Unknown on the first training step. The same
build with the mechanism off runs 10/10 clean, and so does the pre-PR baseline.

Root cause: the grad buffer does key on (param_dtype, grad_dtype) and allocate
one buffer per combination — that part of 6f's claim is true — but the distributed
optimizer on top does not follow; distrib_optimizer.py carries five
assert len(gbuf_range_maps) == 1, "single dtype supported, for now." guards. So
6f's claim that "mixed parameter dtypes are fine downstream" was wrong, and
wrong in a way no test here could catch: single node with PP1 does not reproduce
it, and that is what every unit test and soak run used.

Turning it off costs almost nothing — store_param_remainders keeps FP32 master
params either way, and every consumer already promotes at the use site, so both
update and forward precision are unchanged. What the mark buys is stored
resolution matching the released checkpoint: checkpoint parity, not training. (A
related correction: the FP32 sink assertions are in the gluon / flydsl_v1
kernels, not triton_v2 as 6f said, and the callers' promotion already satisfies
them.) Enable with PRIMUS_V4_KEEP_FP32=1 at PP1.

10. Correct the yaml comment on the distillation loss. Comment-only. The yaml
said 0.0 was "correct when loading an already-trained indexer"; the papers keep
the KL on through the whole sparse stage. Corrected guidance is under item 5.

11. Build the indexer KL target without the dense per-head tensors. Profiling
the 4-node run showed the loss adding 750 ms of GPU kernel time per iteration at
gbs64, and the target branch — not the indexer being trained — was most of it:
pool gather +126 ms, dtype promote / mask / mul ~300 ms, fp32 einsum +45 ms.
The gather produced [B, S, K, head_dim] (2.1 GB per microbatch at Flash widths)
and was then promoted to fp32, so the promotion alone moved more bytes than the
GEMM it fed.

The whole target branch is detached, so nothing has to survive backward. That
makes three things safe: gather and GEMM at the model dtype, promoting only the
result (head_dim times less traffic, and closer to the distribution being
imitated since the main attention computes its logits at the model dtype too);
head-sum inside the loop so [B, H, S, K] never exists in full; and chunk the
query axis so the gather does not either. Measured on 4 nodes, gbs256,
coeff=1e-2: 10721.7 → 10434.1 ms/iter (−2.7%), so the loss now costs
+17.7% over coeff=0 instead of +20.9%; lm loss moves by 1.7e-6. Chunk size
via PRIMUS_V4_DISTILL_TARGET_CHUNK.

Not addressed: ~40% of the remaining overhead is nccl kernels getting slower at
an unchanged call count, i.e. the loss is serialised ahead of the MoE all-to-all
instead of overlapping with it.

12. Rebalance the 4-node PP layout. Et*10|t*11|t*11|t*11mL →
Et*10|t*12|t*12|t*9mL: the last stage carries the MTP layer and the loss on top
of its decoder layers, so giving it the same 11 as the middle stages made it the
bubble (total is still 43). PRIMUS_RECOMPUTE_LAYERS now defaults to 0 on 4
nodes (was hard-coded 3). Both knobs and the 8-node layout now honour an
incoming environment variable, so a sweep can override them without editing the
script. MTP=0 variants unchanged; no separate A/B, the E2E numbers below predate
it.

---

Known gap: the two new auxiliary losses train but are not visible in the log

Both log_seq_balance_loss (7a) and log_indexer_distill_loss (6d) start with an
if not layer_number: return sentinel, and on the V4 path it always fires.
DeepseekV4HybridLayer.__init__ deliberately bypasses
TransformerLayer.__init__, so the upstream self.mlp.set_layer_number(...) call
never runs; the layer computes its own 1-based layer_number but does not forward
it when building the attention and MLP submodules. So
DeepseekV4Attention.layer_number stays at its 0 default and
DeepseekV4MoE.layer_number — hence learned_router.layer_number — stays None.
DeepseekV4MoE.set_layer_number exists but has no production caller; only the
unit tests invoke it, which is why the suite is green.

• The balance loss reads as a hard zero rather than as missing.
  seq_load_balancing_loss is in track_names and force_initialize=True
  creates a zero-filled tracker entry and reduces it, so the field prints 0.0 —
  which looks like perfectly balanced routing.
• The distillation loss never reaches the tracker, and indexer_distill_loss
  is not in track_names either, so fixing only the sentinel would still bypass
  the cross-PP reduction. Both halves need fixing together.
• Gradients are unaffected — neither autoscaler consults layer_number, so
  both objectives train exactly as measured below. This is a reporting defect,
  not a math defect, and it is why the numbers below were read off the loss
  tensors rather than a log.

The fix is small (forward layer_number at the two build_module sites, add the
key to track_names) but touches a path shared with the MTP layer, so it is left
out of this PR rather than bolted on unmeasured. First follow-up.

An earlier revision blamed this on Primus overriding track_moe_metrics with a
version that updates total_loss_dict only inside if writer is not None.
That was wrong — the override exists in the tree but nothing imports it;
training runs Megatron's own, whose update sits outside the writer check.

---

Test plan

Unit tests (8x MI355X, ROCm). Full deepseek_v4 suite: 589 passed, 81
skipped, 0 failed, against a pre-PR baseline of 484 / 81 — the delta is new
cases and the skip count is unchanged, so no gate started skipping. Affected
existing tests: 150 passed, 2 skipped.

Nine new test files. test_v4_parity_suite.py is split into numerical parity
(rebuilds the CSA forward from the module's own parameters with plain matmuls) and
gradient topology (asserts which parameters each objective may and may not
reach) — scale and detach bugs are invisible to a forward-parity test and to a
loss-value assertion alike, and the only thing that catches them is asking where
the gradient went. The rest pin 6a's constant, the Hadamard's orthogonality in
float64, bit-exact FP32 pinning, the balance-loss normalisation, the double-count
guard, and the FP8 knob independence. Two existing inline references encoded the
bugs being fixed and were corrected as part of the fix; the compressed-position
test deliberately uses P > 1, since the two conventions coincide at P == 1 —
which is why the previous HCA test could never have caught it.

End-to-end, 1 node x 8 MI355X, seq=4096, GBS=8, turbo backends. Baseline 2
runs, PR 4 runs with byte-identical argument dumps (so they double as a spread
measurement):

Metric	Baseline	This PR	Delta
iter_time	320.0 ms	324.7 ms	+1.45%
TFLOP/s/GPU	837.6	825.6	−1.43%
peak memory	211.38 GB	209.80 GB	−0.75%

Read the spread before the delta: the four identical runs span 322.8–327.4 ms
(±0.7% around their own mean), so a single pair can land anywhere between +0.8%
and +2.3%. From 2000-iteration runs (314.5 / 312.2 / 314.6 ms, ±0.4%), steady
state is 314.5 ms/iter, 852.3 TFLOP/s/GPU. Items 1–5 measured +0.87% on their
own, so item 6 adds ~half a percent: 6g largely cancels 6e.

4 nodes / PP4 / EP8 — cost of the distillation loss, 10 iterations each:

Metric	coeff=0	coeff=1e-2 (pre-11)	coeff=1e-2 (post-11)
iter_time	8867.8 ms	10721.7 ms (+20.9%)	10434.1 ms (+17.7%)
grad buckets	24	28	unchanged
lm loss @ iter 10	9.976819	9.976658	within 1.7e-6

+17.7% is much more than the +8.7% measured on one node because that run had 3
CSA layers out of 8 while Flash has 21 out of 43 — the cost lands only on CSA
layers. The main loss barely moves, which is 6b working as intended.

Balance-loss magnitude, measured on real GPU shapes because a magnitude error
in a loss formula passes every shape and finiteness check:

Case	Measured	Expected	Ratio
Perfectly balanced routing	1.00000005e-03	alpha = 1e-3	1.000000
Fully collapsed onto K of E	4.26666737e-02	alpha·E/K	1.000000
Real router at Flash widths (bf16)	1.00025884e-03	O(alpha)	1.0003

A missing or extra 1/T would show up immediately as O(100) or O(0.01).

---