NVIDIA / NVIDIA/Megatron-LM

[Bug][DSv4-Hybrid] apply_rope_fusion=True causes NaN at iter 2 in BF16 mock pretrain (DSv4HybridSelfAttention)

Open
#5,317 9 comments 0 reactions 1 assignee Claimed by @guihong-nv View on GitHub
community-request waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.