[Bug][DSv4-Hybrid] apply_rope_fusion=True causes NaN at iter 2 in BF16 mock pretrain (DSv4HybridSelfAttention)
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Description
When training with `experimental_attention_variant="dsv4_hybrid"` and `apply_rope_fusion=True`, the training reliably produces NaN in the **forward loss at iteration 2** under BF16 random-init mock pretrain. Setting `apply_rope_fusion=False` fixes the issue and training proceeds normally.
## Reproduction
**Environment:**
- Megatron-LM commit: `cf081d5df` (current HEAD)
- PyTorch 2.12.0a0 (NV nightly), CUDA 13.2, TransformerEngine
- NCCL 2.29.7
- 8x H800 80GB, EP=8, PP=1, TP=1, launched via `torchrun`
**Minimal config:**
```python
cfg.model.experimental_attention_variant = "dsv4_hybrid"
cfg.model.num_layers = 8
cfg.model.num_moe_experts = 16
cfg.model.csa_compress_ratios = [0, 0, 0, 0, 0, 0, 0, 0]
cfg.model.seq_length = 4096
cfg.model.params_dtype = torch.bfloat16
cfg.model.apply_rope_fusion = True # ← causes NaN
# cfg.model.apply_rope_fusion = False # ← works fine
cfg.model.rope_type = "yarn"
cfg.train.global_batch_size = 8
cfg.train.micro_batch_size = 1
```
Also reproduced with `csa_compress_ratios=[0,0,4,128,4,128,4,0]` and `seq_length=1024`.
## What works vs. what does not
| Config | Result |
|--------|--------|
| `apply_rope_fusion=False` + BF16 | ✅ Trains cleanly (loss 5.8→0.03 over 70 iters) |
| `apply_rope_fusion=True` + BF16, NCCL | ❌ NaN at iter 2 forward |
## Log (standard NCCL, torchrun, apply_rope_fusion=True)
```
Setting rerun_state_machine.current_iteration to 0...
Starting training loop at iteration 0
[rank3]: raise RuntimeError(full_message)
[rank3]: RuntimeError: Rank 3, node aiplatform-wlf3-ge97-39.idchb2az3.hb2.kwaidc.com, device 3, iteration 2: Unexpected result nan (message='found NaN in local forward loss calculation')
[rank2]: raise RuntimeError(full_message)
[rank2]: RuntimeError: Rank 2, node aiplatform-wlf3-ge97-39.idchb2az3.hb2.kwaidc.com, device 2, iteration 2: Unexpected result nan (message='found NaN in local forward loss calculation')
[rank0]: raise RuntimeError(full_message)
[rank0]: RuntimeError: Rank 0, node aiplatform-wlf3-ge97-39.idchb2az3.hb2.kwaidc.com, device 0, iteration 2: Unexpected result nan (message='found NaN in local forward loss calculation')
```
Iteration 1 completes normally (no NaN, grad norm ~18). The crash always occurs at **iteration 2 forward**, consistent with incorrect gradients from iter 1 corrupting weights after optimizer step.
## Investigation
Isolated unit tests for `fused_mla_rope_inplace` (linear → RMSNorm → unsqueeze → fused rope → backward) confirm that **gradient computation is numerically correct in isolation** — both forward values and `weight.grad` match the unfused path within BF16 tolerance. So the issue is specific to the full `DSv4HybridSelfAttention` forward/backward graph.
The fused path differs from the unfused path in:
1. **Forward RoPE on Q/KV**: Triton kernel modifies tensors in-place via direct GPU memory writes (bypasses PyTorch version counter)
2. **Inverse RoPE on `core_attn_out`**: Another in-place modification after attention
3. **`key = value = kv`**: Single-head MQA where K and V are the same tensor — combined gradient accumulation may interact with the in-place modification
Note: The unit test `test_mla_yarn_rope_apply.py` uses `.detach()` for the fused path input, which avoids testing the kernel inside a full autograd graph with prior operations. An end-to-end training test with `apply_rope_fusion=True` on `DSv4HybridSelfAttention` appears to be missing.
## Workaround
```python
cfg.model.apply_rope_fusion = False # must be set for BF16 DSv4-Hybrid training
```
Contributor guide
Assessment
This issue has not been assessed yet.