From 8806f91c05dc7037ab03224e4139bc66c7c197b1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:22:30 +0000 Subject: [PATCH 01/60] Calibrate a distributed model on per-rank data The pieces a calibration loop needs once the model is spread across ranks: - shard_dataloader gives each rank its own data shard: a DistributedSampler for map-style datasets, with drop_last off so every rank runs the same number of steps (a rank that exits early hangs the others on the per-step collectives), and a deterministic strided copy of the stream for an IterableDataset. - distributed_forward_loop builds the forward_loop mtq.quantize calls: it moves each batch to this rank's device and runs the full distributed forward, so the expert-parallel collectives fire inside it, with KV caching off as on the single-process path. It re-plants the torch_a2a forward wrappers that quantizer insertion clobbers on transformers >= 5.16 (on the session's EP mesh), and returns the allocator's reserved pool both before the first forward and after the last -- NCCL's communication buffers live outside the torch allocator, so a full reserved pool starves them. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 191 ++++++++++++++++++++++++++++ 1 file changed, 191 insertions(+) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 26892ef112d..9ac138d42fa 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -16,7 +16,9 @@ """Utility functions for using torch.distributed.""" import functools +import gc import io +import itertools import os import sys import time @@ -32,6 +34,9 @@ import torch.distributed from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard from torch.distributed.tensor import DTensor +from torch.utils.data import DataLoader, IterableDataset +from torch.utils.data.distributed import DistributedSampler +from tqdm import tqdm __all__ = [ "DistributedProcessGroup", @@ -42,6 +47,7 @@ "backend", "barrier", "broadcast_from", + "distributed_forward_loop", "distributed_run", "distribution_plan", "fsdp2_wrap", @@ -55,6 +61,7 @@ "is_master", "rank", "set_distribution_plan", + "shard_dataloader", "size", ] @@ -666,6 +673,190 @@ def distributed_run( _destroy_group(exc) +class _StridedIterableShard(IterableDataset): + """Rank-strided shard of a streaming ``IterableDataset`` (no ``len()`` required). + + ``DistributedSampler`` only works with map-style datasets. Streaming calibration datasets (e.g. the + Nemotron image-text ``NemotronTarPlusJsonlIterable``) are iterable, so each rank instead builds its + own identical stream (deterministic seed) and takes slot ``rank`` from every complete group of + ``world_size`` consecutive samples. Any trailing incomplete group is dropped so every rank yields + exactly the same count -- preserving the equal-per-rank-batch invariant that keeps per-step EP + collectives from hanging when one rank runs short. Assumes single-worker iteration (these loaders use + ``num_workers=0``); with multiple workers the base stream would be duplicated per worker. + """ + + def __init__(self, base: IterableDataset, rank: int, world_size: int): + self.base = base + self.rank = rank + self.world_size = world_size + + def __iter__(self): + it = iter(self.base) + while True: + group = list(itertools.islice(it, self.world_size)) + if len(group) < self.world_size: + return # drop incomplete trailing group -> equal counts across ranks + yield group[self.rank] + + +def shard_dataloader(loader, rank: int, world_size: int): + """Wrap a DataLoader so each rank sees a unique shard. + + For map-style datasets a ``DistributedSampler`` is used. ``drop_last=False`` keeps per-rank batch + counts equal (else a rank exits calibration early and hangs the others on the per-step collectives), + at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. For + streaming ``IterableDataset`` datasets (which have no ``len()`` and so can't use a sampler) each rank + strides its own deterministic copy of the stream via :class:`_StridedIterableShard`. + + Forwards all non-sampler DataLoader settings from ``loader`` (workers, pinning, + prefetch, init fn, generator, ...). + """ + if isinstance(loader.dataset, IterableDataset): + return DataLoader( + _StridedIterableShard(loader.dataset, rank, world_size), + batch_size=loader.batch_size, + collate_fn=loader.collate_fn, + num_workers=loader.num_workers, + pin_memory=loader.pin_memory, + timeout=loader.timeout, + worker_init_fn=loader.worker_init_fn, + multiprocessing_context=loader.multiprocessing_context, + generator=loader.generator, + prefetch_factor=loader.prefetch_factor, + persistent_workers=loader.persistent_workers, + pin_memory_device=getattr(loader, "pin_memory_device", ""), + ) + sampler = DistributedSampler( + loader.dataset, + num_replicas=world_size, + rank=rank, + shuffle=False, + drop_last=False, + ) + return DataLoader( + loader.dataset, + batch_size=loader.batch_size, + sampler=sampler, + collate_fn=loader.collate_fn, + num_workers=loader.num_workers, + pin_memory=loader.pin_memory, + timeout=loader.timeout, + worker_init_fn=loader.worker_init_fn, + multiprocessing_context=loader.multiprocessing_context, + generator=loader.generator, + prefetch_factor=loader.prefetch_factor, + persistent_workers=loader.persistent_workers, + pin_memory_device=getattr(loader, "pin_memory_device", ""), + ) + + +def _reinstall_a2a_forwards(model, mo_session) -> int: + """Re-plant the torch_a2a TP forward wrappers that quantizer insertion clobbered. + + transformers installs its TP input/output transforms by REPLACING ``module.forward`` with an + instance attribute (TensorParallelLayer.install_forward). For torch_a2a_experts that wrapper is + what passes the EP ``process_group`` into the experts forward. ModelOpt's quantizer insertion + rebinds ``module.forward`` on the same instance, dropping the wrapper -- after which + ``torch_a2a_experts_forward`` receives ``process_group=None`` and raises. + + Called at the top of the calibration loop: after conversion (so it wraps the QUANTIZED forward) + and before the first forward. Idempotent via a per-module flag, so repeated calibration passes + do not stack wrappers (which would inject the process_group argument twice). + """ + ep_mesh = getattr(mo_session, "ep_mesh", None) if mo_session is not None else None + if ep_mesh is None: + return 0 + try: + from transformers.distributed.tensor_parallel import ALL_PARALLEL_STYLES + except ImportError: # older transformers + from transformers.integrations.tensor_parallel import ALL_PARALLEL_STYLES + try: + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + except ImportError: + from transformers.integrations.tensor_parallel import _get_parameter_tp_plan + + plan = getattr(model, "tp_plan", None) or {} + if not plan: + return 0 + + # install_forward() -- and the clobber it causes -- only exist on transformers >=5.16. Older + # releases install their TP input/output transforms as forward PRE-HOOKS, which survive + # ModelOpt rebinding ``module.forward``, so there is nothing to re-plant and calling the + # missing method would AttributeError. No-op rather than fail on 5.9/5.13. + if not any( + hasattr(ALL_PARALLEL_STYLES.get(s), "install_forward") + for s in ("torch_a2a_experts", "torch_a2a_router") + ): + return 0 + n = 0 + for name, mod in model.named_modules(): + if getattr(mod, "_modelopt_a2a_forward_reinstalled", False): + continue + style_name = _get_parameter_tp_plan(parameter_name=name, tp_plan=plan, is_weight=False) + if ( + style_name in ("torch_a2a_experts", "torch_a2a_router") + and style_name in ALL_PARALLEL_STYLES + ): + ALL_PARALLEL_STYLES[style_name].install_forward(mod, ep_mesh) + mod._modelopt_a2a_forward_reinstalled = True + n += 1 + if n and (mo_session is None or getattr(mo_session, "rank", 0) == 0): + print( + f"[ep] re-installed {n} torch_a2a forward wrapper(s) clobbered by quantizer insertion", + flush=True, + ) + return n + + +def distributed_forward_loop(model, dataloader, mo_session: "ModelOptSession | None" = None): + """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. + + Moves each batch onto this rank's ``mo_session.device`` and runs the captured ``model`` forward (the + EP all-reduce / all-to-all collectives fire inside that forward). ``mtq.quantize`` hands the + ``forward_loop`` callback an inner module that this closure ignores, calling the captured + top-level ``model`` instead so the full distributed forward (and its collectives) runs. + + KV/state caching is disabled for the duration (via ``_disable_use_cache``), mirroring the + single-node ``_forward_loop``. This distributed path previously omitted the guard, so it ran + calibration at ``config.use_cache`` (default ``True``); for hybrid Mamba/attention models + (e.g. NemotronH) the cache is mutated in-place, which both wastes memory and breaks + calibration correctness. + """ + # Local import to avoid any import cycle within modelopt.torch.utils. + from modelopt.torch.utils.dataset_utils import _disable_use_cache + + device = mo_session.device if mo_session is not None else None + + def calibrate(_inner_model): + _reinstall_a2a_forwards(model, mo_session) + with _disable_use_cache(model): + # Reclaim the caching allocator's reserved-but-unused pool left by the load/quantizer-insertion + # memory peak. Without this, torch hoards ~the whole GPU (e.g. 91 GB live weights vs 275 GB + # reserved), leaving ~0 GB free, so the distributed a2a's NCCL comm buffers cannot be allocated + # ("Failed to CUDA calloc") at the very first all_to_all. Freeing the cache hands that memory + # back to the driver for NCCL + the calibration activations. + gc.collect() + torch.cuda.empty_cache() + for batch in tqdm(dataloader, desc="Calibrating", disable=not is_master()): + if device is not None: + batch = { + k: (v.to(device) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } + model(**batch) + # Symmetric to the pre-forward reclaim above, and for the same reason one stage later: + # the calibration forwards leave their activations in the caching allocator's reserved + # pool, so whatever runs next (amax sync, export collectives) can find the device full + # even though little is live. Measured on Kimi-K3 (EP64, 184 GiB cards) right after the + # single calib forward: alloc 86.0 GiB but reserved 156.6 GiB -> only 0.7 GiB free, and + # the job then hung until the 2 h wall clock instead of exporting. NCCL comm buffers are + # allocated OUTSIDE the torch allocator, so a full reserved pool starves them. + gc.collect() + torch.cuda.empty_cache() + + return calibrate + + def cleanup(): """Cleans up the distributed environment. From 49f76a059a48ef4fef4824082cda5dbc7eef9540 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:19:16 +0000 Subject: [PATCH 02/60] Test the distributed calibration loop Cover how every rank calibrates on its own data and keeps the all-to-all in place: - shard_dataloader: strided shards that together cover the data exactly once, equal batch counts on an uneven split, and every loader setting carried over, for map-style and streaming datasets; - _reinstall_a2a_forwards: re-planting the expert-parallel wrappers that quantizer insertion drops, on the planned modules only and without stacking; - distributed_forward_loop: this rank's shard through the top-level model, use_cache off while it runs, the allocator cache freed around it. Signed-off-by: Shengliang Xu --- .../utils/test_distributed_calibration.py | 475 ++++++++++++++++++ 1 file changed, 475 insertions(+) create mode 100644 tests/unit/torch/utils/test_distributed_calibration.py diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py new file mode 100644 index 00000000000..8d908571a78 --- /dev/null +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -0,0 +1,475 @@ +# 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. + +"""Unit tests for distributed calibration: per-rank data shards and the distributed forward loop. + +Each rank of a distributed (e.g. expert-parallel) PTQ run calibrates on its own shard of the data, +and the EP collectives inside every forward hang unless all ranks run the same number of batches. +The torch_a2a TP forward wrappers come from transformers >= 5.16 and need a process group, so they +are faked here; everything runs in one CPU process. +""" + +import re +import sys +import types +from collections import Counter + +import pytest +import torch +from torch import nn +from torch.utils.data import DataLoader, IterableDataset + +import modelopt.torch.utils.distributed as mto_dist +from modelopt.torch.utils.distributed import ( + ModelOptSession, + _reinstall_a2a_forwards, + distributed_forward_loop, + modelopt_session, + shard_dataloader, +) + +# ------------------------------------------------------------------------------------- data helpers + + +def _collate(ids): + """Calibration-style collate: a dict of batched tensors, as ``model(**batch)`` expects.""" + return {"input_ids": torch.tensor(ids).unsqueeze(-1)} + + +class _Stream(IterableDataset): + """A deterministic stream with no ``len()``, like the streaming image-text calibration sets.""" + + def __init__(self, n): + self.n = n + + def __iter__(self): + return iter(range(self.n)) + + +def _loader(dataset, batch_size=2): + return DataLoader(dataset, batch_size=batch_size, collate_fn=_collate) + + +def _batches(loader): + """The sample ids in each batch the loader yields.""" + return [batch["input_ids"].flatten().tolist() for batch in loader] + + +def _ids(loader): + return [i for batch in _batches(loader) for i in batch] + + +def _seed_worker(worker_id): + torch.manual_seed(worker_id) + + +# --------------------------------------------------------------------------------- shard_dataloader + + +@pytest.mark.parametrize("world_size", [2, 3]) +def test_map_style_shards_are_strided_and_cover_every_sample_once(world_size): + """Rank ``r`` gets samples ``r, r + W, ...``, batched with the original size and collate_fn.""" + n = 12 + shards = [shard_dataloader(_loader(list(range(n))), r, world_size) for r in range(world_size)] + + assert sorted(i for shard in shards for i in _ids(shard)) == list(range(n)) + for r, shard in enumerate(shards): + ids = list(range(r, n, world_size)) + assert _batches(shard) == [ids[i : i + 2] for i in range(0, len(ids), 2)] + + +def test_map_style_uneven_split_pads_every_rank_to_the_same_batch_count(): + """No sample is dropped and no rank runs short (it would hang the others' EP collectives). + + The price of ``drop_last=False`` is that the sampler repeats fewer than ``world_size`` samples. + """ + n, world_size = 10, 4 + shards = [ + _batches(shard_dataloader(_loader(list(range(n))), r, world_size)) + for r in range(world_size) + ] + + assert [len(batches) for batches in shards] == [2] * world_size # 3 samples per rank + seen = Counter(i for batches in shards for batch in batches for i in batch) + assert set(seen) == set(range(n)) + assert seen.total() - n == 2 # padding up to 3 * world_size samples + + +def test_iterable_shards_are_strided_cover_every_sample_once_and_replay(): + """Each rank strides its own copy of a stream; a second pass replays the same shard. + + Multi-pass calibration algorithms (e.g. AWQ) iterate the loader again. + """ + n, world_size = 12, 3 + shards = [shard_dataloader(_loader(_Stream(n)), r, world_size) for r in range(world_size)] + + assert sorted(i for shard in shards for i in _ids(shard)) == list(range(n)) + for r, shard in enumerate(shards): + assert _batches(shard) == _batches(shard) == [[r, r + 3], [r + 6, r + 9]] + + +def test_iterable_uneven_split_drops_the_trailing_partial_group(): + """Equal counts without ``len()``: the last ``n % world_size`` samples go to no rank.""" + n, world_size = 10, 4 + shards = [_ids(shard_dataloader(_loader(_Stream(n)), r, world_size)) for r in range(world_size)] + + assert shards == [[0, 4], [1, 5], [2, 6], [3, 7]] + + +@pytest.mark.parametrize("dataset", [list(range(5)), _Stream(5)], ids=["map-style", "iterable"]) +def test_world_size_one_yields_the_original_batches(dataset): + """A lone rank's shard is the whole (unshuffled) loader, batch for batch.""" + loader = _loader(dataset) + assert _batches(shard_dataloader(loader, 0, 1)) == _batches(loader) == [[0, 1], [2, 3], [4]] + + +@pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) +def test_sharded_loader_keeps_the_original_loader_settings(dataset): + """Workers, pinning, timeout, worker init, RNG and prefetching carry over to the shard.""" + loader = DataLoader( + dataset, + batch_size=4, + collate_fn=_collate, + num_workers=2, + pin_memory=True, + timeout=7.0, + worker_init_fn=_seed_worker, + multiprocessing_context="spawn", + generator=torch.Generator(), + prefetch_factor=3, + persistent_workers=True, + ) + shard = shard_dataloader(loader, 1, 2) + + for setting in ( + "batch_size", + "collate_fn", + "num_workers", + "pin_memory", + "timeout", + "worker_init_fn", + "multiprocessing_context", + "generator", + "prefetch_factor", + "persistent_workers", + ): + assert getattr(shard, setting) == getattr(loader, setting), setting + + +# -------------------------------------------------------------------------- _reinstall_a2a_forwards + +_TP_MODULE = "transformers.distributed.tensor_parallel" +_LEGACY_TP_MODULE = "transformers.integrations.tensor_parallel" + + +class _FakeMesh: + """The EP device mesh; the TP wrappers only need its process group.""" + + def __init__(self): + self.group = object() + + def get_group(self): + return self.group + + +class _WrapperStyle: + """A transformers >= 5.16 TP style: ``install_forward`` wraps the module's current forward.""" + + def __init__(self, inject_group=False): + self.inject_group = inject_group + self.installed = [] + + def install_forward(self, module, mesh): + self.installed.append(module) + inner = module.forward + extra = (mesh.get_group(),) if self.inject_group else () + + def forward(*args, **kwargs): + return inner(*args, *extra, **kwargs) + + module.forward = forward + + +class _PreHookStyle: + """A transformers < 5.16 TP style: its transforms are forward pre-hooks, no ``install_forward``.""" + + +def _get_parameter_tp_plan(parameter_name, tp_plan, is_weight=True): + """transformers' plan lookup: layer indices match ``*``; a weight also takes its module's entry.""" + name = re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), parameter_name) + if name in tp_plan: + return tp_plan[name] + if is_weight and "." in name and (module_name := name.rsplit(".", 1)[0]) in tp_plan: + return tp_plan[module_name] + return None + + +def _install_fake_tp(monkeypatch, styles=None, where=_TP_MODULE): + """Serve ``styles`` as transformers' ``ALL_PARALLEL_STYLES`` from the module ``where``.""" + if styles is None: + styles = { + "torch_a2a_experts": _WrapperStyle(inject_group=True), + "torch_a2a_router": _WrapperStyle(), + "colwise": _WrapperStyle(), + } + module = types.ModuleType(where) + module.ALL_PARALLEL_STYLES = styles + module._get_parameter_tp_plan = _get_parameter_tp_plan + if where != _TP_MODULE: + monkeypatch.setitem(sys.modules, _TP_MODULE, None) # importing it now raises ImportError + monkeypatch.setitem(sys.modules, where, module) + return styles + + +class _Experts(nn.Module): + """A torch_a2a experts module: its forward needs the EP group its TP wrapper injects.""" + + def __init__(self): + super().__init__() + self.weight_quantizer = nn.Identity() # quantizer insertion adds children like this + self.groups = [] + + def forward(self, hidden_states, top_k_index, top_k_weights, process_group=None): + if process_group is None: + raise ValueError("torch_a2a experts forward requires the EP `process_group`") + self.groups.append(process_group) + return hidden_states + + +class _Block(nn.Module): + def __init__(self): + super().__init__() + self.gate = nn.Identity() # the router + self.proj = nn.Identity() # a dense layer under a non-a2a TP style + self.experts = _Experts() + + def forward(self, hidden_states): + return self.experts(self.proj(self.gate(hidden_states)), None, None) + + +class _MoEModel(nn.Module): + """Two MoE layers whose plan sends the experts and routers through torch_a2a.""" + + def __init__(self, with_tp_plan=True): + super().__init__() + self.layers = nn.ModuleList([_Block(), _Block()]) + self.config = types.SimpleNamespace(use_cache=True) + self.tp_plan = ( + { + "layers.*.gate": "torch_a2a_router", + "layers.*.proj": "colwise", + "layers.*.experts": "torch_a2a_experts", + } + if with_tp_plan + else None + ) + + def forward(self, input_ids): + hidden_states = input_ids.float() + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +def _ep_session(rank=0, *, with_mesh=True): + """An EP rank's session: its identity plus the EP mesh the model loader recorded.""" + session = ModelOptSession() + session.initialize( + distributed=True, rank=rank, world_size=2, local_rank=rank, device=None, is_main=rank == 0 + ) + session.ep_mesh = _FakeMesh() if with_mesh else None + return session + + +def _drop_tp_wrapper(module): + """Do what quantizer insertion does to the instance: rebind the class forward over the wrapper. + + (``DynamicModule.convert`` -> ``bind_forward_method``.) + """ + module.forward = types.MethodType(type(module).forward, module) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_reinstall_replants_the_ep_wrapper_on_the_a2a_modules_only(monkeypatch, capsys, rank): + """Every rank re-wraps exactly the planned torch_a2a experts and routers; rank 0 logs the count. + + Other TP styles keep their forward, and so do the quantizers under an experts module: the plan + is looked up per module (``is_weight=False``), so children do not inherit their parent's style. + """ + styles = _install_fake_tp(monkeypatch) + model, session = _MoEModel(), _ep_session(rank) + hidden_states = torch.ones(2, 4) + with pytest.raises(ValueError, match="process_group"): + model(hidden_states) # as quantizer insertion leaves it: no wrapper, no EP group + + assert _reinstall_a2a_forwards(model, session) == 4 + + assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] + assert styles["torch_a2a_router"].installed == [layer.gate for layer in model.layers] + assert styles["colwise"].installed == [] + model(hidden_states) + assert [layer.experts.groups for layer in model.layers] == [[session.ep_mesh.group]] * 2 + assert ("re-installed 4" in capsys.readouterr().out) == (rank == 0) + + +def test_reinstall_does_not_stack_wrappers_across_calibration_passes(monkeypatch): + """A repeat call is a no-op; a second wrapper would pass the experts the EP group twice.""" + styles = _install_fake_tp(monkeypatch) + model, session = _MoEModel(), _ep_session() + assert _reinstall_a2a_forwards(model, session) == 4 + + assert _reinstall_a2a_forwards(model, session) == 0 + assert len(styles["torch_a2a_experts"].installed) == 2 + model(torch.ones(2, 4)) + assert [layer.experts.groups for layer in model.layers] == [[session.ep_mesh.group]] * 2 + + +@pytest.mark.parametrize( + ("has_session", "has_ep_mesh", "has_tp_plan"), + [(False, True, True), (True, False, True), (True, True, False)], + ids=["no-session", "no-ep-mesh", "no-tp-plan"], +) +def test_reinstall_is_a_noop_off_the_expert_parallel_path( + monkeypatch, has_session, has_ep_mesh, has_tp_plan +): + """A single-process or non-EP model is left untouched, even with the torch_a2a styles around.""" + styles = _install_fake_tp(monkeypatch) + model = _MoEModel(with_tp_plan=has_tp_plan) + session = _ep_session(with_mesh=has_ep_mesh) if has_session else None + + assert _reinstall_a2a_forwards(model, session) == 0 + assert not any(style.installed for style in styles.values()) + + +def test_reinstall_is_a_noop_where_tp_styles_are_pre_hooks(monkeypatch): + """transformers < 5.16 has no ``install_forward``; its pre-hooks survive quantizer insertion. + + Those releases keep the styles in ``transformers.integrations``: the fallback import. + """ + _install_fake_tp( + monkeypatch, + styles={"torch_a2a_experts": _PreHookStyle(), "torch_a2a_router": _PreHookStyle()}, + where=_LEGACY_TP_MODULE, + ) + assert _reinstall_a2a_forwards(_MoEModel(), _ep_session()) == 0 + + +# ------------------------------------------------------------------------- distributed_forward_loop + + +class _Recorder(nn.Module): + """Records each forward's kwargs and the ``use_cache`` flag it ran under.""" + + def __init__(self, on_forward=None): + super().__init__() + self.config = types.SimpleNamespace(use_cache=True) + self.calls = [] + self.use_cache_seen = [] + self.on_forward = on_forward + + def forward(self, **batch): + self.calls.append(batch) + self.use_cache_seen.append(self.config.use_cache) + if self.on_forward is not None: + self.on_forward() + + +@pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_forward_loop_runs_this_ranks_shard_through_the_captured_model(dataset, rank): + """Every batch of the rank's shard reaches the top-level model, not the module ``mtq`` passes. + + Only the full model's forward runs the EP collectives; no session means no device move. + """ + model, inner = _Recorder(), _Recorder() + shard = shard_dataloader(_loader(dataset), rank, world_size=2) + + distributed_forward_loop(model, shard)(inner) + + calls = [call["input_ids"].flatten().tolist() for call in model.calls] + assert calls == [[rank, rank + 2], [rank + 4, rank + 6]] + assert inner.calls == [] + + +@pytest.mark.parametrize(("device", "expected"), [(None, "cpu"), ("meta", "meta")]) +def test_forward_loop_moves_tensors_to_the_single_process_session_device( + monkeypatch, device, expected +): + """Batch tensors go to ``mo_session.device`` when it names one; other values pass through. + + The meta device makes the move observable on a CPU-only host. + """ + monkeypatch.delenv("WORLD_SIZE", raising=False) # single process: no process group + batches = [{"input_ids": torch.zeros(2, 3), "logits_to_keep": 1}] * 2 + model = _Recorder() + with modelopt_session(device=device) as session: + distributed_forward_loop(model, batches, session)(model) + + assert [call["input_ids"].device.type for call in model.calls] == [expected] * 2 + assert [call["logits_to_keep"] for call in model.calls] == [1, 1] + + +def test_forward_loop_disables_use_cache_only_while_calibrating(): + """KV/state caching is off for the calibration forwards (it breaks hybrid Mamba models).""" + model = _Recorder() + distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}] * 2)(model) + + assert model.use_cache_seen == [False, False] + assert model.config.use_cache is True + + +def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch): + """Re-plant the a2a wrappers and free the cached pool before the first forward, and again after. + + Building the loop does neither: quantizers are inserted between building and calling it. NCCL + allocates its buffers outside torch's caching allocator, so a pool hoarded by the model load or + by the calibration activations starves the next collective. + """ + events = [] + monkeypatch.setattr( + mto_dist, "_reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args)) + ) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: events.append("empty_cache")) + model, session = _Recorder(on_forward=lambda: events.append("forward")), _ep_session() + forward_loop = distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}] * 2, session) + assert events == [] + + forward_loop(model) + + assert events == [ + ("reinstall", model, session), + "empty_cache", + "forward", + "forward", + "empty_cache", + ] + + +def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypatch): + """hf_ptq's order -- build the loop, insert quantizers, calibrate -- hands the experts the EP group. + + It holds on every pass (multi-pass algorithms call the loop again), with one wrapper per module. + """ + styles = _install_fake_tp(monkeypatch) + model, session = _MoEModel(), _ep_session() + forward_loop = distributed_forward_loop(model, [{"input_ids": torch.ones(2, 4)}] * 3, session) + for layer in model.layers: + _drop_tp_wrapper(layer.experts) + + forward_loop(model) + forward_loop(model) + + assert [layer.experts.groups for layer in model.layers] == [[session.ep_mesh.group] * 6] * 2 + assert len(styles["torch_a2a_experts"].installed) == 2 From 95bd7a704220f32e50246530c25b987f95413c78 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:19:20 +0000 Subject: [PATCH 03/60] Split a streaming shard's groups across DataLoader workers Each worker process iterates its own copy of the stream, and every one of them yielded the rank's whole shard, so with num_workers=2 each sample was calibrated on twice. Have each worker take every num_workers-th group: the shards still cover the data once, with equal counts on every rank. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 13 ++++++++----- .../utils/test_distributed_calibration.py | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 9ac138d42fa..9f3dd0bbf95 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -34,7 +34,7 @@ import torch.distributed from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard from torch.distributed.tensor import DTensor -from torch.utils.data import DataLoader, IterableDataset +from torch.utils.data import DataLoader, IterableDataset, get_worker_info from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm @@ -681,8 +681,8 @@ class _StridedIterableShard(IterableDataset): own identical stream (deterministic seed) and takes slot ``rank`` from every complete group of ``world_size`` consecutive samples. Any trailing incomplete group is dropped so every rank yields exactly the same count -- preserving the equal-per-rank-batch invariant that keeps per-step EP - collectives from hanging when one rank runs short. Assumes single-worker iteration (these loaders use - ``num_workers=0``); with multiple workers the base stream would be duplicated per worker. + collectives from hanging when one rank runs short. Each DataLoader worker process iterates its own + copy of the stream, so the workers split the groups between them rather than each yielding them all. """ def __init__(self, base: IterableDataset, rank: int, world_size: int): @@ -691,12 +691,15 @@ def __init__(self, base: IterableDataset, rank: int, world_size: int): self.world_size = world_size def __iter__(self): + info = get_worker_info() + num_workers, worker = (info.num_workers, info.id) if info is not None else (1, 0) it = iter(self.base) - while True: + for index in itertools.count(): group = list(itertools.islice(it, self.world_size)) if len(group) < self.world_size: return # drop incomplete trailing group -> equal counts across ranks - yield group[self.rank] + if index % num_workers == worker: + yield group[self.rank] def shard_dataloader(loader, rank: int, world_size: int): diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 8d908571a78..9809f2d51d3 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -120,6 +120,24 @@ def test_iterable_shards_are_strided_cover_every_sample_once_and_replay(): assert _batches(shard) == _batches(shard) == [[r, r + 3], [r + 6, r + 9]] +def test_iterable_shards_split_their_groups_across_worker_processes(): + """Each worker process iterates its own copy of the stream. Unless the workers split the groups + between them, every one yields the whole shard and each sample comes out once per worker.""" + n, world_size = 16, 2 + shards = [ + _ids( + shard_dataloader( + DataLoader(_Stream(n), batch_size=2, collate_fn=_collate, num_workers=2), + r, + world_size, + ) + ) + for r in range(world_size) + ] + + assert [sorted(ids) for ids in shards] == [list(range(r, n, world_size)) for r in range(2)] + + def test_iterable_uneven_split_drops_the_trailing_partial_group(): """Equal counts without ``len()``: the last ``n % world_size`` samples go to no rank.""" n, world_size = 10, 4 From 47402e88fd4f1c1b23871d31a9ac8f3c8b2a7814 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 22:21:58 +0000 Subject: [PATCH 04/60] Agree on the calibration batch size across ranks hf_ptq's default --batch_size 0 sizes calibration with get_max_batch_size, which each rank ran on its own: it measured every visible GPU, probed with a forward, and on an out-of-memory error halved the batch and probed again. Under expert parallelism those forwards hold all-to-alls, so a rank that retried ran more of them than the rest, and ranks could settle on different batch sizes. Different sizes mean different numbers of calibration steps, and a rank that finishes early leaves the others waiting in a collective. Inside a distributed session get_max_batch_size is now collective: each rank measures its own GPU only and probes with a single forward, and every rank takes the smallest size any rank allows, through the new min-reduced session.calib_batch_size. Outside a session it is unchanged. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/dataset_utils.py | 48 +++++++++++++++-- tests/unit/torch/utils/test_dataset_utils.py | 57 ++++++++++++++++++++ 2 files changed, 100 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 1b82a045390..8f502f761ec 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1006,16 +1006,54 @@ def get_max_batch_size( ``forward_step(model, inputs)`` optionally probes a custom workload, including during OOM retries. By default, probe a model forward (or generation for encoder-decoder models). + + Inside a distributed :class:`~modelopt.torch.utils.distributed.ModelOptSession` it is collective: + every rank calls it, and every rank gets the smallest batch size any rank's GPU allows + (``session.calib_batch_size``), so that all run the same number of calibration steps. Each rank + then measures its own GPU only and probes with a single forward: retrying a smaller batch after an + out-of-memory error would run a different number of forwards, and of the collectives inside them, + on each rank. """ + from .distributed import get_active_mo_session + + session = get_active_mo_session(required=False) + args = ( + model, + max_sample_length, + sample_memory_usage_ratio, + sample_input_single_batch, + enable_grad, + forward_step, + ) + if session.initialized and session.distributed: + return session.decide( + "calib_batch_size", lambda: _max_batch_size(*args, device=session.device) + ) + return _max_batch_size(*args) + + +def _max_batch_size( + model: torch.nn.Module, + max_sample_length: int, + sample_memory_usage_ratio: float, + sample_input_single_batch: torch.Tensor | None, + enable_grad: bool, + forward_step: Callable[[torch.nn.Module, torch.Tensor], Any] | None, + device: torch.device | None = None, +) -> int: + """:func:`get_max_batch_size` measured on ``device`` with one probe, or on every visible GPU.""" + devices = range(torch.cuda.device_count()) if device is None else [device] def _get_free_gpu_mem(): - min_gpu_free_mem = torch.cuda.get_device_properties(0).total_memory + min_gpu_free_mem = torch.cuda.get_device_properties( + 0 if device is None else device + ).total_memory max_allocated_mem = 0 - for device in range(torch.cuda.device_count()): - free_mem = torch.cuda.mem_get_info(device)[0] + for gpu in devices: + free_mem = torch.cuda.mem_get_info(gpu)[0] if free_mem < min_gpu_free_mem: min_gpu_free_mem = free_mem - max_allocated_mem = torch.cuda.max_memory_allocated(device) + max_allocated_mem = torch.cuda.max_memory_allocated(gpu) return min_gpu_free_mem, max_allocated_mem torch.cuda.empty_cache() @@ -1070,7 +1108,7 @@ def _expand_to(batch: int) -> torch.Tensor: # For some models on multi GPU, we observe the memory per batch is not a constant. # So we just test the target batch size and make sure we do not go OOM. - while target_data_batch > 1: + while target_data_batch > 1 and device is None: with torch.set_grad_enabled(enable_grad): try: infer_method(target_input) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 1666308e1d4..7101214ef39 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -436,6 +436,63 @@ def fake_call(x): assert result == 4 +def test_get_max_batch_size_agrees_on_the_smallest_ranks_in_a_session(monkeypatch): + """Every rank must run the same number of calibration steps, so ranks agree on one batch size. + + Rank 1 of two, on cuda:1, measures its own GPU only, probes with one forward (an OOM retry + would run a different number of collective-bearing forwards per rank), and takes rank 0's + smaller size. + """ + from modelopt.torch.utils.distributed import get_active_mo_session + + session = get_active_mo_session(required=False) + session.initialize( + distributed=True, + rank=1, + world_size=2, + local_rank=1, + device=torch.device("cuda:1"), + is_main=False, + ) + sent = [] + + def all_gather_object(gathered, mine, group=None): + sent.append(mine) + gathered[:] = [("calib_batch_size", 4, None), mine] + + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2) + monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) + seen_batch_sizes: list[int] = [] + model = Mock(spec=torch.nn.Module, side_effect=lambda x: seen_batch_sizes.append(x.shape[0])) + model.__class__.__name__ = "DummyModel" + device_props = Mock() + device_props.total_memory = 10**12 + measured = [] + + def mem_get_info(device): + measured.append(device) + return (1000, 0) if len(measured) == 1 else (900, 0) # 100 bytes a sample -> 10 + + try: + with ( + patch("torch.cuda.empty_cache"), + patch("torch.cuda.get_device_properties", return_value=device_props), + patch("torch.cuda.device_count", return_value=4), + patch("torch.cuda.mem_get_info", side_effect=mem_get_info), + patch("torch.cuda.max_memory_allocated", return_value=0), + ): + result = get_max_batch_size( + model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) + ) + finally: + session._reset() + + assert measured == [torch.device("cuda:1")] * 2 + assert seen_batch_sizes == [1] + assert sent == [("calib_batch_size", 8, None)] # this rank's own 10, regulated to 8 + assert result == 4 + + def test_get_dataset_samples_with_unsupported_dataset(make_toy_hf_dataset): """A dataset not in ``SUPPORTED_DATASET_CONFIG`` loads via the auto-detect path.""" dataset_name = make_toy_hf_dataset() # basename never matches a registered key From 34fac6575020b404f68c7d6f9c23abaa7ed28fa7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 01:01:28 +0000 Subject: [PATCH 05/60] Agree on the calibration batch size without keeping it on the session Every rank takes the value get_max_batch_size returns, and nothing reads the batch size back from the session, so the ranks agree on it (ModelOptSession.agree, combined with min) rather than store it as a synced attribute. Every rank still gets the smallest size any rank's GPU allows. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/dataset_utils.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 8f502f761ec..41bf37b05eb 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1009,10 +1009,10 @@ def get_max_batch_size( Inside a distributed :class:`~modelopt.torch.utils.distributed.ModelOptSession` it is collective: every rank calls it, and every rank gets the smallest batch size any rank's GPU allows - (``session.calib_batch_size``), so that all run the same number of calibration steps. Each rank - then measures its own GPU only and probes with a single forward: retrying a smaller batch after an - out-of-memory error would run a different number of forwards, and of the collectives inside them, - on each rank. + (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.agree`), so that all run the same + number of calibration steps. Each rank then measures its own GPU only and probes with a single + forward: retrying a smaller batch after an out-of-memory error would run a different number of + forwards, and of the collectives inside them, on each rank. """ from .distributed import get_active_mo_session @@ -1026,8 +1026,8 @@ def get_max_batch_size( forward_step, ) if session.initialized and session.distributed: - return session.decide( - "calib_batch_size", lambda: _max_batch_size(*args, device=session.device) + return session.agree( + "calib_batch_size", lambda: _max_batch_size(*args, device=session.device), min ) return _max_batch_size(*args) From 0a2636e8d77e9a26cd1fc0385b7232bc08a2ca59 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 03:09:35 +0000 Subject: [PATCH 06/60] Reduce the calibration batch size through the renamed all_reduce Signed-off-by: Shengliang Xu --- modelopt/torch/utils/dataset_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 41bf37b05eb..4b9c9fe775e 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1009,7 +1009,7 @@ def get_max_batch_size( Inside a distributed :class:`~modelopt.torch.utils.distributed.ModelOptSession` it is collective: every rank calls it, and every rank gets the smallest batch size any rank's GPU allows - (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.agree`), so that all run the same + (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.all_reduce`), so that all run the same number of calibration steps. Each rank then measures its own GPU only and probes with a single forward: retrying a smaller batch after an out-of-memory error would run a different number of forwards, and of the collectives inside them, on each rank. @@ -1026,7 +1026,7 @@ def get_max_batch_size( forward_step, ) if session.initialized and session.distributed: - return session.agree( + return session.all_reduce( "calib_batch_size", lambda: _max_batch_size(*args, device=session.device), min ) return _max_batch_size(*args) From c9e5235ecd49d4791a25f618ced3f8b04374cff0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:23:38 +0000 Subject: [PATCH 07/60] Calibrate distributed without the session: pass the device and EP mesh, reduce over the live group distributed_forward_loop takes the device to move batches to and the expert-parallel mesh whose all-to-all forwards it re-installs, where it took the session that carried both. The calibration batch size is still the smallest any rank allows, reduced with all_reduce over the live process group, each rank measuring its own GPU. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/dataset_utils.py | 17 ++--- modelopt/torch/utils/distributed.py | 20 +++--- tests/unit/torch/utils/test_dataset_utils.py | 38 ++++------- .../utils/test_distributed_calibration.py | 68 +++++++------------ 4 files changed, 56 insertions(+), 87 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 4b9c9fe775e..a0587273d10 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1007,16 +1007,14 @@ def get_max_batch_size( ``forward_step(model, inputs)`` optionally probes a custom workload, including during OOM retries. By default, probe a model forward (or generation for encoder-decoder models). - Inside a distributed :class:`~modelopt.torch.utils.distributed.ModelOptSession` it is collective: - every rank calls it, and every rank gets the smallest batch size any rank's GPU allows - (:meth:`~modelopt.torch.utils.distributed.ModelOptSession.all_reduce`), so that all run the same - number of calibration steps. Each rank then measures its own GPU only and probes with a single + On a distributed run it is collective: every rank calls it, and every rank gets the smallest + batch size any rank's GPU allows (:func:`~modelopt.torch.utils.distributed.all_reduce`), so that + all run the same number of calibration steps. Each rank then measures its own GPU only and probes with a single forward: retrying a smaller batch after an out-of-memory error would run a different number of forwards, and of the collectives inside them, on each rank. """ - from .distributed import get_active_mo_session + from .distributed import all_reduce, is_distributed_run, local_rank - session = get_active_mo_session(required=False) args = ( model, max_sample_length, @@ -1025,10 +1023,9 @@ def get_max_batch_size( enable_grad, forward_step, ) - if session.initialized and session.distributed: - return session.all_reduce( - "calib_batch_size", lambda: _max_batch_size(*args, device=session.device), min - ) + if is_distributed_run(): + device = torch.device("cuda", local_rank()) + return all_reduce("calib_batch_size", lambda: _max_batch_size(*args, device=device), min) return _max_batch_size(*args) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 9f3dd0bbf95..9b799887f66 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -753,7 +753,7 @@ def shard_dataloader(loader, rank: int, world_size: int): ) -def _reinstall_a2a_forwards(model, mo_session) -> int: +def _reinstall_a2a_forwards(model, ep_mesh) -> int: """Re-plant the torch_a2a TP forward wrappers that quantizer insertion clobbered. transformers installs its TP input/output transforms by REPLACING ``module.forward`` with an @@ -764,9 +764,9 @@ def _reinstall_a2a_forwards(model, mo_session) -> int: Called at the top of the calibration loop: after conversion (so it wraps the QUANTIZED forward) and before the first forward. Idempotent via a per-module flag, so repeated calibration passes - do not stack wrappers (which would inject the process_group argument twice). + do not stack wrappers (which would inject the process_group argument twice). ``ep_mesh`` is the + mesh the routed experts are sharded over; without one there is nothing to re-plant. """ - ep_mesh = getattr(mo_session, "ep_mesh", None) if mo_session is not None else None if ep_mesh is None: return 0 try: @@ -803,7 +803,7 @@ def _reinstall_a2a_forwards(model, mo_session) -> int: ALL_PARALLEL_STYLES[style_name].install_forward(mod, ep_mesh) mod._modelopt_a2a_forward_reinstalled = True n += 1 - if n and (mo_session is None or getattr(mo_session, "rank", 0) == 0): + if n and is_master(): print( f"[ep] re-installed {n} torch_a2a forward wrapper(s) clobbered by quantizer insertion", flush=True, @@ -811,11 +811,13 @@ def _reinstall_a2a_forwards(model, mo_session) -> int: return n -def distributed_forward_loop(model, dataloader, mo_session: "ModelOptSession | None" = None): +def distributed_forward_loop(model, dataloader, device=None, ep_mesh=None): """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. - Moves each batch onto this rank's ``mo_session.device`` and runs the captured ``model`` forward (the - EP all-reduce / all-to-all collectives fire inside that forward). ``mtq.quantize`` hands the + Moves each batch onto ``device``, when given, and runs the captured ``model`` forward (the EP + all-reduce / all-to-all collectives fire inside that forward). An expert-parallel model passes + the ``ep_mesh`` its experts are sharded over, for the all-to-all forwards that quantizer insertion + drops to be re-installed before the first forward. ``mtq.quantize`` hands the ``forward_loop`` callback an inner module that this closure ignores, calling the captured top-level ``model`` instead so the full distributed forward (and its collectives) runs. @@ -828,10 +830,8 @@ def distributed_forward_loop(model, dataloader, mo_session: "ModelOptSession | N # Local import to avoid any import cycle within modelopt.torch.utils. from modelopt.torch.utils.dataset_utils import _disable_use_cache - device = mo_session.device if mo_session is not None else None - def calibrate(_inner_model): - _reinstall_a2a_forwards(model, mo_session) + _reinstall_a2a_forwards(model, ep_mesh) with _disable_use_cache(model): # Reclaim the caching allocator's reserved-but-unused pool left by the load/quantizer-insertion # memory peak. Without this, torch hoards ~the whole GPU (e.g. 91 GB live weights vs 275 GB diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 7101214ef39..efd0e239fbf 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -436,24 +436,15 @@ def fake_call(x): assert result == 4 -def test_get_max_batch_size_agrees_on_the_smallest_ranks_in_a_session(monkeypatch): +def test_get_max_batch_size_takes_the_smallest_ranks_in_a_distributed_run(monkeypatch): """Every rank must run the same number of calibration steps, so ranks agree on one batch size. Rank 1 of two, on cuda:1, measures its own GPU only, probes with one forward (an OOM retry would run a different number of collective-bearing forwards per rank), and takes rank 0's smaller size. """ - from modelopt.torch.utils.distributed import get_active_mo_session - - session = get_active_mo_session(required=False) - session.initialize( - distributed=True, - rank=1, - world_size=2, - local_rank=1, - device=torch.device("cuda:1"), - is_main=False, - ) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setenv("LOCAL_RANK", "1") sent = [] def all_gather_object(gathered, mine, group=None): @@ -473,19 +464,16 @@ def mem_get_info(device): measured.append(device) return (1000, 0) if len(measured) == 1 else (900, 0) # 100 bytes a sample -> 10 - try: - with ( - patch("torch.cuda.empty_cache"), - patch("torch.cuda.get_device_properties", return_value=device_props), - patch("torch.cuda.device_count", return_value=4), - patch("torch.cuda.mem_get_info", side_effect=mem_get_info), - patch("torch.cuda.max_memory_allocated", return_value=0), - ): - result = get_max_batch_size( - model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) - ) - finally: - session._reset() + with ( + patch("torch.cuda.empty_cache"), + patch("torch.cuda.get_device_properties", return_value=device_props), + patch("torch.cuda.device_count", return_value=4), + patch("torch.cuda.mem_get_info", side_effect=mem_get_info), + patch("torch.cuda.max_memory_allocated", return_value=0), + ): + result = get_max_batch_size( + model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) + ) assert measured == [torch.device("cuda:1")] * 2 assert seen_batch_sizes == [1] diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 9809f2d51d3..f3b0fdc631c 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -33,10 +33,8 @@ import modelopt.torch.utils.distributed as mto_dist from modelopt.torch.utils.distributed import ( - ModelOptSession, _reinstall_a2a_forwards, distributed_forward_loop, - modelopt_session, shard_dataloader, ) @@ -301,16 +299,6 @@ def forward(self, input_ids): return hidden_states -def _ep_session(rank=0, *, with_mesh=True): - """An EP rank's session: its identity plus the EP mesh the model loader recorded.""" - session = ModelOptSession() - session.initialize( - distributed=True, rank=rank, world_size=2, local_rank=rank, device=None, is_main=rank == 0 - ) - session.ep_mesh = _FakeMesh() if with_mesh else None - return session - - def _drop_tp_wrapper(module): """Do what quantizer insertion does to the instance: rebind the class forward over the wrapper. @@ -327,47 +315,43 @@ def test_reinstall_replants_the_ep_wrapper_on_the_a2a_modules_only(monkeypatch, is looked up per module (``is_weight=False``), so children do not inherit their parent's style. """ styles = _install_fake_tp(monkeypatch) - model, session = _MoEModel(), _ep_session(rank) + monkeypatch.setattr(mto_dist, "is_master", lambda group=None: rank == 0) + model, mesh = _MoEModel(), _FakeMesh() hidden_states = torch.ones(2, 4) with pytest.raises(ValueError, match="process_group"): model(hidden_states) # as quantizer insertion leaves it: no wrapper, no EP group - assert _reinstall_a2a_forwards(model, session) == 4 + assert _reinstall_a2a_forwards(model, mesh) == 4 assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] assert styles["torch_a2a_router"].installed == [layer.gate for layer in model.layers] assert styles["colwise"].installed == [] model(hidden_states) - assert [layer.experts.groups for layer in model.layers] == [[session.ep_mesh.group]] * 2 + assert [layer.experts.groups for layer in model.layers] == [[mesh.group]] * 2 assert ("re-installed 4" in capsys.readouterr().out) == (rank == 0) def test_reinstall_does_not_stack_wrappers_across_calibration_passes(monkeypatch): """A repeat call is a no-op; a second wrapper would pass the experts the EP group twice.""" styles = _install_fake_tp(monkeypatch) - model, session = _MoEModel(), _ep_session() - assert _reinstall_a2a_forwards(model, session) == 4 + model, mesh = _MoEModel(), _FakeMesh() + assert _reinstall_a2a_forwards(model, mesh) == 4 - assert _reinstall_a2a_forwards(model, session) == 0 + assert _reinstall_a2a_forwards(model, mesh) == 0 assert len(styles["torch_a2a_experts"].installed) == 2 model(torch.ones(2, 4)) - assert [layer.experts.groups for layer in model.layers] == [[session.ep_mesh.group]] * 2 + assert [layer.experts.groups for layer in model.layers] == [[mesh.group]] * 2 @pytest.mark.parametrize( - ("has_session", "has_ep_mesh", "has_tp_plan"), - [(False, True, True), (True, False, True), (True, True, False)], - ids=["no-session", "no-ep-mesh", "no-tp-plan"], + ("has_ep_mesh", "has_tp_plan"), [(False, True), (True, False)], ids=["no-ep-mesh", "no-tp-plan"] ) -def test_reinstall_is_a_noop_off_the_expert_parallel_path( - monkeypatch, has_session, has_ep_mesh, has_tp_plan -): +def test_reinstall_is_a_noop_off_the_expert_parallel_path(monkeypatch, has_ep_mesh, has_tp_plan): """A single-process or non-EP model is left untouched, even with the torch_a2a styles around.""" styles = _install_fake_tp(monkeypatch) model = _MoEModel(with_tp_plan=has_tp_plan) - session = _ep_session(with_mesh=has_ep_mesh) if has_session else None - assert _reinstall_a2a_forwards(model, session) == 0 + assert _reinstall_a2a_forwards(model, _FakeMesh() if has_ep_mesh else None) == 0 assert not any(style.installed for style in styles.values()) @@ -381,7 +365,7 @@ def test_reinstall_is_a_noop_where_tp_styles_are_pre_hooks(monkeypatch): styles={"torch_a2a_experts": _PreHookStyle(), "torch_a2a_router": _PreHookStyle()}, where=_LEGACY_TP_MODULE, ) - assert _reinstall_a2a_forwards(_MoEModel(), _ep_session()) == 0 + assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0 # ------------------------------------------------------------------------- distributed_forward_loop @@ -409,7 +393,7 @@ def forward(self, **batch): def test_forward_loop_runs_this_ranks_shard_through_the_captured_model(dataset, rank): """Every batch of the rank's shard reaches the top-level model, not the module ``mtq`` passes. - Only the full model's forward runs the EP collectives; no session means no device move. + Only the full model's forward runs the EP collectives; without a device the batches stay put. """ model, inner = _Recorder(), _Recorder() shard = shard_dataloader(_loader(dataset), rank, world_size=2) @@ -422,18 +406,14 @@ def test_forward_loop_runs_this_ranks_shard_through_the_captured_model(dataset, @pytest.mark.parametrize(("device", "expected"), [(None, "cpu"), ("meta", "meta")]) -def test_forward_loop_moves_tensors_to_the_single_process_session_device( - monkeypatch, device, expected -): - """Batch tensors go to ``mo_session.device`` when it names one; other values pass through. +def test_forward_loop_moves_tensors_to_the_device_it_is_given(device, expected): + """Batch tensors go to ``device`` when one is given; other values pass through. The meta device makes the move observable on a CPU-only host. """ - monkeypatch.delenv("WORLD_SIZE", raising=False) # single process: no process group batches = [{"input_ids": torch.zeros(2, 3), "logits_to_keep": 1}] * 2 model = _Recorder() - with modelopt_session(device=device) as session: - distributed_forward_loop(model, batches, session)(model) + distributed_forward_loop(model, batches, device=device)(model) assert [call["input_ids"].device.type for call in model.calls] == [expected] * 2 assert [call["logits_to_keep"] for call in model.calls] == [1, 1] @@ -460,14 +440,16 @@ def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch) mto_dist, "_reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args)) ) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: events.append("empty_cache")) - model, session = _Recorder(on_forward=lambda: events.append("forward")), _ep_session() - forward_loop = distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}] * 2, session) + model, mesh = _Recorder(on_forward=lambda: events.append("forward")), _FakeMesh() + forward_loop = distributed_forward_loop( + model, [{"input_ids": torch.zeros(1, 1)}] * 2, ep_mesh=mesh + ) assert events == [] forward_loop(model) assert events == [ - ("reinstall", model, session), + ("reinstall", model, mesh), "empty_cache", "forward", "forward", @@ -481,13 +463,15 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat It holds on every pass (multi-pass algorithms call the loop again), with one wrapper per module. """ styles = _install_fake_tp(monkeypatch) - model, session = _MoEModel(), _ep_session() - forward_loop = distributed_forward_loop(model, [{"input_ids": torch.ones(2, 4)}] * 3, session) + model, mesh = _MoEModel(), _FakeMesh() + forward_loop = distributed_forward_loop( + model, [{"input_ids": torch.ones(2, 4)}] * 3, ep_mesh=mesh + ) for layer in model.layers: _drop_tp_wrapper(layer.experts) forward_loop(model) forward_loop(model) - assert [layer.experts.groups for layer in model.layers] == [[session.ep_mesh.group] * 6] * 2 + assert [layer.experts.groups for layer in model.layers] == [[mesh.group] * 6] * 2 assert len(styles["torch_a2a_experts"].installed) == 2 From 08bc8197dc67fb1977eac304533efdc53164311b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 19:53:57 +0000 Subject: [PATCH 08/60] Size the calibration batch and check the EP mesh by the model's distribution plan get_max_batch_size agrees on one batch size across ranks when the model's recorded plan spreads it over several, rather than whenever a multi-rank process group is live. distributed_forward_loop takes the ep_mesh exactly when the plan is expert-parallel: an expert-parallel model without its mesh would run its experts without an EP group at the first step, and a mesh for any other model would be ignored. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/dataset_utils.py | 16 +++++++------- modelopt/torch/utils/distributed.py | 14 ++++++++++--- tests/unit/torch/utils/test_dataset_utils.py | 6 +++++- .../utils/test_distributed_calibration.py | 21 +++++++++++++++++++ 4 files changed, 46 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index a0587273d10..a5c2cb1329d 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1007,13 +1007,15 @@ def get_max_batch_size( ``forward_step(model, inputs)`` optionally probes a custom workload, including during OOM retries. By default, probe a model forward (or generation for encoder-decoder models). - On a distributed run it is collective: every rank calls it, and every rank gets the smallest - batch size any rank's GPU allows (:func:`~modelopt.torch.utils.distributed.all_reduce`), so that - all run the same number of calibration steps. Each rank then measures its own GPU only and probes with a single - forward: retrying a smaller batch after an out-of-memory error would run a different number of - forwards, and of the collectives inside them, on each rank. + For a model distributed over several ranks (its + :func:`~modelopt.torch.utils.distributed.distribution_plan`) it is collective: every rank calls + it, and every rank gets the smallest batch size any rank's GPU allows + (:func:`~modelopt.torch.utils.distributed.all_reduce`), so that all run the same number of + calibration steps. Each rank then measures its own GPU only and probes with a single forward: + retrying a smaller batch after an out-of-memory error would run a different number of forwards, + and of the collectives inside them, on each rank. """ - from .distributed import all_reduce, is_distributed_run, local_rank + from .distributed import all_reduce, distribution_plan, local_rank args = ( model, @@ -1023,7 +1025,7 @@ def get_max_batch_size( enable_grad, forward_step, ) - if is_distributed_run(): + if distribution_plan(model).world_size > 1: device = torch.device("cuda", local_rank()) return all_reduce("calib_batch_size", lambda: _max_batch_size(*args, device=device), min) return _max_batch_size(*args) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 9b799887f66..32cf5fef5d3 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -815,9 +815,10 @@ def distributed_forward_loop(model, dataloader, device=None, ep_mesh=None): """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. Moves each batch onto ``device``, when given, and runs the captured ``model`` forward (the EP - all-reduce / all-to-all collectives fire inside that forward). An expert-parallel model passes - the ``ep_mesh`` its experts are sharded over, for the all-to-all forwards that quantizer insertion - drops to be re-installed before the first forward. ``mtq.quantize`` hands the + all-reduce / all-to-all collectives fire inside that forward). A model whose plan + (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its experts are sharded + over, for the all-to-all forwards that quantizer insertion drops to be re-installed before the + first forward; any other model passes none. ``mtq.quantize`` hands the ``forward_loop`` callback an inner module that this closure ignores, calling the captured top-level ``model`` instead so the full distributed forward (and its collectives) runs. @@ -830,6 +831,13 @@ def distributed_forward_loop(model, dataloader, device=None, ep_mesh=None): # Local import to avoid any import cycle within modelopt.torch.utils. from modelopt.torch.utils.dataset_utils import _disable_use_cache + plan = distribution_plan(model) + if plan.is_expert_parallel != (ep_mesh is not None): + raise ValueError( + f"{plan!r}: an expert-parallel model calibrates with the ep_mesh its routed experts " + "are sharded over, and any other model without one." + ) + def calibrate(_inner_model): _reinstall_a2a_forwards(model, ep_mesh) with _disable_use_cache(model): diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index efd0e239fbf..ed9fbee26f3 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -34,6 +34,7 @@ get_max_batch_size, prepare_messages_for_chat_template, ) +from modelopt.torch.utils.distributed import DistributionPlan def test_prepare_messages_for_chat_template(): @@ -406,6 +407,7 @@ def fake_call(x): raise torch.cuda.OutOfMemoryError model = Mock(spec=torch.nn.Module) + model.modules.return_value = [] # no plan recorded and no FSDP2 module: a single process's # get_max_batch_size calls the module (model(...)) so FSDP2 hooks fire, not model.forward, # so route the mock's __call__ via side_effect. model.side_effect = fake_call @@ -437,7 +439,8 @@ def fake_call(x): def test_get_max_batch_size_takes_the_smallest_ranks_in_a_distributed_run(monkeypatch): - """Every rank must run the same number of calibration steps, so ranks agree on one batch size. + """Every rank must run the same number of calibration steps, so ranks agree on one batch size + for a model whose plan distributes it over them. Rank 1 of two, on cuda:1, measures its own GPU only, probes with one forward (an OOM retry would run a different number of collective-bearing forwards per rank), and takes rank 0's @@ -456,6 +459,7 @@ def all_gather_object(gathered, mine, group=None): seen_batch_sizes: list[int] = [] model = Mock(spec=torch.nn.Module, side_effect=lambda x: seen_batch_sizes.append(x.shape[0])) model.__class__.__name__ = "DummyModel" + model._modelopt_distribution_plan = DistributionPlan(2) device_props = Mock() device_props.total_memory = 10**12 measured = [] diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index f3b0fdc631c..3f2e8b71b01 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -33,11 +33,15 @@ import modelopt.torch.utils.distributed as mto_dist from modelopt.torch.utils.distributed import ( + DistributionPlan, _reinstall_a2a_forwards, distributed_forward_loop, + set_distribution_plan, shard_dataloader, ) +_EXPERT_PARALLEL = DistributionPlan(world_size=2, ep_size=2) + # ------------------------------------------------------------------------------------- data helpers @@ -441,6 +445,7 @@ def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch) ) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: events.append("empty_cache")) model, mesh = _Recorder(on_forward=lambda: events.append("forward")), _FakeMesh() + set_distribution_plan(model, _EXPERT_PARALLEL) forward_loop = distributed_forward_loop( model, [{"input_ids": torch.zeros(1, 1)}] * 2, ep_mesh=mesh ) @@ -464,6 +469,7 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat """ styles = _install_fake_tp(monkeypatch) model, mesh = _MoEModel(), _FakeMesh() + set_distribution_plan(model, _EXPERT_PARALLEL) forward_loop = distributed_forward_loop( model, [{"input_ids": torch.ones(2, 4)}] * 3, ep_mesh=mesh ) @@ -475,3 +481,18 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat assert [layer.experts.groups for layer in model.layers] == [[mesh.group] * 6] * 2 assert len(styles["torch_a2a_experts"].installed) == 2 + + +@pytest.mark.parametrize( + ("plan", "mesh"), + [(_EXPERT_PARALLEL, None), (DistributionPlan(world_size=2), _FakeMesh())], + ids=["expert-parallel-without-mesh", "mesh-for-a-replicated-model"], +) +def test_the_ep_mesh_comes_with_an_expert_parallel_plan_only(plan, mesh): + """Without its mesh, the experts' forward would run with no EP group at the first step; a mesh + for a model whose plan shards no experts would be ignored.""" + model = _Recorder() + set_distribution_plan(model, plan) + + with pytest.raises(ValueError, match="calibrates with the ep_mesh"): + distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}], ep_mesh=mesh) From b6a145f3be2f2dfdff04c62e81ceb27fa7564e70 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:44:20 +0000 Subject: [PATCH 09/60] State that a streaming calibration dataset must give every worker its whole stream The rank shard splits each stream's groups between the DataLoader workers, so a dataset that already splits itself by get_worker_info() is split twice and drops samples. hf_ptq loads with num_workers=0; say so for other callers. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 32cf5fef5d3..799c6f19721 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -709,7 +709,9 @@ def shard_dataloader(loader, rank: int, world_size: int): counts equal (else a rank exits calibration early and hangs the others on the per-step collectives), at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. For streaming ``IterableDataset`` datasets (which have no ``len()`` and so can't use a sampler) each rank - strides its own deterministic copy of the stream via :class:`_StridedIterableShard`. + strides its own deterministic copy of the stream via :class:`_StridedIterableShard`. Every DataLoader + worker must see that whole stream: a dataset that splits itself across workers by + ``get_worker_info()`` would be split twice and lose samples, so load one with ``num_workers=0``. Forwards all non-sampler DataLoader settings from ``loader`` (workers, pinning, prefetch, init fn, generator, ...). From ebfb28155f2d7119e7cb801da7d6393a579ace15 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 06:57:34 +0000 Subject: [PATCH 10/60] Free the CUDA cache around calibration only in a process that started CUDA distributed_forward_loop collects garbage and empties the CUDA cache before the first calibration forward and after the last, so NCCL can allocate its buffers outside the caching allocator. A process that never started CUDA has no pool to free, yet each call still paid a full garbage collection, about a second on a large heap: nearly 12 s of the unit run across this file's forward-loop tests. The release now runs only once CUDA is initialized. The forward-loop test that walked a 4-case shard grid keeps one case; shard_dataloader's own tests cover the striding. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 17 ++++++++++--- .../utils/test_distributed_calibration.py | 25 ++++++++----------- 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 799c6f19721..3963fd0c363 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -813,6 +813,17 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: return n +def _release_cached_cuda_memory(): + """Hand the caching allocator's reserved-but-unused pool back to the driver. + + Only a process that started CUDA has a pool to hand back; elsewhere the collection alone costs + about a second. + """ + if torch.cuda.is_initialized(): + gc.collect() + torch.cuda.empty_cache() + + def distributed_forward_loop(model, dataloader, device=None, ep_mesh=None): """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. @@ -848,8 +859,7 @@ def calibrate(_inner_model): # reserved), leaving ~0 GB free, so the distributed a2a's NCCL comm buffers cannot be allocated # ("Failed to CUDA calloc") at the very first all_to_all. Freeing the cache hands that memory # back to the driver for NCCL + the calibration activations. - gc.collect() - torch.cuda.empty_cache() + _release_cached_cuda_memory() for batch in tqdm(dataloader, desc="Calibrating", disable=not is_master()): if device is not None: batch = { @@ -864,8 +874,7 @@ def calibrate(_inner_model): # single calib forward: alloc 86.0 GiB but reserved 156.6 GiB -> only 0.7 GiB free, and # the job then hung until the 2 h wall clock instead of exporting. NCCL comm buffers are # allocated OUTSIDE the torch allocator, so a full reserved pool starves them. - gc.collect() - torch.cuda.empty_cache() + _release_cached_cuda_memory() return calibrate diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 3f2e8b71b01..9e3d4456c54 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -392,20 +392,18 @@ def forward(self, **batch): self.on_forward() -@pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) -@pytest.mark.parametrize("rank", [0, 1]) -def test_forward_loop_runs_this_ranks_shard_through_the_captured_model(dataset, rank): +def test_forward_loop_runs_this_ranks_shard_through_the_captured_model(): """Every batch of the rank's shard reaches the top-level model, not the module ``mtq`` passes. Only the full model's forward runs the EP collectives; without a device the batches stay put. """ model, inner = _Recorder(), _Recorder() - shard = shard_dataloader(_loader(dataset), rank, world_size=2) + shard = shard_dataloader(_loader(list(range(8))), 0, world_size=2) distributed_forward_loop(model, shard)(inner) calls = [call["input_ids"].flatten().tolist() for call in model.calls] - assert calls == [[rank, rank + 2], [rank + 4, rank + 6]] + assert calls == [[0, 2], [4, 6]] assert inner.calls == [] @@ -432,17 +430,21 @@ def test_forward_loop_disables_use_cache_only_while_calibrating(): assert model.config.use_cache is True -def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch): +@pytest.mark.parametrize("cuda_started", [True, False], ids=["cuda", "cpu-only"]) +def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch, cuda_started): """Re-plant the a2a wrappers and free the cached pool before the first forward, and again after. Building the loop does neither: quantizers are inserted between building and calling it. NCCL allocates its buffers outside torch's caching allocator, so a pool hoarded by the model load or - by the calibration activations starves the next collective. + by the calibration activations starves the next collective. A process that never started CUDA + has no pool to free and skips the garbage collection, which alone takes about a second. """ events = [] monkeypatch.setattr( mto_dist, "_reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args)) ) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: cuda_started) + monkeypatch.setattr(mto_dist.gc, "collect", lambda: events.append("collect")) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: events.append("empty_cache")) model, mesh = _Recorder(on_forward=lambda: events.append("forward")), _FakeMesh() set_distribution_plan(model, _EXPERT_PARALLEL) @@ -453,13 +455,8 @@ def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch) forward_loop(model) - assert events == [ - ("reinstall", model, mesh), - "empty_cache", - "forward", - "forward", - "empty_cache", - ] + release = ["collect", "empty_cache"] if cuda_started else [] + assert events == [("reinstall", model, mesh), *release, "forward", "forward", *release] def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypatch): From f9919008efa90de12068128293a4c901e5557e31 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 06:58:23 +0000 Subject: [PATCH 11/60] Test a streaming shard's split across DataLoader workers in one process The test started two DataLoader workers per rank, forked from pytest's multithreaded process, which warns that fork may deadlock; in a sandbox where the workers could not use shared memory it hung until its timeout. The split it guards is decided in the shard's __iter__ from get_worker_info(), so the test now runs each worker's pass in this process as the worker it is told it is. Signed-off-by: Shengliang Xu --- .../utils/test_distributed_calibration.py | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 9e3d4456c54..58f1d74684c 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -122,22 +122,23 @@ def test_iterable_shards_are_strided_cover_every_sample_once_and_replay(): assert _batches(shard) == _batches(shard) == [[r, r + 3], [r + 6, r + 9]] -def test_iterable_shards_split_their_groups_across_worker_processes(): +def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch): """Each worker process iterates its own copy of the stream. Unless the workers split the groups - between them, every one yields the whole shard and each sample comes out once per worker.""" - n, world_size = 16, 2 - shards = [ - _ids( - shard_dataloader( - DataLoader(_Stream(n), batch_size=2, collate_fn=_collate, num_workers=2), - r, - world_size, - ) - ) - for r in range(world_size) - ] + between them, every one yields the whole shard and each sample comes out once per worker. - assert [sorted(ids) for ids in shards] == [list(range(r, n, world_size)) for r in range(2)] + Each worker's pass runs here, in this process, as the worker it is told it is. + """ + n, world_size, num_workers = 16, 2, 2 + for rank in range(world_size): + loader = DataLoader(_Stream(n), batch_size=2, collate_fn=_collate, num_workers=num_workers) + shard = shard_dataloader(loader, rank, world_size).dataset + yielded = [] + for worker in range(num_workers): + info = types.SimpleNamespace(num_workers=num_workers, id=worker) + monkeypatch.setattr(mto_dist, "get_worker_info", lambda info=info: info) + yielded += list(shard) + + assert sorted(yielded) == list(range(rank, n, world_size)) def test_iterable_uneven_split_drops_the_trailing_partial_group(): From a3608492272fb4cad68bc1ac6d790ee64ac8de49 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 06:58:47 +0000 Subject: [PATCH 12/60] Say why shard_dataloader strides streams with its own shard class vlm_dataset_utils._ShardedIterable already strides an iterable dataset across ranks, but it keeps ranks in step by padding each to a per-rank count the VLM loader knows. shard_dataloader wraps any loader's stream, which may have no length, so it drops the trailing partial group instead, and it must also split across the loader's workers, which the VLM loader never starts. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 3963fd0c363..f487f76ecab 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -683,6 +683,8 @@ class _StridedIterableShard(IterableDataset): exactly the same count -- preserving the equal-per-rank-batch invariant that keeps per-step EP collectives from hanging when one rank runs short. Each DataLoader worker process iterates its own copy of the stream, so the workers split the groups between them rather than each yielding them all. + ``vlm_dataset_utils._ShardedIterable`` cannot stand in: it pads every rank to a count its own loader + knows, while this wraps any loader's stream, which may know none, and splits across its workers. """ def __init__(self, base: IterableDataset, rank: int, world_size: int): From 0849f6113e2de728b7ade295eb85535aa491b8b1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:05:49 +0000 Subject: [PATCH 13/60] Agree on the calibration batch size over two real gloo ranks in its test The test simulated rank 1 of two by stubbing all_gather_object and the world size, and asserted the tuple this rank put on the wire. It now runs get_max_batch_size on both ranks of the shared gloo workers, each with its own GPU's memory faked, and checks that both take the smaller size through the real reduction, measuring only their own GPU with one probe forward. A rank left out of the reduction gives up on a group with a 10 s timeout instead of hanging the run. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_dataset_utils.py | 64 +++++++++++--------- 1 file changed, 35 insertions(+), 29 deletions(-) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index ed9fbee26f3..2658c728fb2 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -14,6 +14,9 @@ # limitations under the License. import json +from datetime import timedelta +from functools import partial +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest @@ -34,7 +37,7 @@ get_max_batch_size, prepare_messages_for_chat_template, ) -from modelopt.torch.utils.distributed import DistributionPlan +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan def test_prepare_messages_for_chat_template(): @@ -438,51 +441,54 @@ def fake_call(x): assert result == 4 -def test_get_max_batch_size_takes_the_smallest_ranks_in_a_distributed_run(monkeypatch): - """Every rank must run the same number of calibration steps, so ranks agree on one batch size - for a model whose plan distributes it over them. +class _BatchSizeProbe(torch.nn.Module): + """Records the batch size of every forward.""" - Rank 1 of two, on cuda:1, measures its own GPU only, probes with one forward (an OOM retry - would run a different number of collective-bearing forwards per rank), and takes rank 0's - smaller size. - """ - monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) - monkeypatch.setenv("LOCAL_RANK", "1") - sent = [] + def __init__(self): + super().__init__() + self.batch_sizes = [] - def all_gather_object(gathered, mine, group=None): - sent.append(mine) - gathered[:] = [("calib_batch_size", 4, None), mine] + def forward(self, x): + self.batch_sizes.append(x.shape[0]) - monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2) - monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) - seen_batch_sizes: list[int] = [] - model = Mock(spec=torch.nn.Module, side_effect=lambda x: seen_batch_sizes.append(x.shape[0])) - model.__class__.__name__ = "DummyModel" - model._modelopt_distribution_plan = DistributionPlan(2) - device_props = Mock() - device_props.total_memory = 10**12 + +def _smallest_batch_size_job(rank, world_size): + model = _BatchSizeProbe() + set_distribution_plan(model, DistributionPlan(world_size)) + # A sample takes 100 bytes: rank 0's GPU has room for 4 of them, rank 1's for 10 (regulated to 8). + free = [400, 1000][rank] measured = [] def mem_get_info(device): measured.append(device) - return (1000, 0) if len(measured) == 1 else (900, 0) # 100 bytes a sample -> 10 + return (free if len(measured) == 1 else free - 100, 0) + # A rank left waiting in the reduction gives up in seconds on this group. + group = torch.distributed.new_group(timeout=timedelta(seconds=10)) + gather = partial(torch.distributed.all_gather_object, group=group) with ( + patch.object(torch.distributed, "all_gather_object", gather), patch("torch.cuda.empty_cache"), - patch("torch.cuda.get_device_properties", return_value=device_props), + patch( + "torch.cuda.get_device_properties", return_value=SimpleNamespace(total_memory=10**12) + ), patch("torch.cuda.device_count", return_value=4), patch("torch.cuda.mem_get_info", side_effect=mem_get_info), patch("torch.cuda.max_memory_allocated", return_value=0), ): - result = get_max_batch_size( + batch_size = get_max_batch_size( model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) ) - assert measured == [torch.device("cuda:1")] * 2 - assert seen_batch_sizes == [1] - assert sent == [("calib_batch_size", 8, None)] # this rank's own 10, regulated to 8 - assert result == 4 + assert measured == [torch.device("cuda", rank)] * 2 # its own GPU (LOCAL_RANK) only + assert model.batch_sizes == [1] # one probe: an OOM retry would run more forwards on some ranks + assert batch_size == 4 # rank 0's + + +def test_get_max_batch_size_takes_the_smallest_ranks_in_a_distributed_run(cpu_dist_workers): + """Every rank must run the same number of calibration steps, so ranks agree on one batch size + for a model whose plan distributes it over them: the smallest any rank's GPU allows.""" + cpu_dist_workers.run(_smallest_batch_size_job) def test_get_dataset_samples_with_unsupported_dataset(make_toy_hf_dataset): From 666d0320f30a7cbbe14930ea1ffe9f3bbf1bc030 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:52:25 +0000 Subject: [PATCH 14/60] Trim the distributed calibration tests: one per behavior Merge the shard tests that differ only in dataset kind or world size: - the strided map-style test, the iterable strided-and-replay test and the world-size-one test become one table over {map-style, iterable} x world size {1, 2, 3}, which also checks the replay for both kinds; - the map-style padding test and the iterable trailing-drop test become one uneven-split test that keeps both kinds' asserts. Fold the _reinstall_a2a_forwards replant test into the forward-loop replant test (its router, colwise and children asserts; the rank-0 log line it also checked is wording). Drop, each checked by breaking the code it guards: - the no-stacking test: the forward-loop replant test runs two passes and fails when a repeat call wraps again; - the no-tp-plan no-op case: removing that early return changes nothing, so the case caught nothing (the no-mesh case stays: it is the only guard of its early return); - the device test's no-device case: the shard-routing test passes no device and reads the values back. Signed-off-by: Shengliang Xu --- .../utils/test_distributed_calibration.py | 147 ++++++------------ 1 file changed, 48 insertions(+), 99 deletions(-) diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 58f1d74684c..ba8d1dff532 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -77,49 +77,52 @@ def _seed_worker(worker_id): torch.manual_seed(worker_id) +def _dataset(kind, n): + """Samples ``0 .. n - 1`` as a list (map-style) or as a stream with no ``len()``.""" + return list(range(n)) if kind == "map-style" else _Stream(n) + + # --------------------------------------------------------------------------------- shard_dataloader -@pytest.mark.parametrize("world_size", [2, 3]) -def test_map_style_shards_are_strided_and_cover_every_sample_once(world_size): - """Rank ``r`` gets samples ``r, r + W, ...``, batched with the original size and collate_fn.""" +@pytest.mark.parametrize("kind", ["map-style", "iterable"]) +@pytest.mark.parametrize("world_size", [1, 2, 3]) +def test_shards_are_strided_cover_every_sample_once_and_replay(kind, world_size): + """Rank ``r`` gets samples ``r, r + W, ...``, batched with the original size and collate_fn, and + a second pass replays the same shard: multi-pass calibration algorithms (e.g. AWQ) iterate the + loader again. A lone rank's shard is the whole (unshuffled) loader, batch for batch. + """ n = 12 - shards = [shard_dataloader(_loader(list(range(n))), r, world_size) for r in range(world_size)] + shards = [ + shard_dataloader(_loader(_dataset(kind, n)), r, world_size) for r in range(world_size) + ] assert sorted(i for shard in shards for i in _ids(shard)) == list(range(n)) for r, shard in enumerate(shards): ids = list(range(r, n, world_size)) - assert _batches(shard) == [ids[i : i + 2] for i in range(0, len(ids), 2)] - + assert _batches(shard) == _batches(shard) == [ids[i : i + 2] for i in range(0, len(ids), 2)] -def test_map_style_uneven_split_pads_every_rank_to_the_same_batch_count(): - """No sample is dropped and no rank runs short (it would hang the others' EP collectives). - The price of ``drop_last=False`` is that the sampler repeats fewer than ``world_size`` samples. +@pytest.mark.parametrize( + ("kind", "padded", "dropped"), + [("map-style", 2, set()), ("iterable", 0, {8, 9})], + ids=["map-style", "iterable"], +) +def test_an_uneven_split_gives_every_rank_the_same_batch_count(kind, padded, dropped): + """No rank runs short: it would hang the others' EP collectives. A map-style loader repeats a + few samples to even the ranks out (``drop_last=False``); a stream, with no ``len()`` to pad to, + drops its trailing partial group instead. """ n, world_size = 10, 4 shards = [ - _batches(shard_dataloader(_loader(list(range(n))), r, world_size)) + _batches(shard_dataloader(_loader(_dataset(kind, n)), r, world_size)) for r in range(world_size) ] - assert [len(batches) for batches in shards] == [2] * world_size # 3 samples per rank + assert len({len(batches) for batches in shards}) == 1 seen = Counter(i for batches in shards for batch in batches for i in batch) - assert set(seen) == set(range(n)) - assert seen.total() - n == 2 # padding up to 3 * world_size samples - - -def test_iterable_shards_are_strided_cover_every_sample_once_and_replay(): - """Each rank strides its own copy of a stream; a second pass replays the same shard. - - Multi-pass calibration algorithms (e.g. AWQ) iterate the loader again. - """ - n, world_size = 12, 3 - shards = [shard_dataloader(_loader(_Stream(n)), r, world_size) for r in range(world_size)] - - assert sorted(i for shard in shards for i in _ids(shard)) == list(range(n)) - for r, shard in enumerate(shards): - assert _batches(shard) == _batches(shard) == [[r, r + 3], [r + 6, r + 9]] + assert set(seen) == set(range(n)) - dropped + assert seen.total() - (n - len(dropped)) == padded def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch): @@ -141,21 +144,6 @@ def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch) assert sorted(yielded) == list(range(rank, n, world_size)) -def test_iterable_uneven_split_drops_the_trailing_partial_group(): - """Equal counts without ``len()``: the last ``n % world_size`` samples go to no rank.""" - n, world_size = 10, 4 - shards = [_ids(shard_dataloader(_loader(_Stream(n)), r, world_size)) for r in range(world_size)] - - assert shards == [[0, 4], [1, 5], [2, 6], [3, 7]] - - -@pytest.mark.parametrize("dataset", [list(range(5)), _Stream(5)], ids=["map-style", "iterable"]) -def test_world_size_one_yields_the_original_batches(dataset): - """A lone rank's shard is the whole (unshuffled) loader, batch for batch.""" - loader = _loader(dataset) - assert _batches(shard_dataloader(loader, 0, 1)) == _batches(loader) == [[0, 1], [2, 3], [4]] - - @pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) def test_sharded_loader_keeps_the_original_loader_settings(dataset): """Workers, pinning, timeout, worker init, RNG and prefetching carry over to the shard.""" @@ -283,19 +271,15 @@ def forward(self, hidden_states): class _MoEModel(nn.Module): """Two MoE layers whose plan sends the experts and routers through torch_a2a.""" - def __init__(self, with_tp_plan=True): + def __init__(self): super().__init__() self.layers = nn.ModuleList([_Block(), _Block()]) self.config = types.SimpleNamespace(use_cache=True) - self.tp_plan = ( - { - "layers.*.gate": "torch_a2a_router", - "layers.*.proj": "colwise", - "layers.*.experts": "torch_a2a_experts", - } - if with_tp_plan - else None - ) + self.tp_plan = { + "layers.*.gate": "torch_a2a_router", + "layers.*.proj": "colwise", + "layers.*.experts": "torch_a2a_experts", + } def forward(self, input_ids): hidden_states = input_ids.float() @@ -312,51 +296,11 @@ def _drop_tp_wrapper(module): module.forward = types.MethodType(type(module).forward, module) -@pytest.mark.parametrize("rank", [0, 1]) -def test_reinstall_replants_the_ep_wrapper_on_the_a2a_modules_only(monkeypatch, capsys, rank): - """Every rank re-wraps exactly the planned torch_a2a experts and routers; rank 0 logs the count. - - Other TP styles keep their forward, and so do the quantizers under an experts module: the plan - is looked up per module (``is_weight=False``), so children do not inherit their parent's style. - """ - styles = _install_fake_tp(monkeypatch) - monkeypatch.setattr(mto_dist, "is_master", lambda group=None: rank == 0) - model, mesh = _MoEModel(), _FakeMesh() - hidden_states = torch.ones(2, 4) - with pytest.raises(ValueError, match="process_group"): - model(hidden_states) # as quantizer insertion leaves it: no wrapper, no EP group - - assert _reinstall_a2a_forwards(model, mesh) == 4 - - assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] - assert styles["torch_a2a_router"].installed == [layer.gate for layer in model.layers] - assert styles["colwise"].installed == [] - model(hidden_states) - assert [layer.experts.groups for layer in model.layers] == [[mesh.group]] * 2 - assert ("re-installed 4" in capsys.readouterr().out) == (rank == 0) - - -def test_reinstall_does_not_stack_wrappers_across_calibration_passes(monkeypatch): - """A repeat call is a no-op; a second wrapper would pass the experts the EP group twice.""" - styles = _install_fake_tp(monkeypatch) - model, mesh = _MoEModel(), _FakeMesh() - assert _reinstall_a2a_forwards(model, mesh) == 4 - - assert _reinstall_a2a_forwards(model, mesh) == 0 - assert len(styles["torch_a2a_experts"].installed) == 2 - model(torch.ones(2, 4)) - assert [layer.experts.groups for layer in model.layers] == [[mesh.group]] * 2 - - -@pytest.mark.parametrize( - ("has_ep_mesh", "has_tp_plan"), [(False, True), (True, False)], ids=["no-ep-mesh", "no-tp-plan"] -) -def test_reinstall_is_a_noop_off_the_expert_parallel_path(monkeypatch, has_ep_mesh, has_tp_plan): +def test_reinstall_is_a_noop_without_an_ep_mesh(monkeypatch): """A single-process or non-EP model is left untouched, even with the torch_a2a styles around.""" styles = _install_fake_tp(monkeypatch) - model = _MoEModel(with_tp_plan=has_tp_plan) - assert _reinstall_a2a_forwards(model, _FakeMesh() if has_ep_mesh else None) == 0 + assert _reinstall_a2a_forwards(_MoEModel(), None) == 0 assert not any(style.installed for style in styles.values()) @@ -408,17 +352,17 @@ def test_forward_loop_runs_this_ranks_shard_through_the_captured_model(): assert inner.calls == [] -@pytest.mark.parametrize(("device", "expected"), [(None, "cpu"), ("meta", "meta")]) -def test_forward_loop_moves_tensors_to_the_device_it_is_given(device, expected): +def test_forward_loop_moves_tensors_to_the_device_it_is_given(): """Batch tensors go to ``device`` when one is given; other values pass through. - The meta device makes the move observable on a CPU-only host. + The meta device makes the move observable on a CPU-only host. (Without a device the batches stay + put: the shard-routing test reads their values back.) """ batches = [{"input_ids": torch.zeros(2, 3), "logits_to_keep": 1}] * 2 model = _Recorder() - distributed_forward_loop(model, batches, device=device)(model) + distributed_forward_loop(model, batches, device="meta")(model) - assert [call["input_ids"].device.type for call in model.calls] == [expected] * 2 + assert [call["input_ids"].device.type for call in model.calls] == ["meta"] * 2 assert [call["logits_to_keep"] for call in model.calls] == [1, 1] @@ -464,6 +408,9 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat """hf_ptq's order -- build the loop, insert quantizers, calibrate -- hands the experts the EP group. It holds on every pass (multi-pass algorithms call the loop again), with one wrapper per module. + Only the planned torch_a2a experts and routers are re-wrapped: other TP styles keep their forward, + and so do the quantizers under an experts module, since the plan is looked up per module + (``is_weight=False``) and children do not inherit their parent's style. """ styles = _install_fake_tp(monkeypatch) model, mesh = _MoEModel(), _FakeMesh() @@ -478,7 +425,9 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat forward_loop(model) assert [layer.experts.groups for layer in model.layers] == [[mesh.group] * 6] * 2 - assert len(styles["torch_a2a_experts"].installed) == 2 + assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] + assert styles["torch_a2a_router"].installed == [layer.gate for layer in model.layers] + assert styles["colwise"].installed == [] @pytest.mark.parametrize( From 90543d13100f7ef8e0ffc98d399d3c009c773b48 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:48:29 +0000 Subject: [PATCH 15/60] Refuse a calibration stream shorter than the world size The strided iterable shard drops a trailing incomplete group so every rank yields the same count. A stream with fewer samples than ranks has no complete group at all, so every rank silently calibrated on nothing. Refuse it instead, as the VLM loader's own sharding already does for num_samples < dp_size. Every rank sees the same stream, so all of them refuse it at the first batch, before any forward. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 12 ++++++++++-- .../unit/torch/utils/test_distributed_calibration.py | 8 ++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index f487f76ecab..d875b90d1a8 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -681,8 +681,10 @@ class _StridedIterableShard(IterableDataset): own identical stream (deterministic seed) and takes slot ``rank`` from every complete group of ``world_size`` consecutive samples. Any trailing incomplete group is dropped so every rank yields exactly the same count -- preserving the equal-per-rank-batch invariant that keeps per-step EP - collectives from hanging when one rank runs short. Each DataLoader worker process iterates its own - copy of the stream, so the workers split the groups between them rather than each yielding them all. + collectives from hanging when one rank runs short. A stream shorter than ``world_size`` has no + complete group, so it is refused rather than leaving every rank with nothing to calibrate on. Each + DataLoader worker process iterates its own copy of the stream, so the workers split the groups + between them rather than each yielding them all. ``vlm_dataset_utils._ShardedIterable`` cannot stand in: it pads every rank to a count its own loader knows, while this wraps any loader's stream, which may know none, and splits across its workers. """ @@ -699,6 +701,12 @@ def __iter__(self): for index in itertools.count(): group = list(itertools.islice(it, self.world_size)) if len(group) < self.world_size: + if index == 0: + # Every rank and worker sees the same stream, so all of them refuse it together. + raise ValueError( + f"The calibration stream has {len(group)} sample(s), fewer than the " + f"{self.world_size} ranks: no rank would get a sample to calibrate on." + ) return # drop incomplete trailing group -> equal counts across ranks if index % num_workers == worker: yield group[self.rank] diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index ba8d1dff532..c9b28d19424 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -125,6 +125,14 @@ def test_an_uneven_split_gives_every_rank_the_same_batch_count(kind, padded, dro assert seen.total() - (n - len(dropped)) == padded +def test_a_stream_shorter_than_the_world_is_refused_on_every_rank(): + """With fewer samples than ranks there is no complete group, so every rank would calibrate on + nothing; each rank refuses the stream instead, before any forward.""" + for rank in range(4): + with pytest.raises(ValueError, match="fewer than the 4 ranks"): + _ids(shard_dataloader(_loader(_Stream(3)), rank, 4)) + + def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch): """Each worker process iterates its own copy of the stream. Unless the workers split the groups between them, every one yields the whole shard and each sample comes out once per worker. From 66a65b04af1835fc77183ea9ee15135fa4d26520 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:49:38 +0000 Subject: [PATCH 16/60] Keep the calibration loader's drop_last in its per-rank shard shard_dataloader forwards the original loader's non-sampler settings to the rebuilt per-rank DataLoader, but not drop_last, so a loader that drops its incomplete final batch got one again. Forward it on both paths. Every rank holds the same number of samples (the sampler pads, the stream keeps whole groups), so ranks still run the same number of batches. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 2 ++ tests/unit/torch/utils/test_distributed_calibration.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index d875b90d1a8..8491600f932 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -730,6 +730,7 @@ def shard_dataloader(loader, rank: int, world_size: int): return DataLoader( _StridedIterableShard(loader.dataset, rank, world_size), batch_size=loader.batch_size, + drop_last=loader.drop_last, collate_fn=loader.collate_fn, num_workers=loader.num_workers, pin_memory=loader.pin_memory, @@ -751,6 +752,7 @@ def shard_dataloader(loader, rank: int, world_size: int): return DataLoader( loader.dataset, batch_size=loader.batch_size, + drop_last=loader.drop_last, sampler=sampler, collate_fn=loader.collate_fn, num_workers=loader.num_workers, diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index c9b28d19424..be5d0ccea65 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -154,10 +154,12 @@ def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch) @pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) def test_sharded_loader_keeps_the_original_loader_settings(dataset): - """Workers, pinning, timeout, worker init, RNG and prefetching carry over to the shard.""" + """Batching, workers, pinning, timeout, worker init, RNG and prefetching carry over to the + shard.""" loader = DataLoader( dataset, batch_size=4, + drop_last=True, collate_fn=_collate, num_workers=2, pin_memory=True, @@ -172,6 +174,7 @@ def test_sharded_loader_keeps_the_original_loader_settings(dataset): for setting in ( "batch_size", + "drop_last", "collate_fn", "num_workers", "pin_memory", From a34bfdee0efeb4ca3598ff8d5dff9952d04eafa5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:04:23 +0000 Subject: [PATCH 17/60] Refuse rank-local calibration on a model distributed over several ranks Each rank of a distributed model calibrates on its own data. AWQ (awq_lite, awq_clip, awq_full, and SVDQuant, which runs awq_lite) averages its activation scales over data parallelism but picks each layer's scale from its own rank's search loss; local_hessian and imatrix share nothing between ranks. The ranks' copies of a weight would come out different, so quantize() and calibrate() refuse these algorithms on a model whose recorded plan spans several ranks, before anything changes. Only a recorded plan counts: Megatron models record none, so their calibration is unchanged. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 32 +++++++++++++++++++ modelopt/torch/utils/distributed.py | 12 ++++++- .../utils/test_distributed_calibration.py | 28 ++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 87c8677985c..81d74431a30 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -72,6 +72,36 @@ "temporarily_fold_weights", ] +# Calibration algorithms whose statistics stay on the rank that gathered them. +_RANK_LOCAL_CALIBRATION = ( + "awq_lite", + "awq_clip", + "awq_full", + "svdquant", + "local_hessian", + "imatrix", +) + + +def _refuse_rank_local_calibration(model: nn.Module, algorithm: QuantizeAlgoCfgType) -> None: + """Refuse a rank-local calibration algorithm on a model distributed over several ranks. + + Each rank of such a model calibrates on its own data, so an algorithm that keeps its statistics + on the rank would leave the ranks' copies of a weight different. Called before anything changes. + """ + plan = dist.recorded_distribution_plan(model) + if plan is None or plan.world_size == 1: + return + methods = [mode.removesuffix("_calibrate") for mode, _ in get_modelike_from_algo_cfg(algorithm)] + refused = [method for method in methods if method in _RANK_LOCAL_CALIBRATION] + if refused: + raise NotImplementedError( + f"Calibrating with {', '.join(refused)} is not supported on a model distributed over " + f"{plan.world_size} ranks. Each rank calibrates on its own data and this calibration " + "keeps what it gathers on the rank, so the ranks' copies of a weight would come out " + "different. Calibrate with max, MSE, GPTQ or SmoothQuant, or in a single process." + ) + def _refuse_awq_on_a_sharded_fsdp2_model( model: nn.Module, algorithm: QuantizeAlgoCfgType, candidate: str | None = None @@ -128,6 +158,7 @@ def calibrate( Returns: The calibrated pytorch model. """ _refuse_awq_on_a_sharded_fsdp2_model(model, algorithm) + _refuse_rank_local_calibration(model, algorithm) if forward_loop is not None: # get the number of arguments of forward_loop num_args = len(inspect.signature(forward_loop).parameters) @@ -502,6 +533,7 @@ def forward_loop(model) -> None: """ quantize_config = QuantizeConfig(**dict(config)) _refuse_awq_on_a_sharded_fsdp2_model(model, config.get("algorithm")) + _refuse_rank_local_calibration(model, config.get("algorithm")) if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 8491600f932..8ae719f66d6 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -60,6 +60,7 @@ "is_initialized", "is_master", "rank", + "recorded_distribution_plan", "set_distribution_plan", "shard_dataloader", "size", @@ -403,7 +404,7 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: a launch sets its group up, each of its processes runs alone with a model of its own, so a model with no plan reads as a single process's there. """ - plan = getattr(model, _PLAN_ATTR, None) + plan = recorded_distribution_plan(model) if plan is not None: return plan sharded = is_fsdp2_model(model) @@ -417,6 +418,15 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: return DistributionPlan() +def recorded_distribution_plan(model: torch.nn.Module) -> DistributionPlan | None: + """The plan :func:`set_distribution_plan` recorded on ``model``, or ``None`` if none was. + + For code that also serves models distributed without a plan, such as Megatron's, which + :func:`distribution_plan` refuses in a run of several ranks. + """ + return getattr(model, _PLAN_ATTR, None) + + _V = TypeVar("_V") diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index be5d0ccea65..d5bec5a71a7 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -454,3 +454,31 @@ def test_the_ep_mesh_comes_with_an_expert_parallel_plan_only(plan, mesh): with pytest.raises(ValueError, match="calibrates with the ep_mesh"): distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}], ep_mesh=mesh) + + +# ------------------------------------------------------------------------ rank-local calibration + + +@pytest.mark.parametrize("entry", ["quantize", "calibrate"]) +def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model(entry): + """Each rank calibrates on its own data, and AWQ keeps what its search gathers on the rank: the + ranks' copies of a weight would come out different. The refusal comes before any change.""" + import modelopt.torch.quantization as mtq + + model = nn.Sequential(nn.Linear(16, 16)) + set_distribution_plan(model, DistributionPlan(world_size=2)) + if entry == "calibrate": + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16))) + modules = [type(module) for module in model.modules()] + state = {name: tensor.clone() for name, tensor in model.state_dict().items()} + + with pytest.raises(NotImplementedError, match=r"awq_lite is not supported .* over 2 ranks"): + if entry == "quantize": + mtq.quantize(model, mtq.INT4_AWQ_CFG, lambda m: m(torch.randn(2, 16))) + else: + mtq.calibrate(model, "awq_lite", lambda m: m(torch.randn(2, 16))) + + assert [type(module) for module in model.modules()] == modules + after = model.state_dict() + assert after.keys() == state.keys() + assert all(torch.equal(after[name], tensor) for name, tensor in state.items()) From 83c6937b6def38f0115c06828938d7f10a098bc6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:18:39 +0000 Subject: [PATCH 18/60] Refuse rank-local calibration that another algorithm runs for its scales nvfp4_act_headroom calibrates its weight scales with its weight_scale_algorithm, and lsq starts from its scale_algorithm's; either may be local_hessian, which keeps its statistics on the rank as well. The refusal now reads the methods those settings name too. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 14 ++++++++-- .../utils/test_distributed_calibration.py | 28 +++++++++++++++---- 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 81d74431a30..09cc1014dfd 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -83,6 +83,17 @@ ) +def _calibration_methods(algorithm: QuantizeAlgoCfgType) -> list[str]: + """The calibration methods ``algorithm`` runs, with those it runs for its scales.""" + methods = [] + for mode, config in get_modelike_from_algo_cfg(algorithm): + methods.append(mode.removesuffix("_calibrate")) + for key in ("weight_scale_algorithm", "scale_algorithm"): + if config.get(key) is not None: + methods += _calibration_methods(config[key]) + return methods + + def _refuse_rank_local_calibration(model: nn.Module, algorithm: QuantizeAlgoCfgType) -> None: """Refuse a rank-local calibration algorithm on a model distributed over several ranks. @@ -92,8 +103,7 @@ def _refuse_rank_local_calibration(model: nn.Module, algorithm: QuantizeAlgoCfgT plan = dist.recorded_distribution_plan(model) if plan is None or plan.world_size == 1: return - methods = [mode.removesuffix("_calibrate") for mode, _ in get_modelike_from_algo_cfg(algorithm)] - refused = [method for method in methods if method in _RANK_LOCAL_CALIBRATION] + refused = [m for m in _calibration_methods(algorithm) if m in _RANK_LOCAL_CALIBRATION] if refused: raise NotImplementedError( f"Calibrating with {', '.join(refused)} is not supported on a model distributed over " diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index d5bec5a71a7..07dff97dbdc 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -459,10 +459,25 @@ def test_the_ep_mesh_comes_with_an_expert_parallel_plan_only(plan, mesh): # ------------------------------------------------------------------------ rank-local calibration -@pytest.mark.parametrize("entry", ["quantize", "calibrate"]) -def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model(entry): +@pytest.mark.parametrize( + ("entry", "algorithm", "refused"), + [ + ("quantize", {"method": "awq_lite"}, "awq_lite"), + ("calibrate", "awq_lite", "awq_lite"), + ( + "calibrate", + {"method": "nvfp4_act_headroom", "weight_scale_algorithm": "local_hessian"}, + "local_hessian", + ), + ], + ids=["quantize", "calibrate", "for-the-weight-scales"], +) +def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model( + entry, algorithm, refused +): """Each rank calibrates on its own data, and AWQ keeps what its search gathers on the rank: the - ranks' copies of a weight would come out different. The refusal comes before any change.""" + ranks' copies of a weight would come out different. So would local_hessian, run by another + algorithm for its weight scales. The refusal comes before any change.""" import modelopt.torch.quantization as mtq model = nn.Sequential(nn.Linear(16, 16)) @@ -472,11 +487,12 @@ def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model modules = [type(module) for module in model.modules()] state = {name: tensor.clone() for name, tensor in model.state_dict().items()} - with pytest.raises(NotImplementedError, match=r"awq_lite is not supported .* over 2 ranks"): + with pytest.raises(NotImplementedError, match=rf"{refused} is not supported .* over 2 ranks"): if entry == "quantize": - mtq.quantize(model, mtq.INT4_AWQ_CFG, lambda m: m(torch.randn(2, 16))) + config = {**mtq.INT4_AWQ_CFG, "algorithm": algorithm} + mtq.quantize(model, config, lambda m: m(torch.randn(2, 16))) else: - mtq.calibrate(model, "awq_lite", lambda m: m(torch.randn(2, 16))) + mtq.calibrate(model, algorithm, lambda m: m(torch.randn(2, 16))) assert [type(module) for module in model.modules()] == modules after = model.state_dict() From f77b7e84217f10f8dedc24ae8b70336c5f2cb3ef Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:56:42 +0000 Subject: [PATCH 19/60] Refuse max and MSE told not to sync, on a model distributed over several ranks Both take distributed_sync. Turned off, max_calibrate returns before its data-, expert- and tensor-parallel amax syncs, and mse passes the option on, so ranks calibrating on their own data keep different scales for the same quantizer. A config that turns it off, nested as another algorithm's scale algorithm included, is refused like the rank-local algorithms. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 24 ++++++++++++------- .../utils/test_distributed_calibration.py | 9 +++++-- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 09cc1014dfd..47cd1e6ea6d 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -72,7 +72,8 @@ "temporarily_fold_weights", ] -# Calibration algorithms whose statistics stay on the rank that gathered them. +# Calibration algorithms whose statistics stay on the rank that gathered them. Any other keeps its +# own there too when its config turns distributed_sync off. _RANK_LOCAL_CALIBRATION = ( "awq_lite", "awq_clip", @@ -83,15 +84,19 @@ ) -def _calibration_methods(algorithm: QuantizeAlgoCfgType) -> list[str]: - """The calibration methods ``algorithm`` runs, with those it runs for its scales.""" - methods = [] +def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: + """The rank-local calibrations ``algorithm`` runs, with those it runs for its scales.""" + found = [] for mode, config in get_modelike_from_algo_cfg(algorithm): - methods.append(mode.removesuffix("_calibrate")) + method = mode.removesuffix("_calibrate") + if method in _RANK_LOCAL_CALIBRATION: + found.append(method) + elif "distributed_sync" in config and not config["distributed_sync"]: + found.append(f"{method} with distributed_sync={config['distributed_sync']}") for key in ("weight_scale_algorithm", "scale_algorithm"): if config.get(key) is not None: - methods += _calibration_methods(config[key]) - return methods + found += _rank_local_calibrations(config[key]) + return found def _refuse_rank_local_calibration(model: nn.Module, algorithm: QuantizeAlgoCfgType) -> None: @@ -103,13 +108,14 @@ def _refuse_rank_local_calibration(model: nn.Module, algorithm: QuantizeAlgoCfgT plan = dist.recorded_distribution_plan(model) if plan is None or plan.world_size == 1: return - refused = [m for m in _calibration_methods(algorithm) if m in _RANK_LOCAL_CALIBRATION] + refused = _rank_local_calibrations(algorithm) if refused: raise NotImplementedError( f"Calibrating with {', '.join(refused)} is not supported on a model distributed over " f"{plan.world_size} ranks. Each rank calibrates on its own data and this calibration " "keeps what it gathers on the rank, so the ranks' copies of a weight would come out " - "different. Calibrate with max, MSE, GPTQ or SmoothQuant, or in a single process." + "different. Calibrate with max, MSE, GPTQ or SmoothQuant, with distributed_sync on, or " + "in a single process." ) diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 07dff97dbdc..e6c124fb94d 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -469,15 +469,20 @@ def test_the_ep_mesh_comes_with_an_expert_parallel_plan_only(plan, mesh): {"method": "nvfp4_act_headroom", "weight_scale_algorithm": "local_hessian"}, "local_hessian", ), + ( + "calibrate", + {"method": "max", "distributed_sync": False}, + "max with distributed_sync=False", + ), ], - ids=["quantize", "calibrate", "for-the-weight-scales"], + ids=["quantize", "calibrate", "for-the-weight-scales", "max-without-its-sync"], ) def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model( entry, algorithm, refused ): """Each rank calibrates on its own data, and AWQ keeps what its search gathers on the rank: the ranks' copies of a weight would come out different. So would local_hessian, run by another - algorithm for its weight scales. The refusal comes before any change.""" + algorithm for its weight scales, and max told not to sync. The refusal comes before any change.""" import modelopt.torch.quantization as mtq model = nn.Sequential(nn.Linear(16, 16)) From 58308c827995ba5fbdbc1f6b5dc948060603b54a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:56:44 +0000 Subject: [PATCH 20/60] Import the quantization package at the top of the distributed calibration tests Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_distributed_calibration.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index e6c124fb94d..6461dcbd5e7 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -31,6 +31,7 @@ from torch import nn from torch.utils.data import DataLoader, IterableDataset +import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as mto_dist from modelopt.torch.utils.distributed import ( DistributionPlan, @@ -483,8 +484,6 @@ def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model """Each rank calibrates on its own data, and AWQ keeps what its search gathers on the rank: the ranks' copies of a weight would come out different. So would local_hessian, run by another algorithm for its weight scales, and max told not to sync. The refusal comes before any change.""" - import modelopt.torch.quantization as mtq - model = nn.Sequential(nn.Linear(16, 16)) set_distribution_plan(model, DistributionPlan(world_size=2)) if entry == "calibrate": From fb6471ee3fa8d1b85e20a44492ab9a2cd71dd466 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:58:41 +0000 Subject: [PATCH 21/60] Record the plan hf_ptq's FSDP2 load distributes the model by From this PR on, get_max_batch_size, which hf_ptq calls with its default --batch_size 0, reads the model's distribution plan, and from the no-gather export on so does export_hf_checkpoint. Each refuses an FSDP2-sharded model with none recorded, and hf_ptq's --use_fsdp2 load, parallel_load_and_prepare_fsdp2, records none: it loads the model whole until hf_ptq moves onto load_pretrained. Record the FSDP2 plan right after that load. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 5 +++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 25 +++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 47d6a8b2c19..c79c25ccf98 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -99,6 +99,7 @@ get_max_batch_size, get_supported_datasets, ) +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan from modelopt.torch.utils.memory_monitor import launch_memory_monitor from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -339,6 +340,10 @@ def load_model(args: argparse.Namespace): attn_implementation=args.attn_implementation, hf_config=hf_config, ) + # Calibration sizes its batch, and the export writes, by how the model is distributed. + set_distribution_plan( + full_model, DistributionPlan(args.dist_state.world_size, dense="fsdp") + ) elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.local_checkpoint_path, diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 2f0b425de40..937868eb7c3 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -480,6 +480,31 @@ def distributed_load(*args, **kwargs): hf_ptq.load_model(args) +def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path): + """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an + FSDP2-sharded model without one is refused rather than taken for a single process's.""" + from transformers import Qwen3Config + + from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan + + Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path) + monkeypatch.setenv("RANK", "0") + hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2") + args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=2) + args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path) # as main() sets them + model = torch.nn.Linear(1, 1) + monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *args, **kwargs: model) + + def loaded(model): + raise RuntimeError("loaded") + + monkeypatch.setattr(hf_ptq, "prepare_model_for_calibration", loaded) + + with pytest.raises(RuntimeError, match="loaded"): + hf_ptq.load_model(args) + assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") + + def test_mlflow_flag_defaults_the_experiment_name(monkeypatch): monkeypatch.setattr(getpass, "getuser", lambda: "tester") hf_ptq, args = _parse_hf_ptq_args( From beb2933cebb8c4cfac12ba2a23ecb6f62de56f59 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 00:33:45 +0000 Subject: [PATCH 22/60] Destroy the group the batch-size test's pooled job creates The job reduced over a group with a 10-second timeout and left it registered on the session's persistent cpu_dist_workers until the pool shut down; it now destroys it. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_dataset_utils.py | 30 +++++++++++--------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 2658c728fb2..5b50b8c263f 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -466,19 +466,23 @@ def mem_get_info(device): # A rank left waiting in the reduction gives up in seconds on this group. group = torch.distributed.new_group(timeout=timedelta(seconds=10)) gather = partial(torch.distributed.all_gather_object, group=group) - with ( - patch.object(torch.distributed, "all_gather_object", gather), - patch("torch.cuda.empty_cache"), - patch( - "torch.cuda.get_device_properties", return_value=SimpleNamespace(total_memory=10**12) - ), - patch("torch.cuda.device_count", return_value=4), - patch("torch.cuda.mem_get_info", side_effect=mem_get_info), - patch("torch.cuda.max_memory_allocated", return_value=0), - ): - batch_size = get_max_batch_size( - model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) - ) + try: + with ( + patch.object(torch.distributed, "all_gather_object", gather), + patch("torch.cuda.empty_cache"), + patch( + "torch.cuda.get_device_properties", + return_value=SimpleNamespace(total_memory=10**12), + ), + patch("torch.cuda.device_count", return_value=4), + patch("torch.cuda.mem_get_info", side_effect=mem_get_info), + patch("torch.cuda.max_memory_allocated", return_value=0), + ): + batch_size = get_max_batch_size( + model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) + ) + finally: + torch.distributed.destroy_process_group(group) assert measured == [torch.device("cuda", rank)] * 2 # its own GPU (LOCAL_RANK) only assert model.batch_sizes == [1] # one probe: an OOM retry would run more forwards on some ranks From 7a63bf4489b393c4f12918063c0c778c7a8415e0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:07:26 +0000 Subject: [PATCH 23/60] Refuse a loader that picks its own samples, and calibrate without gradients shard_dataloader replaces a map-style loader's sampler with the rank's DistributedSampler, so a loader that picked its own samples (a SubsetRandomSampler, or a batch_sampler) silently calibrated on the whole dataset; it is now refused. distributed_forward_loop runs its forwards under no_grad, as the single-node _forward_loop does, and the two new public functions carry type annotations. The comments on releasing the CUDA cache keep their reason and drop the measurements. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 58 ++++++++++++------- .../utils/test_distributed_calibration.py | 17 +++++- 2 files changed, 52 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 8ae719f66d6..6810fd53b55 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -34,7 +34,14 @@ import torch.distributed from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard from torch.distributed.tensor import DTensor -from torch.utils.data import DataLoader, IterableDataset, get_worker_info +from torch.utils.data import ( + BatchSampler, + DataLoader, + IterableDataset, + RandomSampler, + SequentialSampler, + get_worker_info, +) from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm @@ -722,12 +729,14 @@ def __iter__(self): yield group[self.rank] -def shard_dataloader(loader, rank: int, world_size: int): +def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoader: """Wrap a DataLoader so each rank sees a unique shard. For map-style datasets a ``DistributedSampler`` is used. ``drop_last=False`` keeps per-rank batch counts equal (else a rank exits calibration early and hangs the others on the per-step collectives), - at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. For + at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. It + replaces the loader's sampler, so a shuffling one becomes the rank's unshuffled share, and one that + picks samples itself (a subset, or a batch sampler) is refused rather than silently widened. For streaming ``IterableDataset`` datasets (which have no ``len()`` and so can't use a sampler) each rank strides its own deterministic copy of the stream via :class:`_StridedIterableShard`. Every DataLoader worker must see that whole stream: a dataset that splits itself across workers by @@ -736,6 +745,18 @@ def shard_dataloader(loader, rank: int, world_size: int): Forwards all non-sampler DataLoader settings from ``loader`` (workers, pinning, prefetch, init fn, generator, ...). """ + if not isinstance(loader.dataset, IterableDataset): + picker = ( + loader.batch_sampler + if loader.batch_sampler is not None and type(loader.batch_sampler) is not BatchSampler + else loader.sampler + ) + if type(picker) not in (SequentialSampler, RandomSampler, DistributedSampler): + raise ValueError( + f"shard_dataloader: the loader's {type(picker).__name__} picks which samples it " + "reads, and the rank's DistributedSampler would replace it; pass a loader over " + "just the samples to calibrate on (a torch.utils.data.Subset) instead." + ) if isinstance(loader.dataset, IterableDataset): return DataLoader( _StridedIterableShard(loader.dataset, rank, world_size), @@ -846,7 +867,12 @@ def _release_cached_cuda_memory(): torch.cuda.empty_cache() -def distributed_forward_loop(model, dataloader, device=None, ep_mesh=None): +def distributed_forward_loop( + model: torch.nn.Module, + dataloader: DataLoader, + device: torch.device | str | None = None, + ep_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, +) -> Callable[[torch.nn.Module], None]: """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. Moves each batch onto ``device``, when given, and runs the captured ``model`` forward (the EP @@ -857,11 +883,9 @@ def distributed_forward_loop(model, dataloader, device=None, ep_mesh=None): ``forward_loop`` callback an inner module that this closure ignores, calling the captured top-level ``model`` instead so the full distributed forward (and its collectives) runs. - KV/state caching is disabled for the duration (via ``_disable_use_cache``), mirroring the - single-node ``_forward_loop``. This distributed path previously omitted the guard, so it ran - calibration at ``config.use_cache`` (default ``True``); for hybrid Mamba/attention models - (e.g. NemotronH) the cache is mutated in-place, which both wastes memory and breaks - calibration correctness. + The forwards run without gradients and with KV/state caching disabled (``_disable_use_cache``), as + the single-node ``_forward_loop`` runs them: a hybrid Mamba/attention model mutates its cache in + place. """ # Local import to avoid any import cycle within modelopt.torch.utils. from modelopt.torch.utils.dataset_utils import _disable_use_cache @@ -875,12 +899,8 @@ def distributed_forward_loop(model, dataloader, device=None, ep_mesh=None): def calibrate(_inner_model): _reinstall_a2a_forwards(model, ep_mesh) - with _disable_use_cache(model): - # Reclaim the caching allocator's reserved-but-unused pool left by the load/quantizer-insertion - # memory peak. Without this, torch hoards ~the whole GPU (e.g. 91 GB live weights vs 275 GB - # reserved), leaving ~0 GB free, so the distributed a2a's NCCL comm buffers cannot be allocated - # ("Failed to CUDA calloc") at the very first all_to_all. Freeing the cache hands that memory - # back to the driver for NCCL + the calibration activations. + with _disable_use_cache(model), torch.no_grad(): + # Hand the caching allocator's unused pool back: NCCL allocates its buffers outside it. _release_cached_cuda_memory() for batch in tqdm(dataloader, desc="Calibrating", disable=not is_master()): if device is not None: @@ -889,13 +909,7 @@ def calibrate(_inner_model): for k, v in batch.items() } model(**batch) - # Symmetric to the pre-forward reclaim above, and for the same reason one stage later: - # the calibration forwards leave their activations in the caching allocator's reserved - # pool, so whatever runs next (amax sync, export collectives) can find the device full - # even though little is live. Measured on Kimi-K3 (EP64, 184 GiB cards) right after the - # single calib forward: alloc 86.0 GiB but reserved 156.6 GiB -> only 0.7 GiB free, and - # the job then hung until the 2 h wall clock instead of exporting. NCCL comm buffers are - # allocated OUTSIDE the torch allocator, so a full reserved pool starves them. + # Again for what runs next (amax sync, export collectives), past the forwards' pool. _release_cached_cuda_memory() return calibrate diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 6461dcbd5e7..44ef16b31d0 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -29,7 +29,7 @@ import pytest import torch from torch import nn -from torch.utils.data import DataLoader, IterableDataset +from torch.utils.data import DataLoader, IterableDataset, SubsetRandomSampler import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as mto_dist @@ -153,6 +153,21 @@ def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch) assert sorted(yielded) == list(range(rank, n, world_size)) +@pytest.mark.parametrize( + "picker", + [{"sampler": SubsetRandomSampler(range(4))}, {"batch_sampler": [[0, 1], [2, 3]]}], + ids=["subset-sampler", "batch-sampler"], +) +def test_a_loader_that_picks_its_own_samples_is_refused(picker): + """The rank's DistributedSampler would replace what it picks, silently calibrating on the whole + dataset instead.""" + batching = {} if "batch_sampler" in picker else {"batch_size": 2} + loader = DataLoader(list(range(8)), collate_fn=_collate, **batching, **picker) + + with pytest.raises(ValueError, match="picks which samples it reads"): + shard_dataloader(loader, 0, 2) + + @pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) def test_sharded_loader_keeps_the_original_loader_settings(dataset): """Batching, workers, pinning, timeout, worker init, RNG and prefetching carry over to the From bdbc3fe13eb8e5c87d26ac8b62a8a31b5cfb4add Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:14:27 +0000 Subject: [PATCH 24/60] Carry a plain BatchSampler's batching into each rank's calibration shard shard_dataloader took an exactly-BatchSampler batch sampler for the one DataLoader builds from batch_size and checked the loader's sampler instead. A loader given batch_sampler=BatchSampler(SubsetRandomSampler(...), 4, False) still carries a default SequentialSampler, so its subset passed the picker refusal and was widened to the whole dataset, and its batch_size of None rebuilt the shard with auto-batching off. The shard now takes the picker from the plain BatchSampler itself, refusing a subset there too, and batches by that sampler's batch_size and drop_last. A loader with auto-batching off is refused. distributed_forward_loop refuses a batch that is not a mapping of forward keyword arguments by name, rather than failing on batch.items() inside the loop. auto_quantize refuses a rank-local calibration on a distributed model in its format check, before it converts the model, as quantize and calibrate do. _max_batch_size builds its probe input only where it probes. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 1 + modelopt/torch/utils/dataset_utils.py | 7 +- modelopt/torch/utils/distributed.py | 44 +++++++----- .../utils/test_distributed_calibration.py | 67 +++++++++++++++---- 4 files changed, 85 insertions(+), 34 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 47cd1e6ea6d..06643e07849 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -1109,6 +1109,7 @@ def forward_backward_step(model, batch) -> None: ) # Every candidate, before the search calibrates any: calibrate() would refuse mid-search. _refuse_awq_on_a_sharded_fsdp2_model(model, algo, candidate=name) + _refuse_rank_local_calibration(model, algo) # Select the appropriate searcher based on method if method_name not in AUTO_QUANTIZE_SEARCHERS: diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index a5c2cb1329d..ef798596704 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1103,18 +1103,15 @@ def _expand_to(batch: int) -> torch.Tensor: ] ) - target_input = _expand_to(target_data_batch) - # For some models on multi GPU, we observe the memory per batch is not a constant. # So we just test the target batch size and make sure we do not go OOM. - while target_data_batch > 1 and device is None: + while device is None and target_data_batch > 1: with torch.set_grad_enabled(enable_grad): try: - infer_method(target_input) + infer_method(_expand_to(target_data_batch)) break except torch.cuda.OutOfMemoryError: # pragma: no cover - GPU OOM retry path target_data_batch = target_data_batch // 2 # pragma: no cover - target_input = _expand_to(target_data_batch) # pragma: no cover torch.cuda.empty_cache() # pragma: no cover # Regulate the data batch target to be 1, 2, 4, 8, 12, ..., capped at 64 diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 6810fd53b55..deb0488ce7a 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -23,7 +23,7 @@ import sys import time import traceback -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager, suppress from dataclasses import dataclass from datetime import timedelta @@ -736,21 +736,26 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad counts equal (else a rank exits calibration early and hangs the others on the per-step collectives), at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. It replaces the loader's sampler, so a shuffling one becomes the rank's unshuffled share, and one that - picks samples itself (a subset, or a batch sampler) is refused rather than silently widened. For - streaming ``IterableDataset`` datasets (which have no ``len()`` and so can't use a sampler) each rank - strides its own deterministic copy of the stream via :class:`_StridedIterableShard`. Every DataLoader - worker must see that whole stream: a dataset that splits itself across workers by - ``get_worker_info()`` would be split twice and lose samples, so load one with ``num_workers=0``. - - Forwards all non-sampler DataLoader settings from ``loader`` (workers, pinning, - prefetch, init fn, generator, ...). + picks samples itself (a subset, or a batch sampler other than a plain ``BatchSampler``) is refused + rather than silently widened. For streaming ``IterableDataset`` datasets (which have no ``len()`` + and so can't use a sampler) each rank strides its own deterministic copy of the stream via + :class:`_StridedIterableShard`. Every DataLoader worker must see that whole stream: a dataset that + splits itself across workers by ``get_worker_info()`` would be split twice and lose samples, so + load one with ``num_workers=0``. + + Forwards all non-sampler DataLoader settings from ``loader`` (batching, workers, pinning, + prefetch, init fn, generator, ...). A loader with auto-batching off (``batch_size=None``) is + refused. """ - if not isinstance(loader.dataset, IterableDataset): - picker = ( - loader.batch_sampler - if loader.batch_sampler is not None and type(loader.batch_sampler) is not BatchSampler - else loader.sampler + batching = loader.batch_sampler + if batching is None: + raise ValueError( + "shard_dataloader: the loader has auto-batching off (batch_size=None), so each rank " + "would calibrate on one dataset item per step; pass a loader that batches by batch_size." ) + if not isinstance(loader.dataset, IterableDataset): + # A plain BatchSampler, built from batch_size or passed in, batches what its sampler picks. + picker = batching.sampler if type(batching) is BatchSampler else batching if type(picker) not in (SequentialSampler, RandomSampler, DistributedSampler): raise ValueError( f"shard_dataloader: the loader's {type(picker).__name__} picks which samples it " @@ -782,8 +787,8 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad ) return DataLoader( loader.dataset, - batch_size=loader.batch_size, - drop_last=loader.drop_last, + batch_size=batching.batch_size, + drop_last=batching.drop_last, sampler=sampler, collate_fn=loader.collate_fn, num_workers=loader.num_workers, @@ -885,7 +890,7 @@ def distributed_forward_loop( The forwards run without gradients and with KV/state caching disabled (``_disable_use_cache``), as the single-node ``_forward_loop`` runs them: a hybrid Mamba/attention model mutates its cache in - place. + place. Each batch is a mapping of the forward's keyword arguments, as ``model(**batch)`` takes. """ # Local import to avoid any import cycle within modelopt.torch.utils. from modelopt.torch.utils.dataset_utils import _disable_use_cache @@ -903,6 +908,11 @@ def calibrate(_inner_model): # Hand the caching allocator's unused pool back: NCCL allocates its buffers outside it. _release_cached_cuda_memory() for batch in tqdm(dataloader, desc="Calibrating", disable=not is_master()): + if not isinstance(batch, Mapping): + raise TypeError( + f"distributed_forward_loop: the loader yields {type(batch).__name__}, not a " + "mapping of the forward's keyword arguments." + ) if device is not None: batch = { k: (v.to(device) if isinstance(v, torch.Tensor) else v) diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 44ef16b31d0..29ebdf12cfb 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -29,7 +29,13 @@ import pytest import torch from torch import nn -from torch.utils.data import DataLoader, IterableDataset, SubsetRandomSampler +from torch.utils.data import ( + BatchSampler, + DataLoader, + IterableDataset, + SequentialSampler, + SubsetRandomSampler, +) import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as mto_dist @@ -154,20 +160,38 @@ def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch) @pytest.mark.parametrize( - "picker", - [{"sampler": SubsetRandomSampler(range(4))}, {"batch_sampler": [[0, 1], [2, 3]]}], - ids=["subset-sampler", "batch-sampler"], + ("loader_kwargs", "refusal"), + [ + ({"sampler": SubsetRandomSampler(range(4)), "batch_size": 2}, "picks which samples"), + ({"batch_sampler": [[0, 1], [2, 3]]}, "picks which samples"), + ( + {"batch_sampler": BatchSampler(SubsetRandomSampler(range(4)), 2, drop_last=False)}, + "picks which samples", + ), + ({"batch_size": None}, "auto-batching off"), + ], + ids=["subset-sampler", "batch-sampler", "subset-in-a-plain-batch-sampler", "auto-batching-off"], ) -def test_a_loader_that_picks_its_own_samples_is_refused(picker): - """The rank's DistributedSampler would replace what it picks, silently calibrating on the whole - dataset instead.""" - batching = {} if "batch_sampler" in picker else {"batch_size": 2} - loader = DataLoader(list(range(8)), collate_fn=_collate, **batching, **picker) +def test_a_loader_the_shard_cannot_reproduce_is_refused(loader_kwargs, refusal): + """The rank's DistributedSampler would replace what the loader picks, silently calibrating on + the whole dataset instead, also when a plain BatchSampler batches the picks. A loader with + auto-batching off would calibrate one dataset item per step.""" + loader = DataLoader(list(range(8)), collate_fn=_collate, **loader_kwargs) - with pytest.raises(ValueError, match="picks which samples it reads"): + with pytest.raises(ValueError, match=refusal): shard_dataloader(loader, 0, 2) +def test_a_plain_batch_sampler_keeps_its_batching_in_the_shard(): + """A loader given a BatchSampler reports ``batch_size=None``; its shard still batches by the + sampler's size and drop_last.""" + dataset = list(range(10)) + batch_sampler = BatchSampler(SequentialSampler(dataset), 2, drop_last=True) + loader = DataLoader(dataset, batch_sampler=batch_sampler, collate_fn=_collate) + + assert _batches(shard_dataloader(loader, 1, 2)) == [[1, 3], [5, 7]] + + @pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) def test_sharded_loader_keeps_the_original_loader_settings(dataset): """Batching, workers, pinning, timeout, worker init, RNG and prefetching carry over to the @@ -472,6 +496,15 @@ def test_the_ep_mesh_comes_with_an_expert_parallel_plan_only(plan, mesh): distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}], ep_mesh=mesh) +def test_forward_loop_refuses_a_batch_that_is_not_a_mapping(): + """The loop runs ``model(**batch)``; a loader of bare tensors is refused as such at its first + batch, rather than failing on ``batch.items()`` inside the loop.""" + model = _Recorder() + + with pytest.raises(TypeError, match="yields Tensor"): + distributed_forward_loop(model, [torch.zeros(1, 1)], device="meta")(model) + + # ------------------------------------------------------------------------ rank-local calibration @@ -490,8 +523,9 @@ def test_the_ep_mesh_comes_with_an_expert_parallel_plan_only(plan, mesh): {"method": "max", "distributed_sync": False}, "max with distributed_sync=False", ), + ("auto_quantize", {"method": "awq_lite"}, "awq_lite"), ], - ids=["quantize", "calibrate", "for-the-weight-scales", "max-without-its-sync"], + ids=["quantize", "calibrate", "for-the-weight-scales", "max-without-its-sync", "auto_quantize"], ) def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model( entry, algorithm, refused @@ -510,8 +544,17 @@ def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model if entry == "quantize": config = {**mtq.INT4_AWQ_CFG, "algorithm": algorithm} mtq.quantize(model, config, lambda m: m(torch.randn(2, 16))) - else: + elif entry == "calibrate": mtq.calibrate(model, algorithm, lambda m: m(torch.randn(2, 16))) + else: + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[({**mtq.INT4_AWQ_CFG, "algorithm": algorithm}, "awq")], + data_loader=[torch.randn(2, 16)], + forward_step=lambda m, batch: m(batch), + loss_func=lambda output, batch: output.sum(), + ) assert [type(module) for module in model.modules()] == modules after = model.state_dict() From a1a54d93a084bbc941e49bb8ea1141dd6230b8c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:33:56 +0000 Subject: [PATCH 25/60] Shard a loader with auto-batching off instead of refusing it A loader built with batch_size=None has auto-batching off: it reads one dataset item per step, a pre-batched item included, and its rank's shard, rebuilt with batch_size=None, reads the same way. Sharding it was correct, so refusing it only turned such loaders away. shard_dataloader shards it again, on the map-style and the streaming path, and still refuses a sampler there that picks its own samples. The plain-BatchSampler look-through and the custom batch sampler refusal stay. The tests shard a map-style and a streaming loader with auto-batching off one item per step, and refuse a subset sampler under one. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 20 +++---- .../utils/test_distributed_calibration.py | 57 ++++++++++++------- 2 files changed, 44 insertions(+), 33 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index deb0488ce7a..8719bb61a83 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -744,18 +744,16 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad load one with ``num_workers=0``. Forwards all non-sampler DataLoader settings from ``loader`` (batching, workers, pinning, - prefetch, init fn, generator, ...). A loader with auto-batching off (``batch_size=None``) is - refused. + prefetch, init fn, generator, ...). """ batching = loader.batch_sampler - if batching is None: - raise ValueError( - "shard_dataloader: the loader has auto-batching off (batch_size=None), so each rank " - "would calibrate on one dataset item per step; pass a loader that batches by batch_size." - ) if not isinstance(loader.dataset, IterableDataset): - # A plain BatchSampler, built from batch_size or passed in, batches what its sampler picks. - picker = batching.sampler if type(batching) is BatchSampler else batching + # A plain BatchSampler, built from batch_size or passed in, batches what its sampler picks; + # with auto-batching off there is none, and the loader's sampler picks each item. + if type(batching) is BatchSampler: + picker = batching.sampler + else: + picker = loader.sampler if batching is None else batching if type(picker) not in (SequentialSampler, RandomSampler, DistributedSampler): raise ValueError( f"shard_dataloader: the loader's {type(picker).__name__} picks which samples it " @@ -787,8 +785,8 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad ) return DataLoader( loader.dataset, - batch_size=batching.batch_size, - drop_last=batching.drop_last, + batch_size=loader.batch_size if batching is None else batching.batch_size, + drop_last=loader.drop_last if batching is None else batching.drop_last, sampler=sampler, collate_fn=loader.collate_fn, num_workers=loader.num_workers, diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 29ebdf12cfb..38eb1d316a2 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -160,36 +160,49 @@ def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch) @pytest.mark.parametrize( - ("loader_kwargs", "refusal"), + "picker", [ - ({"sampler": SubsetRandomSampler(range(4)), "batch_size": 2}, "picks which samples"), - ({"batch_sampler": [[0, 1], [2, 3]]}, "picks which samples"), - ( - {"batch_sampler": BatchSampler(SubsetRandomSampler(range(4)), 2, drop_last=False)}, - "picks which samples", - ), - ({"batch_size": None}, "auto-batching off"), + {"sampler": SubsetRandomSampler(range(4)), "batch_size": 2}, + {"sampler": SubsetRandomSampler(range(4)), "batch_size": None}, + {"batch_sampler": [[0, 1], [2, 3]]}, + {"batch_sampler": BatchSampler(SubsetRandomSampler(range(4)), 2, drop_last=False)}, + ], + ids=[ + "subset-sampler", + "subset-sampler-unbatched", + "batch-sampler", + "subset-in-a-plain-batch-sampler", ], - ids=["subset-sampler", "batch-sampler", "subset-in-a-plain-batch-sampler", "auto-batching-off"], ) -def test_a_loader_the_shard_cannot_reproduce_is_refused(loader_kwargs, refusal): - """The rank's DistributedSampler would replace what the loader picks, silently calibrating on - the whole dataset instead, also when a plain BatchSampler batches the picks. A loader with - auto-batching off would calibrate one dataset item per step.""" - loader = DataLoader(list(range(8)), collate_fn=_collate, **loader_kwargs) +def test_a_loader_that_picks_its_own_samples_is_refused(picker): + """The rank's DistributedSampler would replace what it picks, silently calibrating on the whole + dataset instead, however the loader batches the picks.""" + loader = DataLoader(list(range(8)), collate_fn=_collate, **picker) - with pytest.raises(ValueError, match=refusal): + with pytest.raises(ValueError, match="picks which samples it reads"): shard_dataloader(loader, 0, 2) -def test_a_plain_batch_sampler_keeps_its_batching_in_the_shard(): - """A loader given a BatchSampler reports ``batch_size=None``; its shard still batches by the - sampler's size and drop_last.""" - dataset = list(range(10)) - batch_sampler = BatchSampler(SequentialSampler(dataset), 2, drop_last=True) - loader = DataLoader(dataset, batch_sampler=batch_sampler, collate_fn=_collate) +@pytest.mark.parametrize( + ("kind", "batching", "batches"), + [ + ( + "map-style", + {"batch_sampler": BatchSampler(SequentialSampler(range(10)), 2, drop_last=True)}, + [[1, 3], [5, 7]], + ), + ("map-style", {"batch_size": None}, [[1], [3], [5], [7], [9]]), + ("iterable", {"batch_size": None}, [[1], [3], [5], [7], [9]]), + ], + ids=["plain-batch-sampler", "auto-batching-off", "auto-batching-off-stream"], +) +def test_the_shard_batches_as_the_loader_does(kind, batching, batches): + """A loader given a plain BatchSampler reports ``batch_size=None``; its shard still batches by the + sampler's size and drop_last. A loader with auto-batching off reads one item per step, and so + does its shard.""" + loader = DataLoader(_dataset(kind, 10), collate_fn=_collate, **batching) - assert _batches(shard_dataloader(loader, 1, 2)) == [[1, 3], [5, 7]] + assert _batches(shard_dataloader(loader, 1, 2)) == batches @pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) From ae57933fe3f729e13c95a128bb7851089f94e23f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:13:45 +0000 Subject: [PATCH 26/60] Refuse expert-parallel calibration that would merge experts' scales Under expert parallelism each rank's per-expert quantizers hold its own experts, and max and GPTQ sync them over the data-parallel group their module's ParallelState names, which defaults to every rank: a slot holding expert 0 on one rank was merged with the slot holding another expert on the next, silently. quantize, calibrate and auto_quantize now refuse any calibration on an expert-parallel plan whose modules carry no expert-parallel group, before converting anything; a model whose routed experts carry their ParallelState calibrates as before. The guard is renamed for what it now refuses. get_max_batch_size reads the recorded plan, as that guard does: a model with no plan recorded in a run of several ranks (a Megatron model, say) is measured on every visible GPU again, as on main, instead of being refused. distributed_forward_loop also refuses an ep_mesh whose size is not the plan's ep_size, and the docstrings say that tensors nested in a batch value are not moved and what each DataLoader worker of a streaming shard decodes. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 34 +++++++++-- modelopt/torch/utils/dataset_utils.py | 14 +++-- modelopt/torch/utils/distributed.py | 10 ++- tests/unit/torch/utils/test_dataset_utils.py | 30 ++++++--- .../utils/test_distributed_calibration.py | 61 ++++++++++++++++--- 5 files changed, 117 insertions(+), 32 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 06643e07849..07a3d8415d6 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -99,11 +99,17 @@ def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: return found -def _refuse_rank_local_calibration(model: nn.Module, algorithm: QuantizeAlgoCfgType) -> None: - """Refuse a rank-local calibration algorithm on a model distributed over several ranks. +def _refuse_unsupported_distributed_calibration( + model: nn.Module, algorithm: QuantizeAlgoCfgType +) -> None: + """Refuse a calibration that would come out wrong on a model distributed over several ranks. Each rank of such a model calibrates on its own data, so an algorithm that keeps its statistics - on the rank would leave the ranks' copies of a weight different. Called before anything changes. + on the rank would leave the ranks' copies of a weight different. Under expert parallelism each + rank's per-expert quantizers hold its own experts, and calibration syncs them over the groups + their module's ``ParallelState`` names, every rank unless set: without the experts' + expert-parallel group, any calibration would merge different experts' scales. Called before + anything changes. """ plan = dist.recorded_distribution_plan(model) if plan is None or plan.world_size == 1: @@ -117,6 +123,22 @@ def _refuse_rank_local_calibration(model: nn.Module, algorithm: QuantizeAlgoCfgT "different. Calibrate with max, MSE, GPTQ or SmoothQuant, with distributed_sync on, or " "in a single process." ) + if ( + algorithm is not None + and plan.is_expert_parallel + and not any( + (state := getattr(module, "_parallel_state", None)) is not None + and state.expert_model_parallel_group.is_initialized() + for module in model.modules() + ) + ): + raise NotImplementedError( + f"Calibrating a model whose experts are sharded over {plan.ep_size} ranks needs their " + "expert-parallel group, which none of its modules carries: each rank's per-expert " + "quantizers would sync over every rank, merging different experts' scales. Give the " + "routed experts a ParallelState with their expert-parallel group, and the ranks " + "holding the same experts as their data-parallel group." + ) def _refuse_awq_on_a_sharded_fsdp2_model( @@ -174,7 +196,7 @@ def calibrate( Returns: The calibrated pytorch model. """ _refuse_awq_on_a_sharded_fsdp2_model(model, algorithm) - _refuse_rank_local_calibration(model, algorithm) + _refuse_unsupported_distributed_calibration(model, algorithm) if forward_loop is not None: # get the number of arguments of forward_loop num_args = len(inspect.signature(forward_loop).parameters) @@ -549,7 +571,7 @@ def forward_loop(model) -> None: """ quantize_config = QuantizeConfig(**dict(config)) _refuse_awq_on_a_sharded_fsdp2_model(model, config.get("algorithm")) - _refuse_rank_local_calibration(model, config.get("algorithm")) + _refuse_unsupported_distributed_calibration(model, config.get("algorithm")) if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: @@ -1109,7 +1131,7 @@ def forward_backward_step(model, batch) -> None: ) # Every candidate, before the search calibrates any: calibrate() would refuse mid-search. _refuse_awq_on_a_sharded_fsdp2_model(model, algo, candidate=name) - _refuse_rank_local_calibration(model, algo) + _refuse_unsupported_distributed_calibration(model, algo) # Select the appropriate searcher based on method if method_name not in AUTO_QUANTIZE_SEARCHERS: diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index ef798596704..7f9a44161ed 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1007,15 +1007,16 @@ def get_max_batch_size( ``forward_step(model, inputs)`` optionally probes a custom workload, including during OOM retries. By default, probe a model forward (or generation for encoder-decoder models). - For a model distributed over several ranks (its - :func:`~modelopt.torch.utils.distributed.distribution_plan`) it is collective: every rank calls - it, and every rank gets the smallest batch size any rank's GPU allows + For a model whose recorded plan distributes it over several ranks + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) it is collective: every rank + calls it, and every rank gets the smallest batch size any rank's GPU allows (:func:`~modelopt.torch.utils.distributed.all_reduce`), so that all run the same number of calibration steps. Each rank then measures its own GPU only and probes with a single forward: retrying a smaller batch after an out-of-memory error would run a different number of forwards, - and of the collectives inside them, on each rank. + and of the collectives inside them, on each rank. A model with no plan recorded, such as a + Megatron model, is measured on every visible GPU, as in a single process. """ - from .distributed import all_reduce, distribution_plan, local_rank + from .distributed import all_reduce, local_rank, recorded_distribution_plan args = ( model, @@ -1025,7 +1026,8 @@ def get_max_batch_size( enable_grad, forward_step, ) - if distribution_plan(model).world_size > 1: + plan = recorded_distribution_plan(model) + if plan is not None and plan.world_size > 1: device = torch.device("cuda", local_rank()) return all_reduce("calib_batch_size", lambda: _max_batch_size(*args, device=device), min) return _max_batch_size(*args) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 8719bb61a83..70052fa8c08 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -701,7 +701,9 @@ class _StridedIterableShard(IterableDataset): collectives from hanging when one rank runs short. A stream shorter than ``world_size`` has no complete group, so it is refused rather than leaving every rank with nothing to calibrate on. Each DataLoader worker process iterates its own copy of the stream, so the workers split the groups - between them rather than each yielding them all. + between them rather than each yielding them all: every worker still decodes the whole stream and + keeps ``1/num_workers`` of it, so ``num_workers > 0`` buys parallelism at ``num_workers`` times + the decode cost. ``vlm_dataset_utils._ShardedIterable`` cannot stand in: it pads every rank to a count its own loader knows, while this wraps any loader's stream, which may know none, and splits across its workers. """ @@ -888,7 +890,9 @@ def distributed_forward_loop( The forwards run without gradients and with KV/state caching disabled (``_disable_use_cache``), as the single-node ``_forward_loop`` runs them: a hybrid Mamba/attention model mutates its cache in - place. Each batch is a mapping of the forward's keyword arguments, as ``model(**batch)`` takes. + place. Each batch is a mapping of the forward's keyword arguments, as ``model(**batch)`` takes; + only its tensor values move to ``device``, so tensors nested in a value (a list of images, say) + must already be there. """ # Local import to avoid any import cycle within modelopt.torch.utils. from modelopt.torch.utils.dataset_utils import _disable_use_cache @@ -899,6 +903,8 @@ def distributed_forward_loop( f"{plan!r}: an expert-parallel model calibrates with the ep_mesh its routed experts " "are sharded over, and any other model without one." ) + if ep_mesh is not None and ep_mesh.size() != plan.ep_size: + raise ValueError(f"{plan!r}: the ep_mesh spans {ep_mesh.size()} ranks, not {plan.ep_size}.") def calibrate(_inner_model): _reinstall_a2a_forwards(model, ep_mesh) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 5b50b8c263f..4d757f5010e 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -410,7 +410,6 @@ def fake_call(x): raise torch.cuda.OutOfMemoryError model = Mock(spec=torch.nn.Module) - model.modules.return_value = [] # no plan recorded and no FSDP2 module: a single process's # get_max_batch_size calls the module (model(...)) so FSDP2 hooks fire, not model.forward, # so route the mock's __call__ via side_effect. model.side_effect = fake_call @@ -452,16 +451,19 @@ def forward(self, x): self.batch_sizes.append(x.shape[0]) -def _smallest_batch_size_job(rank, world_size): +def _batch_size_job(rank, world_size, planned): model = _BatchSizeProbe() - set_distribution_plan(model, DistributionPlan(world_size)) + if planned: + set_distribution_plan(model, DistributionPlan(world_size)) # A sample takes 100 bytes: rank 0's GPU has room for 4 of them, rank 1's for 10 (regulated to 8). free = [400, 1000][rank] + # What one measurement reads: the rank's own GPU, or every visible one. + gpus = 1 if planned else 4 measured = [] def mem_get_info(device): measured.append(device) - return (free if len(measured) == 1 else free - 100, 0) + return (free if len(measured) <= gpus else free - 100, 0) # A rank left waiting in the reduction gives up in seconds on this group. group = torch.distributed.new_group(timeout=timedelta(seconds=10)) @@ -484,15 +486,23 @@ def mem_get_info(device): finally: torch.distributed.destroy_process_group(group) - assert measured == [torch.device("cuda", rank)] * 2 # its own GPU (LOCAL_RANK) only - assert model.batch_sizes == [1] # one probe: an OOM retry would run more forwards on some ranks - assert batch_size == 4 # rank 0's + if planned: + assert measured == [torch.device("cuda", rank)] * 2 # its own GPU (LOCAL_RANK) only + # One probe: an OOM retry would run more forwards on some ranks. + assert model.batch_sizes == [1] + assert batch_size == 4 # rank 0's + else: + assert measured == list(range(4)) * 2 # every visible GPU, as in a single process + assert batch_size == [4, 8][rank] # its own -def test_get_max_batch_size_takes_the_smallest_ranks_in_a_distributed_run(cpu_dist_workers): +@pytest.mark.parametrize("planned", [True, False], ids=["planned", "no-plan"]) +def test_get_max_batch_size_in_a_distributed_run(cpu_dist_workers, planned): """Every rank must run the same number of calibration steps, so ranks agree on one batch size - for a model whose plan distributes it over them: the smallest any rank's GPU allows.""" - cpu_dist_workers.run(_smallest_batch_size_job) + for a model whose plan distributes it over them: the smallest any rank's GPU allows. A model + with no plan recorded, such as a Megatron model, is measured by each rank as a single process + measures it, and is not refused.""" + cpu_dist_workers.run(_batch_size_job, planned=planned) def test_get_dataset_samples_with_unsupported_dataset(make_toy_hf_dataset): diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 38eb1d316a2..7c992de5b8c 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -28,6 +28,7 @@ import pytest import torch +from _test_utils.torch.distributed.utils import bounded_collectives from torch import nn from torch.utils.data import ( BatchSampler, @@ -41,6 +42,7 @@ import modelopt.torch.utils.distributed as mto_dist from modelopt.torch.utils.distributed import ( DistributionPlan, + ParallelState, _reinstall_a2a_forwards, distributed_forward_loop, set_distribution_plan, @@ -248,14 +250,18 @@ def test_sharded_loader_keeps_the_original_loader_settings(dataset): class _FakeMesh: - """The EP device mesh; the TP wrappers only need its process group.""" + """The EP device mesh; the TP wrappers only need its process group, the loop its size.""" - def __init__(self): + def __init__(self, ranks=2): self.group = object() + self.ranks = ranks def get_group(self): return self.group + def size(self): + return self.ranks + class _WrapperStyle: """A transformers >= 5.16 TP style: ``install_forward`` wraps the module's current forward.""" @@ -495,17 +501,22 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat @pytest.mark.parametrize( - ("plan", "mesh"), - [(_EXPERT_PARALLEL, None), (DistributionPlan(world_size=2), _FakeMesh())], - ids=["expert-parallel-without-mesh", "mesh-for-a-replicated-model"], + ("plan", "mesh", "match"), + [ + (_EXPERT_PARALLEL, None, "calibrates with the ep_mesh"), + (DistributionPlan(world_size=2), _FakeMesh(), "calibrates with the ep_mesh"), + (DistributionPlan(world_size=4, ep_size=2), _FakeMesh(ranks=4), "spans 4 ranks"), + ], + ids=["expert-parallel-without-mesh", "mesh-for-a-replicated-model", "mesh-of-the-whole-world"], ) -def test_the_ep_mesh_comes_with_an_expert_parallel_plan_only(plan, mesh): +def test_the_ep_mesh_is_the_plans_expert_parallel_one(plan, mesh, match): """Without its mesh, the experts' forward would run with no EP group at the first step; a mesh - for a model whose plan shards no experts would be ignored.""" + for a model whose plan shards no experts would be ignored; and a mesh over other ranks than an + expert-parallel group would hand the all-to-all the wrong group.""" model = _Recorder() set_distribution_plan(model, plan) - with pytest.raises(ValueError, match="calibrates with the ep_mesh"): + with pytest.raises(ValueError, match=match): distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}], ep_mesh=mesh) @@ -573,3 +584,37 @@ def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model after = model.state_dict() assert after.keys() == state.keys() assert all(torch.equal(after[name], tensor) for name, tensor in state.items()) + + +def _expert_parallel_calibration_job(rank, size, seeded): + model = nn.Sequential(nn.Linear(16, 16)) + if seeded: + # A routed expert's state under expert parallelism: its expert-parallel group (here every + # rank), and no other rank holding the same experts for its data-parallel group. + model[0]._parallel_state = ParallelState( + data_parallel_group=-1, expert_model_parallel_group=None + ) + set_distribution_plan(model, DistributionPlan(world_size=size, ep_size=size)) + raised = None + with bounded_collectives(): + # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take + # the shared worker down instead of failing this test alone. + try: + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16))) + except NotImplementedError as e: + raised = str(e) + if seeded: + assert raised is None, raised + else: + assert raised is not None and "expert-parallel group" in raised, raised + assert type(model[0]) is nn.Linear # refused before converting anything + + +@pytest.mark.parametrize("seeded", [False, True], ids=["without-the-group", "seeded"]) +def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( + cpu_dist_workers, seeded +): + """Each rank's per-expert quantizers hold its own experts, and max syncs them over the groups + their module's ParallelState names, every rank unless set, which would merge different experts' + scales. Calibration is refused before anything changes until the experts carry their group.""" + cpu_dist_workers.run(_expert_parallel_calibration_job, seeded=seeded) From e2137fcbbed72587e3a5687d9406b6933dbd19e9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 02:09:08 +0000 Subject: [PATCH 27/60] Check each routed expert's groups and re-plant a2a forwards when rebound The expert-parallel calibration guard passed once any module carried an expert-parallel group, an expert-parallel group of None (every rank) included. A router seeded with the group passed it while the routed experts kept the conversion default, a data-parallel group of every rank, which merges different experts' scales; so did experts seeded with a data-parallel group over every rank. The guard now checks every module holding routed experts: the modules carrying an expert-parallel group and, once quantization has converted them, the quantized modules under a container named experts, as transformers' expert-parallel plans name it. Each must carry an expert-parallel group and, for data parallelism, just the plan's edp ranks. _reinstall_a2a_forwards flagged a module once it re-planted its all-to-all forward and skipped it from then on, so a forward rebound again, by another conversion or an export restoring an earlier forward, kept no EP process group. It now skips a module only while its forward is still the wrapper planted there. get_max_batch_size's docstring no longer implies that a distributed model's probe retries after an out-of-memory error, and says that the batch size it arrives at is not probed. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 50 ++++++++++----- modelopt/torch/utils/dataset_utils.py | 12 ++-- modelopt/torch/utils/distributed.py | 14 +++-- .../utils/test_distributed_calibration.py | 62 ++++++++++++------- 4 files changed, 88 insertions(+), 50 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 07a3d8415d6..8034d99ff22 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -107,9 +107,11 @@ def _refuse_unsupported_distributed_calibration( Each rank of such a model calibrates on its own data, so an algorithm that keeps its statistics on the rank would leave the ranks' copies of a weight different. Under expert parallelism each rank's per-expert quantizers hold its own experts, and calibration syncs them over the groups - their module's ``ParallelState`` names, every rank unless set: without the experts' - expert-parallel group, any calibration would merge different experts' scales. Called before - anything changes. + their module's ``ParallelState`` names, every rank for a module converted without one: a module + holding routed experts must name the experts' expert-parallel group and, for data parallelism, + only the ranks holding the same experts, or any calibration would merge different experts' + scales. Called before anything changes; a routed expert's quantized module, whose state + calibration reads, is checked once quantization has converted it, before it calibrates. """ plan = dist.recorded_distribution_plan(model) if plan is None or plan.world_size == 1: @@ -123,21 +125,35 @@ def _refuse_unsupported_distributed_calibration( "different. Calibrate with max, MSE, GPTQ or SmoothQuant, with distributed_sync on, or " "in a single process." ) - if ( - algorithm is not None - and plan.is_expert_parallel - and not any( - (state := getattr(module, "_parallel_state", None)) is not None - and state.expert_model_parallel_group.is_initialized() - for module in model.modules() - ) - ): + if algorithm is None or not plan.is_expert_parallel: + return + # Routed experts: the modules carrying an expert-parallel group, and the quantized modules under + # a container named ``experts``, as every expert-parallel plan in transformers names it. + holders = {} + for name, module in model.named_modules(): + state = getattr(module, "_parallel_state", None) + if (state is not None and state.expert_model_parallel_group.is_initialized()) or ( + isinstance(module, QuantModule) and "experts" in name.split(".") + ): + holders[name] = state + wrong = [ + name + for name, state in holders.items() + if state is None + or not state.expert_model_parallel_group.is_initialized() + or max(state.data_parallel_group.world_size(), 1) != plan.edp + ] + if not holders or wrong: raise NotImplementedError( - f"Calibrating a model whose experts are sharded over {plan.ep_size} ranks needs their " - "expert-parallel group, which none of its modules carries: each rank's per-expert " - "quantizers would sync over every rank, merging different experts' scales. Give the " - "routed experts a ParallelState with their expert-parallel group, and the ranks " - "holding the same experts as their data-parallel group." + f"Calibrating a model whose experts are sharded over {plan.ep_size} ranks needs each " + "routed expert's ParallelState to name its expert-parallel group and, as its " + f"data-parallel group, the {plan.edp} rank(s) holding the same experts, or calibration " + "would sync their scales over the wrong ranks. " + + ( + f"These modules do not: {wrong[:3]}." + if wrong + else "None of the model's modules carries an expert-parallel group." + ) ) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 7f9a44161ed..319247ebaeb 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1005,16 +1005,18 @@ def get_max_batch_size( """Get the maximum batch size that can be used for the model. ``forward_step(model, inputs)`` optionally probes a custom workload, including during - OOM retries. By default, probe a model forward (or generation for encoder-decoder models). + OOM retries (which a model distributed over several ranks skips, below). By default, probe a + model forward (or generation for encoder-decoder models). For a model whose recorded plan distributes it over several ranks (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) it is collective: every rank calls it, and every rank gets the smallest batch size any rank's GPU allows (:func:`~modelopt.torch.utils.distributed.all_reduce`), so that all run the same number of - calibration steps. Each rank then measures its own GPU only and probes with a single forward: - retrying a smaller batch after an out-of-memory error would run a different number of forwards, - and of the collectives inside them, on each rank. A model with no plan recorded, such as a - Megatron model, is measured on every visible GPU, as in a single process. + calibration steps. Each rank then measures its own GPU only, from one forward of the + single-sample input, and does not probe the batch size it arrives at: retrying a smaller batch + after an out-of-memory error would run a different number of forwards, and of the collectives + inside them, on each rank. A model with no plan recorded, such as a Megatron model, is measured + on every visible GPU, as in a single process. """ from .distributed import all_reduce, local_rank, recorded_distribution_plan diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 70052fa8c08..b15a4899474 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -813,9 +813,11 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: ``torch_a2a_experts_forward`` receives ``process_group=None`` and raises. Called at the top of the calibration loop: after conversion (so it wraps the QUANTIZED forward) - and before the first forward. Idempotent via a per-module flag, so repeated calibration passes - do not stack wrappers (which would inject the process_group argument twice). ``ep_mesh`` is the - mesh the routed experts are sharded over; without one there is nothing to re-plant. + and before the first forward. A module whose forward is still the wrapper planted here is + skipped, so repeated calibration passes do not stack wrappers; one whose forward was rebound + since (another conversion, or an export restoring an earlier forward) is planted again. + ``ep_mesh`` is the mesh the routed experts are sharded over; without one there is nothing to + re-plant. """ if ep_mesh is None: return 0 @@ -843,15 +845,15 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: return 0 n = 0 for name, mod in model.named_modules(): - if getattr(mod, "_modelopt_a2a_forward_reinstalled", False): - continue + if mod.forward is getattr(mod, "_modelopt_a2a_forward", None): + continue # the wrapper planted here is still in place style_name = _get_parameter_tp_plan(parameter_name=name, tp_plan=plan, is_weight=False) if ( style_name in ("torch_a2a_experts", "torch_a2a_router") and style_name in ALL_PARALLEL_STYLES ): ALL_PARALLEL_STYLES[style_name].install_forward(mod, ep_mesh) - mod._modelopt_a2a_forward_reinstalled = True + mod._modelopt_a2a_forward = mod.forward n += 1 if n and is_master(): print( diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 7c992de5b8c..4ed6c3b34eb 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -24,7 +24,7 @@ import re import sys import types -from collections import Counter +from collections import Counter, OrderedDict import pytest import torch @@ -477,10 +477,11 @@ def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch, def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypatch): """hf_ptq's order -- build the loop, insert quantizers, calibrate -- hands the experts the EP group. - It holds on every pass (multi-pass algorithms call the loop again), with one wrapper per module. - Only the planned torch_a2a experts and routers are re-wrapped: other TP styles keep their forward, - and so do the quantizers under an experts module, since the plan is looked up per module - (``is_weight=False``) and children do not inherit their parent's style. + It holds on every pass (multi-pass algorithms call the loop again), with one wrapper per module, + and again once another conversion drops the wrappers anew. Only the planned torch_a2a experts + and routers are re-wrapped: other TP styles keep their forward, and so do the quantizers under an + experts module, since the plan is looked up per module (``is_weight=False``) and children do not + inherit their parent's style. """ styles = _install_fake_tp(monkeypatch) model, mesh = _MoEModel(), _FakeMesh() @@ -493,9 +494,12 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat forward_loop(model) forward_loop(model) + for layer in model.layers: + _drop_tp_wrapper(layer.experts) + forward_loop(model) - assert [layer.experts.groups for layer in model.layers] == [[mesh.group] * 6] * 2 - assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] + assert [layer.experts.groups for layer in model.layers] == [[mesh.group] * 9] * 2 + assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] * 2 assert styles["torch_a2a_router"].installed == [layer.gate for layer in model.layers] assert styles["colwise"].installed == [] @@ -586,13 +590,23 @@ def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model assert all(torch.equal(after[name], tensor) for name, tensor in state.items()) -def _expert_parallel_calibration_job(rank, size, seeded): - model = nn.Sequential(nn.Linear(16, 16)) - if seeded: - # A routed expert's state under expert parallelism: its expert-parallel group (here every - # rank), and no other rank holding the same experts for its data-parallel group. - model[0]._parallel_state = ParallelState( - data_parallel_group=-1, expert_model_parallel_group=None +# {module: (data_parallel_group, expert_model_parallel_group)} each layout seeds. With the experts +# sharded over every rank, their expert-parallel group is every rank (None), and no other rank holds +# the same experts (-1: no data-parallel group). +_EXPERT_LAYOUTS = { + "seeded": {"experts": (-1, None)}, + "without-the-group": {}, + "on-the-router-only": {"gate": (-1, None)}, + "data-parallel-over-every-rank": {"experts": (None, None)}, +} + + +def _expert_parallel_calibration_job(rank, size, layout): + # A router and a routed expert, named as expert-parallel plans name the experts' container. + model = nn.Sequential(OrderedDict(gate=nn.Linear(16, 16), experts=nn.Linear(16, 16))) + for name, (dp, ep) in _EXPERT_LAYOUTS[layout].items(): + model.get_submodule(name)._parallel_state = ParallelState( + data_parallel_group=dp, expert_model_parallel_group=ep ) set_distribution_plan(model, DistributionPlan(world_size=size, ep_size=size)) raised = None @@ -603,18 +617,22 @@ def _expert_parallel_calibration_job(rank, size, seeded): mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16))) except NotImplementedError as e: raised = str(e) - if seeded: + if layout == "seeded": assert raised is None, raised - else: - assert raised is not None and "expert-parallel group" in raised, raised - assert type(model[0]) is nn.Linear # refused before converting anything + return + assert raised is not None and "expert-parallel group" in raised, raised + # Refused before converting anything, but for unseeded experts beside a seeded router: those + # are told apart once quantization has converted them, and refused before they calibrate. + assert (type(model.experts) is nn.Linear) == (layout != "on-the-router-only") -@pytest.mark.parametrize("seeded", [False, True], ids=["without-the-group", "seeded"]) +@pytest.mark.parametrize("layout", list(_EXPERT_LAYOUTS)) def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( - cpu_dist_workers, seeded + cpu_dist_workers, layout ): """Each rank's per-expert quantizers hold its own experts, and max syncs them over the groups their module's ParallelState names, every rank unless set, which would merge different experts' - scales. Calibration is refused before anything changes until the experts carry their group.""" - cpu_dist_workers.run(_expert_parallel_calibration_job, seeded=seeded) + scales. Calibration is refused until the experts carry their groups: with no expert-parallel + group at all, with one on the router alone, and with a data-parallel group spanning ranks that + hold other experts.""" + cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) From aa060dfd54aba62c5068f4e667def6c7335ce0b3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:38:17 +0000 Subject: [PATCH 28/60] Keep a lone rank's calibration loader as is; document FSDP2 refusals shard_dataloader(loader, rank, 1) now returns the loader itself, before the sampler check, so a single process keeps its own sampler, a subset or a shuffling one included, rather than having it refused or replaced by an unshuffled DistributedSampler. Recording the FSDP2 plan in hf_ptq's --use_fsdp2 load makes the refusal of rank-local calibration apply there on several GPUs: AWQ, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate. The FSDP2 section of the hf_ptq README and the changelog now say so, and the refusal's message and docstring cover the ranks' shards of a weight as well as their copies. shard_dataloader's docstring states that a stream must yield the same samples in the same order on every rank and worker, a shuffling one with a fixed seed of its own; the strided shard's docstring and its short-stream comment read that as the requirement it is. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + examples/hf_ptq/README.md | 4 +- modelopt/torch/quantization/model_quant.py | 21 ++++++----- modelopt/torch/utils/distributed.py | 37 +++++++++++-------- .../utils/test_distributed_calibration.py | 4 +- 5 files changed, 40 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 89fb2a8b8d7..15430999f8b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -46,6 +46,7 @@ Changelog **Backward Breaking Changes** +- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank: AWQ, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate. Each rank calibrates on its own share of the data, so these calibrated the ranks' shards of a weight differently; run them without ``--use_fsdp2``. - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index d108dcbf786..d28ea2c524e 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -610,12 +610,14 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. -> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental.* +> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental, and refused when a candidate format is one of those refused below.* > *Nemotron-H checkpoints containing MTP weights are not supported with `--use_fsdp2` and are rejected before distributed model loading. Remove `--use_fsdp2` to load and calibrate the MTP modules; the distributed loader does not construct them.* > *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* +> *With `--use_fsdp2` on more than one GPU, calibration that keeps its statistics on each rank is refused when quantization starts: AWQ and SVDQuant (`--qformat int4_awq`, `w4a8_awq_beta`, `nvfp4_awq_lite`, `nvfp4_awq_clip`, `nvfp4_awq_full`, `nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`.* + ### Usage #### Slurm (recommended) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 8034d99ff22..90639e8d844 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -105,13 +105,14 @@ def _refuse_unsupported_distributed_calibration( """Refuse a calibration that would come out wrong on a model distributed over several ranks. Each rank of such a model calibrates on its own data, so an algorithm that keeps its statistics - on the rank would leave the ranks' copies of a weight different. Under expert parallelism each - rank's per-expert quantizers hold its own experts, and calibration syncs them over the groups - their module's ``ParallelState`` names, every rank for a module converted without one: a module - holding routed experts must name the experts' expert-parallel group and, for data parallelism, - only the ranks holding the same experts, or any calibration would merge different experts' - scales. Called before anything changes; a routed expert's quantized module, whose state - calibration reads, is checked once quantization has converted it, before it calibrates. + on the rank would calibrate the ranks' copies, or shards, of a weight differently. Under expert + parallelism each rank's per-expert quantizers hold its own experts, and calibration syncs them + over the groups their module's ``ParallelState`` names, every rank for a module converted + without one: a module holding routed experts must name the experts' expert-parallel group and, + for data parallelism, only the ranks holding the same experts, or any calibration would merge + different experts' scales. Called before anything changes; a routed expert's quantized module, + whose state calibration reads, is checked once quantization has converted it, before it + calibrates. """ plan = dist.recorded_distribution_plan(model) if plan is None or plan.world_size == 1: @@ -121,9 +122,9 @@ def _refuse_unsupported_distributed_calibration( raise NotImplementedError( f"Calibrating with {', '.join(refused)} is not supported on a model distributed over " f"{plan.world_size} ranks. Each rank calibrates on its own data and this calibration " - "keeps what it gathers on the rank, so the ranks' copies of a weight would come out " - "different. Calibrate with max, MSE, GPTQ or SmoothQuant, with distributed_sync on, or " - "in a single process." + "keeps what it gathers on the rank, so the ranks would calibrate their copies, or their " + "shards, of a weight differently. Calibrate with max, MSE, GPTQ or SmoothQuant, with " + "distributed_sync on, or in a single process." ) if algorithm is None or not plan.is_expert_parallel: return diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index b15a4899474..7e34b75516e 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -695,15 +695,15 @@ class _StridedIterableShard(IterableDataset): ``DistributedSampler`` only works with map-style datasets. Streaming calibration datasets (e.g. the Nemotron image-text ``NemotronTarPlusJsonlIterable``) are iterable, so each rank instead builds its - own identical stream (deterministic seed) and takes slot ``rank`` from every complete group of - ``world_size`` consecutive samples. Any trailing incomplete group is dropped so every rank yields - exactly the same count -- preserving the equal-per-rank-batch invariant that keeps per-step EP - collectives from hanging when one rank runs short. A stream shorter than ``world_size`` has no - complete group, so it is refused rather than leaving every rank with nothing to calibrate on. Each - DataLoader worker process iterates its own copy of the stream, so the workers split the groups - between them rather than each yielding them all: every worker still decodes the whole stream and - keeps ``1/num_workers`` of it, so ``num_workers > 0`` buys parallelism at ``num_workers`` times - the decode cost. + own copy of the stream, the same on every rank and worker as :func:`shard_dataloader` requires, + and takes slot ``rank`` from every complete group of ``world_size`` consecutive samples. Any + trailing incomplete group is dropped so every rank yields exactly the same count -- preserving + the equal-per-rank-batch invariant that keeps per-step EP collectives from hanging when one rank + runs short. A stream shorter than ``world_size`` has no complete group, so it is refused rather + than leaving every rank with nothing to calibrate on. Each DataLoader worker process iterates + its own copy of the stream, so the workers split the groups between them rather than each + yielding them all: every worker still decodes the whole stream and keeps ``1/num_workers`` of + it, so ``num_workers > 0`` buys parallelism at ``num_workers`` times the decode cost. ``vlm_dataset_utils._ShardedIterable`` cannot stand in: it pads every rank to a count its own loader knows, while this wraps any loader's stream, which may know none, and splits across its workers. """ @@ -721,7 +721,8 @@ def __iter__(self): group = list(itertools.islice(it, self.world_size)) if len(group) < self.world_size: if index == 0: - # Every rank and worker sees the same stream, so all of them refuse it together. + # shard_dataloader requires the same stream on every rank and worker, so all of + # them refuse it together. raise ValueError( f"The calibration stream has {len(group)} sample(s), fewer than the " f"{self.world_size} ranks: no rank would get a sample to calibrate on." @@ -740,14 +741,20 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad replaces the loader's sampler, so a shuffling one becomes the rank's unshuffled share, and one that picks samples itself (a subset, or a batch sampler other than a plain ``BatchSampler``) is refused rather than silently widened. For streaming ``IterableDataset`` datasets (which have no ``len()`` - and so can't use a sampler) each rank strides its own deterministic copy of the stream via - :class:`_StridedIterableShard`. Every DataLoader worker must see that whole stream: a dataset that - splits itself across workers by ``get_worker_info()`` would be split twice and lose samples, so - load one with ``num_workers=0``. + and so can't use a sampler) each rank strides its own copy of the stream via + :class:`_StridedIterableShard`. That stream must yield the same samples in the same order on + every rank and in every DataLoader worker, or the shards overlap and the ranks' step counts can + differ, which hangs the collectives: one that shuffles needs a fixed seed of its own, as torch + seeds each worker's RNG differently and the ranks' RNGs need not agree. Every DataLoader worker + must see that whole stream: a dataset that splits itself across workers by ``get_worker_info()`` + would be split twice and lose samples, so load one with ``num_workers=0``. Forwards all non-sampler DataLoader settings from ``loader`` (batching, workers, pinning, - prefetch, init fn, generator, ...). + prefetch, init fn, generator, ...). With one rank there is nothing to shard: ``loader`` itself + comes back, its sampler and shuffling included. """ + if world_size == 1: + return loader batching = loader.batch_sampler if not isinstance(loader.dataset, IterableDataset): # A plain BatchSampler, built from batch_size or passed in, batches what its sampler picks; diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 4ed6c3b34eb..39adffc2bf0 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -178,11 +178,13 @@ def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch) ) def test_a_loader_that_picks_its_own_samples_is_refused(picker): """The rank's DistributedSampler would replace what it picks, silently calibrating on the whole - dataset instead, however the loader batches the picks.""" + dataset instead, however the loader batches the picks. A lone rank has nothing to shard and + keeps the loader as it is.""" loader = DataLoader(list(range(8)), collate_fn=_collate, **picker) with pytest.raises(ValueError, match="picks which samples it reads"): shard_dataloader(loader, 0, 2) + assert shard_dataloader(loader, 0, 1) is loader @pytest.mark.parametrize( From ce49b0ee419fffc122343c0eb41eec917bd6045c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 04:19:04 +0000 Subject: [PATCH 29/60] Let modules without routed experts calibrate under an EP plan hf_ptq quantizes a VLM's vision tower and its other non-language modules apart from the language model, with every quantizer disabled. Under an expert-parallel plan the calibration check refused each of them, since no module there carries an expert-parallel group, though they hold no routed experts whose scales could merge. A module with no expert-parallel group and no container named experts now passes. Experts that carry no group are still refused before conversion, and converted experts with a missing or wrong group before they calibrate. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 2 ++ .../torch/utils/test_distributed_calibration.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 90639e8d844..c77115943e1 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -137,6 +137,8 @@ def _refuse_unsupported_distributed_calibration( isinstance(module, QuantModule) and "experts" in name.split(".") ): holders[name] = state + if not holders and not any("experts" in name.split(".") for name, _ in model.named_modules()): + return # no routed experts here, as in a VLM's vision tower: no experts' scales to merge wrong = [ name for name, state in holders.items() diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 39adffc2bf0..c1beb81e3be 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -638,3 +638,20 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( group at all, with one on the router alone, and with a data-parallel group spanning ranks that hold other experts.""" cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) + + +@pytest.mark.parametrize( + "config", + [ + {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"}, + mtq.INT8_DEFAULT_CFG, + ], + ids=["all-disabled", "quantized"], +) +def test_a_module_without_routed_experts_calibrates_under_an_expert_parallel_plan(config): + """hf_ptq quantizes a VLM's vision tower apart from its language model, every quantizer + disabled. It holds no routed experts, so it has no experts' scales to merge, quantized or not.""" + model = nn.Sequential(nn.Linear(16, 16)) + set_distribution_plan(model, _EXPERT_PARALLEL) + + mtq.quantize(model, config, lambda m: m(torch.randn(2, 16))) From 49d238d1c7f2cedf70152919b2883dbbc3fd62a4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:01 +0000 Subject: [PATCH 30/60] Re-plant only the a2a forward wrappers quantizer insertion dropped The first calibration pass planted a wrapper on every module the plan sends through torch_a2a, its guard recognising only a wrapper it had planted itself. A planned module quantization leaves alone, such as a router quantizer insertion does not convert, still holds transformers' wrapper from the load, and got a second one around it, running its transforms twice. install_forward stores a plain function as the forward and quantizer insertion a bound method over it, so a module is now planted only where its forward is a bound method, or none is set; the marker attribute goes. The transformers-version probe checked that either style has install_forward, then called it on both; it now checks the style about to be called. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 45 +++++++++---------- .../utils/test_distributed_calibration.py | 39 +++++++++++----- 2 files changed, 51 insertions(+), 33 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7e34b75516e..016e5acb483 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -23,6 +23,7 @@ import sys import time import traceback +import types from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager, suppress from dataclasses import dataclass @@ -820,11 +821,11 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: ``torch_a2a_experts_forward`` receives ``process_group=None`` and raises. Called at the top of the calibration loop: after conversion (so it wraps the QUANTIZED forward) - and before the first forward. A module whose forward is still the wrapper planted here is - skipped, so repeated calibration passes do not stack wrappers; one whose forward was rebound - since (another conversion, or an export restoring an earlier forward) is planted again. - ``ep_mesh`` is the mesh the routed experts are sharded over; without one there is nothing to - re-plant. + and before the first forward. Only a forward quantizer insertion rebound is planted again: a + wrapper still in place -- transformers' from the load, on a module quantization left alone, or + one planted here on an earlier pass -- is kept, since wrapping it again would run its transforms + twice. ``ep_mesh`` is the mesh the routed experts are sharded over; without one there is nothing + to re-plant. """ if ep_mesh is None: return 0 @@ -841,27 +842,25 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: if not plan: return 0 - # install_forward() -- and the clobber it causes -- only exist on transformers >=5.16. Older - # releases install their TP input/output transforms as forward PRE-HOOKS, which survive - # ModelOpt rebinding ``module.forward``, so there is nothing to re-plant and calling the - # missing method would AttributeError. No-op rather than fail on 5.9/5.13. - if not any( - hasattr(ALL_PARALLEL_STYLES.get(s), "install_forward") - for s in ("torch_a2a_experts", "torch_a2a_router") - ): - return 0 n = 0 for name, mod in model.named_modules(): - if mod.forward is getattr(mod, "_modelopt_a2a_forward", None): - continue # the wrapper planted here is still in place style_name = _get_parameter_tp_plan(parameter_name=name, tp_plan=plan, is_weight=False) - if ( - style_name in ("torch_a2a_experts", "torch_a2a_router") - and style_name in ALL_PARALLEL_STYLES - ): - ALL_PARALLEL_STYLES[style_name].install_forward(mod, ep_mesh) - mod._modelopt_a2a_forward = mod.forward - n += 1 + if style_name not in ("torch_a2a_experts", "torch_a2a_router"): + continue + # install_forward() -- and the clobber it causes -- only exist on transformers >=5.16. Older + # releases install their TP input/output transforms as forward PRE-HOOKS, which survive + # ModelOpt rebinding ``module.forward``: nothing to re-plant. Checked per style, the one + # about to be called. + style = ALL_PARALLEL_STYLES.get(style_name) + if not hasattr(style, "install_forward"): + continue + # install_forward stores a plain function as the instance's forward, and quantizer + # insertion (bind_forward_method) a bound method over it: only the latter lost its wrapper. + current = vars(mod).get("forward") + if current is not None and not isinstance(current, types.MethodType): + continue + style.install_forward(mod, ep_mesh) + n += 1 if n and is_master(): print( f"[ep] re-installed {n} torch_a2a forward wrapper(s) clobbered by quantizer insertion", diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index c1beb81e3be..d039b895e87 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -360,6 +360,15 @@ def forward(self, input_ids): return hidden_states +def _load_tp_wrappers(model, styles, mesh): + """Do what transformers' expert-parallel load does: wrap every planned module's forward.""" + for name, module in model.named_modules(): + if (style := _get_parameter_tp_plan(name, model.tp_plan, is_weight=False)) is not None: + styles[style].install_forward(module, mesh) + for style in styles.values(): + style.installed.clear() + + def _drop_tp_wrapper(module): """Do what quantizer insertion does to the instance: rebind the class forward over the wrapper. @@ -376,17 +385,24 @@ def test_reinstall_is_a_noop_without_an_ep_mesh(monkeypatch): assert not any(style.installed for style in styles.values()) -def test_reinstall_is_a_noop_where_tp_styles_are_pre_hooks(monkeypatch): +@pytest.mark.parametrize( + ("experts_style", "planted"), + [(_PreHookStyle(), 0), (_WrapperStyle(inject_group=True), 2)], + ids=["all-pre-hooks", "only-the-experts-wrap"], +) +def test_reinstall_skips_styles_whose_transforms_are_pre_hooks(monkeypatch, experts_style, planted): """transformers < 5.16 has no ``install_forward``; its pre-hooks survive quantizer insertion. - Those releases keep the styles in ``transformers.integrations``: the fallback import. + Each style is checked for it, the one about to be called: a style without one is skipped even + where the other has one. Those releases keep the styles in ``transformers.integrations``: the + fallback import. """ _install_fake_tp( monkeypatch, - styles={"torch_a2a_experts": _PreHookStyle(), "torch_a2a_router": _PreHookStyle()}, + styles={"torch_a2a_experts": experts_style, "torch_a2a_router": _PreHookStyle()}, where=_LEGACY_TP_MODULE, ) - assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0 + assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == planted # ------------------------------------------------------------------------- distributed_forward_loop @@ -477,16 +493,19 @@ def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch, def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypatch): - """hf_ptq's order -- build the loop, insert quantizers, calibrate -- hands the experts the EP group. + """hf_ptq's order -- load, build the loop, insert quantizers, calibrate -- hands the experts the + EP group. It holds on every pass (multi-pass algorithms call the loop again), with one wrapper per module, - and again once another conversion drops the wrappers anew. Only the planned torch_a2a experts - and routers are re-wrapped: other TP styles keep their forward, and so do the quantizers under an - experts module, since the plan is looked up per module (``is_weight=False``) and children do not - inherit their parent's style. + and again once another conversion drops the wrappers anew. Only the wrappers conversion dropped + are planted again: a router quantization left alone keeps the one from the load, since wrapping + it again would run its transforms twice. Other TP styles keep their forward, and so do the + quantizers under an experts module, since the plan is looked up per module (``is_weight=False``) + and children do not inherit their parent's style. """ styles = _install_fake_tp(monkeypatch) model, mesh = _MoEModel(), _FakeMesh() + _load_tp_wrappers(model, styles, mesh) set_distribution_plan(model, _EXPERT_PARALLEL) forward_loop = distributed_forward_loop( model, [{"input_ids": torch.ones(2, 4)}] * 3, ep_mesh=mesh @@ -502,7 +521,7 @@ def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypat assert [layer.experts.groups for layer in model.layers] == [[mesh.group] * 9] * 2 assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] * 2 - assert styles["torch_a2a_router"].installed == [layer.gate for layer in model.layers] + assert styles["torch_a2a_router"].installed == [] assert styles["colwise"].installed == [] From b71bc660a6dfb2e026c0f2aaa605b731a1c25b3c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:37:33 +0000 Subject: [PATCH 31/60] Measure a distributed batch size on the GPU the rank computes on The batch size of a model distributed over several ranks was measured on cuda:LOCAL_RANK, which is not the rank's GPU under a launcher that shows each rank one GPU: there the device lookup failed on every rank but the first. It is measured on the current device now, the one the forward runs on. The docstring and the hf_ptq README also say why that size is not probed for out-of-memory, as a single process's is: a rank that runs out of memory partway through a forward leaves the others waiting in its collectives, so probing in lockstep would hang instead of backing off. Set the batch size explicitly when the estimate does not fit. The single-sample input's expansion now lives in the probing branch, its only user. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 + modelopt/torch/utils/dataset_utils.py | 51 +++++++++++--------- tests/unit/torch/utils/test_dataset_utils.py | 4 +- 3 files changed, 34 insertions(+), 23 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index d28ea2c524e..649b3db5fe6 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -618,6 +618,8 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *With `--use_fsdp2` on more than one GPU, calibration that keeps its statistics on each rank is refused when quantization starts: AWQ and SVDQuant (`--qformat int4_awq`, `w4a8_awq_beta`, `nvfp4_awq_lite`, `nvfp4_awq_clip`, `nvfp4_awq_full`, `nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`.* +> *With `--use_fsdp2` on more than one GPU and `--batch_size 0`, each rank estimates the calibration batch size from one forward and the ranks use the smallest. That size is not probed for out-of-memory, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives: set `--batch_size` if calibration runs out of memory.* + ### Usage #### Slurm (recommended) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 319247ebaeb..5fe326c2f6f 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1012,13 +1012,15 @@ def get_max_batch_size( (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) it is collective: every rank calls it, and every rank gets the smallest batch size any rank's GPU allows (:func:`~modelopt.torch.utils.distributed.all_reduce`), so that all run the same number of - calibration steps. Each rank then measures its own GPU only, from one forward of the + calibration steps. Each rank then measures the GPU it computes on, from one forward of the single-sample input, and does not probe the batch size it arrives at: retrying a smaller batch after an out-of-memory error would run a different number of forwards, and of the collectives - inside them, on each rank. A model with no plan recorded, such as a Megatron model, is measured - on every visible GPU, as in a single process. + inside them, on each rank. Probing in lockstep would not make it safe either: a rank that runs out + of memory partway through a forward has left its peers waiting in that forward's collectives. If + the estimate does not fit, pass an explicit batch size instead. A model with no plan recorded, + such as a Megatron model, is measured on every visible GPU, as in a single process. """ - from .distributed import all_reduce, local_rank, recorded_distribution_plan + from .distributed import all_reduce, recorded_distribution_plan args = ( model, @@ -1030,7 +1032,9 @@ def get_max_batch_size( ) plan = recorded_distribution_plan(model) if plan is not None and plan.world_size > 1: - device = torch.device("cuda", local_rank()) + # The device this rank computes on, which LOCAL_RANK does not name under a launcher that + # shows each rank one GPU. + device = torch.device("cuda", torch.cuda.current_device()) return all_reduce("calib_batch_size", lambda: _max_batch_size(*args, device=device), min) return _max_batch_size(*args) @@ -1099,24 +1103,27 @@ def infer_method(inputs): else: target_data_batch = max(int(free_mem_before / mem_diff_per_data_batch), 1) - def _expand_to(batch: int) -> torch.Tensor: - return sample_input_single_batch.expand( - [ - batch if index == 0 else dim - for index, dim in enumerate(sample_input_single_batch.shape) - ] - ) - # For some models on multi GPU, we observe the memory per batch is not a constant. - # So we just test the target batch size and make sure we do not go OOM. - while device is None and target_data_batch > 1: - with torch.set_grad_enabled(enable_grad): - try: - infer_method(_expand_to(target_data_batch)) - break - except torch.cuda.OutOfMemoryError: # pragma: no cover - GPU OOM retry path - target_data_batch = target_data_batch // 2 # pragma: no cover - torch.cuda.empty_cache() # pragma: no cover + # So we just test the target batch size and make sure we do not go OOM. A model distributed + # over several ranks is measured once, not probed (see get_max_batch_size). + if device is None: + + def _expand_to(batch: int) -> torch.Tensor: + return sample_input_single_batch.expand( + [ + batch if index == 0 else dim + for index, dim in enumerate(sample_input_single_batch.shape) + ] + ) + + while target_data_batch > 1: + with torch.set_grad_enabled(enable_grad): + try: + infer_method(_expand_to(target_data_batch)) + break + except torch.cuda.OutOfMemoryError: # pragma: no cover - GPU OOM retry path + target_data_batch = target_data_batch // 2 # pragma: no cover + torch.cuda.empty_cache() # pragma: no cover # Regulate the data batch target to be 1, 2, 4, 8, 12, ..., capped at 64 if target_data_batch < 2: diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 4d757f5010e..cc725f861de 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -477,6 +477,8 @@ def mem_get_info(device): return_value=SimpleNamespace(total_memory=10**12), ), patch("torch.cuda.device_count", return_value=4), + # A launcher that shows each rank one GPU: every rank computes on cuda:0. + patch("torch.cuda.current_device", return_value=0), patch("torch.cuda.mem_get_info", side_effect=mem_get_info), patch("torch.cuda.max_memory_allocated", return_value=0), ): @@ -487,7 +489,7 @@ def mem_get_info(device): torch.distributed.destroy_process_group(group) if planned: - assert measured == [torch.device("cuda", rank)] * 2 # its own GPU (LOCAL_RANK) only + assert measured == [torch.device("cuda", 0)] * 2 # the GPU it computes on, not LOCAL_RANK # One probe: an OOM retry would run more forwards on some ranks. assert model.batch_sizes == [1] assert batch_size == 4 # rank 0's From 36f54d73f4f203f3be33bb465217e839196a0fbd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:37:33 +0000 Subject: [PATCH 32/60] Refuse a streaming calibration loader that has workers A stream that splits itself across the loader's workers by get_worker_info(), as a Hugging Face streaming dataset does, was split again by rank and silently lost samples: the shards stayed disjoint and the step counts equal, so nothing failed. shard_dataloader now refuses a streaming loader with workers, and the rank's shard no longer splits its groups across workers it cannot have. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 27 ++++++------ .../utils/test_distributed_calibration.py | 42 +++++++++---------- 2 files changed, 34 insertions(+), 35 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 016e5acb483..0f6e14fb702 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -41,7 +41,6 @@ IterableDataset, RandomSampler, SequentialSampler, - get_worker_info, ) from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm @@ -701,12 +700,10 @@ class _StridedIterableShard(IterableDataset): trailing incomplete group is dropped so every rank yields exactly the same count -- preserving the equal-per-rank-batch invariant that keeps per-step EP collectives from hanging when one rank runs short. A stream shorter than ``world_size`` has no complete group, so it is refused rather - than leaving every rank with nothing to calibrate on. Each DataLoader worker process iterates - its own copy of the stream, so the workers split the groups between them rather than each - yielding them all: every worker still decodes the whole stream and keeps ``1/num_workers`` of - it, so ``num_workers > 0`` buys parallelism at ``num_workers`` times the decode cost. + than leaving every rank with nothing to calibrate on. It is iterated in the loader's own process + (:func:`shard_dataloader` refuses workers for a stream). ``vlm_dataset_utils._ShardedIterable`` cannot stand in: it pads every rank to a count its own loader - knows, while this wraps any loader's stream, which may know none, and splits across its workers. + knows, while this wraps any loader's stream, which may know none. """ def __init__(self, base: IterableDataset, rank: int, world_size: int): @@ -715,8 +712,6 @@ def __init__(self, base: IterableDataset, rank: int, world_size: int): self.world_size = world_size def __iter__(self): - info = get_worker_info() - num_workers, worker = (info.num_workers, info.id) if info is not None else (1, 0) it = iter(self.base) for index in itertools.count(): group = list(itertools.islice(it, self.world_size)) @@ -729,8 +724,7 @@ def __iter__(self): f"{self.world_size} ranks: no rank would get a sample to calibrate on." ) return # drop incomplete trailing group -> equal counts across ranks - if index % num_workers == worker: - yield group[self.rank] + yield group[self.rank] def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoader: @@ -746,9 +740,9 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad :class:`_StridedIterableShard`. That stream must yield the same samples in the same order on every rank and in every DataLoader worker, or the shards overlap and the ranks' step counts can differ, which hangs the collectives: one that shuffles needs a fixed seed of its own, as torch - seeds each worker's RNG differently and the ranks' RNGs need not agree. Every DataLoader worker - must see that whole stream: a dataset that splits itself across workers by ``get_worker_info()`` - would be split twice and lose samples, so load one with ``num_workers=0``. + seeds each worker's RNG differently and the ranks' RNGs need not agree. A stream loaded with + workers is refused: a dataset that splits itself across them by ``get_worker_info()``, as a + Hugging Face streaming dataset does, would be split twice and silently lose samples. Forwards all non-sampler DataLoader settings from ``loader`` (batching, workers, pinning, prefetch, init fn, generator, ...). With one rank there is nothing to shard: ``loader`` itself @@ -771,6 +765,13 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad "just the samples to calibrate on (a torch.utils.data.Subset) instead." ) if isinstance(loader.dataset, IterableDataset): + if loader.num_workers: + raise ValueError( + f"shard_dataloader: a streaming loader with num_workers={loader.num_workers} is " + "refused: a dataset that splits itself across the workers by get_worker_info() " + "would be split twice, by worker and by rank, and lose samples. Load it with " + "num_workers=0." + ) return DataLoader( _StridedIterableShard(loader.dataset, rank, world_size), batch_size=loader.batch_size, diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index d039b895e87..1bf9a8dbde4 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -142,23 +142,14 @@ def test_a_stream_shorter_than_the_world_is_refused_on_every_rank(): _ids(shard_dataloader(_loader(_Stream(3)), rank, 4)) -def test_iterable_shards_split_their_groups_across_worker_processes(monkeypatch): - """Each worker process iterates its own copy of the stream. Unless the workers split the groups - between them, every one yields the whole shard and each sample comes out once per worker. +def test_a_stream_loaded_with_workers_is_refused(): + """A stream that splits itself across the loader's workers, as a Hugging Face streaming dataset + does, would be split again by rank and silently lose samples. A lone rank keeps its loader.""" + loader = DataLoader(_Stream(16), batch_size=2, collate_fn=_collate, num_workers=2) - Each worker's pass runs here, in this process, as the worker it is told it is. - """ - n, world_size, num_workers = 16, 2, 2 - for rank in range(world_size): - loader = DataLoader(_Stream(n), batch_size=2, collate_fn=_collate, num_workers=num_workers) - shard = shard_dataloader(loader, rank, world_size).dataset - yielded = [] - for worker in range(num_workers): - info = types.SimpleNamespace(num_workers=num_workers, id=worker) - monkeypatch.setattr(mto_dist, "get_worker_info", lambda info=info: info) - yielded += list(shard) - - assert sorted(yielded) == list(range(rank, n, world_size)) + with pytest.raises(ValueError, match="num_workers=2 is refused"): + shard_dataloader(loader, 0, 2) + assert shard_dataloader(loader, 0, 1) is loader @pytest.mark.parametrize( @@ -211,21 +202,28 @@ def test_the_shard_batches_as_the_loader_does(kind, batching, batches): @pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) def test_sharded_loader_keeps_the_original_loader_settings(dataset): - """Batching, workers, pinning, timeout, worker init, RNG and prefetching carry over to the - shard.""" + """Batching, pinning, timeout, worker init and RNG carry over to the shard, and for a map-style + loader its workers and prefetching too (a stream is refused workers).""" + workers = ( + {} + if isinstance(dataset, _Stream) + else { + "num_workers": 2, + "multiprocessing_context": "spawn", + "prefetch_factor": 3, + "persistent_workers": True, + } + ) loader = DataLoader( dataset, batch_size=4, drop_last=True, collate_fn=_collate, - num_workers=2, pin_memory=True, timeout=7.0, worker_init_fn=_seed_worker, - multiprocessing_context="spawn", generator=torch.Generator(), - prefetch_factor=3, - persistent_workers=True, + **workers, ) shard = shard_dataloader(loader, 1, 2) From 2174d523a0aa5468a6be631d9fa13033ac2406f4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:49:11 +0000 Subject: [PATCH 33/60] Drop the worker clauses a refused streaming loader makes moot shard_dataloader refuses a streaming loader with workers, so asking the stream to be the same in every worker, and explaining that torch seeds each worker differently, described a case that cannot arise. _StridedIterableShard's docstring carried the same clause. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 0f6e14fb702..c3eeca007f1 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -695,7 +695,7 @@ class _StridedIterableShard(IterableDataset): ``DistributedSampler`` only works with map-style datasets. Streaming calibration datasets (e.g. the Nemotron image-text ``NemotronTarPlusJsonlIterable``) are iterable, so each rank instead builds its - own copy of the stream, the same on every rank and worker as :func:`shard_dataloader` requires, + own copy of the stream, the same on every rank as :func:`shard_dataloader` requires, and takes slot ``rank`` from every complete group of ``world_size`` consecutive samples. Any trailing incomplete group is dropped so every rank yields exactly the same count -- preserving the equal-per-rank-batch invariant that keeps per-step EP collectives from hanging when one rank @@ -738,11 +738,11 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad rather than silently widened. For streaming ``IterableDataset`` datasets (which have no ``len()`` and so can't use a sampler) each rank strides its own copy of the stream via :class:`_StridedIterableShard`. That stream must yield the same samples in the same order on - every rank and in every DataLoader worker, or the shards overlap and the ranks' step counts can - differ, which hangs the collectives: one that shuffles needs a fixed seed of its own, as torch - seeds each worker's RNG differently and the ranks' RNGs need not agree. A stream loaded with - workers is refused: a dataset that splits itself across them by ``get_worker_info()``, as a - Hugging Face streaming dataset does, would be split twice and silently lose samples. + every rank, or the shards overlap and the ranks' step counts can differ, which hangs the + collectives: one that shuffles needs a fixed seed of its own, as the ranks' RNGs need not agree. + A stream loaded with workers is refused: a dataset that splits itself across them by + ``get_worker_info()``, as a Hugging Face streaming dataset does, would be split twice and + silently lose samples. Forwards all non-sampler DataLoader settings from ``loader`` (batching, workers, pinning, prefetch, init fn, generator, ...). With one rank there is nothing to shard: ``loader`` itself From 99928330dcd41be3c5d77001dba04e035aa50c6b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:49:10 +0000 Subject: [PATCH 34/60] Refuse a random sampler that draws its own number of samples shard_dataloader let any RandomSampler through, but one drawing fewer samples than the dataset (num_samples) or with replacement picks which and how many samples the loader reads, as a subset does. The rank's DistributedSampler then replaced it over the whole dataset, so calibration silently ran on all of it instead of the few asked for. Such a sampler is now refused like the other pickers; a sampler that shuffles the whole dataset is still sharded as the rank's unshuffled share. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 12 +++++++++--- .../unit/torch/utils/test_distributed_calibration.py | 10 ++++++++-- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index c3eeca007f1..046af6d90ee 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -734,8 +734,9 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad counts equal (else a rank exits calibration early and hangs the others on the per-step collectives), at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. It replaces the loader's sampler, so a shuffling one becomes the rank's unshuffled share, and one that - picks samples itself (a subset, or a batch sampler other than a plain ``BatchSampler``) is refused - rather than silently widened. For streaming ``IterableDataset`` datasets (which have no ``len()`` + picks samples itself (a subset, a random sampler drawing its own number of samples or with + replacement, or a batch sampler other than a plain ``BatchSampler``) is refused rather than + silently widened. For streaming ``IterableDataset`` datasets (which have no ``len()`` and so can't use a sampler) each rank strides its own copy of the stream via :class:`_StridedIterableShard`. That stream must yield the same samples in the same order on every rank, or the shards overlap and the ranks' step counts can differ, which hangs the @@ -758,7 +759,12 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad picker = batching.sampler else: picker = loader.sampler if batching is None else batching - if type(picker) not in (SequentialSampler, RandomSampler, DistributedSampler): + # A RandomSampler drawing its own number of samples, or with replacement, picks which and how + # many it reads, as a subset does; one that shuffles the whole dataset does not. + narrowing = isinstance(picker, RandomSampler) and ( + picker.replacement or picker.num_samples != len(picker.data_source) + ) + if type(picker) not in (SequentialSampler, RandomSampler, DistributedSampler) or narrowing: raise ValueError( f"shard_dataloader: the loader's {type(picker).__name__} picks which samples it " "reads, and the rank's DistributedSampler would replace it; pass a loader over " diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 1bf9a8dbde4..b4a3e47b2c7 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -34,6 +34,7 @@ BatchSampler, DataLoader, IterableDataset, + RandomSampler, SequentialSampler, SubsetRandomSampler, ) @@ -159,12 +160,16 @@ def test_a_stream_loaded_with_workers_is_refused(): {"sampler": SubsetRandomSampler(range(4)), "batch_size": None}, {"batch_sampler": [[0, 1], [2, 3]]}, {"batch_sampler": BatchSampler(SubsetRandomSampler(range(4)), 2, drop_last=False)}, + {"sampler": RandomSampler(range(8), num_samples=4), "batch_size": 2}, + {"sampler": RandomSampler(range(8), replacement=True), "batch_size": 2}, ], ids=[ "subset-sampler", "subset-sampler-unbatched", "batch-sampler", "subset-in-a-plain-batch-sampler", + "random-sampler-drawing-fewer", + "random-sampler-with-replacement", ], ) def test_a_loader_that_picks_its_own_samples_is_refused(picker): @@ -188,13 +193,14 @@ def test_a_loader_that_picks_its_own_samples_is_refused(picker): ), ("map-style", {"batch_size": None}, [[1], [3], [5], [7], [9]]), ("iterable", {"batch_size": None}, [[1], [3], [5], [7], [9]]), + ("map-style", {"batch_size": 2, "shuffle": True}, [[1, 3], [5, 7], [9]]), ], - ids=["plain-batch-sampler", "auto-batching-off", "auto-batching-off-stream"], + ids=["plain-batch-sampler", "auto-batching-off", "auto-batching-off-stream", "shuffled"], ) def test_the_shard_batches_as_the_loader_does(kind, batching, batches): """A loader given a plain BatchSampler reports ``batch_size=None``; its shard still batches by the sampler's size and drop_last. A loader with auto-batching off reads one item per step, and so - does its shard.""" + does its shard. A loader shuffling the whole dataset gives the rank its unshuffled share.""" loader = DataLoader(_dataset(kind, 10), collate_fn=_collate, **batching) assert _batches(shard_dataloader(loader, 1, 2)) == batches From f9aae8c0cc18c41168d22d865ea1750c59a055e5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:03:09 +0000 Subject: [PATCH 35/60] Name the calibration defaults the rank-local refusal relies on The refusal reads the algorithm's options as the caller spelled them, which holds while every registered default syncs over the ranks; the comment names those defaults, so one that stops syncing is known to need refusing here. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index c77115943e1..c446a998f22 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -87,6 +87,10 @@ def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: """The rank-local calibrations ``algorithm`` runs, with those it runs for its scales.""" found = [] + # Reads the options as the caller spelled them, which is enough while every registered default + # syncs over the ranks: max and mse default to distributed_sync=True, nvfp4_act_headroom's + # weight_scale_algorithm to max, and lsq calibrates its scales with mse when it names none. A + # default that keeps its statistics on the rank has to be refused here too. for mode, config in get_modelike_from_algo_cfg(algorithm): method = mode.removesuffix("_calibrate") if method in _RANK_LOCAL_CALIBRATION: From 8d7daa5d1516feb5221e2e3d2cf9a087fa3a9ab4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:31:42 +0000 Subject: [PATCH 36/60] Size the two-rank batch test's group with bounded_collectives The test created its own gloo group and destroyed it at once when the job ran no collective on it, as the no-plan case does. gloo's new_group can return on one rank while a peer is still connecting, so the peer's new_group failed with "Connection closed by peer" (seen once in CI). bounded_collectives makes the same bounded group and passes a barrier on it first. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_dataset_utils.py | 42 +++++++++----------- 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index cc725f861de..138cc35c1da 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -14,13 +14,12 @@ # limitations under the License. import json -from datetime import timedelta -from functools import partial from types import SimpleNamespace from unittest.mock import Mock, patch import pytest import torch +from _test_utils.torch.distributed.utils import bounded_collectives from _test_utils.torch.transformers_models import get_tiny_tokenizer from torch.utils.data import DataLoader @@ -465,28 +464,23 @@ def mem_get_info(device): measured.append(device) return (free if len(measured) <= gpus else free - 100, 0) - # A rank left waiting in the reduction gives up in seconds on this group. - group = torch.distributed.new_group(timeout=timedelta(seconds=10)) - gather = partial(torch.distributed.all_gather_object, group=group) - try: - with ( - patch.object(torch.distributed, "all_gather_object", gather), - patch("torch.cuda.empty_cache"), - patch( - "torch.cuda.get_device_properties", - return_value=SimpleNamespace(total_memory=10**12), - ), - patch("torch.cuda.device_count", return_value=4), - # A launcher that shows each rank one GPU: every rank computes on cuda:0. - patch("torch.cuda.current_device", return_value=0), - patch("torch.cuda.mem_get_info", side_effect=mem_get_info), - patch("torch.cuda.max_memory_allocated", return_value=0), - ): - batch_size = get_max_batch_size( - model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) - ) - finally: - torch.distributed.destroy_process_group(group) + with ( + # A rank left waiting in the reduction gives up in seconds. + bounded_collectives(), + patch("torch.cuda.empty_cache"), + patch( + "torch.cuda.get_device_properties", + return_value=SimpleNamespace(total_memory=10**12), + ), + patch("torch.cuda.device_count", return_value=4), + # A launcher that shows each rank one GPU: every rank computes on cuda:0. + patch("torch.cuda.current_device", return_value=0), + patch("torch.cuda.mem_get_info", side_effect=mem_get_info), + patch("torch.cuda.max_memory_allocated", return_value=0), + ): + batch_size = get_max_batch_size( + model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) + ) if planned: assert measured == [torch.device("cuda", 0)] * 2 # the GPU it computes on, not LOCAL_RANK From 74ddb3d33fa2d853733601f304aa4e178aadc777 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:31:45 +0000 Subject: [PATCH 37/60] Import the batch sizing's distributed helpers at module top Nothing defers them: modelopt.torch.utils.logging, which dataset_utils imports, already imports modelopt.torch.utils.distributed at its top. The hf_ptq argument test's new imports move to the top too. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/dataset_utils.py | 3 +-- tests/examples/hf_ptq/test_hf_ptq_args.py | 7 ++----- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 5fe326c2f6f..124cded7a38 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -30,6 +30,7 @@ from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm +from .distributed import all_reduce, recorded_distribution_plan from .logging import print_rank_0, warn_rank_0 if TYPE_CHECKING: @@ -1020,8 +1021,6 @@ def get_max_batch_size( the estimate does not fit, pass an explicit batch size instead. A model with no plan recorded, such as a Megatron model, is measured on every visible GPU, as in a single process. """ - from .distributed import all_reduce, recorded_distribution_plan - args = ( model, max_sample_length, diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 937868eb7c3..5543e9de902 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +26,7 @@ import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 -from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig +from transformers import Qwen3_5MoeConfig, Qwen3Config, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -38,6 +38,7 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.utils import mlflow as mlflow_lib +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.mlflow import describe_run, run_tags _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -483,10 +484,6 @@ def distributed_load(*args, **kwargs): def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path): """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an FSDP2-sharded model without one is refused rather than taken for a single process's.""" - from transformers import Qwen3Config - - from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan - Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path) monkeypatch.setenv("RANK", "0") hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2") From d4bc40b42b2b524da9e409f08c2279f085ea1987 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:47:54 +0000 Subject: [PATCH 38/60] Read the calibration options through each mode's schema before refusing The rank-local refusal read the algorithm dict as the caller spelled it, so {"method": "max", "distributed_sync": "false"}, which MaxCalibConfig coerces to False, calibrated without the sync. Each mode's options are now validated by its config class first, nested scale algorithms included, so coerced values and defaults count as calibration reads them. The test spells the sync off as "false", and gives nvfp4_act_headroom its local_hessian weight-scale algorithm in the form its schema accepts. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 19 +++++++++---------- .../utils/test_distributed_calibration.py | 7 ++++--- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index c446a998f22..926267f0167 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -53,7 +53,7 @@ get_kv_cache_auto_quantize_config, ) from .kv_cache_auto_quant import _validate_search_inputs as _validate_kv_cache_search_inputs -from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg +from .mode import CalibrateModeRegistry, QuantizeModeRegistry, get_modelike_from_algo_cfg from .nn import AnyQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer from .utils import is_quantized from .utils.core_utils import _get_fsdp2_mesh @@ -87,19 +87,18 @@ def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: """The rank-local calibrations ``algorithm`` runs, with those it runs for its scales.""" found = [] - # Reads the options as the caller spelled them, which is enough while every registered default - # syncs over the ranks: max and mse default to distributed_sync=True, nvfp4_act_headroom's - # weight_scale_algorithm to max, and lsq calibrates its scales with mse when it names none. A - # default that keeps its statistics on the rank has to be refused here too. - for mode, config in get_modelike_from_algo_cfg(algorithm): + for mode, options in get_modelike_from_algo_cfg(algorithm): method = mode.removesuffix("_calibrate") + # The options as the calibration reads them, through its mode's schema: a string such as + # "false" is coerced, and a default counts as if spelled out. + config = CalibrateModeRegistry[mode].config_class.model_validate(options) if method in _RANK_LOCAL_CALIBRATION: found.append(method) - elif "distributed_sync" in config and not config["distributed_sync"]: - found.append(f"{method} with distributed_sync={config['distributed_sync']}") + elif getattr(config, "distributed_sync", True) is False: + found.append(f"{method} with distributed_sync=False") for key in ("weight_scale_algorithm", "scale_algorithm"): - if config.get(key) is not None: - found += _rank_local_calibrations(config[key]) + if getattr(config, key, None) is not None: + found += _rank_local_calibrations(getattr(config, key)) return found diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index b4a3e47b2c7..57304f869dd 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -568,12 +568,12 @@ def test_forward_loop_refuses_a_batch_that_is_not_a_mapping(): ("calibrate", "awq_lite", "awq_lite"), ( "calibrate", - {"method": "nvfp4_act_headroom", "weight_scale_algorithm": "local_hessian"}, + {"method": "nvfp4_act_headroom", "weight_scale_algorithm": {"method": "local_hessian"}}, "local_hessian", ), ( "calibrate", - {"method": "max", "distributed_sync": False}, + {"method": "max", "distributed_sync": "false"}, "max with distributed_sync=False", ), ("auto_quantize", {"method": "awq_lite"}, "awq_lite"), @@ -585,7 +585,8 @@ def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model ): """Each rank calibrates on its own data, and AWQ keeps what its search gathers on the rank: the ranks' copies of a weight would come out different. So would local_hessian, run by another - algorithm for its weight scales, and max told not to sync. The refusal comes before any change.""" + algorithm for its weight scales, and max told not to sync, in any spelling its schema reads as + off. The refusal comes before any change.""" model = nn.Sequential(nn.Linear(16, 16)) set_distribution_plan(model, DistributionPlan(world_size=2)) if entry == "calibrate": From 01731d7fd686e3e8649ea3a79646ab50af58a99d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:47:57 +0000 Subject: [PATCH 39/60] Check the experts' groups by the ranks they span, not only the data-parallel size A data-parallel group of the plan's edp ranks passed the expert-parallel check even when it was the experts' own expert-parallel group, whose ranks hold other experts: at EP 2 over four ranks, max's and GPTQ's syncs would then merge different experts' scales. A routed expert's ParallelState now passes only with an expert-parallel group of ep_size ranks and a data-parallel group of edp ranks that meets it at this rank alone. The two-rank pool cannot build such groups, so the new test answers four ranks' process-group queries. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 40 ++++++++++++----- .../utils/test_distributed_calibration.py | 43 +++++++++++++++++++ 2 files changed, 73 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 926267f0167..aeea86427de 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -102,6 +102,31 @@ def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: return found +def _group_ranks(group: dist.DistributedProcessGroup) -> set[int]: + """The ranks ``group`` spans; just this rank for a group not set.""" + if not group.is_initialized(): + return {dist.rank()} + return set( + torch.distributed.get_process_group_ranks(group.group or torch.distributed.group.WORLD) + ) + + +def _names_the_experts_groups( + state: "dist.ParallelState | None", plan: dist.DistributionPlan +) -> bool: + """Whether ``state`` names the groups a routed expert of ``plan`` syncs its scales over. + + Its expert-parallel group of ``ep_size`` ranks, and as its data-parallel group the ``edp`` + ranks holding the same experts, which meet the expert-parallel group, whose ranks hold the + others, only at this rank. + """ + if state is None or not state.expert_model_parallel_group.is_initialized(): + return False + ep = _group_ranks(state.expert_model_parallel_group) + dp = _group_ranks(state.data_parallel_group) + return len(ep) == plan.ep_size and len(dp) == plan.edp and ep & dp == {dist.rank()} + + def _refuse_unsupported_distributed_calibration( model: nn.Module, algorithm: QuantizeAlgoCfgType ) -> None: @@ -142,19 +167,14 @@ def _refuse_unsupported_distributed_calibration( holders[name] = state if not holders and not any("experts" in name.split(".") for name, _ in model.named_modules()): return # no routed experts here, as in a VLM's vision tower: no experts' scales to merge - wrong = [ - name - for name, state in holders.items() - if state is None - or not state.expert_model_parallel_group.is_initialized() - or max(state.data_parallel_group.world_size(), 1) != plan.edp - ] + wrong = [name for name, state in holders.items() if not _names_the_experts_groups(state, plan)] if not holders or wrong: raise NotImplementedError( f"Calibrating a model whose experts are sharded over {plan.ep_size} ranks needs each " - "routed expert's ParallelState to name its expert-parallel group and, as its " - f"data-parallel group, the {plan.edp} rank(s) holding the same experts, or calibration " - "would sync their scales over the wrong ranks. " + f"routed expert's ParallelState to name its expert-parallel group of {plan.ep_size} " + f"ranks and, as its data-parallel group, the {plan.edp} rank(s) holding the same " + "experts, which share only this rank with that group, or calibration would sync their " + "scales over the wrong ranks. " + ( f"These modules do not: {wrong[:3]}." if wrong diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 57304f869dd..bfd765823a7 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -41,6 +41,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as mto_dist +from modelopt.torch.quantization.model_quant import _refuse_unsupported_distributed_calibration from modelopt.torch.utils.distributed import ( DistributionPlan, ParallelState, @@ -664,6 +665,48 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) +# Rank 0 of four at EP 2: ranks 0 and 1 shard the experts between them, and rank 2 holds the same +# experts as rank 0. The two-rank pool cannot build these groups, so the process group answers as a +# four-rank launch's would. +_FOUR_RANK_GROUPS = {"ep": [0, 1], "edp": [0, 2], "world": [0, 1, 2, 3]} + + +@pytest.mark.parametrize( + ("ep", "dp", "seeded"), + [("ep", "edp", True), ("ep", "ep", False), ("world", "edp", False)], + ids=[ + "seeded", + "data-parallel-over-the-expert-parallel-ranks", + "expert-parallel-over-every-rank", + ], +) +def test_the_experts_groups_are_checked_by_their_ranks(monkeypatch, ep, dp, seeded): + """A data-parallel group of the plan's edp ranks is not enough: over the experts' + expert-parallel ranks, which hold other experts, max's and GPTQ's syncs would merge different + experts' scales. Nor is an expert-parallel group over more ranks than the plan shards over.""" + monkeypatch.setattr(mto_dist, "backend", lambda: "torch") + monkeypatch.setattr(mto_dist, "rank", lambda group=None: 0) + monkeypatch.setattr( + torch.distributed, + "get_world_size", + lambda group=None: 4 if group is None else len(_FOUR_RANK_GROUPS[group]), + ) + monkeypatch.setattr( + torch.distributed, "get_process_group_ranks", lambda group: _FOUR_RANK_GROUPS[group] + ) + model = nn.Sequential(OrderedDict(experts=nn.Linear(16, 16))) + model.experts._parallel_state = ParallelState( + data_parallel_group=dp, expert_model_parallel_group=ep + ) + set_distribution_plan(model, DistributionPlan(world_size=4, ep_size=2)) + + if seeded: + _refuse_unsupported_distributed_calibration(model, "max") + else: + with pytest.raises(NotImplementedError, match="expert-parallel group of 2 ranks"): + _refuse_unsupported_distributed_calibration(model, "max") + + @pytest.mark.parametrize( "config", [ From e1304e0c504a0e00b10fda0f2468f42db834c80f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 01:02:14 +0000 Subject: [PATCH 40/60] Average awq_clip's per-tensor clip losses over the ranks, and allow it distributed awq_clip's static-block branch already all-reduced each candidate clip's loss over the module's data-parallel group; its per-tensor branch, which NVFP4 takes, kept each rank's own, so the ranks' copies of a weight took different clips. Both branches now average through one helper, issued on every rank in the same order, and awq_clip is no longer refused on a model distributed over several ranks. This also makes Megatron data-parallel replicas agree on the per-tensor clip. Tests: a two-rank CPU test of the static-block branch, and a two-GPU NVFP4 test of the per-tensor branch, each with ranks whose own data would pick different clips. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- examples/hf_ptq/README.md | 4 ++- modelopt/torch/quantization/model_calib.py | 17 ++++++----- modelopt/torch/quantization/model_quant.py | 5 ++-- .../gpu/torch/quantization/test_calib_cuda.py | 30 +++++++++++++++++++ tests/unit/torch/quantization/test_calib.py | 28 +++++++++++++++++ 6 files changed, 74 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 15430999f8b..31afcc36e58 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -46,7 +46,7 @@ Changelog **Backward Breaking Changes** -- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank: AWQ, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate. Each rank calibrates on its own share of the data, so these calibrated the ranks' shards of a weight differently; run them without ``--use_fsdp2``. +- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). AWQ-Clip, as in ``nvfp4_awq_clip``, now averages its clip search over the ranks and works with ``--use_fsdp2`` on several GPUs. - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 649b3db5fe6..f76215adc8f 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -616,7 +616,9 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* -> *With `--use_fsdp2` on more than one GPU, calibration that keeps its statistics on each rank is refused when quantization starts: AWQ and SVDQuant (`--qformat int4_awq`, `w4a8_awq_beta`, `nvfp4_awq_lite`, `nvfp4_awq_clip`, `nvfp4_awq_full`, `nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`.* +> *With `--use_fsdp2` on more than one GPU, calibration that keeps its statistics on each rank is refused when quantization starts: AWQ and SVDQuant (`--qformat int4_awq`, `w4a8_awq_beta`, `nvfp4_awq_lite`, `nvfp4_awq_full`, `nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`.* + +> *`nvfp4_awq_clip` works with `--use_fsdp2` on more than one GPU: its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* > *With `--use_fsdp2` on more than one GPU and `--batch_size 0`, each rank estimates the calibration batch size from one forward and the ranks use the smallest. That size is not probed for out-of-memory, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives: set `--batch_size` if calibration runs out of memory.* diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7397d79eb56..6a89f49f348 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2067,6 +2067,14 @@ def update_best_params(self): indices, self.awq_clip.w_amax * shrink, self.awq_clip.best_clip_val ) + def _data_parallel_mean(self, loss): + """``loss`` averaged over the data-parallel ranks, so every replica picks the same clip.""" + parallel_group = self.parallel_state.data_parallel_group + if parallel_group.is_initialized(): + dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=parallel_group.group) + loss /= parallel_group.world_size() + return loss + def _clip_search(self, inputs, co_bsz=256, max_tokens=16): weight = self.weight self.weight_quantizer.enable() @@ -2080,7 +2088,7 @@ def _clip_search(self, inputs, co_bsz=256, max_tokens=16): self.weight_quantizer.amax = original_amax * shrink out = inputs @ self.weight_quantizer(self.weight).T loss = (out - out_actual).float().pow(2).mean() - self.awq_clip.loss[shrink] += loss + self.awq_clip.loss[shrink] += _data_parallel_mean(self, loss) else: # weight [co, ci] -> [co, 1, n_block, block_size] # inputs [..., ci] -> [1, max_tokens, n_block, block_size] @@ -2124,12 +2132,7 @@ def _clip_search(self, inputs, co_bsz=256, max_tokens=16): cur_out = (inputs * cur_w).sum(dim=-1) # co_bsz, max_tokens, n_block # co_bsz, n_block - loss = (cur_out - org_out).float().pow(2).mean(dim=1) - - parallel_group = self.parallel_state.data_parallel_group - if parallel_group.is_initialized(): - dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=parallel_group.group) - loss /= parallel_group.world_size() + loss = _data_parallel_mean(self, (cur_out - org_out).float().pow(2).mean(dim=1)) del cur_out, cur_w self.awq_clip.loss[shrink][ diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index aeea86427de..ed24ee9e29a 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -76,7 +76,6 @@ # own there too when its config turns distributed_sync off. _RANK_LOCAL_CALIBRATION = ( "awq_lite", - "awq_clip", "awq_full", "svdquant", "local_hessian", @@ -151,8 +150,8 @@ def _refuse_unsupported_distributed_calibration( f"Calibrating with {', '.join(refused)} is not supported on a model distributed over " f"{plan.world_size} ranks. Each rank calibrates on its own data and this calibration " "keeps what it gathers on the rank, so the ranks would calibrate their copies, or their " - "shards, of a weight differently. Calibrate with max, MSE, GPTQ or SmoothQuant, with " - "distributed_sync on, or in a single process." + "shards, of a weight differently. Calibrate with max, MSE, GPTQ, SmoothQuant or " + "awq_clip, with distributed_sync on, or in a single process." ) if algorithm is None or not plan.is_expert_parallel: return diff --git a/tests/gpu/torch/quantization/test_calib_cuda.py b/tests/gpu/torch/quantization/test_calib_cuda.py index 677a03dc0e2..a00aae6e59d 100644 --- a/tests/gpu/torch/quantization/test_calib_cuda.py +++ b/tests/gpu/torch/quantization/test_calib_cuda.py @@ -16,9 +16,11 @@ """Calibration tests.""" import torch +import torch.distributed as dist import torch.nn as nn import modelopt.torch.quantization as mtq +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan class _TwoBranchModel(nn.Module): @@ -71,3 +73,31 @@ def _forward_loop(m): if uncal_q.amax is not None: assert uncal_q.axis is None assert uncal_q.amax.numel() == 1 + + +def _nvfp4_awq_clip_on_every_rank_job(rank, size): + device = torch.device(f"cuda:{rank}") + model = nn.Sequential(nn.Linear(64, 16, bias=False)).to(device) + with torch.no_grad(): + # Every rank's copy: tiny weights, whose FP8 block scales a tighter per-tensor amax resolves + # better, and a large one per row on input channel 0, which a tighter amax clips. + model[0].weight.copy_(torch.linspace(-1e-3, 1e-3, 16 * 64, device=device).view(16, 64)) + model[0].weight[:, 0] = 100.0 + set_distribution_plan(model, DistributionPlan(size)) + x = torch.randn(8, 64, device=device, generator=torch.Generator(device).manual_seed(rank)) + if rank % 2 == 0: + x[:, 0] = 0 # never reaches the large weights: the tightest clip suits it + else: + x[:, 1:] = 0 # reaches only them: no clip suits it + mtq.quantize(model, mtq.NVFP4_AWQ_CLIP_CFG, lambda m: m(x)) + amax = model[0].weight_quantizer.amax.float().reshape(1) + clips = [torch.empty_like(amax) for _ in range(size)] + dist.all_gather(clips, amax) + assert all(torch.equal(clip, amax) for clip in clips), clips + + +def test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank(need_2_gpus, dist_workers): + """NVFP4's awq_clip clips the per-tensor amax. Each rank of a model distributed over several + ranks calibrates on its own data, and averaging each candidate clip's loss over the ranks gives + every rank's copy of a weight the same clip, though each rank's own data would pick another.""" + dist_workers.run(_nvfp4_awq_clip_on_every_rank_job) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 27a3dae1776..046ec714680 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -20,6 +20,7 @@ import pytest import torch import torch.nn as nn +from _test_utils.torch.distributed.utils import bounded_collectives from _test_utils.torch.quantization.quantize_common import get_awq_config import modelopt.torch.quantization as mtq @@ -33,6 +34,7 @@ ) from modelopt.torch.quantization.nn import QuantLinear, TensorQuantizer from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan class _SimpleMLP(nn.Module): @@ -266,6 +268,32 @@ def test_awq_clip(): _awq_clip_tester(model, model_ref, ref_data, config) +def _awq_clip_on_every_rank_job(rank, size): + model = nn.Sequential(nn.Linear(16, 4, bias=False)) + with torch.no_grad(): + # Every rank's copy: small weights, and a large one per row on input channel 0. + model[0].weight.copy_(torch.linspace(-0.1, 0.1, 64).view(4, 16)) + model[0].weight[:, 0] = 1.0 + set_distribution_plan(model, DistributionPlan(size)) + x = torch.randn(8, 16, generator=torch.Generator().manual_seed(rank)) + if rank % 2 == 0: + x[:, 0] = 0 # never reaches the large weights: a tighter clip only sharpens the rest + else: + x[:, 1:] = 0 # reaches only them: clipping them costs the most + with bounded_collectives(): + mtq.quantize(model, get_awq_config("awq_clip"), lambda m: m(x)) + clips = [None] * size + torch.distributed.all_gather_object(clips, model[0].weight_quantizer.amax) + assert all(torch.equal(clip, clips[0]) for clip in clips), clips + + +def test_awq_clip_picks_the_same_clips_on_every_rank(cpu_dist_workers): + """Each rank of a model distributed over several ranks calibrates on its own data. awq_clip + averages each candidate clip's loss over the data-parallel ranks, so the ranks' copies of a + weight take the same clips, though each rank's own data would pick others.""" + cpu_dist_workers.run(_awq_clip_on_every_rank_job) + + @pytest.mark.parametrize("use_weight_scale", [None, False]) def test_awq_full(use_weight_scale): """Test awq.""" From 1f6b1e29e9ca4c6f37265c616f6b6ddd62a65290 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 01:26:12 +0000 Subject: [PATCH 41/60] Refuse a calibration whose distributed_sync is None, as when it is False The schema-validated refusal read distributed_sync with "is False", but max and MSE accept None too, which skips their sync as False does, so {"method": "max", "distributed_sync": None} calibrated the ranks apart. Any falsy value is now refused, nested scale algorithms included. The refusal test gains an MSE weight-scale algorithm with distributed_sync None. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 4 ++-- .../utils/test_distributed_calibration.py | 22 ++++++++++++++++--- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index ed24ee9e29a..8607df994c1 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -93,8 +93,8 @@ def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: config = CalibrateModeRegistry[mode].config_class.model_validate(options) if method in _RANK_LOCAL_CALIBRATION: found.append(method) - elif getattr(config, "distributed_sync", True) is False: - found.append(f"{method} with distributed_sync=False") + elif not getattr(config, "distributed_sync", True): # None skips the sync too + found.append(f"{method} with distributed_sync={config.distributed_sync}") for key in ("weight_scale_algorithm", "scale_algorithm"): if getattr(config, key, None) is not None: found += _rank_local_calibrations(getattr(config, key)) diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index bfd765823a7..03a51638438 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -577,17 +577,33 @@ def test_forward_loop_refuses_a_batch_that_is_not_a_mapping(): {"method": "max", "distributed_sync": "false"}, "max with distributed_sync=False", ), + ( + "calibrate", + { + "method": "nvfp4_act_headroom", + "weight_scale_algorithm": {"method": "mse", "distributed_sync": None}, + }, + "mse with distributed_sync=None", + ), ("auto_quantize", {"method": "awq_lite"}, "awq_lite"), ], - ids=["quantize", "calibrate", "for-the-weight-scales", "max-without-its-sync", "auto_quantize"], + ids=[ + "quantize", + "calibrate", + "for-the-weight-scales", + "max-without-its-sync", + "for-the-weight-scales-without-their-sync", + "auto_quantize", + ], ) def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model( entry, algorithm, refused ): """Each rank calibrates on its own data, and AWQ keeps what its search gathers on the rank: the ranks' copies of a weight would come out different. So would local_hessian, run by another - algorithm for its weight scales, and max told not to sync, in any spelling its schema reads as - off. The refusal comes before any change.""" + algorithm for its weight scales, and max or MSE told not to sync, by False in any spelling its + schema reads as False, or by None, which skips the sync too. The refusal comes before any + change.""" model = nn.Sequential(nn.Linear(16, 16)) set_distribution_plan(model, DistributionPlan(world_size=2)) if entry == "calibrate": From cea62b0d6558c4bec4b3fae4f17a7f220ef4ef3f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 01:26:21 +0000 Subject: [PATCH 42/60] Check the experts' groups against the expert-parallel mesh's rows and columns The previous check passed any expert-parallel group of ep_size ranks and data-parallel group of edp ranks meeting it only at this rank, so diagonal data-parallel groups ([0, 3] and [1, 2] at EP 2 over four ranks, where ranks 0 and 2 hold the same shard) still merged different experts' scales. A routed expert's groups must now be this rank's row of the (edp, ep) mesh the expert-parallel load lays the ranks out in, as its expert-parallel group in row order, and its column as its data-parallel group. The four-rank test adds the diagonal data-parallel group and an expert-parallel group in another order. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 33 ++++++++++--------- .../utils/test_distributed_calibration.py | 32 +++++++++++++----- 2 files changed, 41 insertions(+), 24 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 8607df994c1..1ea04a98326 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -101,13 +101,11 @@ def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: return found -def _group_ranks(group: dist.DistributedProcessGroup) -> set[int]: - """The ranks ``group`` spans; just this rank for a group not set.""" +def _group_ranks(group: dist.DistributedProcessGroup) -> list[int]: + """The ranks ``group`` spans, in its rank order; just this rank for a group not set.""" if not group.is_initialized(): - return {dist.rank()} - return set( - torch.distributed.get_process_group_ranks(group.group or torch.distributed.group.WORLD) - ) + return [dist.rank()] + return torch.distributed.get_process_group_ranks(group.group or torch.distributed.group.WORLD) def _names_the_experts_groups( @@ -115,15 +113,18 @@ def _names_the_experts_groups( ) -> bool: """Whether ``state`` names the groups a routed expert of ``plan`` syncs its scales over. - Its expert-parallel group of ``ep_size`` ranks, and as its data-parallel group the ``edp`` - ranks holding the same experts, which meet the expert-parallel group, whose ranks hold the - others, only at this rank. + Those of the ``(edp, ep)`` mesh the expert-parallel load lays the ranks out in: as its + expert-parallel group, this rank's row, whose ranks hold the experts' shards in row order, and + as its data-parallel group, its column, the ranks holding the same shard. """ if state is None or not state.expert_model_parallel_group.is_initialized(): return False - ep = _group_ranks(state.expert_model_parallel_group) - dp = _group_ranks(state.data_parallel_group) - return len(ep) == plan.ep_size and len(dp) == plan.edp and ep & dp == {dist.rank()} + row, column = divmod(dist.rank(), plan.ep_size) + expert_parallel = range(row * plan.ep_size, (row + 1) * plan.ep_size) + data_parallel = range(column, plan.world_size, plan.ep_size) + return _group_ranks(state.expert_model_parallel_group) == list(expert_parallel) and sorted( + _group_ranks(state.data_parallel_group) + ) == list(data_parallel) def _refuse_unsupported_distributed_calibration( @@ -170,10 +171,10 @@ def _refuse_unsupported_distributed_calibration( if not holders or wrong: raise NotImplementedError( f"Calibrating a model whose experts are sharded over {plan.ep_size} ranks needs each " - f"routed expert's ParallelState to name its expert-parallel group of {plan.ep_size} " - f"ranks and, as its data-parallel group, the {plan.edp} rank(s) holding the same " - "experts, which share only this rank with that group, or calibration would sync their " - "scales over the wrong ranks. " + "routed expert's ParallelState to name, as its expert-parallel group, this rank's row " + f"of the ({plan.edp}, {plan.ep_size}) expert-parallel mesh and, as its data-parallel " + f"group, its column, the {plan.edp} rank(s) holding the same experts, or calibration " + "would sync their scales over the wrong ranks. " + ( f"These modules do not: {wrong[:3]}." if wrong diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 03a51638438..a3323c72592 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -681,25 +681,41 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) -# Rank 0 of four at EP 2: ranks 0 and 1 shard the experts between them, and rank 2 holds the same -# experts as rank 0. The two-rank pool cannot build these groups, so the process group answers as a -# four-rank launch's would. -_FOUR_RANK_GROUPS = {"ep": [0, 1], "edp": [0, 2], "world": [0, 1, 2, 3]} +# Rank 0 of four at EP 2, in the (edp, ep) = (2, 2) mesh: ranks 0 and 1 shard the experts between +# them, ranks 2 and 3 do the same, and rank 2 holds the same shard as rank 0. The two-rank pool cannot +# build these groups, so the process group answers as a four-rank launch's would. +_FOUR_RANK_GROUPS = { + "ep": [0, 1], + "ep-reordered": [1, 0], + "edp": [0, 2], + "diagonal": [0, 3], + "world": [0, 1, 2, 3], +} @pytest.mark.parametrize( ("ep", "dp", "seeded"), - [("ep", "edp", True), ("ep", "ep", False), ("world", "edp", False)], + [ + ("ep", "edp", True), + ("ep", "ep", False), + ("ep", "diagonal", False), + ("ep-reordered", "edp", False), + ("world", "edp", False), + ], ids=[ "seeded", "data-parallel-over-the-expert-parallel-ranks", + "data-parallel-over-another-shard", + "expert-parallel-in-another-order", "expert-parallel-over-every-rank", ], ) def test_the_experts_groups_are_checked_by_their_ranks(monkeypatch, ep, dp, seeded): """A data-parallel group of the plan's edp ranks is not enough: over the experts' - expert-parallel ranks, which hold other experts, max's and GPTQ's syncs would merge different - experts' scales. Nor is an expert-parallel group over more ranks than the plan shards over.""" + expert-parallel ranks, or ranks that meet them only at this one but hold another shard (rank 3 + holds rank 1's), max's and GPTQ's syncs would merge different experts' scales. Nor is an + expert-parallel group over more ranks than the plan shards over, or in another order, which + gives this rank another shard.""" monkeypatch.setattr(mto_dist, "backend", lambda: "torch") monkeypatch.setattr(mto_dist, "rank", lambda group=None: 0) monkeypatch.setattr( @@ -719,7 +735,7 @@ def test_the_experts_groups_are_checked_by_their_ranks(monkeypatch, ep, dp, seed if seeded: _refuse_unsupported_distributed_calibration(model, "max") else: - with pytest.raises(NotImplementedError, match="expert-parallel group of 2 ranks"): + with pytest.raises(NotImplementedError, match=r"row of the \(2, 2\) expert-parallel mesh"): _refuse_unsupported_distributed_calibration(model, "max") From b1766f1ae2b47f3f1bfff2d67fccba88461646f4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 01:37:09 +0000 Subject: [PATCH 43/60] Give AWQ-Clip's data-parallel clip its own changelog entry, under Bug Fixes It is not backward breaking and not specific to hf_ptq: it also changes the amax values that data-parallel runs (Megatron included) export. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 31afcc36e58..32013212eb0 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -46,7 +46,7 @@ Changelog **Backward Breaking Changes** -- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). AWQ-Clip, as in ``nvfp4_awq_clip``, now averages its clip search over the ranks and works with ``--use_fsdp2`` on several GPUs. +- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. @@ -69,6 +69,7 @@ Changelog **Bug Fixes** - Megatron-Core fake quantization of Gated DeltaNet (Qwen3.5 / Qwen3.6) with a Transformer Engine ``in_proj`` now keeps its beta / alpha gates in BF16 during evaluation, QAD, ``fold_weight`` and the vLLM fake-quant export, matching deployment. Other ``in_proj`` layers stay fully quantized and warn. +- AWQ-Clip's per-tensor clip, which ``nvfp4_awq_clip`` uses, is now searched over every data-parallel rank's data, as its block-wise clip already was, so the replicas of a weight take the same clip (exported amax values of data-parallel runs may change) and ``nvfp4_awq_clip`` works with ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on several GPUs. - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts. - Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead. - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges. From 4861acb95db47c9295133bb60dcf4894d45b804b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:05:11 +0000 Subject: [PATCH 44/60] Average awq_clip's clip search over the ranks once per module, after the forwards The data-parallel mean ran inside each forward, in both branches (the static-block one on main too), so a module issued its all-reduces once per forward it ran. An MoE expert's forward count is each rank's own, so ranks that ran it a different number of times paired mismatched collectives and hung. Each rank now accumulates its own losses, and after the forward loop every module's candidate losses and token count are reduced once over its data-parallel group, in module order on every rank: the same numbers, as the steps are equal, with a collective count the model fixes. That also replaces up to one all-reduce per candidate, step and output block with one per module. The two-rank test now runs an expert a different number of times on each rank; it fails without the change, on the bounded group's timeout. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 26 ++++++---- tests/unit/torch/quantization/test_calib.py | 55 +++++++++++++++------ 2 files changed, 58 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 6a89f49f348..fd8d0ca7d7d 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2067,13 +2067,20 @@ def update_best_params(self): indices, self.awq_clip.w_amax * shrink, self.awq_clip.best_clip_val ) - def _data_parallel_mean(self, loss): - """``loss`` averaged over the data-parallel ranks, so every replica picks the same clip.""" - parallel_group = self.parallel_state.data_parallel_group - if parallel_group.is_initialized(): - dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=parallel_group.group) - loss /= parallel_group.world_size() - return loss + def sync_across_dp(module): + """Average the clip search over the data-parallel ranks, so every replica picks one clip. + + Once, after the forwards: a module's forward count, as an MoE expert's, can differ by rank. + """ + group = module.parallel_state.data_parallel_group + if not group.is_initialized(): + return + losses = torch.stack(list(module.awq_clip.loss.values())) + tokens = torch.tensor(module.awq_clip.num_tokens, device=losses.device) + for tensor in (losses, tokens): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group.group) + module.awq_clip.loss = dict(zip(module.awq_clip.loss, losses / group.world_size())) + module.awq_clip.num_tokens = int(tokens.item()) def _clip_search(self, inputs, co_bsz=256, max_tokens=16): weight = self.weight @@ -2088,7 +2095,7 @@ def _clip_search(self, inputs, co_bsz=256, max_tokens=16): self.weight_quantizer.amax = original_amax * shrink out = inputs @ self.weight_quantizer(self.weight).T loss = (out - out_actual).float().pow(2).mean() - self.awq_clip.loss[shrink] += _data_parallel_mean(self, loss) + self.awq_clip.loss[shrink] += loss else: # weight [co, ci] -> [co, 1, n_block, block_size] # inputs [..., ci] -> [1, max_tokens, n_block, block_size] @@ -2132,7 +2139,7 @@ def _clip_search(self, inputs, co_bsz=256, max_tokens=16): cur_out = (inputs * cur_w).sum(dim=-1) # co_bsz, max_tokens, n_block # co_bsz, n_block - loss = _data_parallel_mean(self, (cur_out - org_out).float().pow(2).mean(dim=1)) + loss = (cur_out - org_out).float().pow(2).mean(dim=1) del cur_out, cur_w self.awq_clip.loss[shrink][ @@ -2203,6 +2210,7 @@ def postprocess(module): for name, module in model.named_modules(): if is_quantized_linear(module) and hasattr(module, "awq_clip"): + sync_across_dp(module) if module.awq_clip.num_tokens > 0: with enable_weight_access_and_writeback(module, model, names): postprocess(module) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 046ec714680..f30f2ff91c2 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -15,12 +15,13 @@ """Unittests for AWQ and SVDQuant""" +from datetime import timedelta from functools import partial import pytest import torch +import torch.distributed as dist import torch.nn as nn -from _test_utils.torch.distributed.utils import bounded_collectives from _test_utils.torch.quantization.quantize_common import get_awq_config import modelopt.torch.quantization as mtq @@ -34,7 +35,7 @@ ) from modelopt.torch.quantization.nn import QuantLinear, TensorQuantizer from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan +from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan class _SimpleMLP(nn.Module): @@ -268,29 +269,55 @@ def test_awq_clip(): _awq_clip_tester(model, model_ref, ref_data, config) +class _Experts(nn.Module): + """Two experts, each run only for the steps that route tokens to it, as an MoE runs them.""" + + def __init__(self): + super().__init__() + self.experts = nn.ModuleList(nn.Linear(16, 4, bias=False) for _ in range(2)) + + def forward(self, x, hit): + for index in hit: + self.experts[index](x) + + def _awq_clip_on_every_rank_job(rank, size): - model = nn.Sequential(nn.Linear(16, 4, bias=False)) - with torch.no_grad(): - # Every rank's copy: small weights, and a large one per row on input channel 0. - model[0].weight.copy_(torch.linspace(-0.1, 0.1, 64).view(4, 16)) - model[0].weight[:, 0] = 1.0 + model = _Experts() + # A rank left waiting in the clip search's sync gives up in seconds on this group. + group = dist.new_group(timeout=timedelta(seconds=10)) + dist.barrier(group=group) + for expert in model.experts: + with torch.no_grad(): + # Every rank's copy: small weights, and a large one per row on input channel 0. + expert.weight.copy_(torch.linspace(-0.1, 0.1, 64).view(4, 16)) + expert.weight[:, 0] = 1.0 + expert._parallel_state = ParallelState(data_parallel_group=group) set_distribution_plan(model, DistributionPlan(size)) x = torch.randn(8, 16, generator=torch.Generator().manual_seed(rank)) if rank % 2 == 0: x[:, 0] = 0 # never reaches the large weights: a tighter clip only sharpens the rest else: x[:, 1:] = 0 # reaches only them: clipping them costs the most - with bounded_collectives(): - mtq.quantize(model, get_awq_config("awq_clip"), lambda m: m(x)) - clips = [None] * size - torch.distributed.all_gather_object(clips, model[0].weight_quantizer.amax) - assert all(torch.equal(clip, clips[0]) for clip in clips), clips + # The same number of steps on every rank, but expert 1 runs on both of rank 0's and one of + # rank 1's: an expert's forwards, unlike the steps, are each rank's own count. + steps = [(0, 1), (0, 1)] if rank == 0 else [(0, 1), (0,)] + try: + mtq.quantize(model, get_awq_config("awq_clip"), lambda m: [m(x, hit) for hit in steps]) + clips = [expert.weight_quantizer.amax for expert in model.experts] + gathered = [None] * size + dist.all_gather_object(gathered, clips, group=group) + finally: + dist.destroy_process_group(group) + assert all( + torch.equal(mine, theirs) for other in gathered for mine, theirs in zip(clips, other) + ), gathered def test_awq_clip_picks_the_same_clips_on_every_rank(cpu_dist_workers): """Each rank of a model distributed over several ranks calibrates on its own data. awq_clip - averages each candidate clip's loss over the data-parallel ranks, so the ranks' copies of a - weight take the same clips, though each rank's own data would pick others.""" + averages each candidate clip's loss over the data-parallel ranks, once all the forwards have + run, so the ranks' copies of a weight take the same clips, though each rank's own data would + pick others, and though an MoE's expert runs a different number of forwards on each rank.""" cpu_dist_workers.run(_awq_clip_on_every_rank_job) From d87f5af9eca31eec62e2d6448dd3e525a5f7576d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:06:28 +0000 Subject: [PATCH 45/60] Accept the experts' groups of an (edp, ep) mesh with either axis innermost The row and column check took the ep axis as the mesh's innermost, as the expert-parallel load builds it today, and nothing ties the two together. It now accepts this rank's row and column of the mesh with either axis innermost: in both, the data-parallel group holds this rank's shard. Any other grouping is still refused. The four-rank test adds a mesh with ep outermost. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 21 +++++++------ .../utils/test_distributed_calibration.py | 31 ++++++++++--------- 2 files changed, 29 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 1ea04a98326..0b12bea5c09 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -113,18 +113,21 @@ def _names_the_experts_groups( ) -> bool: """Whether ``state`` names the groups a routed expert of ``plan`` syncs its scales over. - Those of the ``(edp, ep)`` mesh the expert-parallel load lays the ranks out in: as its - expert-parallel group, this rank's row, whose ranks hold the experts' shards in row order, and - as its data-parallel group, its column, the ranks holding the same shard. + Those of the ``(edp, ep)`` mesh the expert-parallel load lays the ranks out in, with either axis + innermost: as its expert-parallel group, this rank's row, whose ranks hold the experts' shards + in row order, and as its data-parallel group, its column, the ranks holding the same shard. """ if state is None or not state.expert_model_parallel_group.is_initialized(): return False - row, column = divmod(dist.rank(), plan.ep_size) - expert_parallel = range(row * plan.ep_size, (row + 1) * plan.ep_size) - data_parallel = range(column, plan.world_size, plan.ep_size) - return _group_ranks(state.expert_model_parallel_group) == list(expert_parallel) and sorted( - _group_ranks(state.data_parallel_group) - ) == list(data_parallel) + r, ep, edp, world = dist.rank(), plan.ep_size, plan.edp, plan.world_size + groups = ( + _group_ranks(state.expert_model_parallel_group), + sorted(_group_ranks(state.data_parallel_group)), + ) + return groups in ( # (row, column) with the ep axis innermost, or outermost + (list(range(r - r % ep, r - r % ep + ep)), list(range(r % ep, world, ep))), + (list(range(r % edp, world, edp)), list(range(r - r % edp, r - r % edp + edp))), + ) def _refuse_unsupported_distributed_calibration( diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index a3323c72592..c4c4f1a4aef 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -681,29 +681,32 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) -# Rank 0 of four at EP 2, in the (edp, ep) = (2, 2) mesh: ranks 0 and 1 shard the experts between -# them, ranks 2 and 3 do the same, and rank 2 holds the same shard as rank 0. The two-rank pool cannot -# build these groups, so the process group answers as a four-rank launch's would. +# Rank 0 of four at EP 2, in the (edp, ep) = (2, 2) mesh: with ep innermost, ranks 0 and 1 shard the +# experts between them, ranks 2 and 3 do the same, and rank 2 holds the same shard as rank 0; with ep +# outermost, ranks 0 and 2 shard them and rank 1 holds rank 0's. The two-rank pool cannot build these +# groups, so the process group answers as a four-rank launch's would. _FOUR_RANK_GROUPS = { - "ep": [0, 1], - "ep-reordered": [1, 0], - "edp": [0, 2], - "diagonal": [0, 3], - "world": [0, 1, 2, 3], + "0,1": [0, 1], + "1,0": [1, 0], + "0,2": [0, 2], + "0,3": [0, 3], + "all": [0, 1, 2, 3], } @pytest.mark.parametrize( ("ep", "dp", "seeded"), [ - ("ep", "edp", True), - ("ep", "ep", False), - ("ep", "diagonal", False), - ("ep-reordered", "edp", False), - ("world", "edp", False), + ("0,1", "0,2", True), + ("0,2", "0,1", True), + ("0,1", "0,1", False), + ("0,1", "0,3", False), + ("1,0", "0,2", False), + ("all", "0,2", False), ], ids=[ "seeded", + "seeded-with-ep-outermost", "data-parallel-over-the-expert-parallel-ranks", "data-parallel-over-another-shard", "expert-parallel-in-another-order", @@ -715,7 +718,7 @@ def test_the_experts_groups_are_checked_by_their_ranks(monkeypatch, ep, dp, seed expert-parallel ranks, or ranks that meet them only at this one but hold another shard (rank 3 holds rank 1's), max's and GPTQ's syncs would merge different experts' scales. Nor is an expert-parallel group over more ranks than the plan shards over, or in another order, which - gives this rank another shard.""" + gives this rank another shard. Either axis of the mesh may be the innermost.""" monkeypatch.setattr(mto_dist, "backend", lambda: "torch") monkeypatch.setattr(mto_dist, "rank", lambda group=None: 0) monkeypatch.setattr( From 7a6a2faae67fc3fb9f2ae89789fb529146306bb6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:39:36 +0000 Subject: [PATCH 46/60] Skip awq_clip's clip search for an input with no tokens An MoE expert a rank routes no tokens to can still be called, on an empty input. The per-tensor search then added a NaN loss (a mean over no outputs) for every candidate, which the post-loop sum spread to the replicas that had tokens, leaving them a clip of 0, and the static-block search failed to reshape the empty input. The search now returns at once for an empty input, leaving the losses and token count as they were, and the post-loop reduction still runs on every rank. Both distributed tests now run an expert on no tokens on one rank and check its clips are positive and those a single process picks from the other ranks' tokens: the CPU one (static-block branch) fails without the change on the reshape, the GPU one (NVFP4, per-tensor) would on a clip of 0. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 2 + .../gpu/torch/quantization/test_calib_cuda.py | 51 +++++++++++++----- tests/unit/torch/quantization/test_calib.py | 54 +++++++++++-------- 3 files changed, 73 insertions(+), 34 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index fd8d0ca7d7d..1aacbbb9ac4 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2083,6 +2083,8 @@ def sync_across_dp(module): module.awq_clip.num_tokens = int(tokens.item()) def _clip_search(self, inputs, co_bsz=256, max_tokens=16): + if inputs.numel() == 0: # no tokens to search with; their loss, a mean over none, is NaN + return weight = self.weight self.weight_quantizer.enable() diff --git a/tests/gpu/torch/quantization/test_calib_cuda.py b/tests/gpu/torch/quantization/test_calib_cuda.py index a00aae6e59d..5ee12738cb9 100644 --- a/tests/gpu/torch/quantization/test_calib_cuda.py +++ b/tests/gpu/torch/quantization/test_calib_cuda.py @@ -20,7 +20,7 @@ import torch.nn as nn import modelopt.torch.quantization as mtq -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan +from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan class _TwoBranchModel(nn.Module): @@ -75,29 +75,54 @@ def _forward_loop(m): assert uncal_q.amax.numel() == 1 -def _nvfp4_awq_clip_on_every_rank_job(rank, size): - device = torch.device(f"cuda:{rank}") - model = nn.Sequential(nn.Linear(64, 16, bias=False)).to(device) - with torch.no_grad(): - # Every rank's copy: tiny weights, whose FP8 block scales a tighter per-tensor amax resolves - # better, and a large one per row on input channel 0, which a tighter amax clips. - model[0].weight.copy_(torch.linspace(-1e-3, 1e-3, 16 * 64, device=device).view(16, 64)) - model[0].weight[:, 0] = 100.0 - set_distribution_plan(model, DistributionPlan(size)) +def _nvfp4_linears(device, count): + """``count`` copies of one linear: tiny weights, whose FP8 block scales a tighter per-tensor + amax resolves better, and a large one per row on input channel 0, which a tighter amax clips.""" + model = nn.Sequential(*(nn.Linear(64, 16, bias=False) for _ in range(count))).to(device) + for linear in model: + with torch.no_grad(): + linear.weight.copy_(torch.linspace(-1e-3, 1e-3, 16 * 64, device=device).view(16, 64)) + linear.weight[:, 0] = 100.0 + return model + + +def _nvfp4_rank_input(rank, device): x = torch.randn(8, 64, device=device, generator=torch.Generator(device).manual_seed(rank)) if rank % 2 == 0: x[:, 0] = 0 # never reaches the large weights: the tightest clip suits it else: x[:, 1:] = 0 # reaches only them: no clip suits it - mtq.quantize(model, mtq.NVFP4_AWQ_CLIP_CFG, lambda m: m(x)) - amax = model[0].weight_quantizer.amax.float().reshape(1) + return x + + +def _nvfp4_awq_clip_on_every_rank_job(rank, size): + device = torch.device(f"cuda:{rank}") + model = _nvfp4_linears(device, 2) + set_distribution_plan(model, DistributionPlan(size)) + x = _nvfp4_rank_input(rank, device) + + def forward_loop(m): + m[0](x) + # The second runs on no tokens on the odd ranks, as an MoE expert no token reaches can. + m[1](x if rank % 2 == 0 else x[:0]) + + mtq.quantize(model, mtq.NVFP4_AWQ_CLIP_CFG, forward_loop) + amax = torch.stack([linear.weight_quantizer.amax.float().reshape(()) for linear in model]) clips = [torch.empty_like(amax) for _ in range(size)] dist.all_gather(clips, amax) assert all(torch.equal(clip, amax) for clip in clips), clips + # The second's clip is the one the even ranks' tokens pick, calibrated in one process. + alone = _nvfp4_linears(device, 1) + alone[0]._parallel_state = ParallelState(data_parallel_group=-1) + inputs = [_nvfp4_rank_input(even, device) for even in range(0, size, 2)] + mtq.quantize(alone, mtq.NVFP4_AWQ_CLIP_CFG, lambda m: [m(even) for even in inputs]) + assert amax[1] > 0, amax + assert torch.equal(amax[1], alone[0].weight_quantizer.amax.float().reshape(())), amax def test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank(need_2_gpus, dist_workers): """NVFP4's awq_clip clips the per-tensor amax. Each rank of a model distributed over several ranks calibrates on its own data, and averaging each candidate clip's loss over the ranks gives - every rank's copy of a weight the same clip, though each rank's own data would pick another.""" + every rank's copy of a weight the same clip, though each rank's own data would pick another. A + rank that runs a weight on no tokens leaves it the clip the other ranks' tokens pick.""" dist_workers.run(_nvfp4_awq_clip_on_every_rank_job) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index f30f2ff91c2..69c56ba8770 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -270,39 +270,44 @@ def test_awq_clip(): class _Experts(nn.Module): - """Two experts, each run only for the steps that route tokens to it, as an MoE runs them.""" + """Two experts, each run only on the inputs routed to it, as an MoE runs them.""" - def __init__(self): + def __init__(self, data_parallel_group): super().__init__() self.experts = nn.ModuleList(nn.Linear(16, 4, bias=False) for _ in range(2)) + for expert in self.experts: + with torch.no_grad(): + # Every rank's copy: small weights, and a large one per row on input channel 0. + expert.weight.copy_(torch.linspace(-0.1, 0.1, 64).view(4, 16)) + expert.weight[:, 0] = 1.0 + expert._parallel_state = ParallelState(data_parallel_group=data_parallel_group) - def forward(self, x, hit): - for index in hit: + def forward(self, calls): + for index, x in calls: self.experts[index](x) -def _awq_clip_on_every_rank_job(rank, size): - model = _Experts() - # A rank left waiting in the clip search's sync gives up in seconds on this group. - group = dist.new_group(timeout=timedelta(seconds=10)) - dist.barrier(group=group) - for expert in model.experts: - with torch.no_grad(): - # Every rank's copy: small weights, and a large one per row on input channel 0. - expert.weight.copy_(torch.linspace(-0.1, 0.1, 64).view(4, 16)) - expert.weight[:, 0] = 1.0 - expert._parallel_state = ParallelState(data_parallel_group=group) - set_distribution_plan(model, DistributionPlan(size)) +def _rank_input(rank): x = torch.randn(8, 16, generator=torch.Generator().manual_seed(rank)) if rank % 2 == 0: x[:, 0] = 0 # never reaches the large weights: a tighter clip only sharpens the rest else: x[:, 1:] = 0 # reaches only them: clipping them costs the most - # The same number of steps on every rank, but expert 1 runs on both of rank 0's and one of - # rank 1's: an expert's forwards, unlike the steps, are each rank's own count. - steps = [(0, 1), (0, 1)] if rank == 0 else [(0, 1), (0,)] + return x + + +def _awq_clip_on_every_rank_job(rank, size): + x = _rank_input(rank) + # The same number of steps on every rank, but expert 1 runs on both of rank 0's and on one of + # rank 1's, there on no tokens: an expert's forwards, unlike the steps, are each rank's own. + steps = [[(0, x), (1, x)]] * 2 if rank == 0 else [[(0, x), (1, x[:0])], [(0, x)]] + # A rank left waiting in the clip search's sync gives up in seconds on this group. + group = dist.new_group(timeout=timedelta(seconds=10)) + dist.barrier(group=group) + model = _Experts(group) + set_distribution_plan(model, DistributionPlan(size)) try: - mtq.quantize(model, get_awq_config("awq_clip"), lambda m: [m(x, hit) for hit in steps]) + mtq.quantize(model, get_awq_config("awq_clip"), lambda m: [m(calls) for calls in steps]) clips = [expert.weight_quantizer.amax for expert in model.experts] gathered = [None] * size dist.all_gather_object(gathered, clips, group=group) @@ -311,13 +316,20 @@ def _awq_clip_on_every_rank_job(rank, size): assert all( torch.equal(mine, theirs) for other in gathered for mine, theirs in zip(clips, other) ), gathered + # Expert 1's clips are those rank 0's tokens pick alone, which rank 1's empty call leaves be. + alone = _Experts(data_parallel_group=-1) + x0 = _rank_input(0) + mtq.quantize(alone, get_awq_config("awq_clip"), lambda m: [m([(1, x0)]) for _ in range(2)]) + assert (clips[1] > 0).all(), clips[1] + assert torch.equal(clips[1], alone.experts[1].weight_quantizer.amax) def test_awq_clip_picks_the_same_clips_on_every_rank(cpu_dist_workers): """Each rank of a model distributed over several ranks calibrates on its own data. awq_clip averages each candidate clip's loss over the data-parallel ranks, once all the forwards have run, so the ranks' copies of a weight take the same clips, though each rank's own data would - pick others, and though an MoE's expert runs a different number of forwards on each rank.""" + pick others, and though an MoE's expert runs a different number of forwards on each rank, on + one of them on no tokens, which leaves the clips the other ranks' tokens pick.""" cpu_dist_workers.run(_awq_clip_on_every_rank_job) From cdc0dab35355845a0aad4772a42a4aa53984d125 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 16:21:10 +0000 Subject: [PATCH 47/60] Leave hf_ptq's FSDP2 load recording its plan to the PR whose export reads it hf_ptq --use_fsdp2 recorded the FSDP2 plan after its load, so that calibration sized its batch collectively and refused rank-local calibrations there. That path still retries a batch after a rank-local out-of-memory error, and the fix for that belongs where hf_ptq moves onto the run's abort path. The recording now lands with the no-gather export, the first change whose export needs the plan on that path, together with that fix. The README notes and the changelog entry on the refusal go with it; the AWQ-Clip ones stay, as its sync reads no plan. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 - examples/hf_ptq/README.md | 7 +------ examples/hf_ptq/hf_ptq.py | 5 ----- tests/examples/hf_ptq/test_hf_ptq_args.py | 24 +---------------------- 4 files changed, 2 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 32013212eb0..f6cfb894081 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -46,7 +46,6 @@ Changelog **Backward Breaking Changes** -- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index f76215adc8f..b68621650db 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -610,18 +610,13 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. -> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental, and refused when a candidate format is one of those refused below.* +> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental.* > *Nemotron-H checkpoints containing MTP weights are not supported with `--use_fsdp2` and are rejected before distributed model loading. Remove `--use_fsdp2` to load and calibrate the MTP modules; the distributed loader does not construct them.* > *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* - -> *With `--use_fsdp2` on more than one GPU, calibration that keeps its statistics on each rank is refused when quantization starts: AWQ and SVDQuant (`--qformat int4_awq`, `w4a8_awq_beta`, `nvfp4_awq_lite`, `nvfp4_awq_full`, `nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`.* - > *`nvfp4_awq_clip` works with `--use_fsdp2` on more than one GPU: its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* -> *With `--use_fsdp2` on more than one GPU and `--batch_size 0`, each rank estimates the calibration batch size from one forward and the ranks use the smallest. That size is not probed for out-of-memory, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives: set `--batch_size` if calibration runs out of memory.* - ### Usage #### Slurm (recommended) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index c79c25ccf98..47d6a8b2c19 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -99,7 +99,6 @@ get_max_batch_size, get_supported_datasets, ) -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan from modelopt.torch.utils.memory_monitor import launch_memory_monitor from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -340,10 +339,6 @@ def load_model(args: argparse.Namespace): attn_implementation=args.attn_implementation, hf_config=hf_config, ) - # Calibration sizes its batch, and the export writes, by how the model is distributed. - set_distribution_plan( - full_model, DistributionPlan(args.dist_state.world_size, dense="fsdp") - ) elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.local_checkpoint_path, diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 5543e9de902..2f0b425de40 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +26,7 @@ import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 -from transformers import Qwen3_5MoeConfig, Qwen3Config, Qwen3MoeConfig +from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -38,7 +38,6 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.utils import mlflow as mlflow_lib -from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.mlflow import describe_run, run_tags _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -481,27 +480,6 @@ def distributed_load(*args, **kwargs): hf_ptq.load_model(args) -def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path): - """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an - FSDP2-sharded model without one is refused rather than taken for a single process's.""" - Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path) - monkeypatch.setenv("RANK", "0") - hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2") - args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=2) - args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path) # as main() sets them - model = torch.nn.Linear(1, 1) - monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *args, **kwargs: model) - - def loaded(model): - raise RuntimeError("loaded") - - monkeypatch.setattr(hf_ptq, "prepare_model_for_calibration", loaded) - - with pytest.raises(RuntimeError, match="loaded"): - hf_ptq.load_model(args) - assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") - - def test_mlflow_flag_defaults_the_experiment_name(monkeypatch): monkeypatch.setattr(getpass, "getuser", lambda: "tester") hf_ptq, args = _parse_hf_ptq_args( From 6eeb2fb3855e6a00ae514d0e64edd8fc40f5d11d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 16:24:09 +0000 Subject: [PATCH 48/60] Size an FSDP2 model's calibration batch by its plan, then calibrate it, on GPUs A two-or-more-GPU test of the planned path end to end: a model under fully_shard records its plan, get_max_batch_size measures each rank's GPU in one forward (rank 0 holding a fortieth of its free memory per sample, every other rank a sixth) and every rank takes the smallest size, not its own; then each rank calibrates its shard with max through distributed_forward_loop. Every rank runs the same three steps and ends with the same scales. The batch size's reduction and the test's gathers give up in seconds on a bounded group. Signed-off-by: Shengliang Xu --- tests/gpu/torch/quantization/test_fsdp2.py | 74 +++++++++++++++++++++- 1 file changed, 72 insertions(+), 2 deletions(-) diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index aec8e37f813..632a4ccddac 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -21,10 +21,12 @@ import pytest import torch +import torch.distributed as dist import torch.nn as nn -from _test_utils.torch.distributed.utils import synchronize_state_dict +from _test_utils.torch.distributed.utils import bounded_collectives, synchronize_state_dict from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.tensor import DTensor +from torch.utils.data import DataLoader import modelopt.torch.quantization as mtq import modelopt.torch.quantization.model_calib as model_calib @@ -35,7 +37,13 @@ persistent_materialization, ) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector -from modelopt.torch.utils.dataset_utils import _forward_loop +from modelopt.torch.utils.dataset_utils import _forward_loop, get_max_batch_size +from modelopt.torch.utils.distributed import ( + DistributionPlan, + distributed_forward_loop, + set_distribution_plan, + shard_dataloader, +) def _test_fsdp2_simple_linear(rank, size): @@ -461,3 +469,65 @@ def _test_sharded_root_calibration(rank, size): def test_sharded_root_calibration(dist_workers): dist_workers.run(_test_sharded_root_calibration) + + +def _test_planned_batch_size_then_max_calibration(rank, size): + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64)).cuda(rank) + synchronize_state_dict(model) + fully_shard(model[0]) + fully_shard(model[2]) + model = fully_shard(model) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) + forwards = [] + model.register_forward_pre_hook(lambda module, args: forwards.append(rank)) + # What a sample holds in the measured forward: a fortieth of rank 0's free GPU memory, so its + # own estimate is about 40 samples, and a sixth of every other rank's, about 6 (4 regulated). + held = torch.cuda.mem_get_info()[0] // (40 if rank == 0 else 6) + + def forward_step(model, inputs): + model(inputs) + return torch.empty(held, dtype=torch.uint8, device=inputs.device) + + # A rank left waiting in the batch size's reduction, or in a gather below, gives up in seconds. + with bounded_collectives(seconds=60): + batch_size = get_max_batch_size( + model, sample_input_single_batch=torch.ones(1, 64).cuda(rank), forward_step=forward_step + ) + sizes = [None] * size + dist.all_gather_object(sizes, batch_size) + # Every rank takes the smallest any rank's GPU allows, not its own (rank 0's is about 40). + assert sizes == [batch_size] * size and batch_size <= 4, sizes + + # An uneven split, which the shard pads so that every rank runs the same steps. + data = torch.randn( + 2 * size * batch_size + 1, 64, generator=torch.Generator().manual_seed(1) + ) + loader = DataLoader( + data, batch_size=batch_size, collate_fn=lambda rows: {"input": torch.stack(rows)} + ) + forwards.clear() + forward_loop = distributed_forward_loop( + model, shard_dataloader(loader, rank, size), device=torch.device("cuda", rank) + ) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop) + amaxes = { + name: module.amax.float().cpu() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.is_enabled and module.amax is not None + } + gathered = [None] * size + dist.all_gather_object(gathered, (len(forwards), amaxes)) + + assert amaxes, "no quantizer was calibrated" + for steps, theirs in gathered: + assert steps == len(forwards) == 3, gathered # two whole batches and one of one sample + assert theirs.keys() == amaxes.keys() + assert all(torch.equal(theirs[name], amax) for name, amax in amaxes.items()), theirs + + +def test_planned_batch_size_then_max_calibration(need_2_gpus, dist_workers): + """A model whose plan spreads it over several ranks sizes its calibration batch once for all of + them, the smallest any rank's GPU allows, then calibrates each rank's shard with it: every rank + runs the same steps, and max calibration gives every rank the same scales.""" + dist_workers.run(_test_planned_batch_size_then_max_calibration) From 2c6f0afa68667ba5417884ba145df16159a87f06 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:13:26 +0000 Subject: [PATCH 49/60] Calibrate a real expert-parallel model through the forward loop, on GPUs A two-or-more-GPU test of the expert-parallel calibration lifecycle on stock transformers: a tiny Qwen3-MoE loaded with its experts sharded over every rank by transformers' own expert-parallel load on the torch_a2a styles, seeded with the experts' groups, quantized with FP8 through distributed_forward_loop, each rank on its own batch. Every expert's weight scales and every activation scale match the whole model calibrated on every rank's batch. On transformers >= 5.16 quantizer insertion drops the transport's forward wrappers, which the loop re-plants: the test fails without the re-plant, and without the conversion keeping torch_a2a (#3074). The scales' syncs and the test's gathers give up in seconds on bounded groups. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_moe_a2a.py | 145 ++++++++++++++++++++++++++ 1 file changed, 145 insertions(+) diff --git a/tests/gpu/torch/utils/test_moe_a2a.py b/tests/gpu/torch/utils/test_moe_a2a.py index 9b8a3be53a4..dfd8c704d27 100644 --- a/tests/gpu/torch/utils/test_moe_a2a.py +++ b/tests/gpu/torch/utils/test_moe_a2a.py @@ -16,13 +16,29 @@ """GPU tests for the ``torch_a2a`` expert-parallel MoE transport (``a2a_moe``).""" import copy +import dataclasses +from datetime import timedelta import pytest import torch import torch.distributed as dist +from _test_utils.torch.distributed.utils import bounded_collectives from packaging.version import Version +from torch.distributed.device_mesh import init_device_mesh +import modelopt.torch.quantization as mtq import modelopt.torch.utils.moe as moe +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.utils.distributed import ( + DistributionPlan, + ParallelState, + distributed_forward_loop, + set_distribution_plan, +) + +transformers = pytest.importorskip("transformers") +from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig +from transformers.distributed import DistributedConfig # One MoE layer's expert stage, big enough that the order of each token's top-k adds shows in bf16. TOKENS, HIDDEN, EXPERTS, TOP_K = 4096, 1024, 16, 8 @@ -71,3 +87,132 @@ def test_torch_a2a_gives_grouped_mms_result_bit_for_bit(dist_workers): if Version(transformers.__version__) < Version("5.6"): pytest.skip("expert parallelism needs transformers >= 5.6") dist_workers.run(_grouped_mm_bit_for_bit_job) + + +NUM_EXPERTS = 8 + + +def _save_tiny_qwen3_moe(ckpt_dir): + config = Qwen3MoeConfig( + hidden_size=64, + intermediate_size=128, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=NUM_EXPERTS, + num_experts_per_tok=2, + vocab_size=128, + max_position_embeddings=64, + ) + torch.manual_seed(0) + AutoModelForCausalLM.from_config(config).to(torch.bfloat16).save_pretrained(ckpt_dir) + + +def _load_expert_parallel(ckpt_dir, size): + """The model with its experts sharded over every rank by transformers' own expert-parallel load, + on the ``torch_a2a`` styles and experts forward, and the mesh they are sharded over.""" + moe.register_torch_a2a() + config = AutoConfig.from_pretrained(ckpt_dir) + swap = {"ep_router": "torch_a2a_router", "moe_tp_experts": "torch_a2a_experts"} + config.base_model_ep_plan = {k: swap.get(v, v) for k, v in config.base_model_ep_plan.items()} + mesh = init_device_mesh("cuda", (size,), mesh_dim_names=("ep",)) + # Each option where this release takes it: on DistributedConfig when it has a field for it. + fields = {field.name for field in dataclasses.fields(DistributedConfig)} + options = {"tp_plan": "auto", "tp_size": size} + model = AutoModelForCausalLM.from_pretrained( + ckpt_dir, + config=config, + dtype=torch.bfloat16, + device_mesh=mesh, + distributed_config=DistributedConfig( + enable_expert_parallel=True, **{k: v for k, v in options.items() if k in fields} + ), + experts_implementation="torch_a2a", + **({} if "tp_plan" in fields else {"tp_plan": "auto"}), + ) + for module in model.modules(): # the experts forward reads it at forward time + if hasattr(getattr(module, "config", None), "_experts_implementation"): + module.config._experts_implementation = "torch_a2a" + return model, mesh + + +def _calibrated(model): + return { + name: module.amax.float().cpu() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.is_enabled and module.amax is not None + } + + +def _expert_parallel_calibration_job(rank, size, ckpt_dir): + device = torch.device("cuda", rank) + batches = [ + torch.randint(0, 128, (2, 16), generator=torch.Generator().manual_seed(r)) + for r in range(size) + ] + # The scales' syncs give up in seconds on this group of every rank: the experts' expert-parallel + # group (they span every rank), and the dense modules' data-parallel one. + group = dist.new_group(timeout=timedelta(seconds=120)) + dist.barrier(group=group) + try: + with bounded_collectives(seconds=120): + model, mesh = _load_expert_parallel(ckpt_dir, size) + for name, module in model.named_modules(): + module._parallel_state = ( + ParallelState(data_parallel_group=-1, expert_model_parallel_group=group) + if ".mlp.experts" in name + else ParallelState(data_parallel_group=group) + ) + set_distribution_plan(model, DistributionPlan(size, ep_size=size)) + # Each rank its own batch, through the quantized model's re-planted torch_a2a forwards. + forward_loop = distributed_forward_loop( + model, [{"input_ids": batches[rank]}], device=device, ep_mesh=mesh + ) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop) + + # The whole model, calibrated in one process on every rank's batch. + whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16).to(device) + for module in whole.modules(): + module._parallel_state = ParallelState(data_parallel_group=-1) + mtq.quantize(whole, mtq.FP8_DEFAULT_CFG, lambda m: [m(b.to(device)) for b in batches]) + + sharded, reference = _calibrated(model), _calibrated(whole) + problems, compared = [] if sharded else ["no quantizer was calibrated"], [] + per_rank = NUM_EXPERTS // size + for name, amax in sharded.items(): + if "_weight_quantizers." in name: # this rank's local expert -> the global one + prefix, index = name.rsplit(".", 1) + name = f"{prefix}.{rank * per_rank + int(index)}" + compared.append(name) + want = reference.get(name) + if want is None: + problems.append(f"{name} is calibrated here but not in the whole model") + continue + # bf16 activations, computed in another order: compare by relative norm. + error = ((amax - want).norm() / want.norm().clamp_min(1e-12)).item() + if not error < 2e-2: + problems.append(f"{name}: amax off by {error:.3g} (relative)") + gathered = [None] * size + dist.all_gather_object(gathered, (problems, compared)) + finally: + dist.destroy_process_group(group) + + found = [problem for problems, _ in gathered for problem in problems] + assert not found, "\n".join(found[:20]) + experts = {name for name in reference if "_weight_quantizers." in name} + assert experts and experts <= {name for _, names in gathered for name in names}, experts + + +def test_expert_parallel_calibration_matches_the_whole_model(need_2_gpus, dist_workers, tmp_path): + """transformers' own expert-parallel load on the torch_a2a transport, quantized and calibrated + through distributed_forward_loop, each rank on its own batch: every expert's weight scales and + every activation scale match the whole model calibrated on every rank's batch. Quantizer + insertion drops the transport's forward wrappers on transformers >= 5.16, which the loop + re-plants.""" + if Version(transformers.__version__) < Version("5.16"): + pytest.skip("the torch_a2a forward wrappers are pre-hooks below transformers 5.16") + if NUM_EXPERTS % torch.cuda.device_count(): + pytest.skip(f"{NUM_EXPERTS} experts do not shard over {torch.cuda.device_count()} GPUs") + _save_tiny_qwen3_moe(tmp_path) + dist_workers.run(_expert_parallel_calibration_job, ckpt_dir=str(tmp_path)) From 3b940cb88a94767b12f24142c83292855f861173 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:12:48 +0000 Subject: [PATCH 50/60] Check the experts' groups against the ep-innermost mesh only Accepting the (edp, ep) mesh with either axis innermost left the check open when edp equals ep_size: there the two layouts' (row, column) pairs are each other's transpose, so experts whose expert-parallel and data-parallel groups were swapped passed, and max's and GPTQ's syncs merged different experts' scales. The check is exact again for the layout the expert-parallel load builds, ep innermost, and refuses any other. The four-rank test's transposed groups, which the either-way check accepted, are now refused. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 21 ++++++++----------- .../utils/test_distributed_calibration.py | 15 ++++++------- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 0b12bea5c09..1ea04a98326 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -113,21 +113,18 @@ def _names_the_experts_groups( ) -> bool: """Whether ``state`` names the groups a routed expert of ``plan`` syncs its scales over. - Those of the ``(edp, ep)`` mesh the expert-parallel load lays the ranks out in, with either axis - innermost: as its expert-parallel group, this rank's row, whose ranks hold the experts' shards - in row order, and as its data-parallel group, its column, the ranks holding the same shard. + Those of the ``(edp, ep)`` mesh the expert-parallel load lays the ranks out in: as its + expert-parallel group, this rank's row, whose ranks hold the experts' shards in row order, and + as its data-parallel group, its column, the ranks holding the same shard. """ if state is None or not state.expert_model_parallel_group.is_initialized(): return False - r, ep, edp, world = dist.rank(), plan.ep_size, plan.edp, plan.world_size - groups = ( - _group_ranks(state.expert_model_parallel_group), - sorted(_group_ranks(state.data_parallel_group)), - ) - return groups in ( # (row, column) with the ep axis innermost, or outermost - (list(range(r - r % ep, r - r % ep + ep)), list(range(r % ep, world, ep))), - (list(range(r % edp, world, edp)), list(range(r - r % edp, r - r % edp + edp))), - ) + row, column = divmod(dist.rank(), plan.ep_size) + expert_parallel = range(row * plan.ep_size, (row + 1) * plan.ep_size) + data_parallel = range(column, plan.world_size, plan.ep_size) + return _group_ranks(state.expert_model_parallel_group) == list(expert_parallel) and sorted( + _group_ranks(state.data_parallel_group) + ) == list(data_parallel) def _refuse_unsupported_distributed_calibration( diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index c4c4f1a4aef..172e5eb68e8 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -681,10 +681,9 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) -# Rank 0 of four at EP 2, in the (edp, ep) = (2, 2) mesh: with ep innermost, ranks 0 and 1 shard the -# experts between them, ranks 2 and 3 do the same, and rank 2 holds the same shard as rank 0; with ep -# outermost, ranks 0 and 2 shard them and rank 1 holds rank 0's. The two-rank pool cannot build these -# groups, so the process group answers as a four-rank launch's would. +# Rank 0 of four at EP 2, in the (edp, ep) = (2, 2) mesh: ranks 0 and 1 shard the experts between +# them, ranks 2 and 3 do the same, and rank 2 holds the same shard as rank 0. The two-rank pool cannot +# build these groups, so the process group answers as a four-rank launch's would. _FOUR_RANK_GROUPS = { "0,1": [0, 1], "1,0": [1, 0], @@ -698,7 +697,7 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( ("ep", "dp", "seeded"), [ ("0,1", "0,2", True), - ("0,2", "0,1", True), + ("0,2", "0,1", False), ("0,1", "0,1", False), ("0,1", "0,3", False), ("1,0", "0,2", False), @@ -706,7 +705,7 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( ], ids=[ "seeded", - "seeded-with-ep-outermost", + "groups-transposed", "data-parallel-over-the-expert-parallel-ranks", "data-parallel-over-another-shard", "expert-parallel-in-another-order", @@ -718,7 +717,9 @@ def test_the_experts_groups_are_checked_by_their_ranks(monkeypatch, ep, dp, seed expert-parallel ranks, or ranks that meet them only at this one but hold another shard (rank 3 holds rank 1's), max's and GPTQ's syncs would merge different experts' scales. Nor is an expert-parallel group over more ranks than the plan shards over, or in another order, which - gives this rank another shard. Either axis of the mesh may be the innermost.""" + gives this rank another shard, nor the two groups transposed, the mesh's column as the + expert-parallel group and its row as the data-parallel one, which match in size when edp is + ep_size.""" monkeypatch.setattr(mto_dist, "backend", lambda: "torch") monkeypatch.setattr(mto_dist, "rank", lambda group=None: 0) monkeypatch.setattr( From 367609d7acdc2a5b8eca01a9f7469c8acced87c7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:29:59 +0000 Subject: [PATCH 51/60] Reduce awq_clip's clip search in one all-reduce per module, before any weight gather The token count now rides in the fp32 losses' buffer, exact for any calibration set, so each module issues one all-reduce instead of two. And every module's reduction runs in its own pass ahead of the postprocess pass, whose weight access an FSDP2 model gathers for, instead of alternating with those gathers. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 1aacbbb9ac4..3d1442b32e6 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2076,11 +2076,12 @@ def sync_across_dp(module): if not group.is_initialized(): return losses = torch.stack(list(module.awq_clip.loss.values())) - tokens = torch.tensor(module.awq_clip.num_tokens, device=losses.device) - for tensor in (losses, tokens): - dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group.group) - module.awq_clip.loss = dict(zip(module.awq_clip.loss, losses / group.world_size())) - module.awq_clip.num_tokens = int(tokens.item()) + # One all-reduce: the token count rides in the fp32 losses, exact for any calibration set. + packed = torch.cat([losses.reshape(-1), losses.new_tensor([module.awq_clip.num_tokens])]) + dist.all_reduce(packed, op=dist.ReduceOp.SUM, group=group.group) + reduced = packed[:-1].view_as(losses) / group.world_size() + module.awq_clip.loss = dict(zip(module.awq_clip.loss, reduced)) + module.awq_clip.num_tokens = int(packed[-1].item()) def _clip_search(self, inputs, co_bsz=256, max_tokens=16): if inputs.numel() == 0: # no tokens to search with; their loss, a mean over none, is NaN @@ -2210,9 +2211,13 @@ def postprocess(module): if module.awq_clip.is_input_quantized: module.input_quantizer.enable() - for name, module in model.named_modules(): + # Every module's reduction before any weight gather, which an FSDP2 model's postprocess runs. + for module in model.modules(): if is_quantized_linear(module) and hasattr(module, "awq_clip"): sync_across_dp(module) + + for name, module in model.named_modules(): + if is_quantized_linear(module) and hasattr(module, "awq_clip"): if module.awq_clip.num_tokens > 0: with enable_weight_access_and_writeback(module, model, names): postprocess(module) From 6245fa742477aa8598228bc8884e9e5e7289b613 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:58:42 +0000 Subject: [PATCH 52/60] Reduce awq_clip's clip-search losses in place, with no copy of them The data-parallel sync stacked every candidate clip's loss into a new tensor, packed the token count after it and divided out of place, so while it ran it held up to three copies of the losses beyond the module's own. Each candidate's loss is now a view of one buffer, which the sync reduces and averages in place, and the token count goes in its own int64 all-reduce, exact for any count. Every module's reductions still run ahead of the weight gathers. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 27 ++++++++-------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 3d1442b32e6..1d9ddf63e34 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2029,16 +2029,10 @@ def __init__(self, module): clip_ratios = [ round(float(k), 2) for k in torch.arange(min_clip_ratio, 1.0, shrink_step) ] + [1.0] - if self.is_per_tensor_clip(module): - self.loss = {k: torch.tensor(0.0, device=module.weight.device) for k in clip_ratios} - else: - self.loss = { - k: torch.zeros( - (co, math.ceil(ci / self.block_size)), - device=module.weight.device, - ) - for k in clip_ratios - } + shape = () if self.is_per_tensor_clip(module) else (co, math.ceil(ci / self.block_size)) + # One buffer for every candidate's loss, which the data-parallel sync reduces in place. + self.loss_buffer = torch.zeros(len(clip_ratios), *shape, device=module.weight.device) + self.loss = dict(zip(clip_ratios, self.loss_buffer)) self.best_clip_val = None self.best_loss = None @@ -2075,13 +2069,12 @@ def sync_across_dp(module): group = module.parallel_state.data_parallel_group if not group.is_initialized(): return - losses = torch.stack(list(module.awq_clip.loss.values())) - # One all-reduce: the token count rides in the fp32 losses, exact for any calibration set. - packed = torch.cat([losses.reshape(-1), losses.new_tensor([module.awq_clip.num_tokens])]) - dist.all_reduce(packed, op=dist.ReduceOp.SUM, group=group.group) - reduced = packed[:-1].view_as(losses) / group.world_size() - module.awq_clip.loss = dict(zip(module.awq_clip.loss, reduced)) - module.awq_clip.num_tokens = int(packed[-1].item()) + losses = module.awq_clip.loss_buffer + dist.all_reduce(losses, op=dist.ReduceOp.SUM, group=group.group) + losses /= group.world_size() + tokens = torch.tensor(module.awq_clip.num_tokens, device=losses.device) + dist.all_reduce(tokens, op=dist.ReduceOp.SUM, group=group.group) + module.awq_clip.num_tokens = int(tokens.item()) def _clip_search(self, inputs, co_bsz=256, max_tokens=16): if inputs.numel() == 0: # no tokens to search with; their loss, a mean over none, is NaN From afefa2aa09cfa802b4bc46166bce60fdfbd18529 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:58:51 +0000 Subject: [PATCH 53/60] Name distributed_sync only for the calibrations that have it The refusal advised calibrating with max, MSE, GPTQ, SmoothQuant or awq_clip "with distributed_sync on", but only max and MSE have that option. And local_hessian, refused although it has the option, says why where it is listed: its distributed_sync syncs only the max amax it starts from. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 1ea04a98326..1ef1ab92b23 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -78,7 +78,7 @@ "awq_lite", "awq_full", "svdquant", - "local_hessian", + "local_hessian", # its distributed_sync syncs only the max amax it starts from "imatrix", ) @@ -151,8 +151,8 @@ def _refuse_unsupported_distributed_calibration( f"Calibrating with {', '.join(refused)} is not supported on a model distributed over " f"{plan.world_size} ranks. Each rank calibrates on its own data and this calibration " "keeps what it gathers on the rank, so the ranks would calibrate their copies, or their " - "shards, of a weight differently. Calibrate with max, MSE, GPTQ, SmoothQuant or " - "awq_clip, with distributed_sync on, or in a single process." + "shards, of a weight differently. Calibrate with max or MSE with distributed_sync on, " + "GPTQ, SmoothQuant or awq_clip, or in a single process." ) if algorithm is None or not plan.is_expert_parallel: return From 359a0647dc2c4247bb743635c7c74f287175ab82 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:38:48 +0000 Subject: [PATCH 54/60] Say that an expert-parallel caller seeds the experts' groups before quantizing The calibration guard runs before quantization converts anything, so on an expert-parallel plan it reads the routed experts' ParallelState from the caller, as the expert-parallel load seeds it, and checks their quantized modules again before they calibrate. The docstring now says so instead of implying that only the converted modules are checked. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 1ef1ab92b23..c153040f2af 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -138,9 +138,9 @@ def _refuse_unsupported_distributed_calibration( over the groups their module's ``ParallelState`` names, every rank for a module converted without one: a module holding routed experts must name the experts' expert-parallel group and, for data parallelism, only the ranks holding the same experts, or any calibration would merge - different experts' scales. Called before anything changes; a routed expert's quantized module, - whose state calibration reads, is checked once quantization has converted it, before it - calibrates. + different experts' scales. Called before anything changes, so the caller seeds the routed + experts' ``ParallelState`` before quantizing, as the expert-parallel load does; their quantized + modules, whose state calibration reads, are checked again once converted, before they calibrate. """ plan = dist.recorded_distribution_plan(model) if plan is None or plan.world_size == 1: From f5d5468350859ce6217bac714c84ba2766a9ef78 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:58:32 +0000 Subject: [PATCH 55/60] Say that every rank still decodes every sample of a sharded stream _StridedIterableShard takes one sample in world_size from the stream each rank builds, so every rank still reads and decodes all of them. Its docstring now says so, and that a map-style dataset, or a source that shards itself, saves that work for a stream costly to decode. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 046af6d90ee..083c47ed1ae 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -701,7 +701,9 @@ class _StridedIterableShard(IterableDataset): the equal-per-rank-batch invariant that keeps per-step EP collectives from hanging when one rank runs short. A stream shorter than ``world_size`` has no complete group, so it is refused rather than leaving every rank with nothing to calibrate on. It is iterated in the loader's own process - (:func:`shard_dataloader` refuses workers for a stream). + (:func:`shard_dataloader` refuses workers for a stream). Every rank still reads and decodes every + sample, keeping one in ``world_size``: for a stream costly to decode, a map-style dataset, or a + source that shards itself, saves that work. ``vlm_dataset_utils._ShardedIterable`` cannot stand in: it pads every rank to a count its own loader knows, while this wraps any loader's stream, which may know none. """ From 8e811ae1ed3c63bfa10bc34569f2ccbf6329585c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:18:22 +0000 Subject: [PATCH 56/60] Say that nvfp4_awq_clip still works where multi-rank FSDP2 refuses AWQ Main now refuses AWQ-Lite and AWQ-Full on an FSDP2 model sharded over several ranks, and its README note on that sits right above this PR's note on nvfp4_awq_clip. The two notes are now separate paragraphs, and this PR's says why nvfp4_awq_clip is not refused: AWQ-Clip folds no smoothing scale into the weights, and its clip search is averaged over the ranks. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index b68621650db..b8aba4de503 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -615,7 +615,8 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *Nemotron-H checkpoints containing MTP weights are not supported with `--use_fsdp2` and are rejected before distributed model loading. Remove `--use_fsdp2` to load and calibrate the MTP modules; the distributed loader does not construct them.* > *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* -> *`nvfp4_awq_clip` works with `--use_fsdp2` on more than one GPU: its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* + +> *`nvfp4_awq_clip` still works with `--use_fsdp2` on more than one GPU: it calibrates with AWQ-Clip alone, which folds no smoothing scale into the weights, and its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* ### Usage From 1df1095600c26676ba0492f6150de65ac99b1e1a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:51:32 +0000 Subject: [PATCH 57/60] Refuse auto_quantize on an expert-parallel model before converting it AutoQuantize calibrates and scores through forward_step on the model it has converted, and conversion drops the expert-parallel forwards that distributed_forward_loop restores. On an expert-parallel model the first of those forwards would fail with the experts forward's missing process_group error, after the model was converted. auto_quantize now refuses a model whose recorded plan shards its experts, before it changes anything, as quantize and calibrate refuse what they cannot calibrate. A two-rank test with seeded experts and an INT8 candidate checks the module types and state are left as they were. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 8 +++++ .../utils/test_distributed_calibration.py | 36 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index c153040f2af..54b952e2607 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -1052,6 +1052,14 @@ def forward_backward_step(model, batch) -> None: if isinstance(quantization_formats, str) or not isinstance(quantization_formats, Sequence): raise TypeError("`quantization_formats` must be a sequence of formats.") quantization_formats = list(quantization_formats) + plan = dist.recorded_distribution_plan(model) + if plan is not None and plan.is_expert_parallel: # before anything converts the model + raise NotImplementedError( + f"auto_quantize does not support a model whose experts are sharded over {plan.ep_size} " + "ranks yet: its forwards would run the experts without the expert-parallel forwards " + "that quantization replaces and distributed_forward_loop restores. Quantize it with " + "mtq.quantize through distributed_forward_loop, or AutoQuantize in a single process." + ) method_name, method_config = _parse_auto_quantize_method(method) is_kv_search = constraints is not None and constraints.get("cost_model") == COST_MODEL_KV_CACHE diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 172e5eb68e8..c25bd1928ef 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -681,6 +681,42 @@ def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) +def _expert_parallel_auto_quantize_job(rank, size): + model = nn.Sequential(OrderedDict(gate=nn.Linear(16, 16), experts=nn.Linear(16, 16))) + dp, ep = _EXPERT_LAYOUTS["seeded"]["experts"] + model.experts._parallel_state = ParallelState( + data_parallel_group=dp, expert_model_parallel_group=ep + ) + set_distribution_plan(model, DistributionPlan(world_size=size, ep_size=size)) + modules = [type(module) for module in model.modules()] + state = {name: tensor.clone() for name, tensor in model.state_dict().items()} + raised = None + with bounded_collectives(): + try: # not pytest.raises, as in _expert_parallel_calibration_job + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[(mtq.INT8_DEFAULT_CFG, "int8")], + data_loader=[torch.randn(2, 16)], + forward_step=lambda m, batch: m(batch), + loss_func=lambda output, batch: output.sum(), + ) + except NotImplementedError as e: + raised = str(e) + assert raised is not None and f"experts are sharded over {size} ranks" in raised, raised + assert [type(module) for module in model.modules()] == modules + after = model.state_dict() + assert after.keys() == state.keys() + assert all(torch.equal(after[name], tensor) for name, tensor in state.items()) + + +def test_auto_quantize_refuses_an_expert_parallel_model_before_changing_it(cpu_dist_workers): + """auto_quantize's forwards would run the experts without the expert-parallel forwards that + quantization replaces, which distributed_forward_loop restores. So it refuses such a model, with + its experts seeded and a candidate every rank syncs, before it converts anything.""" + cpu_dist_workers.run(_expert_parallel_auto_quantize_job) + + # Rank 0 of four at EP 2, in the (edp, ep) = (2, 2) mesh: ranks 0 and 1 shard the experts between # them, ranks 2 and 3 do the same, and rank 2 holds the same shard as rank 0. The two-rank pool cannot # build these groups, so the process group answers as a four-rank launch's would. From 14bb439a7fb7aceeafc2af7da7e5207be8c51162 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:51:33 +0000 Subject: [PATCH 58/60] Refuse a sampler over another source than the loader's dataset shard_dataloader replaces a loader's sampler with the rank's DistributedSampler over the loader's dataset. It checked a RandomSampler's sample count against the sampler's own source, so a SequentialSampler, RandomSampler or DistributedSampler over a shorter source than the dataset passed, and the shard then calibrated on samples the loader never read. The coverage is now compared with the loader's dataset; a sampler that orders or shuffles the whole dataset is still replaced as before. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 24 ++++++++++++------- .../utils/test_distributed_calibration.py | 7 ++++++ 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 083c47ed1ae..1edb57148ae 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -736,9 +736,10 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad counts equal (else a rank exits calibration early and hangs the others on the per-step collectives), at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. It replaces the loader's sampler, so a shuffling one becomes the rank's unshuffled share, and one that - picks samples itself (a subset, a random sampler drawing its own number of samples or with - replacement, or a batch sampler other than a plain ``BatchSampler``) is refused rather than - silently widened. For streaming ``IterableDataset`` datasets (which have no ``len()`` + picks samples itself (a subset, a sampler over another source than the loader's dataset, a + random sampler drawing its own number of samples or with replacement, or a batch sampler other + than a plain ``BatchSampler``) is refused rather than silently widened. For streaming + ``IterableDataset`` datasets (which have no ``len()`` and so can't use a sampler) each rank strides its own copy of the stream via :class:`_StridedIterableShard`. That stream must yield the same samples in the same order on every rank, or the shards overlap and the ranks' step counts can differ, which hangs the @@ -761,12 +762,19 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad picker = batching.sampler else: picker = loader.sampler if batching is None else batching - # A RandomSampler drawing its own number of samples, or with replacement, picks which and how - # many it reads, as a subset does; one that shuffles the whole dataset does not. - narrowing = isinstance(picker, RandomSampler) and ( - picker.replacement or picker.num_samples != len(picker.data_source) + known = type(picker) in (SequentialSampler, RandomSampler, DistributedSampler) + # A sampler over another source than the loader's dataset, or a RandomSampler drawing its + # own number of samples or with replacement, picks which and how many it reads, as a subset + # does; one that orders or shuffles the whole dataset does not. + source = getattr(picker, "data_source", getattr(picker, "dataset", None)) + narrowing = known and ( + len(source) != len(loader.dataset) + or ( + isinstance(picker, RandomSampler) + and (picker.replacement or picker.num_samples != len(loader.dataset)) + ) ) - if type(picker) not in (SequentialSampler, RandomSampler, DistributedSampler) or narrowing: + if not known or narrowing: raise ValueError( f"shard_dataloader: the loader's {type(picker).__name__} picks which samples it " "reads, and the rank's DistributedSampler would replace it; pass a loader over " diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index c25bd1928ef..18bde9efd82 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -33,6 +33,7 @@ from torch.utils.data import ( BatchSampler, DataLoader, + DistributedSampler, IterableDataset, RandomSampler, SequentialSampler, @@ -163,6 +164,9 @@ def test_a_stream_loaded_with_workers_is_refused(): {"batch_sampler": BatchSampler(SubsetRandomSampler(range(4)), 2, drop_last=False)}, {"sampler": RandomSampler(range(8), num_samples=4), "batch_size": 2}, {"sampler": RandomSampler(range(8), replacement=True), "batch_size": 2}, + {"sampler": SequentialSampler(range(4)), "batch_size": 2}, + {"sampler": RandomSampler(range(4)), "batch_size": 2}, + {"sampler": DistributedSampler(range(4), num_replicas=2, rank=0), "batch_size": 2}, ], ids=[ "subset-sampler", @@ -171,6 +175,9 @@ def test_a_stream_loaded_with_workers_is_refused(): "subset-in-a-plain-batch-sampler", "random-sampler-drawing-fewer", "random-sampler-with-replacement", + "sequential-sampler-over-fewer", + "random-sampler-over-fewer", + "distributed-sampler-over-fewer", ], ) def test_a_loader_that_picks_its_own_samples_is_refused(picker): From ad6e1404fc63da5a92cd368f9b1be2ddb0a4ff38 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:52:08 +0000 Subject: [PATCH 59/60] Give mypy a sized default for the sampler's source Only the three known samplers reach len(source), and each has a data_source or a dataset; the empty default only keeps the expression typed. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 1edb57148ae..268d89d9e53 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -766,7 +766,7 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad # A sampler over another source than the loader's dataset, or a RandomSampler drawing its # own number of samples or with replacement, picks which and how many it reads, as a subset # does; one that orders or shuffles the whole dataset does not. - source = getattr(picker, "data_source", getattr(picker, "dataset", None)) + source = getattr(picker, "data_source", getattr(picker, "dataset", ())) narrowing = known and ( len(source) != len(loader.dataset) or ( From dea9b793553d2121185890c382eab0c8493d3d9b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 01:10:37 +0000 Subject: [PATCH 60/60] Re-plant only a forward that quantizer insertion rebound _reinstall_a2a_forwards planted the torch_a2a wrapper on any planned module whose instance forward was missing or a bound method. The clobber it repairs always leaves a bound method: a module with no instance forward is one the load never wrapped, and planting there would add the style's transforms to a module that ran without them. It now plants only over a bound method, as its docstring says. The pre-hook test now starts from rebound forwards, and a new test checks that a module the load did not wrap is left alone. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 3 +-- .../torch/utils/test_distributed_calibration.py | 15 ++++++++++++++- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 268d89d9e53..7528af9e420 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -873,8 +873,7 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: continue # install_forward stores a plain function as the instance's forward, and quantizer # insertion (bind_forward_method) a bound method over it: only the latter lost its wrapper. - current = vars(mod).get("forward") - if current is not None and not isinstance(current, types.MethodType): + if not isinstance(vars(mod).get("forward"), types.MethodType): continue style.install_forward(mod, ep_mesh) n += 1 diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 18bde9efd82..2efef9fe86d 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -414,7 +414,20 @@ def test_reinstall_skips_styles_whose_transforms_are_pre_hooks(monkeypatch, expe styles={"torch_a2a_experts": experts_style, "torch_a2a_router": _PreHookStyle()}, where=_LEGACY_TP_MODULE, ) - assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == planted + model = _MoEModel() + for layer in model.layers: # quantizer insertion rebinds the planned modules' forwards + _drop_tp_wrapper(layer.gate) + _drop_tp_wrapper(layer.experts) + assert _reinstall_a2a_forwards(model, _FakeMesh()) == planted + + +def test_reinstall_leaves_a_module_the_load_did_not_wrap(monkeypatch): + """Only a forward quantizer insertion rebound is planted again: a planned module whose forward + the load did not wrap keeps running as it did, without the style's transforms.""" + styles = _install_fake_tp(monkeypatch) + + assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0 + assert not any(style.installed for style in styles.values()) # ------------------------------------------------------------------------- distributed_forward_loop