diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py index 3c5491d39a1..6345991e057 100644 --- a/megatron/core/optimizer/clip_grads.py +++ b/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 diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 4a74328d0d9..318c3790a02 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/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: diff --git a/megatron/core/optimizer/reproducible_norm.py b/megatron/core/optimizer/reproducible_norm.py new file mode 100644 index 00000000000..3380c3dbf84 --- /dev/null +++ b/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 diff --git a/tests/unit_tests/optimizer/test_reproducible_norm.py b/tests/unit_tests/optimizer/test_reproducible_norm.py new file mode 100644 index 00000000000..9550f29879d --- /dev/null +++ b/tests/unit_tests/optimizer/test_reproducible_norm.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +import megatron.core.optimizer as optimizer_module +from megatron.core.optimizer.clip_grads import clip_grad_by_total_norm_fp32 +from megatron.core.optimizer.optimizer import ChainedOptimizer +from megatron.core.optimizer.optimizer_config import OptimizerConfig +from megatron.core.optimizer.reproducible_norm import ReproducibleL2Norm + + +@pytest.mark.parametrize( + 'values,expected', + [ + ([0.0], 0.0), + ([3.0, 4.0], 25.0), + ([4096.0, 1.0], 16777216.0), + ([4096.0, 1.0, 1.0, 1.0], 16777220.0), + ([2.0**-70], 2.0**-140), + ([float('inf')], float('inf')), + ([3e38], float('inf')), + ], +) +def test_sum_squares_rounding(values, expected): + norm = ReproducibleL2Norm() + _, squared = norm.finish(norm.accumulate(norm.zeros(), norm.tensor(values))) + assert squared.item() == expected + + +def test_nan_and_invalid_input(): + norm = ReproducibleL2Norm() + actual, _ = norm.finish( + norm.accumulate(norm.zeros(), norm.tensor([float('inf'), float('nan')])) + ) + assert torch.isnan(actual).item() + with pytest.raises(TypeError, match='FP32'): + norm.accumulate(norm.zeros(), norm.tensor([1.0]).bfloat16()) + bins = norm.zeros() + bins[22] = 2**40 + 1 + with pytest.raises(OverflowError): + norm.finish(bins) + + +def test_layout_and_chunk_invariance(): + norm = ReproducibleL2Norm() + gradient = torch.arange(1, 12289, device='cuda', dtype=torch.float32).reshape(96, 128) / 16384 + whole = norm.accumulate(norm.zeros(), gradient) + transposed = norm.accumulate(norm.zeros(), gradient.T.contiguous(), chunk_size=123) + split = norm.accumulate(norm.zeros(), gradient.flatten()[:1234]) + split += norm.accumulate(norm.zeros(), gradient.flatten()[1234:]) + assert torch.equal(whole, transposed) + assert torch.equal(whole, split) + + +def test_chained_owner_groups_and_clipping(monkeypatch): + monkeypatch.setenv("USE_ACCURACY_COMPATIBLE", "1") + rank = torch.distributed.get_rank() + world = torch.distributed.get_world_size() + singleton = None + for member in range(world): + group = torch.distributed.new_group([member]) + if member == rank: + singleton = group + config = OptimizerConfig(clip_grad=1.0) + # Dense 3 and 4 have distinct owners; the expert 12 is replicated between + # singleton stats groups. Finishing each child separately loses this contract. + dense = [torch.tensor([3.0 if rank == 0 else 4.0], device='cuda')] if rank < 2 else [] + if world == 1: + dense = [torch.tensor([3.0, 4.0], device='cuda')] + expert = [torch.tensor([12.0], device='cuda')] + children = [ + SimpleNamespace( + config=config, + get_grads_for_grad_norm=lambda _group=None: dense, + get_grad_stats_parallel_group=lambda: torch.distributed.group.WORLD, + ), + SimpleNamespace( + config=config, + get_grads_for_grad_norm=lambda _group=None: expert, + get_grad_stats_parallel_group=lambda: singleton, + ), + ] + actual = ChainedOptimizer(children).get_grad_norm() + assert actual.item() == 13.0 + parameter = torch.nn.Parameter(torch.zeros(2, device='cuda')) + parameter.grad = torch.tensor([3.0, 4.0], device='cuda') + monkeypatch.setattr('megatron.core.optimizer.clip_grads.multi_tensor_scale_tensor_impl', None) + clip_grad_by_total_norm_fp32([parameter], 1.0, actual) + expected = torch.tensor([3.0, 4.0], device='cuda') * (1.0 / (actual + 1e-6)) + assert torch.equal(parameter.grad, expected) + + +@pytest.mark.parametrize('enabled,clip', [(False, 1.0), (True, 0.0)]) +def test_stock_norm_when_disabled_or_not_clipping(monkeypatch, enabled, clip): + monkeypatch.setenv("USE_ACCURACY_COMPATIBLE", str(int(enabled))) + config = OptimizerConfig(clip_grad=clip) + optimizer = ChainedOptimizer([SimpleNamespace(config=config, get_grad_norm=lambda: 7.0)]) + + def unexpected(*args, **kwargs): + pytest.fail('Reproducible norm must not run without explicit enabled clipping') + + monkeypatch.setattr(optimizer, '_get_reproducible_grad_norm', unexpected) + assert optimizer.get_grad_norm() == 7.0 + + +@pytest.fixture(scope='module', autouse=True) +def distributed_norm_device(): + import os + + torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', '0'))) + owns_group = not torch.distributed.is_initialized() + if owns_group: + torch.distributed.init_process_group(backend='nccl') + yield + if owns_group: + torch.distributed.destroy_process_group() + + +@pytest.fixture(scope='session') +def ensure_test_data(): + """The norm tests are self-contained and do not consume external datasets.""" + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_existing_accuracy_mode_selects_native_adam(monkeypatch, enabled): + monkeypatch.setenv("USE_ACCURACY_COMPATIBLE", str(int(enabled))) + config = OptimizerConfig(lr=1e-3) + parameter = torch.nn.Parameter(torch.tensor([1.0], device="cuda")) + optimizer, _ = optimizer_module._get_megatron_optimizer_based_on_param_groups( + config, [], [{"params": [parameter]}], skip_megatron_wrapping=True + ) + if enabled: + assert type(optimizer) is torch.optim.AdamW + assert optimizer.defaults["fused"] is True + assert optimizer.defaults["foreach"] in (None, False) + else: + expected = ( + torch.optim.AdamW if optimizer_module.USING_PYTORCH_OPTIMIZER else optimizer_module.Adam + ) + assert type(optimizer) is expected + parameter.grad = torch.ones_like(parameter) + optimizer.step() + assert parameter.item() < 1.0