[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)

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

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:

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:

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.