NVIDIA / NVIDIA/Megatron-LM

Training crash with TE CUDA graphs

Open
#5,966 7 comments 0 reactions 1 assignee Claimed by @buptzyb View on GitHub
bug community-request waiting-on-customer
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

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.