Fused TE RoPE produces NaN on MLA k_pe in THD/packed-sequence mode (Moonlight, GLM-4.7-Flash)

[akoumpa]

Summary

For MLA models trained with backend.attn: te + packed sequences (THD), the fused Transformer-Engine RoPE path produces non-finite values on the single-head k_pe at the very first forward. The resulting NaN loss/grad at step 0 corrupts all weights through the optimizer step, and at step 1 the all-NaN router collapses MoE routing onto one expert-parallel rank, which surfaces as a CUDA OOM in the token-permute. The OOM is a symptom; the root cause is the RoPE NaN.

This was originally reported as an OOM (moonlight_16b_te_packed_sequence), but it is not a memory/capacity problem and not a checkpoint-loading problem (weights load 100% clean).

Affected configs (confirmed)

Recipe	Model	Arch
examples/llm_finetune/moonlight/moonlight_16b_te_packed_sequence.yaml	moonshotai/Moonlight-16B-A3B	DeepSeek-V3 MLA
examples/llm_finetune/glm/glm_4.7_flash_te_packed_sequence.yaml	zai-org/GLM-4.7-Flash	GLM-4.x MLA

Both reproduce loss nan | grad_norm nan at step 0 → step-1 MoE-dispatch OOM. Any other MLA + TE + packed-sequence MoE recipe is expected to hit this too (e.g. glm_moe_dsa, DeepSeek-V3.2).

Root cause

rope_fusion defaults to True when TE+CUDA are present (nemo_automodel/components/models/common/utils.py:199), and these recipes do not override it. The MLA forward then calls:

deepseek_v3/layers.py: apply_rotary_emb_qk(q_pe, k_pe, ...)
  -> rope_utils.apply_rotary_emb_qk(rope_fusion=True)
     -> transformer_engine.pytorch.attention.rope.apply_rotary_pos_emb(tensor_format="thd", interleaved=True, fused=True)

In THD/packed mode this corrupts only the single-head k_pe (shape (T, 1, 64)); q_pe (multi-head) under the identical freqs is fine.

Instrumented repro on the exact failing commit (d94e9110 = r0.5.0 HEAD), 8×H100, ep_size=8:

ROPE call#1 IN   q_pe[nan=0 finmax=1.0e1]  k_pe[nan=0 finmax=1.65e1]  freqs[finmax=4.096e3]
ROPE call#1 OUT  q_pe[nan=0 finmax=1.0e1]  k_pe[nan=30 finmax=3.390e38]   <-- bf16 max + NaN

• Loaded-weight scan: bad_tensors=0 on all ranks (not a loading bug).
• First non-finite in the whole forward: layers.0.self_attn...fused_attention, bad_input=True.
• The THD fused positions are built as a contiguous arange(num_tokens) (rope_utils.freqs_cis_from_position_ids, for_fused_rope=True), so freqs max ≈ the pack length rather than per-sequence positions.

Downstream chain (why it looks like an OOM)

1. k_pe NaN → k NaN → attention NaN → loss nan / grad_norm nan at step 0.
2. NaN grad → Adam (bf16) corrupts every weight.
3. Step 1: embedding/router all-NaN → torch.topk of all-NaN scores returns experts 0..k-1, all on EP-rank 0 → entire global batch routed to one rank → TE permute buffer OOM.

Workaround (this PR)

Set model.backend.rope_fusion: false in the affected recipes. The non-fused path does the rotation as a unit-magnitude complex multiply in fp32 → bounded. Validated on cw-dfw (8×H100, commit d94e9110):

Recipe	rope_fusion=True (current)	rope_fusion=false (fix)
Moonlight-16B	loss nan → OOM	loss 2.57 → 2.47 → 2.16, val 2.07, no NaN, balanced routing, no OOM
GLM-4.7-Flash	loss nan → OOM	loss 3.03 → 2.84 → 2.67, val 2.50, no NaN, balanced routing, no OOM

Proper fix (tracked here, beyond the workaround PR)

Make the fused TE RoPE path numerically correct for THD so fusion can stay on:

• Fix handling of the single-head k_pe (T,1,D) tensor in the fused THD kernel (the layout that overflows), and/or
• Build per-sequence positions from cu_seqlens for for_fused_rope=True THD instead of a global arange(num_tokens).

Possibly related (not confirmed)

qwen3_moe_30b_te_packed_sequence_lora fails with finite loss but grad_norm nan (backward NaN) → same step-1 collapse/OOM. It is packed-specific (non-packed Qwen3-MoE variants have finite grad_norm) but Qwen3-MoE is not MLA, so it is a separate failure (possibly the fused-RoPE backward on THD); needs its own investigation.

Repro pointers

CI pipeline 55312948 (main-mirror): jobs 344336851 (moonlight), 344336811 (glm_4.7_flash).