[OMNIML-6072] [2/21] Make amax synchronization correct under data-parallel calibration

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
                     )

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.
 

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:

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)