From accadf14488d3af010a5d92103923e12264f0d56 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:22:41 +0000 Subject: [PATCH 01/11] Make amax synchronization correct under data-parallel calibration With data-parallel calibration a module can be exercised on some ranks and not others -- a sparsely routed MoE expert, or anything a given rank's batch never reaches -- so _amax ends up present on only a subset of the group. Every rank must still join the collective or the participant set mismatches and NCCL hangs, which makes "what does a rank with no amax contribute" the whole problem. For a per-tensor amax the answer is easy: the shape is known locally, so a have-not contributes a scalar zero, the identity for MAX over non-negative amax. If the all-reduced result is zero then nobody calibrated it, and the buffer is dropped so downstream `amax is None` handling still applies -- keeping it as 0 would quantize the module to nothing. A static per-channel or per-block amax has no locally-known shape, so a have-not cannot fabricate a correctly-shaped zero. Ranks agree first, via one MAX all-reduce of a packed presence flag that every rank joins: nobody calibrated it is fine and skipped; everybody did falls through to the real reduce; a genuine split raises on every rank at once. Raising only where the amax is missing would leave the other ranks blocked in a collective that never completes. A dynamic-block quantizer -- NVFP4 activations -- stores only a scalar global scale; its block_sizes drive runtime scales and do not describe the stored amax. Classifying by "has block_sizes" would push it down the static path and turn the ordinary uneven-calibration case into a hard error, so it is treated as per-tensor. Every rank also contributes the same dtype and device. A calibrated amax keeps the calibrator's dtype while a have-not's zero is float32, and NCCL hangs on a mismatch rather than raising (seen on Kimi-K3, NVFP4 EP64: 8 ranks with an amax, 56 without, and the collective never returned), so the amax is cast to float32 first. The device comes from the group's backend, which every rank shares -- CUDA under NCCL, CPU otherwise -- so a gloo group on a CPU host works. max_calibrate keeps main's sync structure -- DP always, EP filtered by _should_sync_amax_across_ep -- so Megatron's numerics are unchanged, with three changes: a rank joins the DP sync even without an amax (disabled quantizers are skipped, which is symmetric since it is a config property); the input amax of expert-parallel experts is shared over the EP group BEFORE the completeness check, so experts that received no tokens pick up the layer's value instead of failing it; and the allocator's cached pool is returned first, since NCCL allocates its buffers outside it. A failed amax all-reduce now says it is most likely that out-of-memory rather than surfacing as a bare CUDA error. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 52 +++++- .../nn/modules/tensor_quantizer.py | 105 ++++++++++-- .../test_amax_sync_distributed.py | 152 ++++++++++++++++++ tests/unit/torch/quantization/test_dist.py | 17 ++ 4 files changed, 313 insertions(+), 13 deletions(-) create mode 100644 tests/gpu/torch/quantization/test_amax_sync_distributed.py diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index cf14a5b8c3c..6b03249e2eb 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -16,6 +16,7 @@ """Calibration utilities.""" import fnmatch +import gc import math import time import warnings @@ -389,7 +390,42 @@ def max_calibrate( _finalize_with_shared_state(model, weight_patterns) return - # Check MoE calibration completeness before sync + # The calibration forward leaves large activation tensors held in the caching allocator. The + # distributed amax sync below issues NCCL collectives whose comm buffers are allocated lazily via + # cudaMalloc (outside the caching allocator). Under the memory pressure of long calib sequences this + # can OOM on a single rank when NCCL first allocates -- and since sync_amax_across_distributed_group + # re-raises on collective failure, that rank aborts mid-collective and takes down the group. Return the + # cached activation memory to CUDA first so NCCL has room. (gc.collect() first to drop any cycle-held + # activation references the calibration closures/calibrators may still own.) + if torch.cuda.is_available(): + gc.collect() + torch.cuda.empty_cache() + + # 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. + for name, module in model.named_modules(): + if isinstance(module, QuantModule) and _has_expert_parallelism(module): + 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 +438,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) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 482f71f8119..20b0caef905 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -1419,18 +1419,89 @@ 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`` via a single symmetric MAX all-reduce. + + 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 -- so a rank lacking a local amax contributes a zero + of the amax's shape. That shape is known LOCALLY from the quantizer's own config (a per-tensor amax + is scalar), so NO peer metadata exchange is needed to discover it. (Non-per-tensor quantizers -- + e.g. weights -- are calibrated from their weight on every rank, so they are never ``None`` here.) + The MAX picks up the real amax from the ranks that have it (zero is the identity for max over + non-negative amax); if no rank in the group has one, the all-reduced result is zero and the + uncalibrated amax is dropped so downstream ``amax is None`` handling still applies. + """ + 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 -> fall through to the shared MAX all-reduce 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)." ) + elif amax is None: + # Per-tensor (or dynamic-block, whose stored amax is a scalar) amax, so a rank that never + # calibrated it can still join the collective with a locally-known scalar zero (the identity for + # MAX over non-negative amax). + self._set_buffer("_amax", torch.zeros((), dtype=torch.float32, device=device)) + amax = self._amax + # NCCL requires every rank to contribute the same DTYPE as well as the same shape, and the two + # sides of this collective get their amax from different places: a rank that calibrated the + # quantizer keeps the calibrator's dtype (``_amax_setter_helper`` stores + # ``value.clone().detach()``), while a rank that did not synthesises float32 just above. When a + # group splits across the two -- exactly the sparsely-routed MoE case this path exists to serve -- + # the all_reduce pairs mismatched dtypes, which NCCL HANGS on rather than raising. + # Observed on Kimi-K3 (NVFP4, EP64): for the first expert input_quantizer 8 ranks held a + # calibrated amax and 56 held none, and the collective never returned. The device must agree + # for the same reason, so a calibrated amax moves to the device the have-nots synthesise on. + if amax.dtype != torch.float32 or amax.device != device: + self._set_buffer("_amax", amax.to(device=device, dtype=torch.float32)) + amax = self._amax + try: + dist.all_reduce(amax, 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 + if not bool(torch.any(amax)): + # No rank in the group had a (positive) amax -> uncalibrated across the whole group. Drop it so + # downstream ``amax is None`` handling applies (a real amax is strictly positive). The + # all-reduced value is identical on every rank, so this decision is symmetric. + delattr(self, "_amax") @contextlib.contextmanager def disable_pre_quant_scale(self): @@ -1529,6 +1600,18 @@ def _to_local(t: torch.Tensor) -> torch.Tensor: return t +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/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..ed6682e575c --- /dev/null +++ b/tests/gpu/torch/quantization/test_amax_sync_distributed.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 across a data-parallel calibration group. + +Under data-parallel calibration a module can be exercised on some ranks and not others -- a sparsely +routed MoE expert, or any module a given rank's batch never reaches. Every rank must still join the +collective or the participant set mismatches and NCCL hangs, so these tests pin down what each rank +contributes when it has no local amax, and that the failure modes are symmetric across ranks. +""" + +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 _quantizer(cfg, amax=None, rank_device=None): + q = TensorQuantizer(cfg).to(rank_device) + if amax is not None: + q._set_buffer("_amax", torch.as_tensor(amax, dtype=torch.float32, device=rank_device)) + return q + + +def _whole_world(): + return DistributedProcessGroup(dist.group.WORLD) + + +# --------------------------------------------------------------------------------- per-tensor amax + + +def _per_tensor_all_ranks_present(rank, size): + """Every rank calibrated it: the group agrees on the max.""" + device = torch.device(f"cuda:{rank}") + q = _quantizer(QuantizerAttributeConfig(num_bits=8), amax=float(rank + 1), rank_device=device) + q.sync_amax_across_distributed_group(_whole_world()) + assert q.amax.item() == pytest.approx(float(size)) + + +def test_per_tensor_amax_takes_the_group_max(dist_workers): + dist_workers.run(_per_tensor_all_ranks_present) + + +def _per_tensor_some_ranks_missing(rank, size): + """A have-not rank contributes a locally-known scalar zero -- the identity for MAX.""" + device = torch.device(f"cuda:{rank}") + amax = 7.0 if rank == 0 else None + q = _quantizer(QuantizerAttributeConfig(num_bits=8), amax=amax, rank_device=device) + q.sync_amax_across_distributed_group(_whole_world()) + # Both ranks end up with rank 0's value; the have-not must NOT be left at None or zero. + assert q.amax.item() == pytest.approx(7.0) + + +def test_per_tensor_amax_survives_uneven_calibration(dist_workers): + """This is the common case -- a module only some ranks' batches reached.""" + dist_workers.run(_per_tensor_some_ranks_missing) + + +def _per_tensor_nobody_calibrated(rank, size): + device = torch.device(f"cuda:{rank}") + q = _quantizer(QuantizerAttributeConfig(num_bits=8), amax=None, rank_device=device) + q.sync_amax_across_distributed_group(_whole_world()) + assert q.amax is None, "an all-zero sync must leave the quantizer uncalibrated, not amax==0" + + +def test_uncalibrated_everywhere_stays_uncalibrated(dist_workers): + """amax==0 would quantize everything to zero; downstream expects None for 'never calibrated'.""" + dist_workers.run(_per_tensor_nobody_calibrated) + + +# --------------------------------------------------------------- static non-per-tensor (per-channel) + + +def _per_channel_all_present(rank, size): + device = torch.device(f"cuda:{rank}") + amax = torch.full((4,), float(rank + 1)) + q = _quantizer(QuantizerAttributeConfig(num_bits=8, axis=0), amax=amax, rank_device=device) + q.sync_amax_across_distributed_group(_whole_world()) + assert torch.allclose(q.amax.cpu(), torch.full((4,), float(size))) + + +def test_per_channel_amax_takes_the_elementwise_group_max(dist_workers): + dist_workers.run(_per_channel_all_present) + + +def _per_channel_unevenly_calibrated(rank, size): + """Have-nots cannot invent the right shape, so every rank must raise -- together.""" + device = torch.device(f"cuda:{rank}") + amax = torch.full((4,), 3.0) if rank == 0 else None + q = _quantizer(QuantizerAttributeConfig(num_bits=8, axis=0), amax=amax, rank_device=device) + with pytest.raises(RuntimeError, match="unevenly-calibrated"): + q.sync_amax_across_distributed_group(_whole_world()) + + +def test_uneven_per_channel_calibration_fails_on_every_rank(dist_workers): + """A one-sided raise would leave the other rank blocked in the collective forever. + + The rank that HAS the amax cannot know the have-nots exist without asking, so the agreement is a + single MAX all-reduce of a packed presence flag that every rank joins before anyone diverges. + """ + dist_workers.run(_per_channel_unevenly_calibrated) + + +def _per_channel_nobody_calibrated(rank, size): + device = torch.device(f"cuda:{rank}") + q = _quantizer(QuantizerAttributeConfig(num_bits=8, axis=0), amax=None, rank_device=device) + q.sync_amax_across_distributed_group(_whole_world()) + assert q.amax is None + + +def test_per_channel_uncalibrated_everywhere_is_not_an_error(dist_workers): + """A module no rank's calibration data exercised is expected, not a failure.""" + dist_workers.run(_per_channel_nobody_calibrated) + + +# ------------------------------------------------------------------------- dynamic block (NVFP4 act) + + +def _dynamic_block_some_ranks_missing(rank, size): + """A dynamic-block quantizer stores a SCALAR amax, so the per-tensor path applies.""" + device = torch.device(f"cuda:{rank}") + cfg = QuantizerAttributeConfig( + num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + ) + amax = 5.0 if rank == 0 else None + q = _quantizer(cfg, amax=amax, rank_device=device) + q.sync_amax_across_distributed_group(_whole_world()) + assert q.amax.item() == pytest.approx(5.0) + + +def test_dynamic_block_amax_is_treated_as_per_tensor(dist_workers): + """Its block_sizes drive runtime scales; only the per-tensor global scale is stored. + + Classifying it by ``block_sizes is not None`` would send an NVFP4 activation quantizer down the + static path and raise 'unevenly-calibrated' for the ordinary case of a module some ranks missed. + """ + dist_workers.run(_dynamic_block_some_ranks_missing) diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 11288e12b3c..631894a7d12 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -15,12 +15,15 @@ import torch import torch.distributed as dist +from _test_utils.torch.distributed.utils import spawn_multiprocess_job 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.nn import TensorQuantizer from modelopt.torch.utils import unwrap_model +from modelopt.torch.utils.distributed import DistributedProcessGroup def _test_data_parallel_helper(rank, size): @@ -51,3 +54,17 @@ def forward_loop(model): def test_data_parallel(cpu_dist_workers): cpu_dist_workers.run(_test_data_parallel_helper) + + +def _test_uneven_amax_sync_helper(rank, size): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + if rank == 0: + quantizer._set_buffer("_amax", torch.tensor(7.0)) + quantizer.sync_amax_across_distributed_group(DistributedProcessGroup(dist.group.WORLD)) + assert quantizer.amax.item() == 7.0 + dist.destroy_process_group() + + +def test_amax_sync_with_a_rank_that_never_calibrated(skip_on_windows): + """The rank without an amax still joins the sync, contributing a zero on the group's device.""" + spawn_multiprocess_job(2, _test_uneven_amax_sync_helper, backend="gloo") From e0d6ddfee877cf28cd73db7e9dce800c8f281673 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 03:21:54 +0000 Subject: [PATCH 02/11] Assert the amax-sync invariants mypy cannot see Two reads go through an Optional that is set by then, which the mypy run in pre-commit (and so in CI) rejects: the amax every rank holds once it reaches the MAX all-reduce (calibrated, or the zero synthesised just above), and the parallel state _has_expert_parallelism has already checked. Assert both. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 1 + modelopt/torch/quantization/nn/modules/tensor_quantizer.py | 1 + 2 files changed, 2 insertions(+) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 6b03249e2eb..644df6422d5 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -413,6 +413,7 @@ def max_calibrate( # BEFORE the completeness check so shared-in experts count as calibrated. for name, module in model.named_modules(): if isinstance(module, QuantModule) 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 ( diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 20b0caef905..ffc0dad7b1a 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -1476,6 +1476,7 @@ def sync_amax_across_distributed_group(self, parallel_group: DistributedProcessG # MAX over non-negative amax). self._set_buffer("_amax", torch.zeros((), dtype=torch.float32, device=device)) amax = self._amax + assert amax is not None # calibrated here, or the zero just above # NCCL requires every rank to contribute the same DTYPE as well as the same shape, and the two # sides of this collective get their amax from different places: a rank that calibrated the # quantizer keeps the calibrator's dtype (``_amax_setter_helper`` stores From 4d1c4312a7ca4704eafbec0a0686df26da52aa7a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 10:38:33 +0000 Subject: [PATCH 03/11] Tell a never-calibrated quantizer from a zero amax in the amax sync The sync decided that no rank had calibrated a quantizer when the reduced amax was zero, and dropped it. A calibrated amax can legitimately be zero -- an all-zero weight -- and dropping it made the FP8 export fail on the missing _amax for any run with a process group (Megatron, FSDP2, DDP). The ranks now agree through a presence flag carried in the same all-reduce. The sync also promoted every amax to float32 for the collective and left it there, so later bf16 writes (MSE calibration) stopped rounding as before. Each rank now keeps its calibrator's dtype, and a rank that never calibrated the quantizer adopts the group's. Signed-off-by: Shengliang Xu --- .../nn/modules/tensor_quantizer.py | 101 ++++++++++-------- tests/unit/torch/quantization/test_dist.py | 49 +++++++++ 2 files changed, 106 insertions(+), 44 deletions(-) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index ffc0dad7b1a..d42b01b2eb6 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -1419,18 +1419,17 @@ def set_from_modelopt_state(self, modelopt_state, properties_only: bool = False) ) def sync_amax_across_distributed_group(self, parallel_group: DistributedProcessGroup): - """Synchronize amax across all ranks in ``parallel_group`` via a single symmetric MAX all-reduce. + """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 -- so a rank lacking a local amax contributes a zero - of the amax's shape. That shape is known LOCALLY from the quantizer's own config (a per-tensor amax - is scalar), so NO peer metadata exchange is needed to discover it. (Non-per-tensor quantizers -- - e.g. weights -- are calibrated from their weight on every rank, so they are never ``None`` here.) - The MAX picks up the real amax from the ranks that have it (zero is the identity for max over - non-negative amax); if no rank in the group has one, the all-reduced result is zero and the - uncalibrated amax is dropped so downstream ``amax is None`` handling still applies. + 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 @@ -1454,7 +1453,7 @@ def sync_amax_across_distributed_group(self, parallel_group: DistributedProcessG # - 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 -> fall through to the shared MAX all-reduce below. + # - 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 @@ -1470,39 +1469,36 @@ def sync_amax_across_distributed_group(self, parallel_group: DistributedProcessG "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)." ) - elif amax is None: - # Per-tensor (or dynamic-block, whose stored amax is a scalar) amax, so a rank that never - # calibrated it can still join the collective with a locally-known scalar zero (the identity for - # MAX over non-negative amax). - self._set_buffer("_amax", torch.zeros((), dtype=torch.float32, device=device)) - amax = self._amax - assert amax is not None # calibrated here, or the zero just above - # NCCL requires every rank to contribute the same DTYPE as well as the same shape, and the two - # sides of this collective get their amax from different places: a rank that calibrated the - # quantizer keeps the calibrator's dtype (``_amax_setter_helper`` stores - # ``value.clone().detach()``), while a rank that did not synthesises float32 just above. When a - # group splits across the two -- exactly the sparsely-routed MoE case this path exists to serve -- - # the all_reduce pairs mismatched dtypes, which NCCL HANGS on rather than raising. - # Observed on Kimi-K3 (NVFP4, EP64): for the first expert input_quantizer 8 ranks held a - # calibrated amax and 56 held none, and the collective never returned. The device must agree - # for the same reason, so a calibrated amax moves to the device the have-nots synthesise on. - if amax.dtype != torch.float32 or amax.device != device: - self._set_buffer("_amax", amax.to(device=device, dtype=torch.float32)) - amax = self._amax - try: - dist.all_reduce(amax, 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 - if not bool(torch.any(amax)): - # No rank in the group had a (positive) amax -> uncalibrated across the whole group. Drop it so - # downstream ``amax is None`` handling applies (a real amax is strictly positive). The - # all-reduced value is identical on every rank, so this decision is symmetric. - delattr(self, "_amax") + 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): @@ -1601,6 +1597,23 @@ 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. diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 631894a7d12..7587554c3a6 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -68,3 +68,52 @@ def _test_uneven_amax_sync_helper(rank, size): def test_amax_sync_with_a_rank_that_never_calibrated(skip_on_windows): """The rank without an amax still joins the sync, contributing a zero on the group's device.""" spawn_multiprocess_job(2, _test_uneven_amax_sync_helper, backend="gloo") + + +def _sync(quantizer): + quantizer.sync_amax_across_distributed_group(DistributedProcessGroup(dist.group.WORLD)) + + +def _calibrated_zero_amax_helper(rank, size): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + quantizer._set_buffer("_amax", torch.tensor(0.0)) + _sync(quantizer) + assert quantizer.amax is not None + assert quantizer.amax.item() == 0.0 + dist.destroy_process_group() + + +def test_a_calibrated_zero_amax_survives_the_sync(skip_on_windows): + """An all-zero weight calibrates to amax 0, which is not "never calibrated".""" + spawn_multiprocess_job(2, _calibrated_zero_amax_helper, backend="gloo") + + +def _calibrators_dtype_helper(rank, size): + per_tensor = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + if rank == 0: + per_tensor._set_buffer("_amax", torch.tensor(3.0, 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)) + _sync(per_tensor) + _sync(per_channel) + assert per_tensor.amax.dtype == torch.bfloat16 + assert per_tensor.amax.item() == 3.0 + assert per_channel.amax.dtype == torch.bfloat16 + assert per_channel.amax.tolist() == [1.0, 3.0] + dist.destroy_process_group() + + +def test_the_sync_leaves_every_rank_the_calibrators_dtype(skip_on_windows): + """Including a rank that never calibrated the quantizer and so had no amax of its own.""" + spawn_multiprocess_job(2, _calibrators_dtype_helper, backend="gloo") + + +def _never_calibrated_helper(rank, size): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + _sync(quantizer) + assert getattr(quantizer, "_amax", None) is None + dist.destroy_process_group() + + +def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(skip_on_windows): + spawn_multiprocess_job(2, _never_calibrated_helper, backend="gloo") From 824a4966f465d6af06843efcc208f9cc9fb2bf47 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 10:38:41 +0000 Subject: [PATCH 04/11] Share expert input amax early only for HF fused experts; drop the cache release Sharing the routed experts' input amax over the expert-parallel group before the completeness check exists for HF fused experts, whose experts see only the tokens routed to them. It also applied to Megatron's experts, which then stopped failing on incomplete calibration and silently took the group max. Only modules that opt in (_share_input_amax_across_ep, set by the HF fused experts wrapper) share early; Megatron keeps main's check-then-sync. max_calibrate also ran gc.collect() and emptied the CUDA cache on every call, which AWQ and SmoothQuant make once per linear: ~45 s on an 8B model. The distributed forward loop already returns the cache after its last forward, right before the sync this was for. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 21 ++++------- .../torch/quantization/plugins/huggingface.py | 3 ++ tests/unit/torch/quantization/test_dist.py | 35 +++++++++++++++++++ 3 files changed, 45 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 644df6422d5..23580293b74 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -16,7 +16,6 @@ """Calibration utilities.""" import fnmatch -import gc import math import time import warnings @@ -390,17 +389,6 @@ def max_calibrate( _finalize_with_shared_state(model, weight_patterns) return - # The calibration forward leaves large activation tensors held in the caching allocator. The - # distributed amax sync below issues NCCL collectives whose comm buffers are allocated lazily via - # cudaMalloc (outside the caching allocator). Under the memory pressure of long calib sequences this - # can OOM on a single rank when NCCL first allocates -- and since sync_amax_across_distributed_group - # re-raises on collective failure, that rank aborts mid-collective and takes down the group. Return the - # cached activation memory to CUDA first so NCCL has room. (gc.collect() first to drop any cycle-held - # activation references the calibration closures/calibrators may still own.) - if torch.cuda.is_available(): - gc.collect() - torch.cuda.empty_cache() - # 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 @@ -410,9 +398,14 @@ def max_calibrate( # 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. + # 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 _has_expert_parallelism(module): + 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(): 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/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 7587554c3a6..ceb3f486bf4 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from functools import partial + +import pytest import torch import torch.distributed as dist from _test_utils.torch.distributed.utils import spawn_multiprocess_job @@ -117,3 +120,35 @@ def _never_calibrated_helper(rank, size): def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(skip_on_windows): spawn_multiprocess_job(2, _never_calibrated_helper, backend="gloo") + + +def _expert_input_amax_helper(rank, size, share): + from modelopt.torch.quantization.model_calib import max_calibrate + from modelopt.torch.utils.distributed import ParallelState + + 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) + dist.destroy_process_group() + + +@pytest.mark.parametrize("share", [False, True], ids=["megatron", "opted-in"]) +def test_only_modules_that_opt_in_share_expert_input_amax_before_the_check(skip_on_windows, share): + """Megatron's experts keep main's check-then-sync; HF fused experts share the input amax first.""" + spawn_multiprocess_job(2, partial(_expert_input_amax_helper, share=share), backend="gloo") From 33a9d66d0f1244763cb5a07d50da40cdf4a09dcb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 14:43:32 +0000 Subject: [PATCH 05/11] Skip the uneven per-channel amax sync test below two GPUs With one rank nobody lacks the amax, so nothing raises and pytest.raises fails inside the worker. That failure is a BaseException the worker pool does not catch: the worker dies without reporting, and every later test on the session's pool waits on it. CI's PR GPU runner has one GPU. Signed-off-by: Shengliang Xu --- tests/gpu/torch/quantization/test_amax_sync_distributed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gpu/torch/quantization/test_amax_sync_distributed.py b/tests/gpu/torch/quantization/test_amax_sync_distributed.py index ed6682e575c..e908031df25 100644 --- a/tests/gpu/torch/quantization/test_amax_sync_distributed.py +++ b/tests/gpu/torch/quantization/test_amax_sync_distributed.py @@ -107,7 +107,7 @@ def _per_channel_unevenly_calibrated(rank, size): q.sync_amax_across_distributed_group(_whole_world()) -def test_uneven_per_channel_calibration_fails_on_every_rank(dist_workers): +def test_uneven_per_channel_calibration_fails_on_every_rank(need_2_gpus, dist_workers): """A one-sided raise would leave the other rank blocked in the collective forever. The rank that HAS the amax cannot know the have-nots exist without asking, so the agreement is a From 31dc3318d7a0178607ae2c5581a372ab33cacbcb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 22:08:56 +0000 Subject: [PATCH 06/11] Run the amax-sync tests on the shared gloo workers Each of them spawned two ranks, about 14 s of the linux unit job apiece. Signed-off-by: Shengliang Xu --- tests/unit/torch/quantization/test_dist.py | 28 ++++++++-------------- 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index ceb3f486bf4..684a55ec244 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -13,12 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from functools import partial - import pytest import torch import torch.distributed as dist -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from _test_utils.torch.quantization.models import SimpleLinear from torch.nn.parallel import DistributedDataParallel @@ -65,12 +62,11 @@ def _test_uneven_amax_sync_helper(rank, size): quantizer._set_buffer("_amax", torch.tensor(7.0)) quantizer.sync_amax_across_distributed_group(DistributedProcessGroup(dist.group.WORLD)) assert quantizer.amax.item() == 7.0 - dist.destroy_process_group() -def test_amax_sync_with_a_rank_that_never_calibrated(skip_on_windows): +def test_amax_sync_with_a_rank_that_never_calibrated(cpu_dist_workers): """The rank without an amax still joins the sync, contributing a zero on the group's device.""" - spawn_multiprocess_job(2, _test_uneven_amax_sync_helper, backend="gloo") + cpu_dist_workers.run(_test_uneven_amax_sync_helper) def _sync(quantizer): @@ -83,12 +79,11 @@ def _calibrated_zero_amax_helper(rank, size): _sync(quantizer) assert quantizer.amax is not None assert quantizer.amax.item() == 0.0 - dist.destroy_process_group() -def test_a_calibrated_zero_amax_survives_the_sync(skip_on_windows): +def test_a_calibrated_zero_amax_survives_the_sync(cpu_dist_workers): """An all-zero weight calibrates to amax 0, which is not "never calibrated".""" - spawn_multiprocess_job(2, _calibrated_zero_amax_helper, backend="gloo") + cpu_dist_workers.run(_calibrated_zero_amax_helper) def _calibrators_dtype_helper(rank, size): @@ -103,23 +98,21 @@ def _calibrators_dtype_helper(rank, size): assert per_tensor.amax.item() == 3.0 assert per_channel.amax.dtype == torch.bfloat16 assert per_channel.amax.tolist() == [1.0, 3.0] - dist.destroy_process_group() -def test_the_sync_leaves_every_rank_the_calibrators_dtype(skip_on_windows): +def test_the_sync_leaves_every_rank_the_calibrators_dtype(cpu_dist_workers): """Including a rank that never calibrated the quantizer and so had no amax of its own.""" - spawn_multiprocess_job(2, _calibrators_dtype_helper, backend="gloo") + cpu_dist_workers.run(_calibrators_dtype_helper) def _never_calibrated_helper(rank, size): quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) _sync(quantizer) assert getattr(quantizer, "_amax", None) is None - dist.destroy_process_group() -def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(skip_on_windows): - spawn_multiprocess_job(2, _never_calibrated_helper, backend="gloo") +def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(cpu_dist_workers): + cpu_dist_workers.run(_never_calibrated_helper) def _expert_input_amax_helper(rank, size, share): @@ -145,10 +138,9 @@ def forward_loop(m): # only rank 0's experts receive tokens else: with pytest.raises(RuntimeError, match="MoE calibration incomplete"): max_calibrate(model, forward_loop) - dist.destroy_process_group() @pytest.mark.parametrize("share", [False, True], ids=["megatron", "opted-in"]) -def test_only_modules_that_opt_in_share_expert_input_amax_before_the_check(skip_on_windows, share): +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.""" - spawn_multiprocess_job(2, partial(_expert_input_amax_helper, share=share), backend="gloo") + cpu_dist_workers.run(_expert_input_amax_helper, share=share) From c2e1970e43cbd19c1ca6cd3a9f2840bfdcd57584 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 01:48:23 +0000 Subject: [PATCH 07/11] Join the tensor-parallel amax sync on every rank, calibrated or not The amax sync is a collective every rank of the group joins, and it handles a rank with no amax of its own. The tensor-parallel callers still skipped such a rank, for linears and for the KV-cache quantizers alike, so a rank that calibrated the quantizer waited in the sync for a peer that never joined it. They now gate on whether the quantizer is enabled, as the data-parallel and expert-parallel sync does: the config enables the same quantizers on every rank, and a disabled one, which nothing quantizes with, joins no collective. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 6 ++-- tests/unit/torch/quantization/test_dist.py | 40 +++++++++++++++++++--- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 23580293b74..acd0424f8df 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -494,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) @@ -540,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/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 684a55ec244..99e084950ca 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -13,6 +13,8 @@ # 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 @@ -21,9 +23,10 @@ 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 +from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState def _test_data_parallel_helper(rank, size): @@ -116,9 +119,6 @@ def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(cpu_dist_workers): def _expert_input_amax_helper(rank, size, share): - from modelopt.torch.quantization.model_calib import max_calibrate - from modelopt.torch.utils.distributed import ParallelState - model = torch.nn.Sequential(torch.nn.Linear(16, 16)) mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop=None) expert = model[0] @@ -144,3 +144,35 @@ def forward_loop(m): # only rank 0's experts receive tokens 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)) + + max_calibrate(model, forward_loop) + 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) From e1cd9479493428a060612989c8daaf913352ae01 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 06:58:19 +0000 Subject: [PATCH 08/11] Run the amax sync's multi-rank cases on gloo, where every PR runs them The GPU test file sized its world from the GPU count, so on the 1-GPU PR runner each "some ranks missing" case ran on one rank that held the amax, and its one two-rank case needed two GPUs: no PR exercised the multi-rank behavior. The sync logic runs on gloo through cpu_dist_workers instead: - the uneven per-channel raise, which every rank must reach together, on a group that gives up in seconds rather than hang; - a per-channel quantizer no rank calibrated, which is not an error; - an NVFP4 dynamic-block quantizer, which syncs as a per-tensor one; - the group's max over ranks that all calibrated, which no unit test checked: a rank keeping its own amax went unnoticed. The GPU file keeps the check only NCCL can make, where a rank without an amax puts the zero it contributes, and asks for two GPUs: with one, no rank lacks the amax. Signed-off-by: Shengliang Xu --- .../test_amax_sync_distributed.py | 143 +++--------------- tests/unit/torch/quantization/test_dist.py | 69 +++++++-- 2 files changed, 73 insertions(+), 139 deletions(-) diff --git a/tests/gpu/torch/quantization/test_amax_sync_distributed.py b/tests/gpu/torch/quantization/test_amax_sync_distributed.py index e908031df25..f6ce3033843 100644 --- a/tests/gpu/torch/quantization/test_amax_sync_distributed.py +++ b/tests/gpu/torch/quantization/test_amax_sync_distributed.py @@ -13,12 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Amax synchronization across a data-parallel calibration group. +"""Amax synchronization over NCCL. -Under data-parallel calibration a module can be exercised on some ranks and not others -- a sparsely -routed MoE expert, or any module a given rank's batch never reaches. Every rank must still join the -collective or the participant set mismatches and NCCL hangs, so these tests pin down what each rank -contributes when it has no local amax, and that the failure modes are symmetric across ranks. +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 @@ -30,123 +29,19 @@ from modelopt.torch.utils.distributed import DistributedProcessGroup -def _quantizer(cfg, amax=None, rank_device=None): - q = TensorQuantizer(cfg).to(rank_device) - if amax is not None: - q._set_buffer("_amax", torch.as_tensor(amax, dtype=torch.float32, device=rank_device)) - return q - - -def _whole_world(): - return DistributedProcessGroup(dist.group.WORLD) - - -# --------------------------------------------------------------------------------- per-tensor amax - - -def _per_tensor_all_ranks_present(rank, size): - """Every rank calibrated it: the group agrees on the max.""" - device = torch.device(f"cuda:{rank}") - q = _quantizer(QuantizerAttributeConfig(num_bits=8), amax=float(rank + 1), rank_device=device) - q.sync_amax_across_distributed_group(_whole_world()) - assert q.amax.item() == pytest.approx(float(size)) - - -def test_per_tensor_amax_takes_the_group_max(dist_workers): - dist_workers.run(_per_tensor_all_ranks_present) - - -def _per_tensor_some_ranks_missing(rank, size): - """A have-not rank contributes a locally-known scalar zero -- the identity for MAX.""" - device = torch.device(f"cuda:{rank}") - amax = 7.0 if rank == 0 else None - q = _quantizer(QuantizerAttributeConfig(num_bits=8), amax=amax, rank_device=device) - q.sync_amax_across_distributed_group(_whole_world()) - # Both ranks end up with rank 0's value; the have-not must NOT be left at None or zero. - assert q.amax.item() == pytest.approx(7.0) - - -def test_per_tensor_amax_survives_uneven_calibration(dist_workers): - """This is the common case -- a module only some ranks' batches reached.""" - dist_workers.run(_per_tensor_some_ranks_missing) - - -def _per_tensor_nobody_calibrated(rank, size): - device = torch.device(f"cuda:{rank}") - q = _quantizer(QuantizerAttributeConfig(num_bits=8), amax=None, rank_device=device) - q.sync_amax_across_distributed_group(_whole_world()) - assert q.amax is None, "an all-zero sync must leave the quantizer uncalibrated, not amax==0" - - -def test_uncalibrated_everywhere_stays_uncalibrated(dist_workers): - """amax==0 would quantize everything to zero; downstream expects None for 'never calibrated'.""" - dist_workers.run(_per_tensor_nobody_calibrated) - - -# --------------------------------------------------------------- static non-per-tensor (per-channel) - - -def _per_channel_all_present(rank, size): - device = torch.device(f"cuda:{rank}") - amax = torch.full((4,), float(rank + 1)) - q = _quantizer(QuantizerAttributeConfig(num_bits=8, axis=0), amax=amax, rank_device=device) - q.sync_amax_across_distributed_group(_whole_world()) - assert torch.allclose(q.amax.cpu(), torch.full((4,), float(size))) - - -def test_per_channel_amax_takes_the_elementwise_group_max(dist_workers): - dist_workers.run(_per_channel_all_present) - - -def _per_channel_unevenly_calibrated(rank, size): - """Have-nots cannot invent the right shape, so every rank must raise -- together.""" - device = torch.device(f"cuda:{rank}") - amax = torch.full((4,), 3.0) if rank == 0 else None - q = _quantizer(QuantizerAttributeConfig(num_bits=8, axis=0), amax=amax, rank_device=device) - with pytest.raises(RuntimeError, match="unevenly-calibrated"): - q.sync_amax_across_distributed_group(_whole_world()) - - -def test_uneven_per_channel_calibration_fails_on_every_rank(need_2_gpus, dist_workers): - """A one-sided raise would leave the other rank blocked in the collective forever. - - The rank that HAS the amax cannot know the have-nots exist without asking, so the agreement is a - single MAX all-reduce of a packed presence flag that every rank joins before anyone diverges. - """ - dist_workers.run(_per_channel_unevenly_calibrated) - - -def _per_channel_nobody_calibrated(rank, size): +def _per_tensor_job(rank, size, calibrated): device = torch.device(f"cuda:{rank}") - q = _quantizer(QuantizerAttributeConfig(num_bits=8, axis=0), amax=None, rank_device=device) - q.sync_amax_across_distributed_group(_whole_world()) - assert q.amax is None - - -def test_per_channel_uncalibrated_everywhere_is_not_an_error(dist_workers): - """A module no rank's calibration data exercised is expected, not a failure.""" - dist_workers.run(_per_channel_nobody_calibrated) - - -# ------------------------------------------------------------------------- dynamic block (NVFP4 act) - - -def _dynamic_block_some_ranks_missing(rank, size): - """A dynamic-block quantizer stores a SCALAR amax, so the per-tensor path applies.""" - device = torch.device(f"cuda:{rank}") - cfg = QuantizerAttributeConfig( - num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)} - ) - amax = 5.0 if rank == 0 else None - q = _quantizer(cfg, amax=amax, rank_device=device) - q.sync_amax_across_distributed_group(_whole_world()) - assert q.amax.item() == pytest.approx(5.0) - - -def test_dynamic_block_amax_is_treated_as_per_tensor(dist_workers): - """Its block_sizes drive runtime scales; only the per-tensor global scale is stored. - - Classifying it by ``block_sizes is not None`` would send an NVFP4 activation quantizer down the - static path and raise 'unevenly-calibrated' for the ordinary case of a module some ranks missed. - """ - dist_workers.run(_dynamic_block_some_ranks_missing) + 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 99e084950ca..90383eab1a5 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -89,33 +89,72 @@ def test_a_calibrated_zero_amax_survives_the_sync(cpu_dist_workers): cpu_dist_workers.run(_calibrated_zero_amax_helper) -def _calibrators_dtype_helper(rank, size): - per_tensor = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) +# 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: - per_tensor._set_buffer("_amax", torch.tensor(3.0, dtype=torch.bfloat16)) + 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)) - _sync(per_tensor) - _sync(per_channel) - assert per_tensor.amax.dtype == torch.bfloat16 - assert per_tensor.amax.item() == 3.0 - assert per_channel.amax.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] -def test_the_sync_leaves_every_rank_the_calibrators_dtype(cpu_dist_workers): - """Including a rank that never calibrated the quantizer and so had no amax of its own.""" - cpu_dist_workers.run(_calibrators_dtype_helper) +@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 _never_calibrated_helper(rank, size): - quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) +def _never_calibrated_helper(rank, size, axis): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=axis)) _sync(quantizer) assert getattr(quantizer, "_amax", None) is None -def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(cpu_dist_workers): - cpu_dist_workers.run(_never_calibrated_helper) +@pytest.mark.parametrize("axis", [None, 0], ids=["per-tensor", "per-channel"]) +def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(cpu_dist_workers, axis): + """A module no rank's calibration data reached, which is expected rather than an error.""" + cpu_dist_workers.run(_never_calibrated_helper, axis=axis) + + +def _uneven_per_channel_helper(rank, size): + # A rank left waiting in a collective its peer never joins gives up in seconds on this group. + group = DistributedProcessGroup(dist.new_group(timeout=timedelta(seconds=10))) + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=0)) + if rank == 0: + quantizer._set_buffer("_amax", torch.full((4,), 3.0)) + try: + quantizer.sync_amax_across_distributed_group(group) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + 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): From dfe42ab6d47d589b444e1d4cfadc9618e4b0623d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:54:26 +0000 Subject: [PATCH 09/11] Trim the amax-sync tests: one test for a zero amax vs none, one drop - test_amax_sync_with_a_rank_that_never_calibrated is the one-calibrated half of test_the_sync_leaves_every_rank_the_calibrators_dtype, which runs the same case in bf16 and checks the dtype too. - test_a_calibrated_zero_amax_survives_the_sync and test_a_quantizer_no_rank_calibrated_stays_uncalibrated become one test_the_sync_keeps_a_calibrated_zero_apart_from_uncalibrated over calibrated-zero, uncalibrated and uncalibrated-per-channel, each case keeping its assert. A have-not rank that keeps no amax or takes zero, a zero amax read as uncalibrated, nobody's amax becoming zero, and a per-channel quantizer nobody calibrated raising each fail a remaining test. Signed-off-by: Shengliang Xu --- tests/unit/torch/quantization/test_dist.py | 49 +++++++--------------- 1 file changed, 16 insertions(+), 33 deletions(-) diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 90383eab1a5..e7f0289ea1e 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -59,36 +59,10 @@ def test_data_parallel(cpu_dist_workers): cpu_dist_workers.run(_test_data_parallel_helper) -def _test_uneven_amax_sync_helper(rank, size): - quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) - if rank == 0: - quantizer._set_buffer("_amax", torch.tensor(7.0)) - quantizer.sync_amax_across_distributed_group(DistributedProcessGroup(dist.group.WORLD)) - assert quantizer.amax.item() == 7.0 - - -def test_amax_sync_with_a_rank_that_never_calibrated(cpu_dist_workers): - """The rank without an amax still joins the sync, contributing a zero on the group's device.""" - cpu_dist_workers.run(_test_uneven_amax_sync_helper) - - def _sync(quantizer): quantizer.sync_amax_across_distributed_group(DistributedProcessGroup(dist.group.WORLD)) -def _calibrated_zero_amax_helper(rank, size): - quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) - quantizer._set_buffer("_amax", torch.tensor(0.0)) - _sync(quantizer) - assert quantizer.amax is not None - assert quantizer.amax.item() == 0.0 - - -def test_a_calibrated_zero_amax_survives_the_sync(cpu_dist_workers): - """An all-zero weight calibrates to amax 0, which is not "never calibrated".""" - cpu_dist_workers.run(_calibrated_zero_amax_helper) - - # 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 = { @@ -124,16 +98,25 @@ def test_the_sync_leaves_every_rank_the_calibrators_dtype(cpu_dist_workers, per_ cpu_dist_workers.run(_calibrators_dtype_helper, per_tensor_cfg=per_tensor_cfg) -def _never_calibrated_helper(rank, size, axis): - quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=axis)) +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) - assert getattr(quantizer, "_amax", None) is None + 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("axis", [None, 0], ids=["per-tensor", "per-channel"]) -def test_a_quantizer_no_rank_calibrated_stays_uncalibrated(cpu_dist_workers, axis): - """A module no rank's calibration data reached, which is expected rather than an error.""" - cpu_dist_workers.run(_never_calibrated_helper, axis=axis) +@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): From 6e3670cf4d1479aaae149fc9c963e69fc4bec1c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 19:05:25 +0000 Subject: [PATCH 10/11] Date test_amax_sync_distributed 2026, as LICENSE_HEADER does Signed-off-by: Shengliang Xu --- tests/gpu/torch/quantization/test_amax_sync_distributed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gpu/torch/quantization/test_amax_sync_distributed.py b/tests/gpu/torch/quantization/test_amax_sync_distributed.py index f6ce3033843..614ab87e658 100644 --- a/tests/gpu/torch/quantization/test_amax_sync_distributed.py +++ b/tests/gpu/torch/quantization/test_amax_sync_distributed.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); From 67609291d084283a557ab43812a2a28bab871be7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 00:33:42 +0000 Subject: [PATCH 11/11] Destroy the process groups the amax-sync tests' pooled jobs create The two jobs that sync on a group with a 10-second timeout left it, with its gloo backend, registered on the session's persistent cpu_dist_workers until the pool shut down; they now destroy it. Signed-off-by: Shengliang Xu --- tests/unit/torch/quantization/test_dist.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index e7f0289ea1e..07d298f3de1 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -120,17 +120,19 @@ def test_the_sync_keeps_a_calibrated_zero_apart_from_uncalibrated(cpu_dist_worke def _uneven_per_channel_helper(rank, size): - # A rank left waiting in a collective its peer never joins gives up in seconds on this group. - group = DistributedProcessGroup(dist.new_group(timeout=timedelta(seconds=10))) 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(group) + 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 @@ -188,7 +190,10 @@ def forward_loop(m): # only rank 0 calibrates the synced quantizers for quantizer in synced: quantizer(torch.ones(2, 16)) - max_calibrate(model, forward_loop) + 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