[BUG] get_grad_norm_fp32 / clip_grad_by_total_norm_fp32 pass a mixed bf16/fp32 gradient list to one TE multi_tensor launch -> illegal memory access (precision-aware optimizer + natively fp32 params)
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Steps/Code to reproduce bug**
Train any model that has natively fp32 parameters (e.g. DeepSeek-V4 hyper-connection alpha/bias, `ape`, `attn_sink`) with `use_precision_aware_optimizer=True`, bf16 gradient buffers (`grad_reduce_in_fp32=False`, `main_grads_dtype=bf16`) and the default `clip_grad`. The first `get_grad_norm_fp32` call after backward crashes. Details below.
**Repo:** NVIDIA/Megatron-LM **Component:** `megatron/core/optimizer/clip_grads.py`
**Trigger:** `use_precision_aware_optimizer=True` + bf16 gradient buffers (`grad_reduce_in_fp32=False`) + a model with **natively fp32 parameters**. DeepSeek-V4 has such parameters (hyper-connection alpha/bias, `ape`, `attn_sink`, kept in fp32 via `mark_keep_in_fp32`). Reproduced on Megatron-LM `dev` (`43e45e13`), and `main` still builds a single mixed-dtype list.
**Describe the bug**
Tag: @NVIDIA/mcore-oncall
**Symptom**
```text
NCCL WARN Cuda failure 'an illegal memory access was encountered'
torch.distributed.DistBackendError: NCCL error in: .../NCCLUtils.cpp:93, unhandled cuda error
```
With `CUDA_LAUNCH_BLOCKING=1` the real site appears:
```text
RuntimeError: .../transformer_engine/common/multi_tensor/multi_tensor_apply.cuh:92 in function multi_tensor_apply: CUDA Error: an illegal memory access was encountered
File ".../transformer_engine/pytorch/optimizers/multi_tensor_apply.py", line 21, in __call__
File ".../megatron/core/optimizer/clip_grads.py", line 114, in get_grad_norm_fp32
```
**Root cause**
With the precision-aware optimizer, bf16 parameters expose a bf16 `decoupled_grad` (grad buffer dtype = param dtype), while natively fp32 parameters have fp32 grads. `get_grad_norm_fp32` collects all of them into one list and calls `multi_tensor_applier(l2_norm_impl, dummy_overflow_buf, [grads_for_norm], False)`. The TE kernel dispatches on the dtype of the **first** tensor and reads every tensor with that element size, so bf16 tensors read as fp32 run past their allocation (the reverse pairing would silently read garbage). `clip_grad_by_total_norm_fp32` has the same pattern with `multi_tensor_scale`. The existing `assert grad.dtype in (torch.float32, torch.bfloat16)` allows the mix but nothing groups it.
**Expected behavior / fix (verified, semantics unchanged)**
Group by dtype, run the kernel once per group, combine:
```python
groups = defaultdict(list)
for g in grads_for_norm:
groups[g.dtype].append(g)
total_sq = torch.zeros(1, dtype=torch.float32, device=device)
for grads in groups.values():
norm, _ = multi_tensor_applier(l2_norm_impl, dummy_overflow_buf, [grads], False)
total_sq += norm.float() ** 2
# then the same single all-reduce as before
```
and analogously for `clip_grad_by_total_norm_fp32` (`multi_tensor_scale` per group). The collective count and the result are identical to the current code for homogeneous lists. With this patch our DeepSeek-V4-Pro SFT (TP1/PP32/EP8, bf16 grads, precision-aware CPU offload) trains normally on 32 and 64 nodes; every rank reports e.g. `groups={'torch.float32': 7, 'torch.bfloat16': 295}`.
Alternative we tried and rejected: `grad_reduce_in_fp32=True` makes every gradient fp32 and avoids the mix, but doubles the gradient buffer (a 12.7 B-element expert bucket → +23.6 GiB) and OOMs on H100 80 GB at PP32.
Not the same as TransformerEngine #2918 (int32 overflow for tensors with numel > INT_MAX); none of our tensors are that large.
Environment: nemo:26.08, Megatron-LM dev `43e45e13`, TE 2.17.1+4329ff84, H100 80 GB.
Contributor guide
Assessment
This issue has not been assessed yet.