Training crash with TE CUDA graphs
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
# Training crash with TE CUDA graphs (`--cuda-graph-scope attn`) after ~89 iterations on Qwen3-235B
## Environment
| Component | Version |
|---|---|
| Megatron-Core | `0.19.0+ddaa315fe` |
| PyTorch | `2.8.0a0+5228986c39.nv25.6` |
| Transformer Engine | `2.16.0` |
## Configuration
- Qwen3-235B-A22B, TP=2, PP=8, VPP=4, EP=8, ETP=1, DP=16
- BF16
- MoE alltoall dispatcher
- `--overlap-moe-expert-parallel-comm`
- CUDA graph: `--cuda-graph-impl transformer_engine --cuda-graph-scope attn`
## What Happens
Training runs normally for ~80 iterations, then grad norm starts oscillating wildly between normal values and tens of millions, before hitting NaN at iteration 91 and crashing:
| Iteration | Grad Norm |
|---|---|
| 80 | 91 |
| 81 | 463,582 |
| 82 | 50 |
| 83 | 21,210,958 |
| 87 | 127 |
| 88 | 27,622,466 |
| 89 | 11,356,515 |
| 91 | NaN → crash |
The crash message is:
```
ERROR:megatron.core.rerun_state_machine: found NaN in local grad norm for bucket #0
in backward pass before data-parallel communication collective
```
## Workaround
Removing `--cuda-graph-impl transformer_engine --cuda-graph-scope attn` makes the crash disappear.
## Question
Has anyone seen this before, or knows what might be causing it?
Contributor guide
Assessment
This issue has not been assessed yet.