Make gradient clipping reproducible without changing configuration names

megatron/core/optimizer/clip_grads.py

@@ -50,12 +50,37 @@
 from ..tensor_parallel import param_is_not_tensor_parallel_duplicate
 from ..transformer.module import param_is_not_shared
 from ..utils import get_data_parallel_group_if_dtensor, to_local_if_dtensor
+from .reproducible_norm import ReproducibleL2Norm
+
+
+@torch.no_grad()
+def get_reproducible_grad_norm_bins(
+    grads_for_norm: List[torch.Tensor],
+    grad_stats_parallel_group: torch.distributed.ProcessGroup | None,
+) -> torch.Tensor:
+    """Reduce exact FP32-square bins using the existing gradient ownership groups."""
+    grads_for_norm = list(grads_for_norm)
+    data_parallel_group = None
+    for grad in grads_for_norm:
+        data_parallel_group = get_data_parallel_group_if_dtensor(grad, data_parallel_group)
+    grads_for_norm = [to_local_if_dtensor(grad) for grad in grads_for_norm]
+    accumulator = ReproducibleL2Norm(grads_for_norm[0].device if grads_for_norm else None)
+    bins = accumulator.zeros()
+    for grad in grads_for_norm:
+        if grad.layout != torch.strided:
+            raise TypeError("Reproducible clipping requires dense FP32 gradients")
+        bins = accumulator.accumulate(bins, grad)
+    if data_parallel_group:
+        torch.distributed.all_reduce(bins, group=data_parallel_group)
+    torch.distributed.all_reduce(bins, group=grad_stats_parallel_group)
+    return bins
 
 
 def get_grad_norm_fp32(
     grads_for_norm: Union[List[torch.Tensor], torch.Tensor],
     norm_type: Union[int, float] = 2,
     grad_stats_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+    use_accuracy_compatible: bool = False,
 ) -> float:
     """Calculate the p-norm of gradients in FP32 precision.
 
@@ -80,6 +105,12 @@ def get_grad_norm_fp32(
     if isinstance(grads_for_norm, torch.Tensor):
         grads_for_norm = [grads_for_norm]
 
+    if use_accuracy_compatible:
+        if float(norm_type) != 2.0:
+            raise ValueError("Reproducible clipping supports only the L2 norm")
+        bins = get_reproducible_grad_norm_bins(grads_for_norm, grad_stats_parallel_group)
+        return ReproducibleL2Norm(bins.device).finish(bins)[0]
+
     data_parallel_group = None
     for grad in grads_for_norm:
         data_parallel_group = get_data_parallel_group_if_dtensor(grad, data_parallel_group)
@@ -184,12 +215,14 @@ def clip_grad_by_total_norm_fp32(
     dummy_overflow_buf = torch.zeros(1, dtype=torch.int, device='cuda')
     if isinstance(clip_coeff, torch.Tensor):
         clip_coeff.clamp_max_(1.0)
-        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
-        )
+        if multi_tensor_scale_tensor_impl is not None:
+            multi_tensor_applier(
+                multi_tensor_scale_tensor_impl, dummy_overflow_buf, [grads, grads], clip_coeff
+            )
+        else:
+            multi_tensor_applier(
+                multi_tensor_scale_impl, dummy_overflow_buf, [grads, grads], clip_coeff.item()
+            )
     elif clip_coeff < 1.0:
         multi_tensor_applier(
             multi_tensor_scale_impl, dummy_overflow_buf, [grads, grads], clip_coeff

megatron/core/optimizer/optimizer.py

@@ -46,11 +46,17 @@
     optim_state_to_sharding_state,
 )
 from ..dist_checkpointing.utils import add_prefix_for_sharding
-from ..transformer.module import param_is_not_shared
+from ..transformer.module import _use_accuracy_compatible, param_is_not_shared
 from ..utils import log_single_rank
-from .clip_grads import clip_grad_by_total_norm_fp32, count_zeros_fp32, get_grad_norm_fp32
+from .clip_grads import (
+    clip_grad_by_total_norm_fp32,
+    count_zeros_fp32,
+    get_grad_norm_fp32,
+    get_reproducible_grad_norm_bins,
+)
 from .grad_scaler import MegatronGradScaler
 from .optimizer_config import OptimizerConfig
+from .reproducible_norm import ReproducibleL2Norm
 
 logger = getLogger(__name__)
 
@@ -296,7 +302,9 @@ def get_grad_norm(self):
         """Compute and return grad norm."""
         grads_for_norm = self.get_grads_for_grad_norm()
         total_norm = get_grad_norm_fp32(
-            grads_for_norm, grad_stats_parallel_group=self.get_grad_stats_parallel_group()
+            grads_for_norm,
+            grad_stats_parallel_group=self.get_grad_stats_parallel_group(),
+            use_accuracy_compatible=_use_accuracy_compatible() and self.config.clip_grad > 0,
         )
         return total_norm
 
@@ -308,7 +316,10 @@ def _compute_grad_norms_by_group(self) -> Dict[str, float]:
             if self.has_grad_norm_group(grad_norm_group):
                 grouped_grads = self.get_grads_for_grad_norm(grad_norm_group)
                 group_grad_norm = get_grad_norm_fp32(
-                    grouped_grads, grad_stats_parallel_group=self.get_grad_stats_parallel_group()
+                    grouped_grads,
+                    grad_stats_parallel_group=self.get_grad_stats_parallel_group(),
+                    use_accuracy_compatible=_use_accuracy_compatible()
+                    and self.config.clip_grad > 0,
                 )
                 self.grad_norms_by_group[grad_norm_group] = group_grad_norm
         return self.grad_norms_by_group
@@ -326,7 +337,9 @@ def clip_grad_norm(self, clip_grad: float) -> float:
         else:
             grads_for_norm = []
         grad_norm = get_grad_norm_fp32(
-            grads_for_norm, grad_stats_parallel_group=self.get_grad_stats_parallel_group()
+            grads_for_norm,
+            grad_stats_parallel_group=self.get_grad_stats_parallel_group(),
+            use_accuracy_compatible=_use_accuracy_compatible() and self.config.clip_grad > 0,
         )
 
         if clip_grad > 0.0 and params:
@@ -1566,8 +1579,21 @@ def get_grad_stats_parallel_group(self) -> torch.distributed.ProcessGroup:
         )
         return self.chained_optimizers[0].get_grad_stats_parallel_group()
 
+    @torch.no_grad()
+    def _get_reproducible_grad_norm(self, grad_norm_group=None):
+        bins = None
+        for optimizer in self.chained_optimizers:
+            part = get_reproducible_grad_norm_bins(
+                optimizer.get_grads_for_grad_norm(grad_norm_group),
+                optimizer.get_grad_stats_parallel_group(),
+            )
+            bins = part if bins is None else bins + part
+        return ReproducibleL2Norm(bins.device).finish(bins)[0]
+
     @torch.no_grad()
     def get_grad_norm(self):
+        if _use_accuracy_compatible() and self.config.clip_grad > 0:
+            return self._get_reproducible_grad_norm()
         if len(self.chained_optimizers) == 1:
             return self.chained_optimizers[0].get_grad_norm()
         if self.grads_states_parallel_group_is_shared():
@@ -1631,6 +1657,8 @@ def has_grad_norm_group(self, grad_norm_group: str) -> bool:
     def _get_grad_norm_for_group(self, grad_norm_group: str):
         """Compute gradient norm for a named parameter group."""
         _validate_grad_norm_group(grad_norm_group)
+        if _use_accuracy_compatible() and self.config.clip_grad > 0:
+            return self._get_reproducible_grad_norm(grad_norm_group)
         if self.grads_states_parallel_group_is_shared():
             grouped_grads = []
             for optimizer in self.chained_optimizers:

megatron/core/optimizer/reproducible_norm.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Layout- and partition-independent L2 norm of FP32 gradients.
+
+Squares are computed in FP32 on the gradient device. Their mantissas are
+accumulated into base-65536 integer bins, which may be SUM-reduced across
+owners before rounding the total once to FP32 and applying the native sqrt.
+No gradient values or floating-point norm computation leave the device.
+
+The 20 limbs cover FP32 squares and the supported 2**40 global element count.
+Each uncarried limb is bounded by 2**40 * 65535 < 2**56, so arbitrary parameter
+and rank reduction orders cannot overflow int64. The last three bins count
+infinities, NaNs, and elements. This deliberately costs more than a fused norm
+and is intended only for explicitly enabled accuracy compatibility.
+"""
+
+import torch
+
+
+class ReproducibleL2Norm:
+    """Accumulate on a single device; reduce bins before calling ``finish``."""
+
+    def __init__(self, device: torch.device | None = None) -> None:
+        self.device = (
+            device if device is not None else torch.device("cuda", torch.cuda.current_device())
+        )
+
+    def cast(self, value: torch.Tensor, dtype: str) -> torch.Tensor:
+        """Convert to the named Torch dtype on the tensor's current device."""
+        return value.to(getattr(torch, dtype))
+
+    def zeros(self, count: int = 23) -> torch.Tensor:
+        """Allocate zeroed integer bins on the accumulator device."""
+        return torch.zeros(count, dtype=torch.int64, device=self.device)
+
+    def tensor(
+        self, value: list[int | float] | int | float, dtype: str = "float32"
+    ) -> torch.Tensor:
+        """Create a typed constant on the accumulator device."""
+        return torch.tensor(value, dtype=getattr(torch, dtype), device=self.device)
+
+    def view(self, value: torch.Tensor, dtype: str) -> torch.Tensor:
+        """Interpret bits as the named Torch dtype without value conversion."""
+        return value.view(getattr(torch, dtype))
+
+    def add(self, bins: torch.Tensor, indices: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+        """Add integer contributions to the selected bins in place."""
+        return bins.scatter_add_(0, indices, values)
+
+    def accumulate(
+        self, bins: torch.Tensor, gradient: torch.Tensor, chunk_size: int = 1048576
+    ) -> torch.Tensor:
+        """Add device-computed FP32 gradient squares to exact integer bins."""
+        if gradient.layout != torch.strided or gradient.dtype != torch.float32:
+            raise TypeError("Reproducible clipping requires FP32 gradients")
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be positive")
+        flat = gradient.reshape([-1])
+        if flat.shape[0] > 2**40:
+            raise OverflowError("Reproducible norm supports at most 2**40 global elements")
+        bins = self.add(bins, self.tensor([22], "int64"), self.tensor([flat.shape[0]], "int64"))
+        for start in range(0, flat.shape[0], chunk_size):
+            value = flat[start : start + chunk_size]
+            bits = self.cast(self.view(value * value, "int32"), "int64")
+            exponent = (bits >> 23) & self.tensor(255, "int64")
+            fraction = bits & self.tensor(0x7FFFFF, "int64")
+            finite = exponent != 255
+            mantissa = fraction | torch.where(
+                exponent > 0, torch.full_like(exponent, 0x800000), torch.zeros_like(exponent)
+            )
+            shift = torch.maximum(exponent - 1, torch.zeros_like(exponent))
+            limb = shift // 16
+            shifted = (mantissa << (shift % 16)) * self.cast(finite, "int64")
+            for offset in range(3):
+                bins = self.add(
+                    bins, limb + offset, (shifted >> (16 * offset)) & self.tensor(65535, "int64")
+                )
+            flags = torch.stack(
+                [
+                    self.cast((exponent == 255) & (fraction == 0), "int64").sum(),
+                    self.cast((exponent == 255) & (fraction != 0), "int64").sum(),
+                ]
+            )
+            bins = self.add(bins, self.tensor([20, 21], "int64"), flags)
+        return bins
+
+    def finish(self, bins: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return norm and square sum after one FP32 rounding of reduced bins."""
+        if int(bins[22].item()) > 2**40:
+            raise OverflowError("Reproducible norm supports at most 2**40 global elements")
+        carry = self.zeros(1)[0]
+        digits = []
+        for i in range(20):
+            value = bins[i] + carry
+            digits.append(value & self.tensor(65535, "int64"))
+            carry = value >> 16
+        digits = torch.stack(digits)
+        index = self.tensor(list(range(20)), "int64")
+        top = torch.where(digits != 0, index, torch.full_like(index, -1)).max()
+
+        def get(i):
+            return torch.where(index == i, digits, torch.zeros_like(digits)).sum()
+
+        word = get(top)
+        leading = self.zeros(1)[0]
+        for width in [8, 4, 2, 1]:
+            take = word >= (1 << width)
+            word = torch.where(take, word >> width, word)
+            leading = leading + self.cast(take, "int64") * width
+        highest = top * 16 + leading
+        cut = torch.maximum(highest - 23, self.zeros(1)[0])
+        limb, shift = cut // 16, cut % 16
+        significand = (
+            (get(limb) >> shift) | (get(limb + 1) << (16 - shift)) | (get(limb + 2) << (32 - shift))
+        ) & self.tensor(0xFFFFFF, "int64")
+        round_position = torch.maximum(cut - 1, self.zeros(1)[0])
+        round_limb, round_shift = round_position // 16, round_position % 16
+        round_word = get(round_limb)
+        round_bit = ((round_word >> round_shift) & self.tensor(1, "int64")) * self.cast(
+            cut > 0, "int64"
+        )
+        sticky = torch.where(index < round_limb, digits, torch.zeros_like(digits)).sum() != 0
+        sticky = sticky | ((round_word & ((self.tensor(1, "int64") << round_shift) - 1)) != 0)
+        significand = significand + round_bit * self.cast(
+            sticky | ((significand & self.tensor(1, "int64")) != 0), "int64"
+        )
+        exponent = highest - 22 + self.cast(significand == 0x1000000, "int64")
+        raw = (exponent << 23) | (significand & self.tensor(0x7FFFFF, "int64"))
+        raw = torch.where(exponent >= 255, torch.full_like(raw, 0x7F800000), raw)
+        raw = torch.where(highest < 23, get(0) | (get(1) << 16), raw)
+        raw = torch.where(top < 0, torch.zeros_like(raw), raw)
+        raw = torch.where(bins[20] != 0, torch.full_like(raw, 0x7F800000), raw)
+        raw = torch.where(bins[21] != 0, torch.full_like(raw, 0x7FC00000), raw)
+        square_sum = self.view(self.cast(raw.reshape([1]), "int32"), "float32")
+        return torch.sqrt(square_sum), square_sum