Use one accuracy-compatible mode for model and optimizer numerics

megatron/core/distributed/finalize_model_grads.py

@@ -468,7 +468,7 @@ def finalize_model_grads(
     from ..transformer.module import _use_accuracy_compatible
 
     loss_normalized_in_graph = (
-        _use_accuracy_compatible() and not config.dsa_accuracy_compatible and num_tokens is not None
+        _use_accuracy_compatible() and not config.uses_dsa_reference and num_tokens is not None
     )
     if loss_normalized_in_graph:
         num_tokens = None
@@ -575,7 +575,7 @@ def finalize_model_grads(
             for model_chunk in model:
                 model_chunk.scale_gradients(1.0 / dp_size)
 
-    if loss_normalized_in_graph or config.dsa_accuracy_compatible:
+    if loss_normalized_in_graph or config.uses_dsa_reference:
         for model_chunk in model:
             for param in model_chunk.parameters():
                 gate_wgrad = getattr(param, "_run_torch_gate_fp32_wgrad", None)

megatron/core/models/common/language_module/language_module.py

@@ -206,7 +206,7 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:
             elif self.config.cross_entropy_fusion_impl == 'native':
                 loss = fused_vocab_parallel_cross_entropy(logits, labels, self.pg_collection.tp)
         else:
-            if _use_accuracy_compatible() and not self.config.dsa_accuracy_compatible:
+            if _use_accuracy_compatible() and not self.config.uses_dsa_reference:
                 s, b = labels.shape
                 loss = torch.nn.functional.cross_entropy(
                     logits.float().reshape(s * b, -1),  # [s*b, vocab]
@@ -231,7 +231,7 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:
             import hashlib as _hashlib
 
             _l = loss.detach().float().contiguous()
-            print(
+            print(  # pylint: disable=bad-builtin
                 f"\nper_token_loss: rank={torch.distributed.get_rank()} "
                 f"shape={list(_l.shape)} "
                 f"md5={_hashlib.md5(_l.cpu().numpy().tobytes()).hexdigest()}",

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -60,7 +60,7 @@
 
 def _get_standalone_norm(config: TransformerConfig, backend: BackendSpecProvider, *, for_qk=False):
     rms_norm = config.normalization == "RMSNorm"
-    if rms_norm and config.norm_accuracy_compatible:
+    if rms_norm and config.uses_dsa_reference:
         return WrappedTorchNorm
     return backend.layer_norm(rms_norm=rms_norm, for_qk=for_qk)
 

megatron/core/optimizer/__init__.py

@@ -554,11 +554,11 @@ def _get_megatron_optimizer_based_on_param_groups(
 
             # set Adam class and weight decay mode depending
             # on source of optimizer (Torch or TE/Apex)
-            if USING_PYTORCH_OPTIMIZER:
+            if config.use_accuracy_compatible and not config.use_precision_aware_optimizer:
+                adam_cls = torch.optim.AdamW
+                kwargs.update({"foreach": False, "fused": True})
+            elif USING_PYTORCH_OPTIMIZER:
                 adam_cls = torch.optim.AdamW if config.decoupled_weight_decay else torch.optim.Adam
-            elif config.native_unfused_adamw:
-                adam_cls = torch.optim.AdamW if config.decoupled_weight_decay else torch.optim.Adam
-                kwargs.update({"foreach": False, "fused": False})
             else:
                 kwargs["adam_w_mode"] = config.decoupled_weight_decay
                 adam_cls = Adam
@@ -587,20 +587,7 @@ def _get_megatron_optimizer_based_on_param_groups(
                 if is_te_min_version("2.1.0.dev0"):
                     kwargs.update({"store_param_remainders": config.store_param_remainders})
 
-            # [对齐修复] use_accuracy_compatible=1 时强制 torch.optim.AdamW(fused=True), 与
-            # PaddleFleet `paddle.optimizer.AdamW` 公式逐位对齐, 绕开 TE.FusedAdam / apex.FusedAdam
-            # 在 bias correction 顺序 / weight decay 时机 / eps 位置上的差异。
-            # 仅在非 precision_aware_optimizer 分支启用, 避免触碰 TE 特有 kwargs。
-            from ..transformer.module import _use_accuracy_compatible
-            if _use_accuracy_compatible() and not config.use_precision_aware_optimizer:
-                # torch.optim.AdamW 无 adam_w_mode 形参（AdamW 即 decoupled weight
-                # decay）；若上游 TE/Apex 分支已注入该 kwarg，这里必须剔除，否则
-                # torch.optim.AdamW 会抛 unexpected keyword argument。
-                kwargs.pop("adam_w_mode", None)
-                kwargs["fused"] = True
-                optimizer = torch.optim.AdamW(**kwargs)
-            else:
-                optimizer = adam_cls(**kwargs)
+            optimizer = adam_cls(**kwargs)
 
             def init_state_fn(opt, config=None):
                 for group in opt.param_groups:

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

@@ -48,9 +48,15 @@
 from ..dist_checkpointing.utils import add_prefix_for_sharding
 from ..transformer.module import 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,10 @@ 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=self.config.use_accuracy_compatible
+            and self.config.clip_grad > 0,
         )
         return total_norm
 
@@ -308,7 +317,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=self.config.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 +338,10 @@ 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=self.config.use_accuracy_compatible
+            and self.config.clip_grad > 0,
         )
 
         if clip_grad > 0.0 and params:
@@ -1566,8 +1581,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 self.config.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 +1659,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 self.config.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/optimizer_config.py

@@ -247,9 +247,6 @@ class OptimizerConfig:
     adam_eps: float = 1e-08
     """Term added to the denominator to improve numerical stability in Adam optimizer."""
 
-    native_unfused_adamw: bool = False
-    """Use torch.optim.AdamW with foreach=False and fused=False instead of TE/Apex Adam."""
-
     decoupled_weight_decay: bool = True
     """If true, decouples weight decay from the gradient update, equivalent to AdamW. If false,
     original Adam update rule will be used. Defaults to True.
@@ -371,6 +368,9 @@ class OptimizerConfig:
     ################
     # Miscellaneous
     ################
+    use_accuracy_compatible: bool = False
+    """Select reference optimizer numerics and reproducible gradient clipping."""
+
     clip_grad: float = 1.0
     """Gradient clipping based on global L2 norm."""
 

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