NVIDIA / NVIDIA/Megatron-LM

Qwen3-Next-80B-A3B with MTP NaN issue on Blackwell

Open
#2,597 0 comments 0 reactions 1 assignee Claimed by @yuzhongw-nvidia View on GitHub
bug module: moe
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.

Image

**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

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.