diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index cf14a5b8c3c..acd0424f8df 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -389,7 +389,37 @@ def max_calibrate( _finalize_with_shared_state(model, weight_patterns) return - # Check MoE calibration completeness before sync + # Share the input (activation) amax across EP-sharded routed experts. In a sparse a2a/EP calibration + # each expert sees ONLY the tokens routed to it, so an expert (or a whole rank's experts) that received + # none has an uncalibrated input_quantizer -- while its weight_quantizer is always calibrated from the + # weight directly (above). Experts in a layer conceptually share one input amax, so MAX the + # input_quantizer amax over the expert-model-parallel group: cold experts pick up the shared value + # (sync_amax_across_distributed_group fills a None rank with the group MAX; if NO rank saw tokens the + # amax is dropped group-wide). Expert WEIGHT amax is intentionally NOT shared over EP (per-rank shards + # differ). This mirrors main's _should_sync_amax_across_ep / sync_moe_expert_amax input-amax sharing, + # which the a2a fused-experts path would otherwise miss -> some experts left with no scale. Must run + # BEFORE the completeness check so shared-in experts count as calibrated. Only modules that opt in + # (``_share_input_amax_across_ep``): Megatron's experts keep main's check-then-sync behavior. + for name, module in model.named_modules(): + if ( + isinstance(module, QuantModule) + and getattr(module, "_share_input_amax_across_ep", False) + and _has_expert_parallelism(module) + ): + assert module.parallel_state is not None # _has_expert_parallelism checked it + ep_group = module.parallel_state.expert_model_parallel_group + for child_name, child in module.named_children(): + if isinstance(child, (TensorQuantizer, SequentialQuantizer)) and ( + "input_quantizer" in child_name + ): + for leaf in _iter_leaf_quantizers(child): + # Only enabled quantizers, and every rank must agree on which: a quantizer one + # rank skips and another syncs pairs mismatched tensors in every later + # all_reduce, which NCCL hangs on rather than reporting. + if leaf.is_enabled: + leaf.sync_amax_across_distributed_group(ep_group) + + # Check MoE calibration completeness (after the EP input-amax share above). for name, module in model.named_modules(): if isinstance(module, QuantModule) and _has_expert_parallelism(module): for child in module.children(): @@ -402,8 +432,20 @@ def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, chi for _q in quantizer: sync_quantizer_amax_across_dp_ep(_q, parallel_state, parent_name, child_name) return - if getattr(quantizer, "_amax", None) is None: + # Skip DISABLED quantizers: their amax is never used for quantization, so there is nothing to sync. + # This matters for a VLM whose non-language modules (vision tower / mm-projector) are disabled via a + # ``{"quantizer_name": "*", "enable": False}`` config with ``forward_loop=None`` + # (extract_and_prepare_language_model_from_vl): those quantizers are uncalibrated, and a disabled + # NON-per-tensor one would hit the "amax shape not locally known" error below. Unlike the _amax-None + # case in the NOTE, disabling is a config property applied identically on every rank, so skipping is + # symmetric across the group -> it cannot desync the collective. + if not quantizer.is_enabled: return + # NOTE: do NOT early-return when this rank's `_amax` is None. Under data-parallel calibration a + # module may be activated on some ranks but not others, so `_amax` is present on only a subset of + # the group; every rank must still join the collective or the participant set mismatches -> NCCL + # deadlock. sync_amax_across_distributed_group handles this symmetrically (a rank without a local + # amax contributes a zero of the locally-known amax shape; the MAX picks up the real value). quantizer.sync_amax_across_distributed_group(parallel_state.data_parallel_group) if _should_sync_amax_across_ep(parent_name, child_name, sync_expert_weight_amax): quantizer.sync_amax_across_distributed_group(parallel_state.expert_model_parallel_group) @@ -452,7 +494,9 @@ def sync_quantizer_amax_across_tp( if getattr(quantizer.block_sizes, "type", None) == "dynamic": return - if quantizer.axis in axes_for_sync and quantizer.amax is not None: + # As in the DP/EP sync above: an enabled quantizer is synced on every rank, with or without + # a local amax, and a disabled one on none, which the config makes the same on every rank. + if quantizer.axis in axes_for_sync and quantizer.is_enabled: quantizer.sync_amax_across_distributed_group(parallel_state.tensor_parallel_group) # Step 2: Sync amax across relevant parallelism (such as TP / EP) @@ -498,7 +542,7 @@ def sync_quantizer_amax_across_tp( # We only support KVCache quantization with scalar per-tensor states for now (NVFP4 & FP8 KV cache) # So we should sync amax across DP and TP for these quantizers (DP is already synced from above) for quantizer in [module.k_bmm_quantizer, module.v_bmm_quantizer]: - if isinstance(quantizer, TensorQuantizer) and quantizer.amax is not None: + if isinstance(quantizer, TensorQuantizer) and quantizer.is_enabled: # as above quantizer.sync_amax_across_distributed_group( module.parallel_state.tensor_parallel_group ) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 482f71f8119..d42b01b2eb6 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -1419,18 +1419,86 @@ def set_from_modelopt_state(self, modelopt_state, properties_only: bool = False) ) def sync_amax_across_distributed_group(self, parallel_group: DistributedProcessGroup): - """Synchronize the amax across all ranks in the given group.""" - if parallel_group.is_initialized() and getattr(self, "_amax", None) is not None: - try: - dist.all_reduce(self._amax, op=dist.ReduceOp.MAX, group=parallel_group.group) - except RuntimeError as e: - # This error happens if the distributed backend is using GPU and - # the tensor is not on GPU (or vice versa). - warnings.warn( - f"Failed to synchronize amax: {e}, probably because the tensor is on a device which is not" - "supported by the current distributed backend. This warning can be ignored" - "if happening during modelopt restore." + """Synchronize amax across all ranks in ``parallel_group`` with a MAX all-reduce every rank joins. + + With data-parallel calibration a module can be activated on some ranks but not others (a + sparsely-routed MoE expert, or any module a rank's calib batch never exercises), so ``_amax`` + ends up present on only a SUBSET of the group's ranks. Every rank must still join the + collective or the participant set mismatches and NCCL deadlocks. + + The ranks agree on whether anyone calibrated the quantizer from a presence flag, never from + the amax value: a calibrated amax can legitimately be zero (an all-zero weight). A quantizer + that no rank calibrated keeps no amax, so downstream ``amax is None`` handling still applies. + A calibrated amax keeps its own dtype and device, exactly as an in-place reduce leaves it. + """ + if not parallel_group.is_initialized(): + return + group = parallel_group.group + amax = getattr(self, "_amax", None) + device = _collective_device(group) + # A DYNAMIC block quantizer (e.g. NVFP4 activation: per-block scales computed at runtime) stores only + # a per-TENSOR scalar amax (the global scale -> ``input_scale``); its block_sizes drive the runtime + # scales, not the stored amax shape. So treat it as per-tensor here -- only STATIC per-channel/ + # per-block amax has a shape a have-not rank can't reconstruct. + _block = self.block_sizes + _static_non_per_tensor = self.axis is not None or ( + _block is not None and _block.get("type", "static") != "dynamic" + ) + if _static_non_per_tensor: + # A non-per-tensor amax (per-channel / static per-block) has a shape that is NOT known locally + # when this rank never calibrated the quantizer, so a have-not rank cannot join the MAX all-reduce + # with a correctly-shaped zero. Before diverging, every rank symmetrically agrees -- via a single MAX + # all-reduce of the packed flag ``[present, -present]`` (which yields ``[any_present, + # -all_present]``) that all ranks join -- on the group-wide calibration state: + # - none present -> the quantizer is uncalibrated group-wide (a module that no rank's calib + # data exercised, e.g. a VLM path never hit by the image-text batches); skip the sync and + # leave ``_amax`` None so downstream ``amax is None`` handling applies. + # - all present -> reduce the amax itself below. + # - some present -> genuinely uneven calibration whose amax shape the have-nots cannot recover; + # every rank raises together (symmetric, no deadlock). + present = 1.0 if amax is not None else 0.0 + flag = torch.tensor([present, -present], dtype=torch.float32, device=device) + dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=group) + any_present = bool(flag[0] > 0) + all_present = bool(-flag[1] > 0) + if not any_present: + return + if not all_present: + raise RuntimeError( + f"Cannot amax-sync unevenly-calibrated non-per-tensor quantizer {type(self).__name__}: " + "some ranks calibrated it but others did not, and its amax shape is not locally known on " + "the have-nots. Ensure every rank calibrates it (e.g. increase --calib-size)." ) + assert amax is not None # all_present + # Every rank calibrated it with the same calibrator, so shapes and dtypes agree. + reduced = amax.to(device=device) # the backend's device; a no-op when already there + _all_reduce_amax(reduced, group) + if reduced is not amax: + amax.copy_(reduced) + return + + # Per-tensor (or dynamic-block, whose stored amax is a scalar): one all-reduce carries the value, + # whether this rank calibrated the quantizer, and the dtype it did so in, so a rank that never + # calibrated it joins with a locally known scalar and still ends up with the group's amax. + # float32 throughout: NCCL hangs rather than raising when ranks pair mismatched dtypes, which a + # group split between calibrated and uncalibrated ranks otherwise would (Kimi-K3, NVFP4, EP64: + # 8 ranks held an expert input amax and 56 did not). + packed = torch.zeros(3, dtype=torch.float32, device=device) + if amax is not None: + packed[0] = amax.detach().reshape(()).to(device=device, dtype=torch.float32) + packed[1] = 1.0 + packed[2] = ( + float(_AMAX_DTYPES.index(amax.dtype) + 1) if amax.dtype in _AMAX_DTYPES else 0.0 + ) + _all_reduce_amax(packed, group) + if not bool(packed[1] > 0): + return # no rank calibrated it + if amax is not None: + amax.copy_(packed[0]) + else: + dtype_code = int(packed[2].item()) + dtype = _AMAX_DTYPES[dtype_code - 1] if dtype_code > 0 else torch.float32 + self._set_buffer("_amax", packed[0].to(dtype)) @contextlib.contextmanager def disable_pre_quant_scale(self): @@ -1529,6 +1597,35 @@ def _to_local(t: torch.Tensor) -> torch.Tensor: return t +# The dtypes a calibrated amax can have, so a rank that never calibrated a quantizer adopts the same one. +_AMAX_DTYPES = (torch.float32, torch.bfloat16, torch.float16, torch.float64) + + +def _all_reduce_amax(tensor: torch.Tensor, group) -> None: + """MAX all-reduce ``tensor`` in place, naming the likely cause of a failure.""" + try: + dist.all_reduce(tensor, op=dist.ReduceOp.MAX, group=group) + except RuntimeError as e: + raise RuntimeError( + f"Failed to synchronize amax across the calibration group: {e}. This is commonly a CUDA " + "OOM when NCCL lazily allocates communication buffers under calibration memory pressure; " + "free cached activations (gc.collect()/torch.cuda.empty_cache()) before the sync or reduce " + "the calibration sequence length / batch size." + ) from e + + +def _collective_device(group) -> torch.device: + """The device every rank of ``group`` hands its collectives: CUDA when NCCL serves it, else CPU. + + Derived from the group's backend, which all ranks share, rather than from a local tensor: a rank + that never calibrated a quantizer has no amax to ask, and ranks that disagree on the device type + reach different backends and hang. A gloo-only group has no CUDA backend at all. + """ + if torch.cuda.is_available() and "nccl" in str(dist.get_backend(group)): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + class StaticBlockScaleQuantizer(TensorQuantizer): """TensorQuantizer for static block quantization with two-level scaling. diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 6a6c5d18bfc..a6aaba5263b 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -812,6 +812,9 @@ class _QuantFusedExperts(_QuantFunctionalMixin): _first_proj_attr = "gate_up_proj" # Whether the first projection packs a gate half that must be split on export. _is_gated = True + # Under expert parallelism each rank's experts see only the tokens routed to them, so calibration + # shares the layer's input amax over the expert-parallel group (see max_calibrate). + _share_input_amax_across_ep = True @property def _first_proj_input_quantizer_attr(self) -> str: diff --git a/tests/gpu/torch/quantization/test_amax_sync_distributed.py b/tests/gpu/torch/quantization/test_amax_sync_distributed.py new file mode 100644 index 00000000000..614ab87e658 --- /dev/null +++ b/tests/gpu/torch/quantization/test_amax_sync_distributed.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Amax synchronization over NCCL. + +What the sync decides is tested on gloo, in tests/unit/torch/quantization/test_dist.py. What only +NCCL shows is where a rank with no amax puts the zero it contributes: on the collective's CUDA +device. +""" + +import pytest +import torch +import torch.distributed as dist + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer +from modelopt.torch.utils.distributed import DistributedProcessGroup + + +def _per_tensor_job(rank, size, calibrated): + device = torch.device(f"cuda:{rank}") + q = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)).to(device) + if calibrated and rank == 0: + q._set_buffer("_amax", torch.tensor(7.0, device=device)) + q.sync_amax_across_distributed_group(DistributedProcessGroup(dist.group.WORLD)) + if calibrated: + assert q.amax.item() == pytest.approx(7.0) + else: + assert q.amax is None, "an all-zero sync must leave the quantizer uncalibrated, not amax==0" + + +@pytest.mark.parametrize("calibrated", [True, False], ids=["rank_0_only", "nobody"]) +def test_per_tensor_amax_survives_uneven_calibration(need_2_gpus, dist_workers, calibrated): + """A module only some ranks' batches reached. With one GPU no rank lacks the amax.""" + dist_workers.run(_per_tensor_job, calibrated=calibrated) diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 11288e12b3c..07d298f3de1 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -13,14 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from datetime import timedelta + +import pytest import torch import torch.distributed as dist from _test_utils.torch.quantization.models import SimpleLinear from torch.nn.parallel import DistributedDataParallel import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.model_calib import max_calibrate from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils import unwrap_model +from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState def _test_data_parallel_helper(rank, size): @@ -51,3 +57,149 @@ def forward_loop(model): def test_data_parallel(cpu_dist_workers): cpu_dist_workers.run(_test_data_parallel_helper) + + +def _sync(quantizer): + quantizer.sync_amax_across_distributed_group(DistributedProcessGroup(dist.group.WORLD)) + + +# An NVFP4 activation quantizer's block sizes drive its runtime scales: it stores only a per-tensor +# amax, the global scale, so it syncs as a per-tensor quantizer does. +_PER_TENSOR_CFGS = { + "int8": {"num_bits": 8}, + "nvfp4-dynamic": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + }, +} + + +def _calibrators_dtype_helper(rank, size, per_tensor_cfg): + cfg = QuantizerAttributeConfig(**_PER_TENSOR_CFGS[per_tensor_cfg]) + one_calibrated, both_calibrated = TensorQuantizer(cfg), TensorQuantizer(cfg) + if rank == 0: + one_calibrated._set_buffer("_amax", torch.tensor(3.0, dtype=torch.bfloat16)) + both_calibrated._set_buffer("_amax", torch.tensor(2.0 + rank, dtype=torch.bfloat16)) + per_channel = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=0)) + per_channel._set_buffer("_amax", torch.tensor([1.0, 2.0 + rank], dtype=torch.bfloat16)) + for quantizer in (one_calibrated, both_calibrated, per_channel): + _sync(quantizer) + assert quantizer.amax.dtype == torch.bfloat16 + assert one_calibrated.amax.item() == 3.0 + assert both_calibrated.amax.item() == 3.0 + assert per_channel.amax.tolist() == [1.0, 3.0] + + +@pytest.mark.parametrize("per_tensor_cfg", list(_PER_TENSOR_CFGS)) +def test_the_sync_leaves_every_rank_the_calibrators_dtype(cpu_dist_workers, per_tensor_cfg): + """And the group's max, including on a rank that never calibrated the quantizer and so had no + amax of its own. Synced as a per-block quantizer, an NVFP4 activation quantizer's uncalibrated + rank would raise instead.""" + cpu_dist_workers.run(_calibrators_dtype_helper, per_tensor_cfg=per_tensor_cfg) + + +def _zero_or_uncalibrated_helper(rank, size, case): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=8, axis=0 if "channel" in case else None) + ) + if case == "calibrated-zero": + quantizer._set_buffer("_amax", torch.tensor(0.0)) + _sync(quantizer) + if case == "calibrated-zero": + assert quantizer.amax is not None + assert quantizer.amax.item() == 0.0 + else: + assert getattr(quantizer, "_amax", None) is None + + +@pytest.mark.parametrize("case", ["calibrated-zero", "uncalibrated", "uncalibrated-per-channel"]) +def test_the_sync_keeps_a_calibrated_zero_apart_from_uncalibrated(cpu_dist_workers, case): + """An all-zero weight calibrates to amax 0, which is not "never calibrated"; a quantizer no + rank's calibration data reached, which is expected rather than an error, stays uncalibrated.""" + cpu_dist_workers.run(_zero_or_uncalibrated_helper, case=case) + + +def _uneven_per_channel_helper(rank, size): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=0)) + if rank == 0: + quantizer._set_buffer("_amax", torch.full((4,), 3.0)) + # A rank left waiting in a collective its peer never joins gives up in seconds on this group. + group = dist.new_group(timeout=timedelta(seconds=10)) + try: + quantizer.sync_amax_across_distributed_group(DistributedProcessGroup(group)) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + finally: + dist.destroy_process_group(group) + assert "unevenly-calibrated" in raised, raised + + +def test_uneven_per_channel_calibration_fails_on_every_rank(cpu_dist_workers): + """A rank without the amax cannot rebuild its per-channel shape, so the sync fails, on every + rank together: one rank raising alone would leave the other waiting in the collective.""" + cpu_dist_workers.run(_uneven_per_channel_helper) + + +def _expert_input_amax_helper(rank, size, share): + model = torch.nn.Sequential(torch.nn.Linear(16, 16)) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop=None) + expert = model[0] + expert._parallel_state = ParallelState( + data_parallel_group=-1, expert_model_parallel_group=dist.group.WORLD + ) + if share: + expert._share_input_amax_across_ep = True + + def forward_loop(m): # only rank 0's experts receive tokens + if rank == 0: + m(torch.ones(2, 16)) + + if share: + max_calibrate(model, forward_loop) + assert expert.input_quantizer.amax.item() == 1.0 + else: + with pytest.raises(RuntimeError, match="MoE calibration incomplete"): + max_calibrate(model, forward_loop) + + +@pytest.mark.parametrize("share", [False, True], ids=["megatron", "opted-in"]) +def test_only_modules_that_opt_in_share_expert_input_amax_before_the_check(cpu_dist_workers, share): + """Megatron's experts keep main's check-then-sync; HF fused experts share the input amax first.""" + cpu_dist_workers.run(_expert_input_amax_helper, share=share) + + +def _tensor_parallel_sync_helper(rank, size, kv_cache): + model = torch.nn.Sequential(torch.nn.Linear(16, 16)) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop=None) + linear = model[0] + if kv_cache: + linear.k_bmm_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + linear.v_bmm_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + synced = [linear.k_bmm_quantizer, linear.v_bmm_quantizer] + else: + linear._is_column_parallel = True + synced = [linear.input_quantizer] + # A rank left waiting for a peer that skips the sync gives up in seconds on this group. + tp_group = dist.new_group(timeout=timedelta(seconds=10)) + linear._parallel_state = ParallelState(data_parallel_group=-1, tensor_parallel_group=tp_group) + + def forward_loop(m): # only rank 0 calibrates the synced quantizers + if rank == 0: + for quantizer in synced: + quantizer(torch.ones(2, 16)) + + try: + max_calibrate(model, forward_loop) + finally: + dist.destroy_process_group(tp_group) + for quantizer in synced: + assert quantizer.amax is not None + assert quantizer.amax.item() == 1.0 + + +@pytest.mark.parametrize("kv_cache", [False, True], ids=["column-parallel", "kv-cache"]) +def test_a_rank_without_amax_still_joins_the_tensor_parallel_sync(cpu_dist_workers, kv_cache): + """Rather than leave the rank that has one waiting in the sync for it.""" + cpu_dist_workers.run(_tensor_parallel_sync_helper, kv_cache=kv_cache)