Qwen3-Next-80B-A3B with MTP NaN issue on Blackwell
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
When running dev branch to train Qwen3-Next-80B-A3B with MTP, sometimes meet grad norm NaN or loss NaN on Blackwell.
We find that it is caused by [this TE bug](https://github.com/NVIDIA/TransformerEngine/issues/2489). Specifically, we find that when using Nvjet GEMM and cuDNN RMSNorm on GB200 / B200, it seems that overlapping would happen with previous kernel. It has data dependencies so the input of RMSNorm should relay on the output of the last kernel. It sometimes work well but in some cases it will cause NaN issue or numerical issue.
**Expected behavior**
```
ERROR:megatron.core.rerun_state_machine:Unexpected result nan on rank 13 at iteration #2 invokation #1 (message='found NaN in local grad norm for bucket #0 in backward pass before data-parallel communication collective')
```
Environment details
- PyTorch version 25.10
- Transformer Engine version 2.10.0.dev0
- CUDA version 13.0.r13.0
- CUDNN version 9.16
Contributor guide
Assessment
This issue has not been assessed yet.