Fix mixed-dtype gradient norms and clipping

megatron/core/optimizer/clip_grads.py

@@ -56,6 +56,14 @@
 from ..utils import get_data_parallel_group_if_dtensor, to_local_if_dtensor
 
 
+def _group_grads_by_dtype(grads: List[torch.Tensor]) -> List[List[torch.Tensor]]:
+    # Fused multi-tensor kernels dispatch one dtype for each tensor list.
+    groups: dict[torch.dtype, List[torch.Tensor]] = {}
+    for grad in grads:
+        groups.setdefault(grad.dtype, []).append(grad)
+    return list(groups.values())
+
+
 def get_grad_norm_fp32(
     grads_for_norm: Union[List[torch.Tensor], torch.Tensor],
     norm_type: Union[int, float] = 2,
@@ -117,12 +125,15 @@ def get_grad_norm_fp32(
             # Use apex's multi-tensor applier for efficiency reasons.
             # Multi-tensor applier takes a function and a list of list
             # and performs the operation on that list all in one kernel.
-            grad_norm, _ = multi_tensor_applier(
-                l2_norm_impl, dummy_overflow_buf, [grads_for_norm], False  # no per-parameter norm
-            )
-            # Since we will be summing across data parallel groups,
-            # we need the pow(norm-type).
-            total_norm = grad_norm**norm_type
+            for index, dtype_grads in enumerate(_group_grads_by_dtype(grads_for_norm)):
+                grad_norm, _ = multi_tensor_applier(
+                    l2_norm_impl, dummy_overflow_buf, [dtype_grads], False
+                )
+                # Combine local squared norms before the existing collectives.
+                if index == 0:
+                    total_norm = grad_norm**norm_type
+                else:
+                    total_norm += grad_norm**norm_type
         else:
             for grad in grads_for_norm:
                 grad_norm = torch.norm(grad, norm_type)
@@ -187,13 +198,18 @@ def clip_grad_by_total_norm_fp32(
         assert (
             multi_tensor_scale_tensor_impl is not None
         ), "clip_coeff is tensor type. But multi_tensor_scale_tensor not available."
-        multi_tensor_applier(
-            multi_tensor_scale_tensor_impl, dummy_overflow_buf, [grads, grads], clip_coeff
-        )
+        for dtype_grads in _group_grads_by_dtype(grads):
+            multi_tensor_applier(
+                multi_tensor_scale_tensor_impl,
+                dummy_overflow_buf,
+                [dtype_grads, dtype_grads],
+                clip_coeff,
+            )
     elif clip_coeff < 1.0:
-        multi_tensor_applier(
-            multi_tensor_scale_impl, dummy_overflow_buf, [grads, grads], clip_coeff
-        )
+        for dtype_grads in _group_grads_by_dtype(grads):
+            multi_tensor_applier(
+                multi_tensor_scale_impl, dummy_overflow_buf, [dtype_grads, dtype_grads], clip_coeff
+            )
 
 
 def _gtp_pad_zero_count(param: torch.Tensor, grad: torch.Tensor) -> int:

tests/unit_tests/determinism/kernels/manifest.py

@@ -419,7 +419,8 @@ class KernelEntry:
         ),
         tests=(K + "test_optimizer_kernels.py",),
         kind="external-lib",
-        notes="multi_tensor l2norm / scale (TE, apex or local fallback) and fused Adam; "
+        notes="multi_tensor l2norm / scale (TE, apex or local fallback), including mixed-dtype "
+        "caller bucketing and padded-storage replay, and fused Adam; "
         "optimizer.py (gradient unscaling) and training/utils/common_utils.py (param / grad norm "
         "logging) launch the same multi_tensor kernels through multi_tensor_applier.",
     ),

tests/unit_tests/determinism/kernels/test_optimizer_kernels.py

@@ -73,6 +73,36 @@ def test_clip_grad_by_total_norm_replays(self):
             for j, (a, b) in enumerate(zip(ref, got)):
                 assert bytes_equal(a, b), f"clipped grad {j} differs on replay {i}"
 
+    @pytest.mark.parametrize(
+        'dtypes', [(torch.bfloat16, torch.float32), (torch.float32, torch.bfloat16)]
+    )
+    def test_mixed_dtype_norm_and_clip_replays(self, dtypes):
+        seeded()
+        numel = 1_048_578  # Multiple chunks per dtype, including a partial final chunk.
+        bases = [torch.randn(4 * numel, device='cuda', dtype=dtype) for dtype in dtypes]
+
+        def norm_and_clip(*storage):
+            grads = [base[:numel] for base in storage]
+            params = [torch.zeros_like(grad) for grad in grads]
+            for param, grad in zip(params, grads):
+                param.grad = grad
+            norm = get_grad_norm_fp32(
+                grads, grad_stats_parallel_group=torch.distributed.group.WORLD
+            )
+            clip_grad_by_total_norm_fp32(params, 1.0, norm)
+            return torch.as_tensor(norm, device='cuda'), *storage
+
+        # The harness clones whole bases, preserving padding for regressed mixed dispatch.
+        # Compare the norm and all backing bytes under side-stream scheduling pressure.
+        assert_replays_bit_exact(
+            norm_and_clip,
+            tuple(bases),
+            replays=4,
+            backward=False,
+            contention=True,
+            what='mixed_dtype_norm_and_clip',
+        )
+
 
 def test_fused_adam_step_replays():
     """Same params, grads and optimizer state -> identical updated params and moments."""

tests/unit_tests/optimizer/test_clip_grads.py

@@ -1,8 +1,10 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import pytest
 import torch
 
 from megatron.core import parallel_state
+from megatron.core.optimizer import clip_grads
 from megatron.core.optimizer.clip_grads import count_zeros_fp32
 from megatron.core.optimizer.optimizer_config import OptimizerConfig
 from tests.unit_tests.test_utilities import Utils
@@ -59,3 +61,117 @@ def test_no_gtp_pad_zeros_attribute_counts_all_zeros(self):
         )
 
         assert num_zeros == grad.numel() - 1
+
+
+_DTYPES = [
+    (torch.bfloat16, torch.float32),
+    (torch.float32, torch.bfloat16),
+    (torch.bfloat16, torch.bfloat16),
+    (torch.float32, torch.float32),
+]
+
+
+def _padded_grads(dtypes, numel, value=1.0 / 64):
+    # Keep a regressed fused kernel's wider dtype interpretation inside its allocation.
+    # Padding also exposes writes beyond the logical gradient view.
+    bases = [torch.ones(4 * numel, device='cuda', dtype=dtype) for dtype in dtypes]
+    grads = [base[:numel].fill_(value) for base in bases]
+    return bases, grads
+
+
+@pytest.mark.skipif(
+    clip_grads.l2_norm_impl.__name__ == 'local_multi_tensor_l2_norm',
+    reason='Requires fused multi-tensor kernels; the Python fallback already accepts mixed lists',
+)
+class TestMixedDtypeGradNormAndClip:
+    def setup_method(self):
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, pipeline_model_parallel_size=1
+        )
+
+    def teardown_method(self):
+        Utils.destroy_model_parallel()
+
+    @pytest.mark.parametrize('dtypes', _DTYPES)
+    @pytest.mark.parametrize('numel', [128, 130])
+    @pytest.mark.parametrize('layout', ['both', 'rank0_empty', 'disjoint', 'all_empty'])
+    def test_grad_norm(self, dtypes, numel, layout):
+        bases, grads = _padded_grads(dtypes, numel)
+        rank = torch.distributed.get_rank()
+        if layout == 'all_empty' or (layout == 'rank0_empty' and rank == 0):
+            grads = []
+        elif layout == 'disjoint':
+            grads = [grads[rank % 2]]
+        squared_norm = torch.zeros(1, device='cuda', dtype=torch.float64)
+        for grad in grads:
+            squared_norm += grad.double().square().sum()
+        torch.distributed.all_reduce(squared_norm)
+
+        actual = clip_grads.get_grad_norm_fp32(
+            grads, grad_stats_parallel_group=torch.distributed.group.WORLD
+        )
+
+        assert float(actual) == pytest.approx(float(squared_norm.sqrt()), rel=2e-6, abs=1e-8)
+        for base in bases:
+            assert torch.equal(base[numel:], torch.ones_like(base[numel:]))
+
+    @pytest.mark.parametrize('dtypes', _DTYPES[:2])
+    def test_zero_grad_norm_ignores_padding(self, dtypes):
+        bases, grads = _padded_grads(dtypes, 128, value=0.0)
+        actual = clip_grads.get_grad_norm_fp32(
+            grads, grad_stats_parallel_group=torch.distributed.group.WORLD
+        )
+        assert float(actual) == 0.0
+        assert all(torch.all(base[128:] == 1) for base in bases)
+
+    @pytest.mark.parametrize('dtypes', _DTYPES)
+    @pytest.mark.parametrize('numel', [127, 130])
+    @pytest.mark.parametrize('use_decoupled_grad', [False, True])
+    @pytest.mark.parametrize('tensor_coefficient', [False, True])
+    @pytest.mark.parametrize('clip', [False, True])
+    def test_clip_gradients(self, dtypes, numel, use_decoupled_grad, tensor_coefficient, clip):
+        if tensor_coefficient and clip_grads.multi_tensor_scale_tensor_impl is None:
+            pytest.skip('Backend has no tensor-coefficient scaling API')
+        bases, grads = _padded_grads(dtypes, numel)
+        params = []
+        for grad in grads:
+            # The dtype of a decoupled gradient need not match its parameter's dtype.
+            dtype = torch.bfloat16 if use_decoupled_grad else grad.dtype
+            param = torch.nn.Parameter(torch.zeros_like(grad, dtype=dtype))
+            if use_decoupled_grad:
+                param.decoupled_grad = grad
+            else:
+                param.grad = grad
+            params.append(param)
+        params.append(torch.nn.Parameter(torch.zeros(1, device='cuda')))  # No gradient.
+        total_norm = torch.tensor([2.0], device='cuda') if tensor_coefficient else 2.0
+        max_norm = 1.0000005 if clip else 4.0
+        coefficient = min(float(max_norm / (total_norm + 1e-6)), 1.0)
+        expected = [base.clone() for base in bases]
+        for base in expected:
+            base[:numel].mul_(coefficient)
+
+        clip_grads.clip_grad_by_total_norm_fp32(
+            params, max_norm, total_norm, use_decoupled_grad=use_decoupled_grad
+        )
+
+        for actual, reference in zip(bases, expected):
+            torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+
+    @pytest.mark.parametrize('use_decoupled_grad', [False, True])
+    @pytest.mark.parametrize('tensor_coefficient', [False, True])
+    def test_empty_gradients_do_not_launch_kernel(
+        self, monkeypatch, use_decoupled_grad, tensor_coefficient
+    ):
+        if tensor_coefficient and clip_grads.multi_tensor_scale_tensor_impl is None:
+            pytest.skip('Backend has no tensor-coefficient scaling API')
+
+        def unexpected_launch(*args, **kwargs):
+            pytest.fail('An empty gradient list must not reach a fused kernel')
+
+        # Fail before dispatch: some native backends dereference the first list entry.
+        monkeypatch.setattr(clip_grads, 'multi_tensor_applier', unexpected_launch)
+        total_norm = torch.tensor([2.0], device='cuda') if tensor_coefficient else 2.0
+        clip_grads.clip_grad_by_total_norm_fp32(
+            [], 1.0, total_norm, use_decoupled_grad=use_decoupled_grad
+        )