From d4068884ab7db7b1f853fbd9936cafbb5638e890 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:23:21 +0000 Subject: [PATCH 01/57] Run expert-parallel PTQ from hf_ptq.py hf_ptq already had a distributed path, but only one: --use_fsdp2, which shards a dense model and replicates every expert of a MoE, so the documented CLI -- recipes, export formats, the flag deprecations -- did not reach an expert-parallel run. A distributed launch now runs inside a ModelOptSession, which owns the process group and is what the expert-parallel loader reads its rank and device from. args.dist_state becomes a view of that session rather than a second, independently derived one. load_model decides the plan after applying the model's spec -- for several models the spec is what supplies the base_model_ep_plan, so deciding first would see an expert-parallel-capable MoE as dense -- and gains an expert-parallel arm ahead of the FSDP2 one. It is an arm of the same chain rather than an early return, so the shared post-processing (tokenizer, VLM language-model extraction) still runs. --use_fsdp2 skips the decision entirely, so an explicit request for the FSDP2 loader is honored. MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE override the automatic degrees. A model that CAN be expert-parallel but resolves to no sharding at all raises: that configuration runs and quietly replicates the whole MoE on every rank. Every distributed path calibrates on its own 1/world shard of the data, and through distributed_forward_loop, which re-plants the torch_a2a forward wrappers quantizer insertion clobbers on transformers >= 5.16 (without it the experts forward runs with no EP process group and raises at the first step). A dense model under torchrun is placed on this rank's GPU rather than by accelerate's "auto" heuristic, which picks the same card on every rank of an idle node and aborts the first collective with "Duplicate GPU detected". Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + examples/hf_ptq/example_utils.py | 28 ++++++- examples/hf_ptq/hf_ptq.py | 132 ++++++++++++++++++++++++++++--- 3 files changed, 150 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1b9783d8dec..952b58e4f20 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -329,6 +329,7 @@ Changelog - Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: the invocation, the ModelOpt version, the run log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled. - Add ``--mlflow `` to ``examples/hf_ptq/hf_ptq.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). A tracked run records the invocation, the resolved recipe (``$import``\ s expanded), the run log and the quantization summaries, with every command-line argument as a searchable param; failed runs are recorded with their traceback. The experiment defaults to ``$USER/hf_ptq/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. - Add ``--mlflow `` to ``examples/vllm_serve/vllm_serve_fakequant.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too), so a fake-quant serve records what it quantized and an evaluation of that endpoint can be traced back to a recipe. A tracked run uploads the launcher command, the resolved ``RECIPE_PATH`` (or the merged ``QUANT_CFG``/``KV_QUANT_CFG`` when presets are used), the worker log and the quantizer summary; the experiment defaults to ``$USER/vllm_serve_fakequant/-`` and can be overridden with ``--mlflow-experiment`` / ``--mlflow-run-name``. +- Add expert-parallel PTQ for HuggingFace MoE models to ``examples/hf_ptq/hf_ptq.py``: under ``torchrun``, each rank holds only its share of the routed experts, tokens reach them by all-to-all, and every rank calibrates on its own slice of the data. ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` override the automatic choice of expert-parallel degree and of FSDP2-sharding the dense layers, and the quantized checkpoint is written by every rank in parallel without gathering the model. **Backward Breaking Changes** diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index e90b034960c..c86053fce90 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -72,8 +72,22 @@ class DistributedState: def setup_distributed_args(args): - """Initialize and attach ``args.dist_state`` (single-process if FSDP2 off).""" - if getattr(args, "use_fsdp2", False): + """Initialize and attach ``args.dist_state`` (single-process if not distributed). + + When the run already entered a :func:`modelopt_session` -- which it does for any distributed + launch, because the expert-parallel loaders need the session object -- the state is a view of + that session rather than a second, independently-derived one. Two sources of rank/world would + be one too many. + """ + mo_session = getattr(args, "mo_session", None) + if mo_session is not None and mo_session.distributed: + args.dist_state = DistributedState( + rank=mo_session.rank, + world_size=mo_session.world_size, + device=mo_session.device, + is_main=mo_session.is_main, + ) + elif getattr(args, "use_fsdp2", False): # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can # exceed it, and PyTorch 2.8 has no per-call barrier() timeout (must be set at PG creation). dist_utils.setup(timeout=timedelta(hours=2)) @@ -582,6 +596,16 @@ def get_model( device_map = "auto" if device == "cpu": device_map = "cpu" + elif ":" in str(device): + # A SPECIFIC device was asked for -- "cuda:" under torchrun. Do not fall through to + # "auto": accelerate runs its placement heuristic INDEPENDENTLY in every rank, and on a node + # whose GPUs are all idle every rank scores them identically and picks the SAME card. All ranks' + # weights, and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL + # collective binds the whole data-parallel group to that single device and aborts with + # "Duplicate GPU detected : rank 1 and rank 0 both on CUDA device ...". Honour the caller. + # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under + # data-parallel calibration each rank must own its own GPU. + device_map = {"": str(device)} # Prepare config kwargs for loading config_kwargs = {"trust_remote_code": trust_remote_code} if trust_remote_code else {} diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 78e175b95a3..ceba5bca233 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -14,10 +14,12 @@ # limitations under the License. import argparse +import contextlib import os import random import time import warnings +from datetime import timedelta from pathlib import Path from typing import Any @@ -92,16 +94,30 @@ EagleOfflineDataCollator, OfflineSupervisedDataset, ) -from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils import env, print_rank_0 from modelopt.torch.utils.dataset_utils import ( create_forward_loop, get_dataset_dataloader, get_max_batch_size, get_supported_datasets, ) -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan +from modelopt.torch.utils.distributed import ( + DistributionPlan, + ModelOptSession, + distributed_forward_loop, + is_distributed_launch, + modelopt_session, + set_distribution_plan, + shard_dataloader, +) 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.moe import ( + apply_distributed_support, + decide_distribution_plan, + is_expert_parallel_capable, + load_and_prepare_ep, +) from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_non_model_files, ensure_local_checkpoint, @@ -195,6 +211,7 @@ def make_calib_dataloader( tokenizer: PreTrainedTokenizerBase | None, device: torch.device, model_type: str | None, + mo_session: ModelOptSession | None, autoquant_gradient_recipe: bool = False, ) -> tuple[DataLoader | _DeviceDataLoader, str | None]: calib_dataloader = None @@ -269,12 +286,19 @@ def make_calib_dataloader( max_sample_length=args.calib_seq, device=device, include_labels=include_labels, - distributed=args.use_fsdp2, - sampler_kwargs=( - {"num_replicas": args.dist_state.world_size, "rank": args.dist_state.rank} - if args.use_fsdp2 - else None - ), + ) + if ( + mo_session is not None + and mo_session.distributed + and calib_dataloader is not None + and isinstance(calib_dataloader, DataLoader) + ): + # Every distributed calibration path is data-parallel over the world, so each rank calibrates + # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel + # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2 + # shards the params but still feeds each rank its own data. + calib_dataloader = shard_dataloader( + calib_dataloader, mo_session.rank, mo_session.world_size ) return calib_dataloader, first_text_speech_dataset @@ -317,10 +341,71 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: ) +def _decide_expert_parallel(args: argparse.Namespace): + """Return ``(plan, hf_config)`` when this run should be expert-parallel, else ``(None, None)``. + + The model's own expert-parallel facts are applied to the config BEFORE anything is decided: for + several models that is what SUPPLIES ``base_model_ep_plan`` -- stock transformers ships one for + some MoEs and not others -- so deciding first would read an EP-capable model as dense and + replicate every expert on every rank. + """ + mo_session = getattr(args, "mo_session", None) + if mo_session is None or not mo_session.distributed: + return None, None + # --use_fsdp2 asks for the standalone FSDP2 loader explicitly; the EP + # auto-decision below has no way to know that, so without this it can pick EP anyway + # (MODELOPT_EP_DEGREE set explicitly, or "auto" resolving non-trivially for the model) + # and load_model's `elif args.use_fsdp2:` branch is never reached. + if args.use_fsdp2: + return None, None + + hf_config = AutoConfig.from_pretrained( + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code + ) + hf_config = apply_distributed_support( + hf_config.model_type, + hf_config, + ckpt_path=args.local_checkpoint_path, + trust_remote_code=args.trust_remote_code, + ) + plan = decide_distribution_plan( + hf_config, + mo_session.world_size, + ep_size=env.ep_degree(), + fsdp_size=env.fsdp_size(), + ) + if mo_session.is_main: + print(f"[ep] decided {plan!r}", flush=True) + + if plan.is_expert_parallel: + return plan, hf_config + if is_expert_parallel_capable(hf_config) and not plan.is_fsdp_dense and not args.use_fsdp2: + # Neither sharding applies, so every rank would hold the whole MoE. That runs, and wastes the + # cluster, so say so rather than let it look like it worked. + raise RuntimeError( + f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 with " + f"no dense sharding at world_size={mo_session.world_size} -- every rank would hold the " + "whole MoE. Set MODELOPT_EP_DEGREE>1, or MODELOPT_FSDP_SIZE=world / --use_fsdp2 to shard " + "it with FSDP2 instead." + ) + return None, None + + def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False - if args.use_fsdp2: + _ep_plan, _ep_config = _decide_expert_parallel(args) + if _ep_plan is not None: + full_model = load_and_prepare_ep( + args.local_checkpoint_path, + args.mo_session, + _ep_plan.ep_size, + trust_remote_code=args.trust_remote_code, + attn_implementation=args.attn_implementation, + hf_config=_ep_config, + fsdp_dense=_ep_plan.is_fsdp_dense, + ) + elif args.use_fsdp2: hf_config = AutoConfig.from_pretrained( args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) @@ -556,6 +641,19 @@ def mono_quantize( # quantization target; the full forward still exercises that component's quantizers. if args.calib_with_images: calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) + elif getattr(args, "mo_session", None) is not None and args.mo_session.distributed: + # A distributed run has to calibrate through distributed_forward_loop, not the + # plain one: it is the only caller of _reinstall_a2a_forwards, which re-plants the + # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on + # transformers >= 5.16 (where the styles install a forward rather than a pre-hook). + # Without it the experts forward runs with process_group=None and raises + # "torch_a2a experts forward requires the EP `process_group`" at calibration step 0. + # On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op -- + # which is why this is invisible on older transformers and fatal on newer. + # It also restores the use_cache guard the plain loop does not apply. + calibrate_loop = distributed_forward_loop( + language_model, calib_dataloader, args.mo_session + ) else: calibrate_loop = create_forward_loop( dataloader=calib_dataloader, @@ -1064,6 +1162,7 @@ def quantize_main( tokenizer, device, model_type, + args.mo_session, autoquant_gradient_recipe=( aq_config is not None and aq_config.auto_quantize_method == "gradient" ), @@ -1513,6 +1612,21 @@ def main(args: argparse.Namespace): random.seed(RAND_SEED) np.random.seed(RAND_SEED) + # A distributed launch runs inside a ModelOptSession: it owns the process group and is what the + # expert-parallel loader reads its rank and device from. Entered here rather than inside + # load_model so cleanup still runs if loading raises -- otherwise the surviving ranks block on + # the first collective until the NCCL timeout. + with contextlib.ExitStack() as _stack: + if is_distributed_launch() or args.use_fsdp2: + args.mo_session = _stack.enter_context( + modelopt_session(timeout=timedelta(hours=3), device=args.device) + ) + else: + args.mo_session = None + _main_inner(args) + + +def _main_inner(args: argparse.Namespace): setup_distributed_args(args) try: From 03402881129c997b09716bbc8174d6c60e558494 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 02:11:46 +0000 Subject: [PATCH 02/57] Move the expert-parallel PTQ changelog entry into 0.49.0 Rebasing onto a main that had opened a new release section left the entry where its context matched: in the Misc list of 0.47.0, a release that has already shipped. It belongs with the other new quantization features of 0.49.0, the section main opened for what merges after 0.48.0's cut. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 952b58e4f20..6753acb6ab5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -20,6 +20,7 @@ Changelog - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2. - Add an opt-in no-gather HuggingFace export for FSDP2-sharded models: with ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` sends each tensor to a single owner rank and every rank writes its own shards in parallel, instead of gathering the full state dict on one host. MoE experts stay sharded through the export, so each rank only materializes its own experts. +- Add expert-parallel PTQ for HuggingFace MoE models to ``examples/hf_ptq/hf_ptq.py``: under ``torchrun``, each rank holds only its share of the routed experts, tokens reach them by all-to-all, and every rank calibrates on its own slice of the data. ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` override the automatic choice of expert-parallel degree and of FSDP2-sharding the dense layers, and the quantized checkpoint is written by every rank in parallel without gathering the model. *Speculative Decoding* @@ -329,7 +330,6 @@ Changelog - Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: the invocation, the ModelOpt version, the run log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled. - Add ``--mlflow `` to ``examples/hf_ptq/hf_ptq.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). A tracked run records the invocation, the resolved recipe (``$import``\ s expanded), the run log and the quantization summaries, with every command-line argument as a searchable param; failed runs are recorded with their traceback. The experiment defaults to ``$USER/hf_ptq/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. - Add ``--mlflow `` to ``examples/vllm_serve/vllm_serve_fakequant.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too), so a fake-quant serve records what it quantized and an evaluation of that endpoint can be traced back to a recipe. A tracked run uploads the launcher command, the resolved ``RECIPE_PATH`` (or the merged ``QUANT_CFG``/``KV_QUANT_CFG`` when presets are used), the worker log and the quantizer summary; the experiment defaults to ``$USER/vllm_serve_fakequant/-`` and can be overridden with ``--mlflow-experiment`` / ``--mlflow-run-name``. -- Add expert-parallel PTQ for HuggingFace MoE models to ``examples/hf_ptq/hf_ptq.py``: under ``torchrun``, each rank holds only its share of the routed experts, tokens reach them by all-to-all, and every rank calibrates on its own slice of the data. ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` override the automatic choice of expert-parallel degree and of FSDP2-sharding the dense layers, and the quantized checkpoint is written by every rank in parallel without gathering the model. **Backward Breaking Changes** From c2914d1f765c2448897aac444f84dd7e694f34e3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:07:59 +0000 Subject: [PATCH 03/57] Test hf_ptq.py's expert-parallel wiring Cover how hf_ptq.py runs under torchrun: - the MODELOPT_EP_DEGREE / MODELOPT_FSDP_SIZE knobs, read at call time; - applying a model's spec before the EP decision; - the decision itself: the auto degree, the knobs overriding it, --use_fsdp2 and single-process runs opting out, an unsharded MoE refused; - load_model handing an EP plan to load_and_prepare_ep and to no other loader; - each rank pinning a model to its own GPU; - dist_state following the torchrun session; - main() closing the session even when the run raises; - calibration on per-rank data through distributed_forward_loop. Signed-off-by: Shengliang Xu --- .../hf_ptq/test_hf_ptq_expert_parallel.py | 462 ++++++++++++++++++ 1 file changed, 462 insertions(+) create mode 100644 tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py new file mode 100644 index 00000000000..94425b8d64f --- /dev/null +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -0,0 +1,462 @@ +# 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. + +"""Expert-parallel PTQ wiring in ``examples/hf_ptq``. + +Under torchrun, hf_ptq.py applies the model's spec to its config, decides a distribution plan from +the ``MODELOPT_EP_DEGREE`` / ``MODELOPT_FSDP_SIZE`` knobs, and loads the model through +``load_and_prepare_ep`` when that plan shards the experts. Every rank then calibrates on its own +shard of the data, and a dense model is pinned to the rank's own GPU instead of ``"auto"``. + +Nothing here needs a process group or a checkpoint: an initialized ``ModelOptSession`` stands in +for the torchrun session, the configs are tiny real transformers configs served from memory, and +the loaders are replaced by stand-ins that record how they were called. +""" + +import contextlib +import importlib +import inspect +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import transformers +from _test_utils.examples.hf_ptq_example_utils import example_utils +from _test_utils.torch.transformers_models import get_tiny_tokenizer +from packaging.version import Version +from torch.utils.data import DataLoader +from transformers import DeepseekV3Config, LlamaConfig, Qwen3MoeConfig + +from modelopt.torch.models import get_spec +from modelopt.torch.utils import env +from modelopt.torch.utils.distributed import ModelOptSession +from modelopt.torch.utils.moe import is_expert_parallel_capable, load_and_prepare_ep + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" + +# Small enough that the dense part fits any GPU, so "auto" dense sharding always resolves to +# replication whether or not the machine running the tests has one. +_TINY = { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "vocab_size": 32, +} + +_CKPT = "/ckpt/model" + + +class _LoaderReachedError(Exception): + """Raised by a loader stand-in once the call under test has been recorded.""" + + +class _RunFailedError(Exception): + """Raised by a stand-in for the body of a run that fails.""" + + +@pytest.fixture +def hf_ptq(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + return importlib.import_module("hf_ptq") + + +@pytest.fixture(autouse=True) +def _no_distribution_knobs(monkeypatch): + """A shell that exports the knobs would otherwise steer every decision below.""" + monkeypatch.delenv("MODELOPT_EP_DEGREE", raising=False) + monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False) + + +def _session(world_size, rank=0, distributed=True): + """The session modelopt_session() yields on ``rank`` of a ``world_size``-rank launch.""" + session = ModelOptSession() + session.initialize( + distributed=distributed, + rank=rank, + world_size=world_size, + local_rank=rank, + device=torch.device(f"cuda:{rank}"), + is_main=rank == 0, + ) + return session + + +def _args(session, **overrides): + """What load_model reads before it picks a loader, as main() leaves it.""" + args = SimpleNamespace( + mo_session=session, + use_fsdp2=False, + recipe=None, + local_checkpoint_path=_CKPT, + trust_remote_code=False, + attn_implementation=None, + ) + vars(args).update(overrides) + return args + + +def _serve(monkeypatch, hf_ptq, config): + """Answer hf_ptq's config read with ``config`` rather than a checkpoint on disk.""" + monkeypatch.setattr(hf_ptq.AutoConfig, "from_pretrained", lambda *args, **kwargs: config) + return config + + +def _qwen3_moe(num_experts): + """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself.""" + return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) + + +def _deepseek_v3(n_routed_experts): + """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable.""" + config = DeepseekV3Config(n_routed_experts=n_routed_experts, **_TINY) + config.base_model_ep_plan = None # holds even if a later transformers ships a plan + return config + + +# ----------------------------------------------------------------------------------- env knobs + + +@pytest.mark.parametrize( + ("knob", "name"), + [(env.ep_degree, "MODELOPT_EP_DEGREE"), (env.fsdp_size, "MODELOPT_FSDP_SIZE")], +) +def test_distribution_knobs_default_to_auto(monkeypatch, knob, name): + """Unset and exported-but-empty both mean "let the plan decide".""" + assert knob() == "auto" + monkeypatch.setenv(name, "") + assert knob() == "auto" + + +@pytest.mark.parametrize( + ("knob", "name", "other"), + [ + (env.ep_degree, "MODELOPT_EP_DEGREE", env.fsdp_size), + (env.fsdp_size, "MODELOPT_FSDP_SIZE", env.ep_degree), + ], +) +def test_distribution_knobs_are_read_when_called(monkeypatch, knob, name, other): + """A value snapshotted at import would miss what the launcher sets later; and each knob reads + only its own variable.""" + monkeypatch.setenv(name, "2") + assert knob() == "2" + monkeypatch.setenv(name, "8") + assert knob() == "8" + assert other() == "auto" + + +# -------------------------------------------------------------------------------- the decision + + +@pytest.mark.parametrize("single_process_session", [False, True]) +def test_a_single_process_run_is_never_expert_parallel(hf_ptq, monkeypatch, single_process_session): + """With no distributed session there is nothing to shard over, even for an EP-capable MoE + (which, decided at world_size=1, would otherwise be rejected as an unsharded MoE).""" + session = _session(world_size=1, distributed=False) if single_process_session else None + _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + + assert hf_ptq._decide_expert_parallel(_args(session)) == (None, None) + + +def test_use_fsdp2_keeps_the_standalone_fsdp2_loader(hf_ptq, monkeypatch): + """--use_fsdp2 names its loader, so even an explicit EP degree for an EP-capable MoE must not + divert the run to the expert-parallel one.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") + _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + args = _args(_session(world_size=2), use_fsdp2=True) + + assert hf_ptq._decide_expert_parallel(args) == (None, None) + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), + reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", +) +@pytest.mark.parametrize("rank", [0, 1]) +def test_an_ep_capable_moe_is_planned_expert_parallel(hf_ptq, monkeypatch, capsys, rank): + """Auto picks the largest EP degree dividing both the world and the expert count, hands back a + config the EP loader can shard, and only rank 0 reports the decision.""" + _serve(monkeypatch, hf_ptq, _qwen3_moe(num_experts=8)) + + plan, hf_config = hf_ptq._decide_expert_parallel(_args(_session(world_size=4, rank=rank))) + + assert (plan.world_size, plan.ep_size, plan.edp) == (4, 4, 1) + assert not plan.is_fsdp_dense + assert is_expert_parallel_capable(hf_config) + assert ("[ep] decided" in capsys.readouterr().out) == (rank == 0) + + +def test_the_spec_is_applied_before_the_decision(hf_ptq, monkeypatch): + """DeepSeek-V3's plan comes only from its spec. Deciding before applying it would read the MoE + as dense; applied first, it shards, and the config handed on carries the spec's plan.""" + _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + + plan, hf_config = hf_ptq._decide_expert_parallel(_args(_session(world_size=4))) + + assert plan.ep_size == 4 + assert hf_config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + + +@pytest.mark.parametrize( + ("ep_degree", "experts"), + [(None, 3), ("1", 8)], + ids=["auto_finds_no_degree", "explicit_degree_one"], +) +def test_an_ep_capable_moe_left_unsharded_is_an_error(hf_ptq, monkeypatch, ep_degree, experts): + """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With + the dense part replicated, every rank would hold the whole MoE: that must fail, not run.""" + if ep_degree is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) + _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=experts)) + + with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + hf_ptq._decide_expert_parallel(_args(_session(world_size=2))) + + +def test_a_dense_model_under_torchrun_runs_data_parallel(hf_ptq, monkeypatch): + """A dense model has no experts to shard, so replicating it per rank is not an error.""" + _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) + + assert hf_ptq._decide_expert_parallel(_args(_session(world_size=4))) == (None, None) + + +def test_the_env_knobs_override_the_automatic_plan(hf_ptq, monkeypatch): + """The decision honours both knobs; left on auto it would pick ep_size=4, dense replicated.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + + plan, _ = hf_ptq._decide_expert_parallel(_args(_session(world_size=4))) + + assert (plan.ep_size, plan.edp, plan.fsdp_size) == (2, 2, 4) + + +# ---------------------------------------------------------------------------------- load_model + + +def test_load_model_hands_an_expert_parallel_plan_to_the_ep_loader(hf_ptq, monkeypatch): + """With an EP plan, load_model loads through load_and_prepare_ep -- with the session, the + decided degree and dense sharding, and the spec-prepared config -- and through no other + loader.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + served = _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + session = _session(world_size=4) + calls = [] + + def load_and_prepare_ep_stub(*args, **kwargs): + calls.append(inspect.signature(load_and_prepare_ep).bind(*args, **kwargs).arguments) + raise _LoaderReachedError + + monkeypatch.setattr(hf_ptq, "load_and_prepare_ep", load_and_prepare_ep_stub) + monkeypatch.setattr(hf_ptq, "get_model", lambda *a, **k: pytest.fail("dense loader ran")) + monkeypatch.setattr( + hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *a, **k: pytest.fail("FSDP2 loader ran") + ) + + with pytest.raises(_LoaderReachedError): + hf_ptq.load_model(_args(session, trust_remote_code=True, attn_implementation="eager")) + + (call,) = calls + assert call["ckpt_path"] == _CKPT + assert call["mo_session"] is session + assert call["ep_degree"] == 4 + assert call["fsdp_dense"] is True + assert call["trust_remote_code"] is True + assert call["attn_implementation"] == "eager" + assert call["hf_config"] is served + assert served.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + + +# ------------------------------------------------------------------------- per-rank placement + + +@pytest.mark.parametrize( + ("device", "device_map"), + [ + ("cuda:1", {"": "cuda:1"}), + (torch.device("cuda:1"), {"": "cuda:1"}), + ("cuda", "auto"), + ], + ids=["str", "torch_device", "unindexed"], +) +def test_a_specific_device_pins_the_whole_model_to_it(monkeypatch, tmp_path, device, device_map): + """Under torchrun each rank asks for its own cuda:; "auto" would let every rank's + accelerate pick the same idle GPU. An unindexed "cuda" still spreads the model over the GPUs.""" + LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) + placed = [] + + def from_pretrained(auto_class, ckpt_path, **kwargs): + placed.append(kwargs["device_map"]) + return torch.nn.Linear(1, 1) + + monkeypatch.setattr(example_utils, "_from_pretrained_recording", from_pretrained) + + example_utils.get_model(str(tmp_path), device=device) + + assert placed == [device_map] + + +# -------------------------------------------------------------------- session and dist_state + + +@pytest.mark.parametrize("use_fsdp2", [False, True]) +def test_dist_state_mirrors_the_distributed_session(monkeypatch, use_fsdp2): + """Under torchrun dist_state is read from the session that owns the process group, even with + --use_fsdp2, whose own branch would initialize a second group.""" + monkeypatch.setattr( + example_utils.dist_utils, "setup", lambda *a, **k: pytest.fail("second process group") + ) + session = _session(world_size=4, rank=1) + args = SimpleNamespace(mo_session=session, use_fsdp2=use_fsdp2, device="cuda") + + example_utils.setup_distributed_args(args) + + assert args.dist_state == example_utils.DistributedState( + rank=1, world_size=4, device=torch.device("cuda:1"), is_main=False + ) + + +def test_a_single_process_session_leaves_fsdp2_to_set_up_its_group(monkeypatch): + """A one-rank --use_fsdp2 launch enters a session that owns no process group, so the FSDP2 + branch must still create one.""" + setups = [] + monkeypatch.setattr(example_utils.dist_utils, "setup", lambda timeout=None: setups.append(0)) + monkeypatch.setattr(example_utils.dist_utils, "rank", lambda group=None: 0) + monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: 1) + monkeypatch.setattr(example_utils.dist_utils, "local_rank", lambda: 0) + args = SimpleNamespace(mo_session=_session(world_size=1, distributed=False), use_fsdp2=True) + + example_utils.setup_distributed_args(args) + + assert len(setups) == 1 + assert args.dist_state.world_size == 1 + + +@pytest.mark.parametrize( + ("world_size", "use_fsdp2", "enters_session"), + [("2", False, True), ("1", True, True), (None, False, False)], + ids=["torchrun", "fsdp2", "single_process"], +) +def test_main_runs_a_distributed_launch_inside_a_session( + hf_ptq, monkeypatch, world_size, use_fsdp2, enters_session +): + """main() opens the session around the whole run, with a collective timeout above NCCL's + 30-minute default, and closes it even when the run raises -- or the other ranks would wait in + their next collective until that timeout.""" + if world_size is None: + monkeypatch.delenv("WORLD_SIZE", raising=False) + else: + monkeypatch.setenv("WORLD_SIZE", world_size) + session = _session(world_size=2) + events, timeouts = [], [] + + @contextlib.contextmanager + def modelopt_session(timeout=None, device=None): + timeouts.append(timeout) + events.append(("enter", device)) + try: + yield session + finally: + events.append(("exit",)) + + def main_inner(args): + events.append(("run", args.mo_session)) + raise _RunFailedError + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(hf_ptq, "modelopt_session", modelopt_session) + monkeypatch.setattr(hf_ptq, "_main_inner", main_inner) + + with pytest.raises(_RunFailedError): + hf_ptq.main(SimpleNamespace(use_fsdp2=use_fsdp2, device="cuda")) + + if enters_session: + assert events == [("enter", "cuda"), ("run", session), ("exit",)] + assert timeouts[0] > timedelta(minutes=30) + else: + assert events == [("run", None)] + + +# --------------------------------------------------------------------------------- calibration + + +def test_distributed_calibration_gives_each_rank_its_own_shard(hf_ptq, monkeypatch): + """Every distributed path calibrates data-parallel: the ranks get equal, disjoint shards that + together cover the calibration set. A single process keeps the loader as built.""" + loader = DataLoader(list(range(8)), batch_size=1) + monkeypatch.setattr(hf_ptq, "get_dataset_dataloader", lambda **kwargs: loader) + args = SimpleNamespace( + specdec_offline_dataset=None, + calib_with_images=False, + dataset=["cnn_dailymail"], + batch_size=1, + calib_size=[8], + calib_seq=16, + ) + tokenizer = get_tiny_tokenizer() + + def calib_loader(session): + built, _ = hf_ptq.make_calib_dataloader( + args, None, None, tokenizer, torch.device("cpu"), "llama", session + ) + return built + + shards = [[int(b) for b in calib_loader(_session(world_size=2, rank=r))] for r in range(2)] + + assert len(shards[0]) == len(shards[1]) == 4 + assert sorted(shards[0] + shards[1]) == list(range(8)) + assert calib_loader(None) is loader + + +def test_distributed_calibration_runs_through_the_distributed_forward_loop(hf_ptq, monkeypatch): + """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion + clobbers, so a distributed run must calibrate through it rather than the plain loop.""" + session = _session(world_size=2) + model = torch.nn.Linear(2, 2) + calib_dataloader, loop = object(), object() + built, quantized = [], [] + monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False) + monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: True) + monkeypatch.setattr(hf_ptq, "distributed_forward_loop", lambda *a: built.append(a) or loop) + monkeypatch.setattr( + hf_ptq, "create_forward_loop", lambda **k: pytest.fail("plain calibration loop selected") + ) + monkeypatch.setattr( + hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: quantized.append(forward_loop) or m + ) + args = SimpleNamespace( + qformat="fp8", + calib_with_images=False, + mo_session=session, + layerwise_export=False, + specdec_offline_dataset=None, + ) + + hf_ptq.mono_quantize( + args, + {"algorithm": "max"}, + model, + model, + model_type="qwen3_moe", + calibration_only=False, + calib_dataloader=calib_dataloader, + is_nemotron_vl_model=False, + ) + + assert built == [(model, calib_dataloader, session)] + assert quantized == [loop] From dcfa2f061d0a311b68ded0578c99f080baf4b41b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:08:13 +0000 Subject: [PATCH 04/57] Refuse a plan whose dense sharding nothing would apply Only the expert-parallel loader shards a dense part by itself. When the plan asked for dense sharding (MODELOPT_FSDP_SIZE, or "auto" finding the dense part too big for one GPU) but resolved to ep_size=1, the decision returned no plan without complaint, and load_model replicated the whole model -- MoE or dense -- on every rank. The error for an unsharded MoE even recommended MODELOPT_FSDP_SIZE=world, which only silenced it. Raise instead, pointing at --use_fsdp2 (and at MODELOPT_EP_DEGREE for an MoE), and drop the bad advice and the unreachable --use_fsdp2 check. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 24 +++++++++++++------ .../hf_ptq/test_hf_ptq_expert_parallel.py | 19 +++++++++++++++ 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index ceba5bca233..7955fa2cb18 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -379,14 +379,24 @@ def _decide_expert_parallel(args: argparse.Namespace): if plan.is_expert_parallel: return plan, hf_config - if is_expert_parallel_capable(hf_config) and not plan.is_fsdp_dense and not args.use_fsdp2: - # Neither sharding applies, so every rank would hold the whole MoE. That runs, and wastes the - # cluster, so say so rather than let it look like it worked. + # Past this point load_model replicates the model on every rank, so any sharding the plan asked + # for would be dropped. That runs, and wastes the cluster, so say so rather than let it look like + # it worked. + ep_capable = is_expert_parallel_capable(hf_config) + if plan.is_fsdp_dense: + # Only the expert-parallel loader shards a dense part by itself. raise RuntimeError( - f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 with " - f"no dense sharding at world_size={mo_session.world_size} -- every rank would hold the " - "whole MoE. Set MODELOPT_EP_DEGREE>1, or MODELOPT_FSDP_SIZE=world / --use_fsdp2 to shard " - "it with FSDP2 instead." + f"The plan shards the dense part over {plan.fsdp_size} ranks but leaves ep_size=1, and " + "without expert parallelism nothing shards it -- every rank would hold the whole model. " + "Pass --use_fsdp2 to shard it with FSDP2" + + (", or set MODELOPT_EP_DEGREE>1" if ep_capable else "") + + ", or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU." + ) + if ep_capable: + raise RuntimeError( + f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at " + f"world_size={mo_session.world_size} -- every rank would hold the whole MoE. Set " + "MODELOPT_EP_DEGREE>1, or pass --use_fsdp2 to shard it with FSDP2 instead." ) return None, None diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 94425b8d64f..49df0859e8a 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -228,6 +228,25 @@ def test_an_ep_capable_moe_left_unsharded_is_an_error(hf_ptq, monkeypatch, ep_de hf_ptq._decide_expert_parallel(_args(_session(world_size=2))) +@pytest.mark.parametrize( + ("config", "ep_capable"), + [(_deepseek_v3(n_routed_experts=3), True), (LlamaConfig(**_TINY), False)], + ids=["moe_no_ep_degree", "dense"], +) +def test_dense_sharding_without_expert_parallelism_is_an_error( + hf_ptq, monkeypatch, config, ep_capable +): + """Only the EP loader shards a dense part by itself. Without EP the run would load the whole + model on every rank and drop the sharding the plan asked for, so it must fail and point at the + loader that can shard it.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + _serve(monkeypatch, hf_ptq, config) + + with pytest.raises(RuntimeError, match="--use_fsdp2") as raised: + hf_ptq._decide_expert_parallel(_args(_session(world_size=2))) + assert ("MODELOPT_EP_DEGREE" in str(raised.value)) == ep_capable + + def test_a_dense_model_under_torchrun_runs_data_parallel(hf_ptq, monkeypatch): """A dense model has no experts to shard, so replicating it per rank is not an error.""" _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) From b204c5232500f12168cb2249887c647a636f8c3d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:08:21 +0000 Subject: [PATCH 05/57] Pin a model to one GPU only when torchrun runs several ranks get_model pinned every indexed device to that one GPU so that torchrun's ranks stop sharing a card, but it did so for a single process too: a run passing --device cuda:0 no longer spread a model too big for one GPU across the others, as it did before. Pin only when the run has more than one rank. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 6 +++--- .../hf_ptq/test_hf_ptq_expert_parallel.py | 19 ++++++++++++------- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index c86053fce90..eefba5fa75a 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -596,15 +596,15 @@ def get_model( device_map = "auto" if device == "cpu": device_map = "cpu" - elif ":" in str(device): - # A SPECIFIC device was asked for -- "cuda:" under torchrun. Do not fall through to + elif ":" in str(device) and dist_utils.size() > 1: + # Under torchrun each rank asks for its own "cuda:". Do not fall through to # "auto": accelerate runs its placement heuristic INDEPENDENTLY in every rank, and on a node # whose GPUs are all idle every rank scores them identically and picks the SAME card. All ranks' # weights, and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL # collective binds the whole data-parallel group to that single device and aborts with # "Duplicate GPU detected : rank 1 and rank 0 both on CUDA device ...". Honour the caller. # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under - # data-parallel calibration each rank must own its own GPU. + # data-parallel calibration each rank must own its own GPU. A single process keeps "auto". device_map = {"": str(device)} # Prepare config kwargs for loading diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 49df0859e8a..6923d6df99c 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -305,17 +305,21 @@ def load_and_prepare_ep_stub(*args, **kwargs): @pytest.mark.parametrize( - ("device", "device_map"), + ("device", "world_size", "device_map"), [ - ("cuda:1", {"": "cuda:1"}), - (torch.device("cuda:1"), {"": "cuda:1"}), - ("cuda", "auto"), + ("cuda:1", 4, {"": "cuda:1"}), + (torch.device("cuda:1"), 4, {"": "cuda:1"}), + ("cuda", 4, "auto"), + ("cuda:1", 1, "auto"), ], - ids=["str", "torch_device", "unindexed"], + ids=["str", "torch_device", "unindexed", "single_process"], ) -def test_a_specific_device_pins_the_whole_model_to_it(monkeypatch, tmp_path, device, device_map): +def test_a_rank_device_pins_the_whole_model_to_it( + monkeypatch, tmp_path, device, world_size, device_map +): """Under torchrun each rank asks for its own cuda:; "auto" would let every rank's - accelerate pick the same idle GPU. An unindexed "cuda" still spreads the model over the GPUs.""" + accelerate pick the same idle GPU. An unindexed "cuda" still spreads the model over the GPUs, + and so does a single process, whatever device it names.""" LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) placed = [] @@ -324,6 +328,7 @@ def from_pretrained(auto_class, ckpt_path, **kwargs): return torch.nn.Linear(1, 1) monkeypatch.setattr(example_utils, "_from_pretrained_recording", from_pretrained) + monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: world_size) example_utils.get_model(str(tmp_path), device=device) From d5fc7c68a7e2e79b0796104229612669baec69c1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:43:38 +0000 Subject: [PATCH 06/57] Document expert-parallel PTQ in the hf_ptq README Say what running hf_ptq.py under torchrun does to an MoE, how to launch it on one node or several, what MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE override, and which models it refuses. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 4fe27160db7..89c5da7442b 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -654,6 +654,32 @@ See [Recipe-based Quantization](#recipe-based-quantization) for the recipe forma > *Set `MODELOPT_FSDP2_EXPORT=nogather` to have every rank write its own shards of the checkpoint in parallel instead of gathering the full state dict on one host.* +## Expert-Parallel Post-Training Quantization for MoE Models + +Under `torchrun`, `hf_ptq.py` shards an MoE model's routed experts across the ranks (expert parallelism) instead of loading the whole model on every GPU. Each rank holds `num_experts / ep_degree` experts and calibrates on its own shard of the calibration data. An all-to-all sends every token to the rank that owns its expert, so each expert's quantizers see the tokens routed to it from all ranks. The dense (non-expert) part is replicated on every rank, or FSDP2-sharded when it does not fit one GPU. Each rank then writes its share of the unified HF checkpoint, without gathering the model onto one GPU. + +A model can be expert-parallel when its config declares a `base_model_ep_plan`. transformers ships one for several MoEs, and ModelOpt supplies the plan or an adapter for others (see the `ep_spec` entries in `modelopt/torch/models/*/specs.py`). Expert parallelism needs transformers 5.6 or newer; a model that needs a newer release says so before loading. + +```bash +torchrun --nproc_per_node= \ + hf_ptq.py \ + --pyt_ckpt_path \ + --recipe general/ptq/nvfp4_default-kv_fp8 \ + --calib_size \ + --export_path +``` + +To run across nodes, add the `--nnodes`, `--node_rank`, `--master_addr` and `--master_port` flags shown in [Manual (run on each node)](#manual-run-on-each-node). `--calib_size` is the total over all ranks. + +Two environment variables override the automatic distribution plan: + +| Variable | Default | Effect | +| :--- | :--- | :--- | +| `MODELOPT_EP_DEGREE` | `auto` | Expert-parallel degree. `auto` picks the largest degree that divides both the routed-expert count and the world size. Below the world size, each group of `MODELOPT_EP_DEGREE` ranks holds a full set of experts, and the groups calibrate on different data. | +| `MODELOPT_FSDP_SIZE` | `auto` | Whether to FSDP2-shard the dense part. `auto` shards it when it would take more than 60% of one GPU's memory; `world` always shards it; `1` replicates it. | + +A run that cannot shard the model the way the plan asks fails before loading, rather than loading the whole model on every rank. `--use_fsdp2` skips expert parallelism and uses the [FSDP2 path](#multi-node-post-training-quantization-with-fsdp2) instead. That is also the way to quantize GPT-OSS and Llama-4 across GPUs: their experts are stored transposed, which the expert-parallel export cannot write yet, so expert parallelism refuses them. + ## Evaluate Accuracy ### TensorRT-LLM Validation From 1b6d046a330406f4c9453fbba19dff8035450603 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 10:46:13 +0000 Subject: [PATCH 07/57] Calibrate through the distributed loop only for expert-parallel runs Every distributed run calibrated through distributed_forward_loop, which moved --use_fsdp2 off main's create_forward_loop and with it the split-on-OOM retry and the encoder-decoder generate path. Only expert parallelism needs the distributed loop -- it re-plants the all-to-all forwards quantizer insertion clobbers -- so it is now chosen from the session's ep_mesh, which load_and_prepare_ep records, and FSDP2 keeps the plain loop. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 7 ++-- .../hf_ptq/test_hf_ptq_expert_parallel.py | 36 +++++++++++++------ 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 7955fa2cb18..b42ea08364c 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -651,8 +651,8 @@ def mono_quantize( # quantization target; the full forward still exercises that component's quantizers. if args.calib_with_images: calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) - elif getattr(args, "mo_session", None) is not None and args.mo_session.distributed: - # A distributed run has to calibrate through distributed_forward_loop, not the + elif getattr(getattr(args, "mo_session", None), "ep_mesh", None) is not None: + # An expert-parallel run has to calibrate through distributed_forward_loop, not the # plain one: it is the only caller of _reinstall_a2a_forwards, which re-plants the # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on # transformers >= 5.16 (where the styles install a forward rather than a pre-hook). @@ -660,7 +660,8 @@ def mono_quantize( # "torch_a2a experts forward requires the EP `process_group`" at calibration step 0. # On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op -- # which is why this is invisible on older transformers and fatal on newer. - # It also restores the use_cache guard the plain loop does not apply. + # It also restores the use_cache guard the plain loop does not apply. Other + # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. calibrate_loop = distributed_forward_loop( language_model, calib_dataloader, args.mo_session ) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 6923d6df99c..c8eead2c73b 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -447,18 +447,17 @@ def calib_loader(session): assert calib_loader(None) is loader -def test_distributed_calibration_runs_through_the_distributed_forward_loop(hf_ptq, monkeypatch): - """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion - clobbers, so a distributed run must calibrate through it rather than the plain loop.""" - session = _session(world_size=2) +def _calibration_loop_choice(hf_ptq, monkeypatch, session): + """Run mono_quantize under ``session`` and return which calibration loop it built.""" model = torch.nn.Linear(2, 2) - calib_dataloader, loop = object(), object() - built, quantized = [], [] + chosen, quantized = [], [] monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False) monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: True) - monkeypatch.setattr(hf_ptq, "distributed_forward_loop", lambda *a: built.append(a) or loop) monkeypatch.setattr( - hf_ptq, "create_forward_loop", lambda **k: pytest.fail("plain calibration loop selected") + hf_ptq, "distributed_forward_loop", lambda *a: chosen.append("distributed") or object() + ) + monkeypatch.setattr( + hf_ptq, "create_forward_loop", lambda **k: chosen.append("plain") or object() ) monkeypatch.setattr( hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: quantized.append(forward_loop) or m @@ -478,9 +477,24 @@ def test_distributed_calibration_runs_through_the_distributed_forward_loop(hf_pt model, model_type="qwen3_moe", calibration_only=False, - calib_dataloader=calib_dataloader, + calib_dataloader=object(), is_nemotron_vl_model=False, ) - assert built == [(model, calib_dataloader, session)] - assert quantized == [loop] + assert len(quantized) == 1 + return chosen + + +def test_expert_parallel_calibration_runs_through_the_distributed_forward_loop(hf_ptq, monkeypatch): + """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion + clobbers, so an expert-parallel run must calibrate through it rather than the plain loop.""" + session = _session(world_size=2) + session.ep_mesh = object() # what load_and_prepare_ep records + + assert _calibration_loop_choice(hf_ptq, monkeypatch, session) == ["distributed"] + + +def test_fsdp2_calibration_keeps_the_plain_forward_loop(hf_ptq, monkeypatch): + """Without expert parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its + split-on-OOM retry, as on main.""" + assert _calibration_loop_choice(hf_ptq, monkeypatch, _session(world_size=2)) == ["plain"] From 6b0e1eac339a5f88845985209a3aed005f3249fc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 01:23:22 +0000 Subject: [PATCH 08/57] Initialize NVML before probing it in launch_memory_monitor launch_memory_monitor checks that NVML can report memory before starting the monitor, so hosts where it cannot (#374) skip monitoring instead of crashing. But the probe ran before nvmlInit, which only GPUMemoryMonitor's constructor calls, so every NVML call in it raised NVMLError_Uninitialized: the monitor printed "Failed to get GPU memory info: Uninitialized" and never started on any host. hf_ptq.py launches it on every run, so no run has reported its per-GPU peak memory since. The probe now initializes NVML for itself and shuts it down afterwards, so the init count stays balanced with the monitor's own. Hosts where memory info is unsupported still skip the monitor. Unit tests run launch_memory_monitor against a fake, reference-counted NVML: the monitor starts when memory info is available (this fails on the old code), and is skipped with NVML left shut down when it is not. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/memory_monitor.py | 9 +- tests/unit/torch/utils/test_memory_monitor.py | 89 +++++++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 tests/unit/torch/utils/test_memory_monitor.py diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py index 94ed1cab850..0b1fbbb1fd1 100644 --- a/modelopt/torch/utils/memory_monitor.py +++ b/modelopt/torch/utils/memory_monitor.py @@ -140,8 +140,15 @@ def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor | N Returns: GPUMemoryMonitor: The monitor instance that was launched """ + # Probe that NVML can report memory on this host before starting the monitor. NVML must be + # initialized for the probe itself; without it every call raises NVMLError_Uninitialized, and + # the monitor would never start anywhere. try: - nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) + nvmlInit() + try: + nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) + finally: + nvmlShutdown() except Exception as e: print(f"Failed to get GPU memory info: {e}. Stopping GPU memory monitor.") return None diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py new file mode 100644 index 00000000000..c90052a48fb --- /dev/null +++ b/tests/unit/torch/utils/test_memory_monitor.py @@ -0,0 +1,89 @@ +# 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 launch_memory_monitor's NVML availability probe, against a fake NVML.""" + +import atexit +from types import SimpleNamespace + +import pytest + +import modelopt.torch.utils.memory_monitor as mm + + +class _FakeNVML: + """Reference-counted like NVML: every call but init raises while uninitialized.""" + + def __init__(self, memory_info_supported: bool = True): + self.inits = 0 + self.memory_info_supported = memory_info_supported + + def install(self, monkeypatch): + for name in ("Init", "Shutdown", "DeviceGetCount", "DeviceGetHandleByIndex"): + monkeypatch.setattr(mm, f"nvml{name}", getattr(self, name)) + monkeypatch.setattr(mm, "nvmlDeviceGetMemoryInfo", self.DeviceGetMemoryInfo) + + def _require_init(self): + if self.inits <= 0: + raise RuntimeError("Uninitialized") + + def Init(self): # noqa: N802 + self.inits += 1 + + def Shutdown(self): # noqa: N802 + self._require_init() + self.inits -= 1 + + def DeviceGetCount(self): # noqa: N802 + self._require_init() + return 1 + + def DeviceGetHandleByIndex(self, index): # noqa: N802 + self._require_init() + return index + + def DeviceGetMemoryInfo(self, handle): # noqa: N802 + self._require_init() + if not self.memory_info_supported: + raise RuntimeError("Not Supported") + return SimpleNamespace(used=2 * 1024**3) + + +@pytest.fixture +def launched(): + monitors = [] + yield monitors + for monitor in monitors: + atexit.unregister(monitor.stop) + monitor.stop() + + +def test_the_monitor_starts_when_nvml_can_report_memory(monkeypatch, launched): + """The probe used to run before NVML was initialized, so it always failed and the monitor + never started on any host.""" + nvml = _FakeNVML() + nvml.install(monkeypatch) + monitor = mm.launch_memory_monitor(monitor_interval=0.01) + assert monitor is not None + launched.append(monitor) + assert nvml.inits == 1 # the probe's own init is balanced; the monitor holds one + + +def test_the_monitor_is_skipped_where_nvml_cannot_report_memory(monkeypatch): + """Hosts where NVML cannot report memory still skip the monitor, and leave NVML shut down.""" + nvml = _FakeNVML(memory_info_supported=False) + nvml.install(monkeypatch) + assert mm.launch_memory_monitor() is None + assert nvml.inits == 0 From e787aeb1cccdd08a8aa6564788878172a07b1f33 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 00:32:56 +0000 Subject: [PATCH 09/57] Stub load_pretrained in the device-pinning test get_model now loads through load_pretrained(ckpt_path, auto_class, ...), which replaced example_utils' own recording loader. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index c8eead2c73b..a4dc8797f00 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -323,11 +323,11 @@ def test_a_rank_device_pins_the_whole_model_to_it( LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) placed = [] - def from_pretrained(auto_class, ckpt_path, **kwargs): + def load_pretrained(ckpt_path, auto_class, **kwargs): placed.append(kwargs["device_map"]) return torch.nn.Linear(1, 1) - monkeypatch.setattr(example_utils, "_from_pretrained_recording", from_pretrained) + monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained) monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: world_size) example_utils.get_model(str(tmp_path), device=device) From 73a4e4fa448523bbf3ebc01810f5f7aef1f0d682 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 17:09:14 +0000 Subject: [PATCH 10/57] Load hf_ptq's models through load_pretrained's routing The expert-parallel and --use_fsdp2 runs now load through load_pretrained, with the plan or "fsdp2", instead of calling load_and_prepare_ep and parallel_load_and_prepare_fsdp2 themselves, so they get its checks: the expert-parallel run now refuses a Nemotron-H MTP checkpoint it would have exported with the head unquantized. resolve_distribution_plan takes over the decision, including the refusal to replicate an expert-parallel MoE on every rank; this script keeps only its own rule that FSDP2 runs under --use_fsdp2. get_model and the low-memory load ask for "sequential": their run is already decided, and they place the model themselves. The tests of the decision itself moved with it to test_load_pretrained_distribution.py. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 4 + examples/hf_ptq/hf_ptq.py | 76 ++---- .../hf_ptq/test_hf_ptq_expert_parallel.py | 232 ++++++++---------- 3 files changed, 119 insertions(+), 193 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index eefba5fa75a..3da7a745e48 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -694,6 +694,7 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, + distribution="sequential", device_map=device_map, **model_kwargs, ) @@ -705,6 +706,7 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, + distribution="sequential", device_map="auto", trust_remote_code=trust_remote_code, dtype="auto", @@ -731,6 +733,7 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, + distribution="sequential", device_map="cpu" if device == "cpu" else "sequential", **model_kwargs, ) @@ -839,6 +842,7 @@ def has_pack_quantized_config(config): ckpt_path, auto_model_module, model_type=hf_config.model_type, + distribution="sequential", device_map=device_map, **model_kwargs2, ) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index b42ea08364c..4b472f1f604 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -94,7 +94,7 @@ EagleOfflineDataCollator, OfflineSupervisedDataset, ) -from modelopt.torch.utils import env, print_rank_0 +from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import ( create_forward_loop, get_dataset_dataloader, @@ -112,20 +112,12 @@ ) 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.moe import ( - apply_distributed_support, - decide_distribution_plan, - is_expert_parallel_capable, - load_and_prepare_ep, -) +from modelopt.torch.utils.moe import is_expert_parallel_capable from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_non_model_files, ensure_local_checkpoint, ) -from modelopt.torch.utils.plugins.model_load_utils import ( - load_pretrained, - parallel_load_and_prepare_fsdp2, -) +from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader @@ -344,14 +336,9 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: def _decide_expert_parallel(args: argparse.Namespace): """Return ``(plan, hf_config)`` when this run should be expert-parallel, else ``(None, None)``. - The model's own expert-parallel facts are applied to the config BEFORE anything is decided: for - several models that is what SUPPLIES ``base_model_ep_plan`` -- stock transformers ships one for - some MoEs and not others -- so deciding first would read an EP-capable model as dense and - replicate every expert on every rank. + ``resolve_distribution_plan`` decides, as ``load_pretrained`` would on its own; deciding first + lets a run that loads the model whole go through ``get_model``'s placement. """ - mo_session = getattr(args, "mo_session", None) - if mo_session is None or not mo_session.distributed: - return None, None # --use_fsdp2 asks for the standalone FSDP2 loader explicitly; the EP # auto-decision below has no way to know that, so without this it can pick EP anyway # (MODELOPT_EP_DEGREE set explicitly, or "auto" resolving non-trivially for the model) @@ -359,45 +346,19 @@ def _decide_expert_parallel(args: argparse.Namespace): if args.use_fsdp2: return None, None - hf_config = AutoConfig.from_pretrained( + plan, hf_config = resolve_distribution_plan( args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) - hf_config = apply_distributed_support( - hf_config.model_type, - hf_config, - ckpt_path=args.local_checkpoint_path, - trust_remote_code=args.trust_remote_code, - ) - plan = decide_distribution_plan( - hf_config, - mo_session.world_size, - ep_size=env.ep_degree(), - fsdp_size=env.fsdp_size(), - ) - if mo_session.is_main: - print(f"[ep] decided {plan!r}", flush=True) - - if plan.is_expert_parallel: + if plan is None or plan.is_expert_parallel: # None: not a distributed run return plan, hf_config - # Past this point load_model replicates the model on every rank, so any sharding the plan asked - # for would be dropped. That runs, and wastes the cluster, so say so rather than let it look like - # it worked. - ep_capable = is_expert_parallel_capable(hf_config) if plan.is_fsdp_dense: - # Only the expert-parallel loader shards a dense part by itself. + # FSDP2 alone runs only under --use_fsdp2, which checks what this script cannot do with it. raise RuntimeError( - f"The plan shards the dense part over {plan.fsdp_size} ranks but leaves ep_size=1, and " - "without expert parallelism nothing shards it -- every rank would hold the whole model. " + f"The plan shards the dense part over {plan.fsdp_size} ranks but leaves ep_size=1. " "Pass --use_fsdp2 to shard it with FSDP2" - + (", or set MODELOPT_EP_DEGREE>1" if ep_capable else "") + + (", or set MODELOPT_EP_DEGREE>1" if is_expert_parallel_capable(hf_config) else "") + ", or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU." ) - if ep_capable: - raise RuntimeError( - f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at " - f"world_size={mo_session.world_size} -- every rank would hold the whole MoE. Set " - "MODELOPT_EP_DEGREE>1, or pass --use_fsdp2 to shard it with FSDP2 instead." - ) return None, None @@ -406,29 +367,25 @@ def load_model(args: argparse.Namespace): calibration_only = False _ep_plan, _ep_config = _decide_expert_parallel(args) if _ep_plan is not None: - full_model = load_and_prepare_ep( + full_model = load_pretrained( args.local_checkpoint_path, - args.mo_session, - _ep_plan.ep_size, + distribution=_ep_plan, + hf_config=_ep_config, trust_remote_code=args.trust_remote_code, attn_implementation=args.attn_implementation, - hf_config=_ep_config, - fsdp_dense=_ep_plan.is_fsdp_dense, ) elif args.use_fsdp2: hf_config = AutoConfig.from_pretrained( args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) validate_fsdp2_supported(args, hf_config) - full_model = parallel_load_and_prepare_fsdp2( + full_model = load_pretrained( args.local_checkpoint_path, - args.dist_state.device, - args.dist_state.rank, - args.dist_state.world_size, + distribution="fsdp2", + hf_config=hf_config, trust_remote_code=args.trust_remote_code, cpu_offload=args.cpu_offload, 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( @@ -471,6 +428,7 @@ def load_model(args: argparse.Namespace): args.local_checkpoint_path, AutoModelForCausalLM, model_type=hf_config.model_type, + distribution="sequential", **model_kwargs, ) calibration_only = True diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index a4dc8797f00..852fee55316 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -15,10 +15,10 @@ """Expert-parallel PTQ wiring in ``examples/hf_ptq``. -Under torchrun, hf_ptq.py applies the model's spec to its config, decides a distribution plan from -the ``MODELOPT_EP_DEGREE`` / ``MODELOPT_FSDP_SIZE`` knobs, and loads the model through -``load_and_prepare_ep`` when that plan shards the experts. Every rank then calibrates on its own -shard of the data, and a dense model is pinned to the rank's own GPU instead of ``"auto"``. +Under torchrun, hf_ptq.py decides a distribution plan with ``resolve_distribution_plan`` and loads +the model through ``load_pretrained``: with the plan when it shards the experts, as FSDP2 under +``--use_fsdp2``, and otherwise whole through ``get_model``. Every rank then calibrates on its own +shard of the data, and a whole model is pinned to the rank's own GPU instead of ``"auto"``. Nothing here needs a process group or a checkpoint: an initialized ``ModelOptSession`` stands in for the torchrun session, the configs are tiny real transformers configs served from memory, and @@ -34,17 +34,15 @@ import pytest import torch -import transformers from _test_utils.examples.hf_ptq_example_utils import example_utils from _test_utils.torch.transformers_models import get_tiny_tokenizer -from packaging.version import Version from torch.utils.data import DataLoader -from transformers import DeepseekV3Config, LlamaConfig, Qwen3MoeConfig +from transformers import DeepseekV3Config, LlamaConfig from modelopt.torch.models import get_spec -from modelopt.torch.utils import env -from modelopt.torch.utils.distributed import ModelOptSession -from modelopt.torch.utils.moe import is_expert_parallel_capable, load_and_prepare_ep +from modelopt.torch.utils.distributed import ModelOptSession, get_active_mo_session +from modelopt.torch.utils.plugins import model_load_utils +from modelopt.torch.utils.plugins.model_load_utils import load_pretrained _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -83,6 +81,34 @@ def _no_distribution_knobs(monkeypatch): monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False) +@pytest.fixture(autouse=True) +def _checkpoint_is_local(monkeypatch): + """The checkpoint path names no files; resolving it must not reach the Hub.""" + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda path, group=None: (None, path) + ) + + +@pytest.fixture +def torchrun(): + """Enter the active session as rank ``rank`` of a ``world_size``-rank torchrun launch.""" + active = get_active_mo_session(required=False) + + def enter(world_size, rank=0): + active.initialize( + distributed=world_size > 1, + rank=rank, + world_size=world_size, + local_rank=rank, + device=torch.device(f"cuda:{rank}"), + is_main=rank == 0, + ) + return active + + yield enter + active._reset() + + def _session(world_size, rank=0, distributed=True): """The session modelopt_session() yields on ``rank`` of a ``world_size``-rank launch.""" session = ModelOptSession() @@ -117,11 +143,6 @@ def _serve(monkeypatch, hf_ptq, config): return config -def _qwen3_moe(num_experts): - """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself.""" - return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) - - def _deepseek_v3(n_routed_experts): """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable.""" config = DeepseekV3Config(n_routed_experts=n_routed_experts, **_TINY) @@ -129,176 +150,119 @@ def _deepseek_v3(n_routed_experts): return config -# ----------------------------------------------------------------------------------- env knobs - - -@pytest.mark.parametrize( - ("knob", "name"), - [(env.ep_degree, "MODELOPT_EP_DEGREE"), (env.fsdp_size, "MODELOPT_FSDP_SIZE")], -) -def test_distribution_knobs_default_to_auto(monkeypatch, knob, name): - """Unset and exported-but-empty both mean "let the plan decide".""" - assert knob() == "auto" - monkeypatch.setenv(name, "") - assert knob() == "auto" - - -@pytest.mark.parametrize( - ("knob", "name", "other"), - [ - (env.ep_degree, "MODELOPT_EP_DEGREE", env.fsdp_size), - (env.fsdp_size, "MODELOPT_FSDP_SIZE", env.ep_degree), - ], -) -def test_distribution_knobs_are_read_when_called(monkeypatch, knob, name, other): - """A value snapshotted at import would miss what the launcher sets later; and each knob reads - only its own variable.""" - monkeypatch.setenv(name, "2") - assert knob() == "2" - monkeypatch.setenv(name, "8") - assert knob() == "8" - assert other() == "auto" - - # -------------------------------------------------------------------------------- the decision @pytest.mark.parametrize("single_process_session", [False, True]) def test_a_single_process_run_is_never_expert_parallel(hf_ptq, monkeypatch, single_process_session): - """With no distributed session there is nothing to shard over, even for an EP-capable MoE - (which, decided at world_size=1, would otherwise be rejected as an unsharded MoE).""" + """With no distributed session there is nothing to shard over, even for an EP-capable MoE.""" session = _session(world_size=1, distributed=False) if single_process_session else None _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) assert hf_ptq._decide_expert_parallel(_args(session)) == (None, None) -def test_use_fsdp2_keeps_the_standalone_fsdp2_loader(hf_ptq, monkeypatch): +def test_use_fsdp2_keeps_the_standalone_fsdp2_loader(hf_ptq, monkeypatch, torchrun): """--use_fsdp2 names its loader, so even an explicit EP degree for an EP-capable MoE must not divert the run to the expert-parallel one.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) - args = _args(_session(world_size=2), use_fsdp2=True) + args = _args(torchrun(world_size=2), use_fsdp2=True) assert hf_ptq._decide_expert_parallel(args) == (None, None) -@pytest.mark.skipif( - Version(transformers.__version__) < Version("5.7"), - reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", -) -@pytest.mark.parametrize("rank", [0, 1]) -def test_an_ep_capable_moe_is_planned_expert_parallel(hf_ptq, monkeypatch, capsys, rank): - """Auto picks the largest EP degree dividing both the world and the expert count, hands back a - config the EP loader can shard, and only rank 0 reports the decision.""" - _serve(monkeypatch, hf_ptq, _qwen3_moe(num_experts=8)) - - plan, hf_config = hf_ptq._decide_expert_parallel(_args(_session(world_size=4, rank=rank))) - - assert (plan.world_size, plan.ep_size, plan.edp) == (4, 4, 1) - assert not plan.is_fsdp_dense - assert is_expert_parallel_capable(hf_config) - assert ("[ep] decided" in capsys.readouterr().out) == (rank == 0) - - -def test_the_spec_is_applied_before_the_decision(hf_ptq, monkeypatch): - """DeepSeek-V3's plan comes only from its spec. Deciding before applying it would read the MoE - as dense; applied first, it shards, and the config handed on carries the spec's plan.""" - _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) - - plan, hf_config = hf_ptq._decide_expert_parallel(_args(_session(world_size=4))) - - assert plan.ep_size == 4 - assert hf_config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan - - -@pytest.mark.parametrize( - ("ep_degree", "experts"), - [(None, 3), ("1", 8)], - ids=["auto_finds_no_degree", "explicit_degree_one"], -) -def test_an_ep_capable_moe_left_unsharded_is_an_error(hf_ptq, monkeypatch, ep_degree, experts): - """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With - the dense part replicated, every rank would hold the whole MoE: that must fail, not run.""" - if ep_degree is not None: - monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) - _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=experts)) - - with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): - hf_ptq._decide_expert_parallel(_args(_session(world_size=2))) - - @pytest.mark.parametrize( ("config", "ep_capable"), [(_deepseek_v3(n_routed_experts=3), True), (LlamaConfig(**_TINY), False)], ids=["moe_no_ep_degree", "dense"], ) -def test_dense_sharding_without_expert_parallelism_is_an_error( - hf_ptq, monkeypatch, config, ep_capable +def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2( + hf_ptq, monkeypatch, torchrun, config, ep_capable ): - """Only the EP loader shards a dense part by itself. Without EP the run would load the whole - model on every rank and drop the sharding the plan asked for, so it must fail and point at the - loader that can shard it.""" + """Without expert parallelism a plan that shards the dense part is an FSDP2 run, which this + script runs only under --use_fsdp2, the flag its FSDP2 checks hang off; so it must fail and point + there rather than load the model whole on every rank.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") _serve(monkeypatch, hf_ptq, config) with pytest.raises(RuntimeError, match="--use_fsdp2") as raised: - hf_ptq._decide_expert_parallel(_args(_session(world_size=2))) + hf_ptq._decide_expert_parallel(_args(torchrun(world_size=2))) assert ("MODELOPT_EP_DEGREE" in str(raised.value)) == ep_capable -def test_a_dense_model_under_torchrun_runs_data_parallel(hf_ptq, monkeypatch): - """A dense model has no experts to shard, so replicating it per rank is not an error.""" +def test_a_dense_model_under_torchrun_runs_data_parallel(hf_ptq, monkeypatch, torchrun): + """A dense model has no experts to shard, so it loads whole on every rank, through get_model.""" _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) - assert hf_ptq._decide_expert_parallel(_args(_session(world_size=4))) == (None, None) + assert hf_ptq._decide_expert_parallel(_args(torchrun(world_size=4))) == (None, None) -def test_the_env_knobs_override_the_automatic_plan(hf_ptq, monkeypatch): - """The decision honours both knobs; left on auto it would pick ep_size=4, dense replicated.""" - monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") - monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") - _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) +# ---------------------------------------------------------------------------------- load_model - plan, _ = hf_ptq._decide_expert_parallel(_args(_session(world_size=4))) - assert (plan.ep_size, plan.edp, plan.fsdp_size) == (2, 2, 4) +def _record_load_pretrained(hf_ptq, monkeypatch): + """Replace hf_ptq's load_pretrained and every other loader; return the recorded call.""" + calls = [] + def load_pretrained_stub(*args, **kwargs): + calls.append(inspect.signature(load_pretrained).bind(*args, **kwargs).arguments) + raise _LoaderReachedError -# ---------------------------------------------------------------------------------- load_model + monkeypatch.setattr(hf_ptq, "load_pretrained", load_pretrained_stub) + monkeypatch.setattr(hf_ptq, "get_model", lambda *a, **k: pytest.fail("whole-model loader ran")) + return calls -def test_load_model_hands_an_expert_parallel_plan_to_the_ep_loader(hf_ptq, monkeypatch): - """With an EP plan, load_model loads through load_and_prepare_ep -- with the session, the - decided degree and dense sharding, and the spec-prepared config -- and through no other - loader.""" +def test_load_model_loads_an_expert_parallel_plan_through_load_pretrained( + hf_ptq, monkeypatch, torchrun +): + """With an EP plan, load_model hands load_pretrained the plan -- the decided degree and dense + sharding -- and the spec-prepared config it was decided on.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") served = _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) - session = _session(world_size=4) - calls = [] + calls = _record_load_pretrained(hf_ptq, monkeypatch) + args = _args(torchrun(world_size=4), trust_remote_code=True, attn_implementation="eager") - def load_and_prepare_ep_stub(*args, **kwargs): - calls.append(inspect.signature(load_and_prepare_ep).bind(*args, **kwargs).arguments) - raise _LoaderReachedError + with pytest.raises(_LoaderReachedError): + hf_ptq.load_model(args) - monkeypatch.setattr(hf_ptq, "load_and_prepare_ep", load_and_prepare_ep_stub) - monkeypatch.setattr(hf_ptq, "get_model", lambda *a, **k: pytest.fail("dense loader ran")) - monkeypatch.setattr( - hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *a, **k: pytest.fail("FSDP2 loader ran") + (call,) = calls + plan = call["distribution"] + assert call["ckpt_path"] == _CKPT + assert (plan.ep_size, plan.is_fsdp_dense) == (4, True) + assert call["kwargs"] == { + "hf_config": served, + "trust_remote_code": True, + "attn_implementation": "eager", + } + assert served.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + + +def test_use_fsdp2_loads_through_load_pretrained(hf_ptq, monkeypatch, torchrun): + """--use_fsdp2 asks load_pretrained for its FSDP2 load, with the CPU-offload choice and the + config its checks already read.""" + served = _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY)) + calls = _record_load_pretrained(hf_ptq, monkeypatch) + args = _args( + torchrun(world_size=2), + use_fsdp2=True, + cpu_offload=True, + specdec_offline_dataset=None, + low_memory_mode=False, ) with pytest.raises(_LoaderReachedError): - hf_ptq.load_model(_args(session, trust_remote_code=True, attn_implementation="eager")) + hf_ptq.load_model(args) (call,) = calls - assert call["ckpt_path"] == _CKPT - assert call["mo_session"] is session - assert call["ep_degree"] == 4 - assert call["fsdp_dense"] is True - assert call["trust_remote_code"] is True - assert call["attn_implementation"] == "eager" - assert call["hf_config"] is served - assert served.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + assert call["distribution"] == "fsdp2" + assert call["kwargs"] == { + "hf_config": served, + "trust_remote_code": False, + "cpu_offload": True, + "attn_implementation": None, + } # ------------------------------------------------------------------------- per-rank placement From d6af08eb677b43b981204b10f0f24b8d1f6bacd0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 17:25:21 +0000 Subject: [PATCH 11/57] Stub load_pretrained in the FSDP2 MTP-refusal test hf_ptq's FSDP2 load now goes through load_pretrained, so the test that checks the MTP refusal comes before it stubs that instead of parallel_load_and_prepare_fsdp2, which hf_ptq no longer imports. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 608a03e1beb..020067b2001 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -435,7 +435,7 @@ def test_fsdp2_rejects_mtp_before_distributed_loading(monkeypatch, tmp_path, has def distributed_load(*args, **kwargs): raise RuntimeError("distributed loader reached") - monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load) + monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load) error, message = ( (NotImplementedError, "Nemotron-H MTP") if has_mtp From 8d4d028fe8e094cf15a9b3d8bd0d76af4592c042 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 18:03:19 +0000 Subject: [PATCH 12/57] Answer the session's collectives in the hf_ptq wiring tests The plan is now decided through the session, whose decide is collective. These tests are about what gets decided, so each simulated rank acts as its group's first. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 852fee55316..0aebca7c3dd 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -90,8 +90,14 @@ def _checkpoint_is_local(monkeypatch): @pytest.fixture -def torchrun(): - """Enter the active session as rank ``rank`` of a ``world_size``-rank torchrun launch.""" +def torchrun(monkeypatch): + """Enter the active session as rank ``rank`` of a ``world_size``-rank torchrun launch. + + Every rank acts as its group's first, deciding for itself; how the ranks agree on what they + decide is tested with the session. + """ + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) active = get_active_mo_session(required=False) def enter(world_size, rank=0): From e83fe9a6d02df42dfadf545e5b2efac4b4a0d8fd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:45:57 +0000 Subject: [PATCH 13/57] Run hf_ptq without the session: own the process group, read rank and EP mesh where needed main() runs inside distributed_run, which sets the process group up on a distributed launch and tears it down even when the run raises. The example's DistributedState reads rank and world from that group, the calibration data is sharded by it, and an expert-parallel model calibrates through distributed_forward_loop with the mesh its experts are sharded over (expert_parallel_mesh). Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 18 +- examples/hf_ptq/hf_ptq.py | 33 ++-- .../hf_ptq/test_hf_ptq_expert_parallel.py | 157 ++++++++---------- 3 files changed, 87 insertions(+), 121 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 3da7a745e48..b7bacc4707f 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -74,18 +74,16 @@ class DistributedState: def setup_distributed_args(args): """Initialize and attach ``args.dist_state`` (single-process if not distributed). - When the run already entered a :func:`modelopt_session` -- which it does for any distributed - launch, because the expert-parallel loaders need the session object -- the state is a view of - that session rather than a second, independently-derived one. Two sources of rank/world would - be one too many. + A distributed launch has its process group set up by then (hf_ptq's ``distributed_run``), and the + state reads its rank and world from that group rather than deriving them a second time. """ - mo_session = getattr(args, "mo_session", None) - if mo_session is not None and mo_session.distributed: + if dist_utils.is_distributed_run(): + rank = dist_utils.rank() args.dist_state = DistributedState( - rank=mo_session.rank, - world_size=mo_session.world_size, - device=mo_session.device, - is_main=mo_session.is_main, + rank=rank, + world_size=dist_utils.size(), + device=torch.device(f"cuda:{dist_utils.local_rank()}"), + is_main=rank == 0, ) elif getattr(args, "use_fsdp2", False): # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 4b472f1f604..3cc00dd4993 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -14,7 +14,6 @@ # limitations under the License. import argparse -import contextlib import os import random import time @@ -103,16 +102,14 @@ ) from modelopt.torch.utils.distributed import ( DistributionPlan, - ModelOptSession, distributed_forward_loop, - is_distributed_launch, - modelopt_session, + distributed_run, set_distribution_plan, shard_dataloader, ) 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.moe import is_expert_parallel_capable +from modelopt.torch.utils.moe import expert_parallel_mesh, is_expert_parallel_capable from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_non_model_files, ensure_local_checkpoint, @@ -203,7 +200,6 @@ def make_calib_dataloader( tokenizer: PreTrainedTokenizerBase | None, device: torch.device, model_type: str | None, - mo_session: ModelOptSession | None, autoquant_gradient_recipe: bool = False, ) -> tuple[DataLoader | _DeviceDataLoader, str | None]: calib_dataloader = None @@ -280,8 +276,7 @@ def make_calib_dataloader( include_labels=include_labels, ) if ( - mo_session is not None - and mo_session.distributed + args.dist_state.world_size > 1 and calib_dataloader is not None and isinstance(calib_dataloader, DataLoader) ): @@ -290,7 +285,7 @@ def make_calib_dataloader( # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2 # shards the params but still feeds each rank its own data. calib_dataloader = shard_dataloader( - calib_dataloader, mo_session.rank, mo_session.world_size + calib_dataloader, args.dist_state.rank, args.dist_state.world_size ) return calib_dataloader, first_text_speech_dataset @@ -609,7 +604,7 @@ def mono_quantize( # quantization target; the full forward still exercises that component's quantizers. if args.calib_with_images: calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) - elif getattr(getattr(args, "mo_session", None), "ep_mesh", None) is not None: + elif (ep_mesh := expert_parallel_mesh(full_model)) is not None: # An expert-parallel run has to calibrate through distributed_forward_loop, not the # plain one: it is the only caller of _reinstall_a2a_forwards, which re-plants the # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on @@ -621,7 +616,7 @@ def mono_quantize( # It also restores the use_cache guard the plain loop does not apply. Other # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. calibrate_loop = distributed_forward_loop( - language_model, calib_dataloader, args.mo_session + language_model, calib_dataloader, device=args.dist_state.device, ep_mesh=ep_mesh ) else: calibrate_loop = create_forward_loop( @@ -1131,7 +1126,6 @@ def quantize_main( tokenizer, device, model_type, - args.mo_session, autoquant_gradient_recipe=( aq_config is not None and aq_config.auto_quantize_method == "gradient" ), @@ -1581,17 +1575,10 @@ def main(args: argparse.Namespace): random.seed(RAND_SEED) np.random.seed(RAND_SEED) - # A distributed launch runs inside a ModelOptSession: it owns the process group and is what the - # expert-parallel loader reads its rank and device from. Entered here rather than inside - # load_model so cleanup still runs if loading raises -- otherwise the surviving ranks block on - # the first collective until the NCCL timeout. - with contextlib.ExitStack() as _stack: - if is_distributed_launch() or args.use_fsdp2: - args.mo_session = _stack.enter_context( - modelopt_session(timeout=timedelta(hours=3), device=args.device) - ) - else: - args.mo_session = None + # A distributed launch owns its process group for the whole run. Entered here rather than + # inside load_model so cleanup still runs if loading raises -- otherwise the surviving ranks + # block on the first collective until the NCCL timeout. + with distributed_run(timeout=timedelta(hours=3)): _main_inner(args) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 0aebca7c3dd..f69fbdfa541 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -20,9 +20,9 @@ ``--use_fsdp2``, and otherwise whole through ``get_model``. Every rank then calibrates on its own shard of the data, and a whole model is pinned to the rank's own GPU instead of ``"auto"``. -Nothing here needs a process group or a checkpoint: an initialized ``ModelOptSession`` stands in -for the torchrun session, the configs are tiny real transformers configs served from memory, and -the loaders are replaced by stand-ins that record how they were called. +Nothing here needs a process group or a checkpoint: the process group's rank and size are answered +the way a torchrun launch would, the configs are tiny real transformers configs served from memory, +and the loaders are replaced by stand-ins that record how they were called. """ import contextlib @@ -39,8 +39,8 @@ from torch.utils.data import DataLoader from transformers import DeepseekV3Config, LlamaConfig +import modelopt.torch.utils.distributed as dist from modelopt.torch.models import get_spec -from modelopt.torch.utils.distributed import ModelOptSession, get_active_mo_session from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -91,48 +91,35 @@ def _checkpoint_is_local(monkeypatch): @pytest.fixture def torchrun(monkeypatch): - """Enter the active session as rank ``rank`` of a ``world_size``-rank torchrun launch. + """Run as rank ``rank`` of a ``world_size``-rank torchrun launch; returns its ``dist_state``. Every rank acts as its group's first, deciding for itself; how the ranks agree on what they - decide is tested with the session. + decide is tested with the collectives. """ monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) - active = get_active_mo_session(required=False) def enter(world_size, rank=0): - active.initialize( - distributed=world_size > 1, - rank=rank, - world_size=world_size, - local_rank=rank, - device=torch.device(f"cuda:{rank}"), - is_main=rank == 0, - ) - return active - - yield enter - active._reset() - - -def _session(world_size, rank=0, distributed=True): - """The session modelopt_session() yields on ``rank`` of a ``world_size``-rank launch.""" - session = ModelOptSession() - session.initialize( - distributed=distributed, - rank=rank, - world_size=world_size, - local_rank=rank, - device=torch.device(f"cuda:{rank}"), - is_main=rank == 0, + monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) + monkeypatch.setattr(dist, "size", lambda group=None: world_size) + monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) + monkeypatch.setattr(dist, "local_rank", lambda: rank) + return _dist_state(world_size, rank) + + return enter + + +def _dist_state(world_size, rank=0): + """What setup_distributed_args leaves on ``rank`` of a ``world_size``-rank launch.""" + return example_utils.DistributedState( + rank=rank, world_size=world_size, device=torch.device(f"cuda:{rank}"), is_main=rank == 0 ) - return session -def _args(session, **overrides): +def _args(dist_state, **overrides): """What load_model reads before it picks a loader, as main() leaves it.""" args = SimpleNamespace( - mo_session=session, + dist_state=dist_state, use_fsdp2=False, recipe=None, local_checkpoint_path=_CKPT, @@ -159,13 +146,15 @@ def _deepseek_v3(n_routed_experts): # -------------------------------------------------------------------------------- the decision -@pytest.mark.parametrize("single_process_session", [False, True]) -def test_a_single_process_run_is_never_expert_parallel(hf_ptq, monkeypatch, single_process_session): - """With no distributed session there is nothing to shard over, even for an EP-capable MoE.""" - session = _session(world_size=1, distributed=False) if single_process_session else None +@pytest.mark.parametrize("one_rank_group", [False, True], ids=["no_process_group", "one_rank"]) +def test_a_single_process_run_is_never_expert_parallel( + hf_ptq, monkeypatch, torchrun, one_rank_group +): + """With no distributed run there is nothing to shard over, even for an EP-capable MoE.""" + dist_state = torchrun(world_size=1) if one_rank_group else _dist_state(world_size=1) _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) - assert hf_ptq._decide_expert_parallel(_args(session)) == (None, None) + assert hf_ptq._decide_expert_parallel(_args(dist_state)) == (None, None) def test_use_fsdp2_keeps_the_standalone_fsdp2_loader(hf_ptq, monkeypatch, torchrun): @@ -305,18 +294,21 @@ def load_pretrained(ckpt_path, auto_class, **kwargs): assert placed == [device_map] -# -------------------------------------------------------------------- session and dist_state +# ------------------------------------------------------------------ process group and dist_state @pytest.mark.parametrize("use_fsdp2", [False, True]) -def test_dist_state_mirrors_the_distributed_session(monkeypatch, use_fsdp2): - """Under torchrun dist_state is read from the session that owns the process group, even with +def test_dist_state_reads_the_live_process_group(monkeypatch, use_fsdp2): + """Under torchrun dist_state is read from the process group main() set up, even with --use_fsdp2, whose own branch would initialize a second group.""" monkeypatch.setattr( example_utils.dist_utils, "setup", lambda *a, **k: pytest.fail("second process group") ) - session = _session(world_size=4, rank=1) - args = SimpleNamespace(mo_session=session, use_fsdp2=use_fsdp2, device="cuda") + monkeypatch.setattr(example_utils.dist_utils, "is_distributed_run", lambda: True) + monkeypatch.setattr(example_utils.dist_utils, "rank", lambda group=None: 1) + monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: 4) + monkeypatch.setattr(example_utils.dist_utils, "local_rank", lambda: 1) + args = SimpleNamespace(use_fsdp2=use_fsdp2, device="cuda") example_utils.setup_distributed_args(args) @@ -325,15 +317,15 @@ def test_dist_state_mirrors_the_distributed_session(monkeypatch, use_fsdp2): ) -def test_a_single_process_session_leaves_fsdp2_to_set_up_its_group(monkeypatch): - """A one-rank --use_fsdp2 launch enters a session that owns no process group, so the FSDP2 - branch must still create one.""" +def test_a_single_process_fsdp2_run_sets_up_its_own_group(monkeypatch): + """A one-rank --use_fsdp2 launch has no process group from main(), so the FSDP2 branch must + still create one.""" setups = [] monkeypatch.setattr(example_utils.dist_utils, "setup", lambda timeout=None: setups.append(0)) monkeypatch.setattr(example_utils.dist_utils, "rank", lambda group=None: 0) monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: 1) monkeypatch.setattr(example_utils.dist_utils, "local_rank", lambda: 0) - args = SimpleNamespace(mo_session=_session(world_size=1, distributed=False), use_fsdp2=True) + args = SimpleNamespace(use_fsdp2=True) example_utils.setup_distributed_args(args) @@ -341,49 +333,34 @@ def test_a_single_process_session_leaves_fsdp2_to_set_up_its_group(monkeypatch): assert args.dist_state.world_size == 1 -@pytest.mark.parametrize( - ("world_size", "use_fsdp2", "enters_session"), - [("2", False, True), ("1", True, True), (None, False, False)], - ids=["torchrun", "fsdp2", "single_process"], -) -def test_main_runs_a_distributed_launch_inside_a_session( - hf_ptq, monkeypatch, world_size, use_fsdp2, enters_session -): - """main() opens the session around the whole run, with a collective timeout above NCCL's - 30-minute default, and closes it even when the run raises -- or the other ranks would wait in - their next collective until that timeout.""" - if world_size is None: - monkeypatch.delenv("WORLD_SIZE", raising=False) - else: - monkeypatch.setenv("WORLD_SIZE", world_size) - session = _session(world_size=2) +def test_main_owns_the_process_group_around_the_whole_run(hf_ptq, monkeypatch): + """main() runs inside distributed_run, with a collective timeout above NCCL's 30-minute + default, and leaves it even when the run raises -- or the other ranks would wait in their next + collective until that timeout. distributed_run itself does nothing in a single process.""" events, timeouts = [], [] @contextlib.contextmanager - def modelopt_session(timeout=None, device=None): + def distributed_run(timeout=None): timeouts.append(timeout) - events.append(("enter", device)) + events.append("enter") try: - yield session + yield finally: - events.append(("exit",)) + events.append("exit") def main_inner(args): - events.append(("run", args.mo_session)) + events.append("run") raise _RunFailedError monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(hf_ptq, "modelopt_session", modelopt_session) + monkeypatch.setattr(hf_ptq, "distributed_run", distributed_run) monkeypatch.setattr(hf_ptq, "_main_inner", main_inner) with pytest.raises(_RunFailedError): - hf_ptq.main(SimpleNamespace(use_fsdp2=use_fsdp2, device="cuda")) + hf_ptq.main(SimpleNamespace(use_fsdp2=False, device="cuda")) - if enters_session: - assert events == [("enter", "cuda"), ("run", session), ("exit",)] - assert timeouts[0] > timedelta(minutes=30) - else: - assert events == [("run", None)] + assert events == ["enter", "run", "exit"] + assert timeouts[0] > timedelta(minutes=30) # --------------------------------------------------------------------------------- calibration @@ -404,27 +381,32 @@ def test_distributed_calibration_gives_each_rank_its_own_shard(hf_ptq, monkeypat ) tokenizer = get_tiny_tokenizer() - def calib_loader(session): + def calib_loader(world_size, rank=0): + args.dist_state = _dist_state(world_size, rank) built, _ = hf_ptq.make_calib_dataloader( - args, None, None, tokenizer, torch.device("cpu"), "llama", session + args, None, None, tokenizer, torch.device("cpu"), "llama" ) return built - shards = [[int(b) for b in calib_loader(_session(world_size=2, rank=r))] for r in range(2)] + shards = [[int(b) for b in calib_loader(world_size=2, rank=r)] for r in range(2)] assert len(shards[0]) == len(shards[1]) == 4 assert sorted(shards[0] + shards[1]) == list(range(8)) - assert calib_loader(None) is loader + assert calib_loader(world_size=1) is loader -def _calibration_loop_choice(hf_ptq, monkeypatch, session): - """Run mono_quantize under ``session`` and return which calibration loop it built.""" +def _calibration_loop_choice(hf_ptq, monkeypatch, ep_mesh): + """Run mono_quantize on a model whose experts are sharded over ``ep_mesh`` (``None``: not + expert-parallel) and return which calibration loop it built.""" model = torch.nn.Linear(2, 2) chosen, quantized = [], [] monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False) monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: True) + monkeypatch.setattr(hf_ptq, "expert_parallel_mesh", lambda full_model: ep_mesh) monkeypatch.setattr( - hf_ptq, "distributed_forward_loop", lambda *a: chosen.append("distributed") or object() + hf_ptq, + "distributed_forward_loop", + lambda *a, **k: chosen.append(("distributed", k["ep_mesh"])) or object(), ) monkeypatch.setattr( hf_ptq, "create_forward_loop", lambda **k: chosen.append("plain") or object() @@ -435,7 +417,7 @@ def _calibration_loop_choice(hf_ptq, monkeypatch, session): args = SimpleNamespace( qformat="fp8", calib_with_images=False, - mo_session=session, + dist_state=_dist_state(world_size=2), layerwise_export=False, specdec_offline_dataset=None, ) @@ -458,13 +440,12 @@ def _calibration_loop_choice(hf_ptq, monkeypatch, session): def test_expert_parallel_calibration_runs_through_the_distributed_forward_loop(hf_ptq, monkeypatch): """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion clobbers, so an expert-parallel run must calibrate through it rather than the plain loop.""" - session = _session(world_size=2) - session.ep_mesh = object() # what load_and_prepare_ep records + ep_mesh = object() # the mesh load_and_prepare_ep loaded the experts over - assert _calibration_loop_choice(hf_ptq, monkeypatch, session) == ["distributed"] + assert _calibration_loop_choice(hf_ptq, monkeypatch, ep_mesh) == [("distributed", ep_mesh)] def test_fsdp2_calibration_keeps_the_plain_forward_loop(hf_ptq, monkeypatch): """Without expert parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry, as on main.""" - assert _calibration_loop_choice(hf_ptq, monkeypatch, _session(world_size=2)) == ["plain"] + assert _calibration_loop_choice(hf_ptq, monkeypatch, ep_mesh=None) == ["plain"] From fab91733cac21a648cb9087fd52247b841ab5089 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:54:31 +0000 Subject: [PATCH 14/57] Decide hf_ptq's distribution once and read the plan wherever it distributes hf_ptq.py decides the plan up front with resolve_distribution_plan, --use_fsdp2 asking it for FSDP2 without expert parallelism, and loads every model by it through load_pretrained, get_model's whole-model loads and the low-memory load included. A plan that FSDP2-shards the whole model still needs --use_fsdp2, whose checks now run before anything is set up. Everything after reads the plan back from the model instead of the flag or the process group: the calibration device, the data sharding, the expert-parallel calibration loop, and get_model's pinning of each rank's whole model to its own GPU. The one-rank process group FSDP2 shards through is set up for the plan, right before the load, and torn down with the run; a distributed launch's group stays distributed_run's. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 44 +++-- examples/hf_ptq/hf_ptq.py | 93 ++++----- tests/examples/hf_ptq/test_hf_ptq_args.py | 28 +-- .../hf_ptq/test_hf_ptq_expert_parallel.py | 181 +++++++++++++----- .../hf_ptq/test_hf_ptq_vision_quantization.py | 6 + 5 files changed, 205 insertions(+), 147 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index b7bacc4707f..3bf31dcca3c 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -53,6 +53,7 @@ prepare_model_for_loading, ) from modelopt.torch.utils import distributed as dist_utils +from modelopt.torch.utils.distributed import DistributionPlan from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -85,24 +86,28 @@ def setup_distributed_args(args): device=torch.device(f"cuda:{dist_utils.local_rank()}"), is_main=rank == 0, ) - elif getattr(args, "use_fsdp2", False): + else: + args.dist_state = DistributedState(rank=0, world_size=1, device=args.device, is_main=True) + + +def setup_fsdp2_process_group(plan: DistributionPlan) -> None: + """Set up the one-rank process group a plan that FSDP2-shards the model on one rank needs. + + FSDP2 shards through a process group even on one rank; a distributed launch already has its + group (hf_ptq's ``distributed_run``). + """ + if plan.is_fsdp_dense and not dist_utils.is_initialized(): # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can # exceed it, and PyTorch 2.8 has no per-call barrier() timeout (must be set at PG creation). dist_utils.setup(timeout=timedelta(hours=2)) - rank = dist_utils.rank() - args.dist_state = DistributedState( - rank=rank, - world_size=dist_utils.size(), - device=torch.device(f"cuda:{dist_utils.local_rank()}"), - is_main=rank == 0, - ) - else: - args.dist_state = DistributedState(rank=0, world_size=1, device=args.device, is_main=True) -def cleanup_distributed(args): - """Destroy the process group if ``--use_fsdp2`` set it up.""" - if getattr(args, "use_fsdp2", False): +def cleanup_distributed(): + """Destroy the one-rank process group :func:`setup_fsdp2_process_group` set up, if any. + + A distributed launch's group is ``distributed_run``'s to destroy. + """ + if not dist_utils.is_distributed_launch(): dist_utils.cleanup() @@ -572,6 +577,7 @@ def _fmt_max_memory(max_memory: dict) -> str: def get_model( ckpt_path, device="cuda", + distribution: DistributionPlan | None = None, gpu_mem_percentage=0.8, trust_remote_code=False, use_seq_device_map=False, @@ -581,6 +587,8 @@ def get_model( max_gpu_memory_gb=None, ): print(f"Initializing model from {ckpt_path}") + # The plan to load the whole model by on every rank; a single process's by default. + plan = distribution or DistributionPlan() _disk_offload = offload_folder is not None if _disk_offload and max_cpu_memory_gb is None: @@ -594,7 +602,7 @@ def get_model( device_map = "auto" if device == "cpu": device_map = "cpu" - elif ":" in str(device) and dist_utils.size() > 1: + elif ":" in str(device) and plan.world_size > 1: # Under torchrun each rank asks for its own "cuda:". Do not fall through to # "auto": accelerate runs its placement heuristic INDEPENDENTLY in every rank, and on a node # whose GPUs are all idle every rank scores them identically and picks the SAME card. All ranks' @@ -692,7 +700,7 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, - distribution="sequential", + distribution=plan, device_map=device_map, **model_kwargs, ) @@ -704,7 +712,7 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, - distribution="sequential", + distribution=plan, device_map="auto", trust_remote_code=trust_remote_code, dtype="auto", @@ -731,7 +739,7 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, - distribution="sequential", + distribution=plan, device_map="cpu" if device == "cpu" else "sequential", **model_kwargs, ) @@ -840,7 +848,7 @@ def has_pack_quantized_config(config): ckpt_path, auto_model_module, model_type=hf_config.model_type, - distribution="sequential", + distribution=plan, device_map=device_map, **model_kwargs2, ) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 3cc00dd4993..3d5990fc5aa 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -49,11 +49,11 @@ run_nemotron_vl_preview, save_source_config, setup_distributed_args, + setup_fsdp2_process_group, validate_fsdp2_supported, ) from torch.utils.data import DataLoader from transformers import ( - AutoConfig, AutoModelForCausalLM, AutoProcessor, PreTrainedTokenizer, @@ -101,15 +101,14 @@ get_supported_datasets, ) from modelopt.torch.utils.distributed import ( - DistributionPlan, distributed_forward_loop, distributed_run, - set_distribution_plan, + distribution_plan, shard_dataloader, ) 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.moe import expert_parallel_mesh, is_expert_parallel_capable +from modelopt.torch.utils.moe import expert_parallel_mesh from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_non_model_files, ensure_local_checkpoint, @@ -275,18 +274,13 @@ def make_calib_dataloader( device=device, include_labels=include_labels, ) - if ( - args.dist_state.world_size > 1 - and calib_dataloader is not None - and isinstance(calib_dataloader, DataLoader) - ): + world_size = distribution_plan(language_model).world_size + if world_size > 1 and calib_dataloader is not None and isinstance(calib_dataloader, DataLoader): # Every distributed calibration path is data-parallel over the world, so each rank calibrates # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2 # shards the params but still feeds each rank its own data. - calib_dataloader = shard_dataloader( - calib_dataloader, args.dist_state.rank, args.dist_state.world_size - ) + calib_dataloader = shard_dataloader(calib_dataloader, args.dist_state.rank, world_size) return calib_dataloader, first_text_speech_dataset @@ -328,68 +322,57 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: ) -def _decide_expert_parallel(args: argparse.Namespace): - """Return ``(plan, hf_config)`` when this run should be expert-parallel, else ``(None, None)``. +def _resolve_distribution(args: argparse.Namespace): + """Decide how this run distributes the model: ``(plan, hf_config)``. - ``resolve_distribution_plan`` decides, as ``load_pretrained`` would on its own; deciding first - lets a run that loads the model whole go through ``get_model``'s placement. + ``resolve_distribution_plan`` decides, and --use_fsdp2 asks it for FSDP2 and no expert + parallelism. Without --use_fsdp2 a plan that FSDP2-shards the whole model is refused: the + FSDP2 run is the flag's, which checks what this script cannot do with it. """ - # --use_fsdp2 asks for the standalone FSDP2 loader explicitly; the EP - # auto-decision below has no way to know that, so without this it can pick EP anyway - # (MODELOPT_EP_DEGREE set explicitly, or "auto" resolving non-trivially for the model) - # and load_model's `elif args.use_fsdp2:` branch is never reached. - if args.use_fsdp2: - return None, None - plan, hf_config = resolve_distribution_plan( - args.local_checkpoint_path, trust_remote_code=args.trust_remote_code + args.local_checkpoint_path, + trust_remote_code=args.trust_remote_code, + ep_size=1 if args.use_fsdp2 else None, + dense="fsdp" if args.use_fsdp2 else None, ) - if plan is None or plan.is_expert_parallel: # None: not a distributed run - return plan, hf_config - if plan.is_fsdp_dense: - # FSDP2 alone runs only under --use_fsdp2, which checks what this script cannot do with it. + if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2: raise RuntimeError( - f"The plan shards the dense part over {plan.fsdp_size} ranks but leaves ep_size=1. " - "Pass --use_fsdp2 to shard it with FSDP2" - + (", or set MODELOPT_EP_DEGREE>1" if is_expert_parallel_capable(hf_config) else "") - + ", or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU." + f"{plan!r} FSDP2-shards the model, which this script does under --use_fsdp2. Pass it, " + "or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU; an MoE can instead shard its " + "experts (MODELOPT_EP_DEGREE above 1)." ) - return None, None + return plan, hf_config def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False - _ep_plan, _ep_config = _decide_expert_parallel(args) - if _ep_plan is not None: + plan, hf_config = _resolve_distribution(args) + if plan.is_fsdp_dense and not plan.is_expert_parallel: + validate_fsdp2_supported(args, hf_config) + setup_fsdp2_process_group(plan) + if plan.is_expert_parallel: full_model = load_pretrained( args.local_checkpoint_path, - distribution=_ep_plan, - hf_config=_ep_config, + distribution=plan, + hf_config=hf_config, trust_remote_code=args.trust_remote_code, attn_implementation=args.attn_implementation, ) - elif args.use_fsdp2: - hf_config = AutoConfig.from_pretrained( - args.local_checkpoint_path, trust_remote_code=args.trust_remote_code - ) - validate_fsdp2_supported(args, hf_config) + elif plan.is_fsdp_dense: full_model = load_pretrained( args.local_checkpoint_path, - distribution="fsdp2", + distribution=plan, hf_config=hf_config, trust_remote_code=args.trust_remote_code, cpu_offload=args.cpu_offload, attn_implementation=args.attn_implementation, ) - # 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, args.dist_state.device, + distribution=plan, gpu_mem_percentage=args.gpu_max_mem_percentage, trust_remote_code=args.trust_remote_code, use_seq_device_map=args.use_seq_device_map, @@ -413,9 +396,6 @@ def load_model(args: argparse.Namespace): with init_quantized_weights( quant_cfg, gpu_mem_percentage=args.gpu_max_mem_percentage, quant_gemm=False ): - hf_config = AutoConfig.from_pretrained( - args.local_checkpoint_path, trust_remote_code=args.trust_remote_code - ) model_kwargs = {"trust_remote_code": args.trust_remote_code} if args.attn_implementation is not None: model_kwargs["attn_implementation"] = args.attn_implementation @@ -423,7 +403,7 @@ def load_model(args: argparse.Namespace): args.local_checkpoint_path, AutoModelForCausalLM, model_type=hf_config.model_type, - distribution="sequential", + distribution=plan, **model_kwargs, ) calibration_only = True @@ -435,7 +415,7 @@ def load_model(args: argparse.Namespace): # model reports its own sub-config's type, which varies across transformers versions. model_type = hf_model_type(full_model) - if args.use_fsdp2: + if plan.is_fsdp_dense: device = args.dist_state.device else: device = full_model.device @@ -604,7 +584,7 @@ def mono_quantize( # quantization target; the full forward still exercises that component's quantizers. if args.calib_with_images: calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) - elif (ep_mesh := expert_parallel_mesh(full_model)) is not None: + elif distribution_plan(full_model).is_expert_parallel: # An expert-parallel run has to calibrate through distributed_forward_loop, not the # plain one: it is the only caller of _reinstall_a2a_forwards, which re-plants the # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on @@ -616,7 +596,10 @@ def mono_quantize( # It also restores the use_cache guard the plain loop does not apply. Other # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. calibrate_loop = distributed_forward_loop( - language_model, calib_dataloader, device=args.dist_state.device, ep_mesh=ep_mesh + language_model, + calib_dataloader, + device=args.dist_state.device, + ep_mesh=expert_parallel_mesh(full_model), ) else: calibrate_loop = create_forward_loop( @@ -1629,7 +1612,7 @@ def _main_inner(args: argparse.Namespace): device, ) finally: - cleanup_distributed(args) + cleanup_distributed() if __name__ == "__main__": diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 020067b2001..7adc190eece 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -23,6 +23,7 @@ import pytest import torch +import transformers import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 @@ -436,6 +437,8 @@ def distributed_load(*args, **kwargs): raise RuntimeError("distributed loader reached") monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load) + # A one-rank FSDP2 run sets up its process group of one right before the load. + monkeypatch.setattr(hf_ptq, "setup_fsdp2_process_group", lambda plan: None) error, message = ( (NotImplementedError, "Nemotron-H MTP") if has_mtp @@ -445,31 +448,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.""" - 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( diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index f69fbdfa541..434e50dfa55 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -15,10 +15,11 @@ """Expert-parallel PTQ wiring in ``examples/hf_ptq``. -Under torchrun, hf_ptq.py decides a distribution plan with ``resolve_distribution_plan`` and loads -the model through ``load_pretrained``: with the plan when it shards the experts, as FSDP2 under -``--use_fsdp2``, and otherwise whole through ``get_model``. Every rank then calibrates on its own -shard of the data, and a whole model is pinned to the rank's own GPU instead of ``"auto"``. +hf_ptq.py decides a distribution plan once, with ``resolve_distribution_plan``, and loads the model +by it through ``load_pretrained``: expert-parallel, FSDP2-sharded under ``--use_fsdp2``, or whole +through ``get_model``. Everything after reads the plan back from the model: every rank of a +distributed plan calibrates on its own shard of the data, an expert-parallel one through the +distributed forward loop, and a whole model is pinned to the rank's own GPU instead of ``"auto"``. Nothing here needs a process group or a checkpoint: the process group's rank and size are answered the way a torchrun launch would, the configs are tiny real transformers configs served from memory, @@ -37,10 +38,11 @@ from _test_utils.examples.hf_ptq_example_utils import example_utils from _test_utils.torch.transformers_models import get_tiny_tokenizer from torch.utils.data import DataLoader -from transformers import DeepseekV3Config, LlamaConfig +from transformers import AutoConfig, DeepseekV3Config, LlamaConfig import modelopt.torch.utils.distributed as dist from modelopt.torch.models import get_spec +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -100,6 +102,7 @@ def torchrun(monkeypatch): monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) def enter(world_size, rank=0): + monkeypatch.setattr(dist, "is_initialized", lambda: world_size > 1) monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) monkeypatch.setattr(dist, "size", lambda group=None: world_size) monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) @@ -131,14 +134,25 @@ def _args(dist_state, **overrides): def _serve(monkeypatch, hf_ptq, config): - """Answer hf_ptq's config read with ``config`` rather than a checkpoint on disk.""" - monkeypatch.setattr(hf_ptq.AutoConfig, "from_pretrained", lambda *args, **kwargs: config) + """Answer every config read with ``config`` rather than a checkpoint on disk.""" + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *args, **kwargs: config) return config def _deepseek_v3(n_routed_experts): - """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable.""" - config = DeepseekV3Config(n_routed_experts=n_routed_experts, **_TINY) + """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable. + + Its one layer is an MoE layer: DeepSeek-V3's first layers are dense by default. + """ + config = DeepseekV3Config( + n_routed_experts=n_routed_experts, + first_k_dense_replace=0, + moe_intermediate_size=8, + n_group=1, + topk_group=1, + num_experts_per_tok=1, + **_TINY, + ) config.base_model_ep_plan = None # holds even if a later transformers ships a plan return config @@ -154,43 +168,35 @@ def test_a_single_process_run_is_never_expert_parallel( dist_state = torchrun(world_size=1) if one_rank_group else _dist_state(world_size=1) _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) - assert hf_ptq._decide_expert_parallel(_args(dist_state)) == (None, None) + assert hf_ptq._resolve_distribution(_args(dist_state))[0] == DistributionPlan() -def test_use_fsdp2_keeps_the_standalone_fsdp2_loader(hf_ptq, monkeypatch, torchrun): - """--use_fsdp2 names its loader, so even an explicit EP degree for an EP-capable MoE must not +def test_use_fsdp2_plans_fsdp2_without_expert_parallelism(hf_ptq, monkeypatch, torchrun): + """--use_fsdp2 names its load, so even an explicit EP degree for an EP-capable MoE must not divert the run to the expert-parallel one.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) args = _args(torchrun(world_size=2), use_fsdp2=True) - assert hf_ptq._decide_expert_parallel(args) == (None, None) + assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(2, dense="fsdp") -@pytest.mark.parametrize( - ("config", "ep_capable"), - [(_deepseek_v3(n_routed_experts=3), True), (LlamaConfig(**_TINY), False)], - ids=["moe_no_ep_degree", "dense"], -) -def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2( - hf_ptq, monkeypatch, torchrun, config, ep_capable -): +def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monkeypatch, torchrun): """Without expert parallelism a plan that shards the dense part is an FSDP2 run, which this script runs only under --use_fsdp2, the flag its FSDP2 checks hang off; so it must fail and point there rather than load the model whole on every rank.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") - _serve(monkeypatch, hf_ptq, config) + _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) - with pytest.raises(RuntimeError, match="--use_fsdp2") as raised: - hf_ptq._decide_expert_parallel(_args(torchrun(world_size=2))) - assert ("MODELOPT_EP_DEGREE" in str(raised.value)) == ep_capable + with pytest.raises(RuntimeError, match="--use_fsdp2"): + hf_ptq._resolve_distribution(_args(torchrun(world_size=2))) def test_a_dense_model_under_torchrun_runs_data_parallel(hf_ptq, monkeypatch, torchrun): """A dense model has no experts to shard, so it loads whole on every rank, through get_model.""" _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) - assert hf_ptq._decide_expert_parallel(_args(torchrun(world_size=4))) == (None, None) + assert hf_ptq._resolve_distribution(_args(torchrun(world_size=4)))[0] == DistributionPlan(4) # ---------------------------------------------------------------------------------- load_model @@ -223,9 +229,8 @@ def test_load_model_loads_an_expert_parallel_plan_through_load_pretrained( hf_ptq.load_model(args) (call,) = calls - plan = call["distribution"] assert call["ckpt_path"] == _CKPT - assert (plan.ep_size, plan.is_fsdp_dense) == (4, True) + assert call["distribution"] == DistributionPlan(4, ep_size=4, dense="fsdp") assert call["kwargs"] == { "hf_config": served, "trust_remote_code": True, @@ -251,7 +256,7 @@ def test_use_fsdp2_loads_through_load_pretrained(hf_ptq, monkeypatch, torchrun): hf_ptq.load_model(args) (call,) = calls - assert call["distribution"] == "fsdp2" + assert call["distribution"] == DistributionPlan(2, dense="fsdp") assert call["kwargs"] == { "hf_config": served, "trust_remote_code": False, @@ -260,6 +265,57 @@ def test_use_fsdp2_loads_through_load_pretrained(hf_ptq, monkeypatch, torchrun): } +def test_a_whole_model_loads_by_the_plan_through_get_model(hf_ptq, monkeypatch, torchrun): + """Data-parallel, every rank loads the whole model, by the plan, onto its own GPU.""" + _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) + loads = [] + + def get_model(ckpt_path, device, **kwargs): + loads.append((device, kwargs["distribution"])) + raise _LoaderReachedError + + monkeypatch.setattr(hf_ptq, "get_model", get_model) + args = _args( + torchrun(world_size=4, rank=1), + specdec_offline_dataset=None, + low_memory_mode=False, + gpu_max_mem_percentage=0.8, + use_seq_device_map=False, + offload_folder=None, + max_cpu_memory_gb=None, + max_gpu_memory_gb=None, + ) + + with pytest.raises(_LoaderReachedError): + hf_ptq.load_model(args) + + assert loads == [(torch.device("cuda:1"), DistributionPlan(4))] + + +def test_a_one_rank_fsdp2_run_sets_up_its_group_before_loading(hf_ptq, monkeypatch): + """FSDP2 on one rank shards through a process group, of one, which main() did not set up.""" + _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY)) + calls = _record_load_pretrained(hf_ptq, monkeypatch) + setups = [] + monkeypatch.setattr( + hf_ptq, "setup_fsdp2_process_group", lambda plan: setups.append((plan, len(calls))) + ) + args = _args( + _dist_state(world_size=1), + use_fsdp2=True, + cpu_offload=False, + specdec_offline_dataset=None, + low_memory_mode=False, + ) + + with pytest.raises(_LoaderReachedError): + hf_ptq.load_model(args) + + plan = DistributionPlan(dense="fsdp") + assert setups == [(plan, 0)], "set up before the load" + assert calls[0]["distribution"] == plan + + # ------------------------------------------------------------------------- per-rank placement @@ -282,14 +338,16 @@ def test_a_rank_device_pins_the_whole_model_to_it( LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) placed = [] + plan = DistributionPlan(world_size) + def load_pretrained(ckpt_path, auto_class, **kwargs): + assert kwargs["distribution"] == plan placed.append(kwargs["device_map"]) return torch.nn.Linear(1, 1) monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained) - monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: world_size) - example_utils.get_model(str(tmp_path), device=device) + example_utils.get_model(str(tmp_path), device=device, distribution=plan) assert placed == [device_map] @@ -317,20 +375,37 @@ def test_dist_state_reads_the_live_process_group(monkeypatch, use_fsdp2): ) -def test_a_single_process_fsdp2_run_sets_up_its_own_group(monkeypatch): - """A one-rank --use_fsdp2 launch has no process group from main(), so the FSDP2 branch must - still create one.""" - setups = [] - monkeypatch.setattr(example_utils.dist_utils, "setup", lambda timeout=None: setups.append(0)) - monkeypatch.setattr(example_utils.dist_utils, "rank", lambda group=None: 0) - monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: 1) - monkeypatch.setattr(example_utils.dist_utils, "local_rank", lambda: 0) - args = SimpleNamespace(use_fsdp2=True) +@pytest.mark.parametrize( + ("plan", "live_group", "setups"), + [ + (DistributionPlan(dense="fsdp"), False, 1), + (DistributionPlan(2, dense="fsdp"), True, 0), + (DistributionPlan(), False, 0), + ], + ids=["fsdp2_on_one_rank", "launched_group", "single_process"], +) +def test_fsdp2_on_one_rank_sets_up_a_group_of_its_own(monkeypatch, plan, live_group, setups): + """FSDP2 shards through a process group even on one rank, where main() set up none; a + distributed launch has its group, and a single process's plan needs none.""" + calls = [] + monkeypatch.setattr(example_utils.dist_utils, "is_initialized", lambda: live_group) + monkeypatch.setattr(example_utils.dist_utils, "setup", lambda timeout=None: calls.append(0)) + + example_utils.setup_fsdp2_process_group(plan) + + assert len(calls) == setups - example_utils.setup_distributed_args(args) - assert len(setups) == 1 - assert args.dist_state.world_size == 1 +@pytest.mark.parametrize("launch_world", ["1", "2"], ids=["single_process", "distributed_launch"]) +def test_cleanup_destroys_only_the_group_this_script_set_up(monkeypatch, launch_world): + """A distributed launch's group is distributed_run's to destroy, after the run.""" + destroyed = [] + monkeypatch.setenv("WORLD_SIZE", launch_world) + monkeypatch.setattr(example_utils.dist_utils, "cleanup", lambda: destroyed.append(0)) + + example_utils.cleanup_distributed() + + assert len(destroyed) == (launch_world == "1") def test_main_owns_the_process_group_around_the_whole_run(hf_ptq, monkeypatch): @@ -383,8 +458,10 @@ def test_distributed_calibration_gives_each_rank_its_own_shard(hf_ptq, monkeypat def calib_loader(world_size, rank=0): args.dist_state = _dist_state(world_size, rank) + model = torch.nn.Linear(2, 2) + set_distribution_plan(model, DistributionPlan(world_size)) built, _ = hf_ptq.make_calib_dataloader( - args, None, None, tokenizer, torch.device("cpu"), "llama" + args, model, None, tokenizer, torch.device("cpu"), "llama" ) return built @@ -395,10 +472,11 @@ def calib_loader(world_size, rank=0): assert calib_loader(world_size=1) is loader -def _calibration_loop_choice(hf_ptq, monkeypatch, ep_mesh): - """Run mono_quantize on a model whose experts are sharded over ``ep_mesh`` (``None``: not - expert-parallel) and return which calibration loop it built.""" +def _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh=None): + """Run mono_quantize on a model distributed as ``plan``, its experts loaded over ``ep_mesh``, and + return which calibration loop it built.""" model = torch.nn.Linear(2, 2) + set_distribution_plan(model, plan) chosen, quantized = [], [] monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False) monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: True) @@ -441,11 +519,16 @@ def test_expert_parallel_calibration_runs_through_the_distributed_forward_loop(h """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion clobbers, so an expert-parallel run must calibrate through it rather than the plain loop.""" ep_mesh = object() # the mesh load_and_prepare_ep loaded the experts over + plan = DistributionPlan(2, ep_size=2) - assert _calibration_loop_choice(hf_ptq, monkeypatch, ep_mesh) == [("distributed", ep_mesh)] + assert _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh) == [ + ("distributed", ep_mesh) + ] def test_fsdp2_calibration_keeps_the_plain_forward_loop(hf_ptq, monkeypatch): """Without expert parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry, as on main.""" - assert _calibration_loop_choice(hf_ptq, monkeypatch, ep_mesh=None) == ["plain"] + plan = DistributionPlan(2, dense="fsdp") + + assert _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh=object()) == ["plain"] diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py index d37d3aba3b1..a0a56fbd4b4 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py +++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py @@ -20,6 +20,8 @@ import pytest import torch +from modelopt.torch.utils.distributed import DistributionPlan + _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -64,6 +66,10 @@ def test_image_calibration_model_target_follows_recipe( ) monkeypatch.setattr(hf_ptq, "get_model", lambda *args, **kwargs: full_model) + # A single process's plan, decided without reading the "dummy" checkpoint. + monkeypatch.setattr( + hf_ptq, "resolve_distribution_plan", lambda *args, **kwargs: (DistributionPlan(), None) + ) # A VLM's extracted language model reports its own sub-config type; hf_ptq must keep the root's. monkeypatch.setattr( hf_ptq, From 1fef889dbb988e6274efd2031c4e0440b2ec46cf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 22:20:43 +0000 Subject: [PATCH 15/57] Take hf_ptq's DistributedState from distributed_run hf_ptq built its DistributedState itself, in setup_distributed_args, from the process group distributed_run had just set up. distributed_run yields that state now, so main() hands it to the run, with --device for a single process, and the example's own class and setup are gone. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 29 ------------ examples/hf_ptq/hf_ptq.py | 13 +++-- .../hf_ptq/test_hf_ptq_expert_parallel.py | 47 +++++++------------ 3 files changed, 23 insertions(+), 66 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 3bf31dcca3c..73466cbe25c 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -23,7 +23,6 @@ import warnings from collections.abc import Iterator from contextlib import contextmanager -from dataclasses import dataclass from datetime import timedelta from pathlib import Path from typing import Any @@ -62,34 +61,6 @@ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] -@dataclass -class DistributedState: - """Example-local distributed state for model loading, dataloader sharding, and rank-0 output.""" - - rank: int - world_size: int - device: torch.device | str - is_main: bool - - -def setup_distributed_args(args): - """Initialize and attach ``args.dist_state`` (single-process if not distributed). - - A distributed launch has its process group set up by then (hf_ptq's ``distributed_run``), and the - state reads its rank and world from that group rather than deriving them a second time. - """ - if dist_utils.is_distributed_run(): - rank = dist_utils.rank() - args.dist_state = DistributedState( - rank=rank, - world_size=dist_utils.size(), - device=torch.device(f"cuda:{dist_utils.local_rank()}"), - is_main=rank == 0, - ) - else: - args.dist_state = DistributedState(rank=0, world_size=1, device=args.device, is_main=True) - - def setup_fsdp2_process_group(plan: DistributionPlan) -> None: """Set up the one-rank process group a plan that FSDP2-shards the model on one rank needs. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 3d5990fc5aa..1b659907238 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -48,7 +48,6 @@ recipe_layerwise_blocks, run_nemotron_vl_preview, save_source_config, - setup_distributed_args, setup_fsdp2_process_group, validate_fsdp2_supported, ) @@ -1558,16 +1557,16 @@ def main(args: argparse.Namespace): random.seed(RAND_SEED) np.random.seed(RAND_SEED) - # A distributed launch owns its process group for the whole run. Entered here rather than - # inside load_model so cleanup still runs if loading raises -- otherwise the surviving ranks - # block on the first collective until the NCCL timeout. - with distributed_run(timeout=timedelta(hours=3)): + # A distributed launch owns its process group for the whole run, and the run takes this + # process's rank, size and device from it. Entered here rather than inside load_model so + # cleanup still runs if loading raises -- otherwise the surviving ranks block on the first + # collective until the NCCL timeout. + with distributed_run(timeout=timedelta(hours=3), device=args.device) as dist_state: + args.dist_state = dist_state _main_inner(args) def _main_inner(args: argparse.Namespace): - setup_distributed_args(args) - try: # Entered inside the try: opening the run is fatal by design, and skipping # cleanup_distributed would leave the other ranks blocked on the first collective diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 434e50dfa55..9730c5995d7 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -113,9 +113,9 @@ def enter(world_size, rank=0): def _dist_state(world_size, rank=0): - """What setup_distributed_args leaves on ``rank`` of a ``world_size``-rank launch.""" - return example_utils.DistributedState( - rank=rank, world_size=world_size, device=torch.device(f"cuda:{rank}"), is_main=rank == 0 + """What distributed_run yields to ``rank`` of a ``world_size``-rank launch.""" + return dist.DistributedState( + rank=rank, world_size=world_size, device=torch.device(f"cuda:{rank}") ) @@ -355,26 +355,6 @@ def load_pretrained(ckpt_path, auto_class, **kwargs): # ------------------------------------------------------------------ process group and dist_state -@pytest.mark.parametrize("use_fsdp2", [False, True]) -def test_dist_state_reads_the_live_process_group(monkeypatch, use_fsdp2): - """Under torchrun dist_state is read from the process group main() set up, even with - --use_fsdp2, whose own branch would initialize a second group.""" - monkeypatch.setattr( - example_utils.dist_utils, "setup", lambda *a, **k: pytest.fail("second process group") - ) - monkeypatch.setattr(example_utils.dist_utils, "is_distributed_run", lambda: True) - monkeypatch.setattr(example_utils.dist_utils, "rank", lambda group=None: 1) - monkeypatch.setattr(example_utils.dist_utils, "size", lambda group=None: 4) - monkeypatch.setattr(example_utils.dist_utils, "local_rank", lambda: 1) - args = SimpleNamespace(use_fsdp2=use_fsdp2, device="cuda") - - example_utils.setup_distributed_args(args) - - assert args.dist_state == example_utils.DistributedState( - rank=1, world_size=4, device=torch.device("cuda:1"), is_main=False - ) - - @pytest.mark.parametrize( ("plan", "live_group", "setups"), [ @@ -411,20 +391,23 @@ def test_cleanup_destroys_only_the_group_this_script_set_up(monkeypatch, launch_ def test_main_owns_the_process_group_around_the_whole_run(hf_ptq, monkeypatch): """main() runs inside distributed_run, with a collective timeout above NCCL's 30-minute default, and leaves it even when the run raises -- or the other ranks would wait in their next - collective until that timeout. distributed_run itself does nothing in a single process.""" - events, timeouts = [], [] + collective until that timeout. The run reads its rank, size and device from the state + distributed_run yields, which puts a single process on --device.""" + events, calls, seen = [], [], [] + state = dist.DistributedState(rank=0, world_size=1, device="cpu") @contextlib.contextmanager - def distributed_run(timeout=None): - timeouts.append(timeout) + def distributed_run(timeout=None, device="cuda"): + calls.append((timeout, device)) events.append("enter") try: - yield + yield state finally: events.append("exit") def main_inner(args): events.append("run") + seen.append(args.dist_state) raise _RunFailedError monkeypatch.setattr(torch.cuda, "is_available", lambda: True) @@ -432,10 +415,14 @@ def main_inner(args): monkeypatch.setattr(hf_ptq, "_main_inner", main_inner) with pytest.raises(_RunFailedError): - hf_ptq.main(SimpleNamespace(use_fsdp2=False, device="cuda")) + hf_ptq.main(SimpleNamespace(use_fsdp2=False, device="cpu")) assert events == ["enter", "run", "exit"] - assert timeouts[0] > timedelta(minutes=30) + timeout, device = calls[0] + assert timeout > timedelta(minutes=30) + assert device == "cpu" + assert len(seen) == 1 + assert seen[0] is state # --------------------------------------------------------------------------------- calibration From 2667f5dfd048daa89e64202856f7f6f73372b57d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 22:37:36 +0000 Subject: [PATCH 16/57] Let distributed_run own hf_ptq's one-rank FSDP2 group hf_ptq set up the one-rank group FSDP2 needs in a single process and tore it down in its own finally. Both live in modelopt's distributed module now: load_model sets the group up with setup_fsdp2_process_group, keeping its two-hour collective timeout, and distributed_run destroys it with the run, as it does a launch's. The example's own helpers go. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 23 ------- examples/hf_ptq/hf_ptq.py | 69 +++++++++---------- tests/examples/hf_ptq/test_hf_ptq_args.py | 2 +- .../hf_ptq/test_hf_ptq_expert_parallel.py | 45 +++--------- 4 files changed, 41 insertions(+), 98 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 73466cbe25c..fa81ff7af8b 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -23,7 +23,6 @@ import warnings from collections.abc import Iterator from contextlib import contextmanager -from datetime import timedelta from pathlib import Path from typing import Any @@ -51,7 +50,6 @@ prefers_bundled_code, prepare_model_for_loading, ) -from modelopt.torch.utils import distributed as dist_utils from modelopt.torch.utils.distributed import DistributionPlan from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -61,27 +59,6 @@ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] -def setup_fsdp2_process_group(plan: DistributionPlan) -> None: - """Set up the one-rank process group a plan that FSDP2-shards the model on one rank needs. - - FSDP2 shards through a process group even on one rank; a distributed launch already has its - group (hf_ptq's ``distributed_run``). - """ - if plan.is_fsdp_dense and not dist_utils.is_initialized(): - # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can - # exceed it, and PyTorch 2.8 has no per-call barrier() timeout (must be set at PG creation). - dist_utils.setup(timeout=timedelta(hours=2)) - - -def cleanup_distributed(): - """Destroy the one-rank process group :func:`setup_fsdp2_process_group` set up, if any. - - A distributed launch's group is ``distributed_run``'s to destroy. - """ - if not dist_utils.is_distributed_launch(): - dist_utils.cleanup() - - def validate_fsdp2_supported(args, config): """Raise ``NotImplementedError`` for model/CLI combos the FSDP2 path doesn't support yet.""" issues = [] diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 1b659907238..030c41554ee 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -36,7 +36,6 @@ HF_PTQ, _prepare_quant_cfg, build_quant_cfg, - cleanup_distributed, create_vlm_calibration_loop, generate_excludes_prompt, get_model, @@ -48,7 +47,6 @@ recipe_layerwise_blocks, run_nemotron_vl_preview, save_source_config, - setup_fsdp2_process_group, validate_fsdp2_supported, ) from torch.utils.data import DataLoader @@ -103,6 +101,7 @@ distributed_forward_loop, distributed_run, distribution_plan, + setup_fsdp2_process_group, shard_dataloader, ) from modelopt.torch.utils.memory_monitor import launch_memory_monitor @@ -349,7 +348,9 @@ def load_model(args: argparse.Namespace): plan, hf_config = _resolve_distribution(args) if plan.is_fsdp_dense and not plan.is_expert_parallel: validate_fsdp2_supported(args, hf_config) - setup_fsdp2_process_group(plan) + # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can + # exceed it, and PyTorch 2.8 has no per-call barrier() timeout (must be set at PG creation). + setup_fsdp2_process_group(plan, timeout=timedelta(hours=2)) if plan.is_expert_parallel: full_model = load_pretrained( args.local_checkpoint_path, @@ -1567,25 +1568,35 @@ def main(args: argparse.Namespace): def _main_inner(args: argparse.Namespace): - try: - # Entered inside the try: opening the run is fatal by design, and skipping - # cleanup_distributed would leave the other ranks blocked on the first collective - # until the NCCL timeout. - with mlflow_run(args): - # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole - # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local - # path. - args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint( - args.pyt_ckpt_path - ) + with mlflow_run(args): + # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole + # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local + # path. + args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path) - # launch a memory monitor to read the currently used GPU memory. - launch_memory_monitor() + # launch a memory monitor to read the currently used GPU memory. + launch_memory_monitor() - # Force eager execution for all model types. - torch.compiler.set_stance("force_eager") + # Force eager execution for all model types. + torch.compiler.set_stance("force_eager") + + ( + full_model, + language_model, + model_type, + calibration_only, + processor, + tokenizer, + device, + ) = load_model(args) - ( + if args.sparsity_fmt != "dense": + # Sparse + sparsity_main(args, full_model, tokenizer, device) + else: + # Quantize + quantize_main( + args, full_model, language_model, model_type, @@ -1593,25 +1604,7 @@ def _main_inner(args: argparse.Namespace): processor, tokenizer, device, - ) = load_model(args) - - if args.sparsity_fmt != "dense": - # Sparse - sparsity_main(args, full_model, tokenizer, device) - else: - # Quantize - quantize_main( - args, - full_model, - language_model, - model_type, - calibration_only, - processor, - tokenizer, - device, - ) - finally: - cleanup_distributed() + ) if __name__ == "__main__": diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 7adc190eece..175dca176e6 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -438,7 +438,7 @@ def distributed_load(*args, **kwargs): monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load) # A one-rank FSDP2 run sets up its process group of one right before the load. - monkeypatch.setattr(hf_ptq, "setup_fsdp2_process_group", lambda plan: None) + monkeypatch.setattr(hf_ptq, "setup_fsdp2_process_group", lambda plan, timeout=None: None) error, message = ( (NotImplementedError, "Nemotron-H MTP") if has_mtp diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 9730c5995d7..6a2767d7483 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -293,12 +293,15 @@ def get_model(ckpt_path, device, **kwargs): def test_a_one_rank_fsdp2_run_sets_up_its_group_before_loading(hf_ptq, monkeypatch): - """FSDP2 on one rank shards through a process group, of one, which main() did not set up.""" + """FSDP2 on one rank shards through a process group, of one, which main() did not set up. Its + collective timeout is raised above NCCL's 30-minute default, for rank 0's checkpoint write.""" _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY)) calls = _record_load_pretrained(hf_ptq, monkeypatch) setups = [] monkeypatch.setattr( - hf_ptq, "setup_fsdp2_process_group", lambda plan: setups.append((plan, len(calls))) + hf_ptq, + "setup_fsdp2_process_group", + lambda plan, timeout=None: setups.append((plan, len(calls), timeout)), ) args = _args( _dist_state(world_size=1), @@ -312,7 +315,10 @@ def test_a_one_rank_fsdp2_run_sets_up_its_group_before_loading(hf_ptq, monkeypat hf_ptq.load_model(args) plan = DistributionPlan(dense="fsdp") - assert setups == [(plan, 0)], "set up before the load" + [(set_up_for, loads_before, timeout)] = setups + assert set_up_for == plan + assert loads_before == 0, "set up before the load" + assert timeout > timedelta(minutes=30) assert calls[0]["distribution"] == plan @@ -355,39 +361,6 @@ def load_pretrained(ckpt_path, auto_class, **kwargs): # ------------------------------------------------------------------ process group and dist_state -@pytest.mark.parametrize( - ("plan", "live_group", "setups"), - [ - (DistributionPlan(dense="fsdp"), False, 1), - (DistributionPlan(2, dense="fsdp"), True, 0), - (DistributionPlan(), False, 0), - ], - ids=["fsdp2_on_one_rank", "launched_group", "single_process"], -) -def test_fsdp2_on_one_rank_sets_up_a_group_of_its_own(monkeypatch, plan, live_group, setups): - """FSDP2 shards through a process group even on one rank, where main() set up none; a - distributed launch has its group, and a single process's plan needs none.""" - calls = [] - monkeypatch.setattr(example_utils.dist_utils, "is_initialized", lambda: live_group) - monkeypatch.setattr(example_utils.dist_utils, "setup", lambda timeout=None: calls.append(0)) - - example_utils.setup_fsdp2_process_group(plan) - - assert len(calls) == setups - - -@pytest.mark.parametrize("launch_world", ["1", "2"], ids=["single_process", "distributed_launch"]) -def test_cleanup_destroys_only_the_group_this_script_set_up(monkeypatch, launch_world): - """A distributed launch's group is distributed_run's to destroy, after the run.""" - destroyed = [] - monkeypatch.setenv("WORLD_SIZE", launch_world) - monkeypatch.setattr(example_utils.dist_utils, "cleanup", lambda: destroyed.append(0)) - - example_utils.cleanup_distributed() - - assert len(destroyed) == (launch_world == "1") - - def test_main_owns_the_process_group_around_the_whole_run(hf_ptq, monkeypatch): """main() runs inside distributed_run, with a collective timeout above NCCL's 30-minute default, and leaves it even when the run raises -- or the other ranks would wait in their next From b1b8bd8af3ea8e7e3849b0f5e4e141cb9b2ca95e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:02:27 +0000 Subject: [PATCH 17/57] Leave a single process's FSDP2 group to the load load_pretrained sets up the one-rank group its FSDP2 load shards through ([3/21]), and distributed_run destroys it with the run, so hf_ptq no longer sets it up before loading. The group keeps PyTorch's default collective timeout: on one rank no collective waits for a peer, which is what the two hours were for. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 4 --- tests/examples/hf_ptq/test_hf_ptq_args.py | 2 -- .../hf_ptq/test_hf_ptq_expert_parallel.py | 30 ------------------- 3 files changed, 36 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 030c41554ee..22501541a26 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -101,7 +101,6 @@ distributed_forward_loop, distributed_run, distribution_plan, - setup_fsdp2_process_group, shard_dataloader, ) from modelopt.torch.utils.memory_monitor import launch_memory_monitor @@ -348,9 +347,6 @@ def load_model(args: argparse.Namespace): plan, hf_config = _resolve_distribution(args) if plan.is_fsdp_dense and not plan.is_expert_parallel: validate_fsdp2_supported(args, hf_config) - # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can - # exceed it, and PyTorch 2.8 has no per-call barrier() timeout (must be set at PG creation). - setup_fsdp2_process_group(plan, timeout=timedelta(hours=2)) if plan.is_expert_parallel: full_model = load_pretrained( 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 175dca176e6..fa687017724 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -437,8 +437,6 @@ def distributed_load(*args, **kwargs): raise RuntimeError("distributed loader reached") monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load) - # A one-rank FSDP2 run sets up its process group of one right before the load. - monkeypatch.setattr(hf_ptq, "setup_fsdp2_process_group", lambda plan, timeout=None: None) error, message = ( (NotImplementedError, "Nemotron-H MTP") if has_mtp diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 6a2767d7483..8625be4669f 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -292,36 +292,6 @@ def get_model(ckpt_path, device, **kwargs): assert loads == [(torch.device("cuda:1"), DistributionPlan(4))] -def test_a_one_rank_fsdp2_run_sets_up_its_group_before_loading(hf_ptq, monkeypatch): - """FSDP2 on one rank shards through a process group, of one, which main() did not set up. Its - collective timeout is raised above NCCL's 30-minute default, for rank 0's checkpoint write.""" - _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY)) - calls = _record_load_pretrained(hf_ptq, monkeypatch) - setups = [] - monkeypatch.setattr( - hf_ptq, - "setup_fsdp2_process_group", - lambda plan, timeout=None: setups.append((plan, len(calls), timeout)), - ) - args = _args( - _dist_state(world_size=1), - use_fsdp2=True, - cpu_offload=False, - specdec_offline_dataset=None, - low_memory_mode=False, - ) - - with pytest.raises(_LoaderReachedError): - hf_ptq.load_model(args) - - plan = DistributionPlan(dense="fsdp") - [(set_up_for, loads_before, timeout)] = setups - assert set_up_for == plan - assert loads_before == 0, "set up before the load" - assert timeout > timedelta(minutes=30) - assert calls[0]["distribution"] == plan - - # ------------------------------------------------------------------------- per-rank placement From c8514f757c540206c8085812e3904fd7d007ad46 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 18:20:55 +0000 Subject: [PATCH 18/57] Keep every rank decoding hf_ptq's preview of a sharded model Each rank previews the first sample of its own calibration shard, so the ranks reach EOS at different steps. Under expert parallelism, or with the dense part FSDP2-sharded, every forward runs a collective (the experts' all-to-all, FSDP2's all-gathers), and a rank that stopped decoding left the others blocked in it: Qwen3-30B-A3B at NVFP4 and EP4 hung in the post-PTQ preview, while the same run with --skip_generate finished. The previews, Nemotron-VL's included, now generate with synced_gpus when the model's plan shards it, so every rank keeps stepping until all have finished. The preview tests' stand-in models are modules now, since the previews read their distribution plan. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 23 ++++- examples/hf_ptq/hf_ptq.py | 3 + examples/hf_ptq/vlm_utils.py | 6 +- .../hf_ptq/test_hf_ptq_expert_parallel.py | 91 +++++++++++++++++++ .../examples/hf_ptq/test_hf_ptq_model_type.py | 18 +++- 5 files changed, 132 insertions(+), 9 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index fa81ff7af8b..2071c7db655 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -50,7 +50,7 @@ prefers_bundled_code, prepare_model_for_loading, ) -from modelopt.torch.utils.distributed import DistributionPlan +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -83,6 +83,19 @@ def validate_fsdp2_supported(args, config): ) +def preview_generate_kwargs(model) -> dict[str, Any]: + """Extra ``generate`` kwargs that keep every rank of a sharded model decoding a preview. + + Each rank previews the first sample of its own calibration shard, so the ranks finish at + different steps. A model its plan shards runs a collective in every forward (the routed experts' + all-to-all, FSDP2's all-gathers): a rank that stopped decoding would leave the others blocked in + it. ``synced_gpus`` keeps every rank stepping until all of them have finished. + """ + plan = distribution_plan(model) + sharded = plan.world_size > 1 and (plan.is_expert_parallel or plan.is_fsdp_dense) + return {"synced_gpus": True} if sharded else {} + + def run_nemotron_vl_preview( full_model, tokenizer, @@ -109,10 +122,12 @@ def run_nemotron_vl_preview( print(f"Running text-only preview generation for Nemotron VL model ({stage_name})...") question = tokenizer.decode(input_ids[0], skip_special_tokens=True) + lockstep = preview_generate_kwargs(full_model) generation_config = { "max_new_tokens": 100, "do_sample": False, "eos_token_id": tokenizer.eos_token_id, + **lockstep, } # Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse) @@ -126,11 +141,13 @@ def run_nemotron_vl_preview( generated_ids = text_response elif allow_fallback: print("Text-only generation failed, falling back to standard generate...") - generated_ids = full_model.generate(input_ids, max_new_tokens=100) + generated_ids = full_model.generate(input_ids, max_new_tokens=100, **lockstep) # Run additional VL test with images print(f"Running additional VL test with images ({stage_name})...") - run_vl_preview_generation(full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code) + run_vl_preview_generation( + full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code, lockstep + ) return generated_ids diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 22501541a26..8804283c08c 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -44,6 +44,7 @@ is_nemotron_vl, is_trtllm_enc_dec_export, mlflow_run, + preview_generate_kwargs, recipe_layerwise_blocks, run_nemotron_vl_preview, save_source_config, @@ -846,6 +847,7 @@ def pre_quantize( preview_input_ids, attention_mask=preview_attention_mask, max_new_tokens=100, + **preview_generate_kwargs(full_model), ) return preview_input_ids, preview_attention_mask, generated_ids_before_ptq @@ -904,6 +906,7 @@ def post_quantize( preview_input_ids, attention_mask=preview_attention_mask, max_new_tokens=100, + **preview_generate_kwargs(full_model), ) except Exception as e: warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") diff --git a/examples/hf_ptq/vlm_utils.py b/examples/hf_ptq/vlm_utils.py index abfebbd4f01..16c779c0d71 100644 --- a/examples/hf_ptq/vlm_utils.py +++ b/examples/hf_ptq/vlm_utils.py @@ -21,7 +21,9 @@ from transformers import AutoImageProcessor, AutoProcessor -def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_remote_code=False): +def run_vl_preview_generation( + model, tokenizer, model_path, stage_name, trust_remote_code=False, generation_kwargs=None +): """Run preview generation for VL models using sample images. Args: @@ -30,6 +32,7 @@ def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_re model_path: Path to the model (for loading image processor) stage_name: Description of the stage (e.g., "before quantization") trust_remote_code: Whether to trust remote code for Huggingface models and tokenizers + generation_kwargs: Extra ``generate`` kwargs, such as ``synced_gpus`` for a sharded model Returns: Generated response text for logging/comparison """ @@ -79,6 +82,7 @@ def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_re "max_new_tokens": 50, "do_sample": False, "eos_token_id": tokenizer.eos_token_id, + **(generation_kwargs or {}), } print(f"Generating VL response ({stage_name})...") diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 8625be4669f..8fe2c9c3087 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -29,6 +29,7 @@ import contextlib import importlib import inspect +import sys from datetime import timedelta from pathlib import Path from types import SimpleNamespace @@ -462,3 +463,93 @@ def test_fsdp2_calibration_keeps_the_plain_forward_loop(hf_ptq, monkeypatch): plan = DistributionPlan(2, dense="fsdp") assert _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh=object()) == ["plain"] + + +class _PreviewModel(torch.nn.Module): + """A model whose ``generate`` records the kwargs each preview passes it.""" + + def __init__(self, plan=None): + super().__init__() + self.config = SimpleNamespace(is_encoder_decoder=False, model_type="qwen3_moe") + self.calls = [] + if plan is not None: + set_distribution_plan(self, plan) + + def generate(self, input_ids, **kwargs): + self.calls.append(kwargs) + return torch.cat([input_ids, input_ids[:, :1]], dim=1) + + +_PREVIEW_PLANS = pytest.mark.parametrize( + ("plan", "lockstep"), + [ + (None, False), + (DistributionPlan(4), False), + (DistributionPlan(4, ep_size=4), True), + (DistributionPlan(4, dense="fsdp"), True), + ], + ids=["single_process", "replicated", "expert_parallel", "fsdp2"], +) + + +@_PREVIEW_PLANS +def test_a_sharded_model_previews_in_lockstep(hf_ptq, monkeypatch, plan, lockstep): + """Each rank previews its own calibration shard's first sample, so the ranks finish decoding at + different steps. A sharded model runs a collective in every forward, and a rank that stopped + early left the others blocked in it: Qwen3-30B-A3B NVFP4 at EP4 hung in the post-PTQ preview. + Both previews keep every rank of a sharded model decoding until all have finished.""" + model = _PreviewModel(plan) + loader = [{"input_ids": torch.tensor([[1, 2, 3]]), "attention_mask": torch.ones(1, 3)}] + args = SimpleNamespace( + specdec_offline_dataset=None, + skip_generate=False, + verbose=False, + local_checkpoint_path=_CKPT, + trust_remote_code=False, + ) + monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None) + + ids, mask, before = hf_ptq.pre_quantize(args, model, "qwen3_moe", None, loader, False) + hf_ptq.post_quantize( + args, + model, + model, + "qwen3_moe", + get_tiny_tokenizer(), + None, + ids, + mask, + before, + False, + None, + loader, + ) + + synced = [call.get("synced_gpus", False) for call in model.calls] + assert synced == [lockstep, lockstep] # before and after quantization + + +@_PREVIEW_PLANS +def test_a_sharded_nemotron_vl_model_previews_in_lockstep(monkeypatch, plan, lockstep): + """The Nemotron-VL preview's text, fallback and image generations keep a sharded model's ranks + decoding together too.""" + recorded = {} + vlm_utils = SimpleNamespace( + run_text_only_generation=lambda model, tokenizer, question, config, *args: recorded.update( + text=config + ), + run_vl_preview_generation=lambda model, tokenizer, path, stage, trust, kwargs: ( + recorded.update(image=kwargs) + ), + ) + monkeypatch.setitem(sys.modules, "vlm_utils", vlm_utils) + model = _PreviewModel(plan) + + example_utils.run_nemotron_vl_preview( + model, get_tiny_tokenizer(), torch.tensor([[1, 2, 3]]), _CKPT, "stage", allow_fallback=True + ) + + expected = {"synced_gpus": True} if lockstep else {} + assert recorded["text"].get("synced_gpus", False) == lockstep + assert recorded["image"] == expected + assert model.calls[0].get("synced_gpus", False) == lockstep # the fallback generate diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py index 328dcd5e564..34b5194e1c2 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -24,6 +24,16 @@ _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" +class _Model(torch.nn.Module): + """A model whose ``generate`` is ``generate``: the previews read the module's distribution plan.""" + + def __init__(self, generate, config=None): + super().__init__() + self.generate = generate + if config is not None: + self.config = config + + @pytest.fixture def hf_ptq(monkeypatch: pytest.MonkeyPatch): monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) @@ -43,7 +53,7 @@ def hf_ptq(monkeypatch: pytest.MonkeyPatch): ) def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: bool): generate_calls = [] - full_model = SimpleNamespace(generate=lambda *args, **kwargs: generate_calls.append(1) or "out") + full_model = _Model(lambda *args, **kwargs: generate_calls.append(1) or "out") args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) batch = {"input_ids": torch.ones(2, 4, dtype=torch.long)} @@ -58,7 +68,7 @@ def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: ) def test_pre_quantize_preview_removes_padding(hf_ptq, ids, mask): calls = [] - model = SimpleNamespace(generate=lambda ids, **kwargs: calls.append((ids, kwargs)) or ids) + model = _Model(lambda ids, **kwargs: calls.append((ids, kwargs)) or ids) args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) batch = {"input_ids": torch.tensor([ids]), "attention_mask": torch.tensor([mask])} prompt, attention_mask, _ = hf_ptq.pre_quantize(args, model, "llama", None, [batch], False) @@ -122,9 +132,7 @@ def test_post_quantize_preview_keeps_whole_answer( ): """The preview must decode the answer, not slice a prompt-length prefix off it.""" monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None) - full_model = SimpleNamespace( - config=config, generate=lambda *args, **kwargs: torch.tensor(generated) - ) + full_model = _Model(lambda *args, **kwargs: torch.tensor(generated), config) tokenizer = SimpleNamespace( batch_decode=lambda ids, **kwargs: [" ".join(map(str, row.tolist())) for row in ids] ) From ac5bc1b6e0cebb520fe05b47c652a03403f2a21d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 18:54:56 +0000 Subject: [PATCH 19/57] Test the preview lockstep for a one-rank FSDP2 plan and in the VL image preview A one-rank FSDP2 plan needs no lockstep, but no case pinned it, so a check dropping world_size went unnoticed. The Nemotron-VL test stubbed vlm_utils whole, so nothing checked that the image preview merges the kwargs it is handed into the config its chat() gets. Both gaps found by the epdp port. Signed-off-by: Shengliang Xu --- .../hf_ptq/test_hf_ptq_expert_parallel.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 8fe2c9c3087..a31803b149e 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -38,6 +38,7 @@ import torch from _test_utils.examples.hf_ptq_example_utils import example_utils from _test_utils.torch.transformers_models import get_tiny_tokenizer +from PIL import Image from torch.utils.data import DataLoader from transformers import AutoConfig, DeepseekV3Config, LlamaConfig @@ -487,8 +488,9 @@ def generate(self, input_ids, **kwargs): (DistributionPlan(4), False), (DistributionPlan(4, ep_size=4), True), (DistributionPlan(4, dense="fsdp"), True), + (DistributionPlan(1, dense="fsdp"), False), ], - ids=["single_process", "replicated", "expert_parallel", "fsdp2"], + ids=["single_process", "replicated", "expert_parallel", "fsdp2", "one_rank_fsdp2"], ) @@ -553,3 +555,33 @@ def test_a_sharded_nemotron_vl_model_previews_in_lockstep(monkeypatch, plan, loc assert recorded["text"].get("synced_gpus", False) == lockstep assert recorded["image"] == expected assert model.calls[0].get("synced_gpus", False) == lockstep # the fallback generate + + +@pytest.mark.parametrize("generation_kwargs", [None, {"synced_gpus": True}]) +def test_the_vl_image_preview_generates_with_the_kwargs_it_is_handed( + monkeypatch, tmp_path, generation_kwargs +): + """run_vl_preview_generation hands the kwargs it is given to the model's chat(), in its config.""" + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + vlm_utils = importlib.import_module("vlm_utils") + (tmp_path / "images").mkdir() + Image.new("RGB", (4, 4)).save(tmp_path / "images" / "example1a.jpeg") + monkeypatch.setattr(vlm_utils, "__file__", str(tmp_path / "vlm_utils.py")) + monkeypatch.setattr( + vlm_utils.AutoImageProcessor, "from_pretrained", lambda *args, **kwargs: lambda images: {} + ) + received = [] + + class _ChatModel: + device = torch.device("cpu") + + def chat(self, tokenizer, question, generation_config, **image_features): + received.append(generation_config) + return "a picture" + + vlm_utils.run_vl_preview_generation( + _ChatModel(), SimpleNamespace(eos_token_id=0), _CKPT, "stage", False, generation_kwargs + ) + + assert len(received) == 1 + assert received[0].get("synced_gpus", False) == (generation_kwargs is not None) From e475b1623179cccda49f120e46158528f934b899 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:49:02 +0000 Subject: [PATCH 20/57] Write a replicated model's checkpoint from rank 0 only A dense model under torchrun gets a replicated plan, and every rank then ran the export into the same directory, racing each other's files and the in-place config.json rewrite. Rank 0 alone writes a replicated model; a sharded one still needs every rank. The source config of a multimodal model is rank 0's in every plan: another rank's, written after rank 0's export, would replace the config.json carrying the quantization config. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 10 ++++- .../hf_ptq/test_hf_ptq_expert_parallel.py | 42 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 8804283c08c..cae45104f37 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -699,6 +699,11 @@ def export_quantized( language_model: torch.nn.Module, model_type: str | None, ): + plan = distribution_plan(full_model) + if not (args.dist_state.is_main or plan.is_fsdp_dense or plan.is_expert_parallel): + # Replicated over the ranks, the model is whole on each: rank 0 writes it, and the rest + # would race it into the same directory. A sharded model needs every rank for its part. + return # Not inference_mode: the FSDP2 path gathers full params in this context and # inference tensors break the subsequent state_dict() -> param.detach(). with torch.no_grad(): @@ -714,8 +719,9 @@ def export_quantized( if is_multimodal_model(full_model): # Per-layer export writes its own config.json with quantization_config, which - # the source config would replace; it never writes a processor config. - if not args.layerwise_export: + # the source config would replace; it never writes a processor config. Rank 0's + # export writes config.json, so no other rank may write the source config after it. + if not args.layerwise_export and args.dist_state.is_main: save_source_config(args, export_path) start_time = time.time() diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index a31803b149e..c83a596a869 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -585,3 +585,45 @@ def chat(self, tokenizer, question, generation_config, **image_features): assert len(received) == 1 assert received[0].get("synced_gpus", False) == (generation_kwargs is not None) + + +# ------------------------------------------------------------------------------------- export + + +@pytest.mark.parametrize( + ("plan", "rank", "exports"), + [ + (DistributionPlan(4), 0, True), + (DistributionPlan(4), 1, False), + (DistributionPlan(4, dense="fsdp"), 1, True), + (DistributionPlan(4, ep_size=4), 1, True), + ], + ids=["replicated_rank0", "replicated_rank1", "fsdp2_rank1", "expert_parallel_rank1"], +) +def test_only_rank_0_writes_a_replicated_model(hf_ptq, monkeypatch, plan, rank, exports): + """A replicated model is whole on every rank, whose exports would write one checkpoint into the + same directory over each other; a sharded one needs every rank for its part. The source config + is rank 0's either way: another rank's, written after rank 0's export, would drop the + quantization config.""" + written = [] + monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: False) + monkeypatch.setattr(hf_ptq, "is_multimodal_model", lambda model: True) + monkeypatch.setattr(hf_ptq, "save_source_config", lambda args, path: written.append("config")) + monkeypatch.setattr( + hf_ptq, "export_hf_checkpoint", lambda model, export_dir: written.append("checkpoint") + ) + args = SimpleNamespace( + dist_state=_dist_state(plan.world_size, rank), + export_path="/export", + layerwise_export=False, + sparsity_fmt="dense", + qformat="fp8", + inference_tensor_parallel=1, + inference_pipeline_parallel=1, + vllm_fakequant_export=False, + ) + + model = _PreviewModel(plan) + hf_ptq.export_quantized(args, model, model, "qwen3_moe") + + assert written == ["config"] * (rank == 0) + ["checkpoint"] * exports From 4e13b9818b4925235d1ef0da41889be2ee6e2fdc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:50:28 +0000 Subject: [PATCH 21/57] Refuse dense FSDP2 within expert-data-parallel replicas before loading With fewer expert-parallel ranks than the world (the automatic degree when the world size does not divide the expert count), FSDP2 shards the dense part within each replica. The no-gather export cannot write that layout, and found out only after the load and the whole calibration. Refuse the plan up front and name the settings that avoid it. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 +- examples/hf_ptq/hf_ptq.py | 8 ++++++++ .../hf_ptq/test_hf_ptq_expert_parallel.py | 19 +++++++++++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 89c5da7442b..7d28f3c8b72 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -676,7 +676,7 @@ Two environment variables override the automatic distribution plan: | Variable | Default | Effect | | :--- | :--- | :--- | | `MODELOPT_EP_DEGREE` | `auto` | Expert-parallel degree. `auto` picks the largest degree that divides both the routed-expert count and the world size. Below the world size, each group of `MODELOPT_EP_DEGREE` ranks holds a full set of experts, and the groups calibrate on different data. | -| `MODELOPT_FSDP_SIZE` | `auto` | Whether to FSDP2-shard the dense part. `auto` shards it when it would take more than 60% of one GPU's memory; `world` always shards it; `1` replicates it. | +| `MODELOPT_FSDP_SIZE` | `auto` | Whether to FSDP2-shard the dense part. `auto` shards it when it would take more than 60% of one GPU's memory; `world` always shards it; `1` replicates it. With `MODELOPT_EP_DEGREE` below the world size the export cannot write a sharded dense part yet, so such a run fails before loading. | A run that cannot shard the model the way the plan asks fails before loading, rather than loading the whole model on every rank. `--use_fsdp2` skips expert parallelism and uses the [FSDP2 path](#multi-node-post-training-quantization-with-fsdp2) instead. That is also the way to quantize GPT-OSS and Llama-4 across GPUs: their experts are stored transposed, which the expert-parallel export cannot write yet, so expert parallelism refuses them. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index cae45104f37..d66f76495a1 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -339,6 +339,14 @@ def _resolve_distribution(args: argparse.Namespace): "or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU; an MoE can instead shard its " "experts (MODELOPT_EP_DEGREE above 1)." ) + if plan.is_expert_parallel and plan.is_fsdp_dense and plan.edp > 1: + # Refused before the load: the export, after the whole calibration, is what cannot do it. + raise NotImplementedError( + f"{plan!r} FSDP2-shards the dense part within each of {plan.edp} expert-data-parallel " + "replicas, a layout the export cannot write yet. Set MODELOPT_EP_DEGREE to the world " + "size if it divides the expert count, or MODELOPT_FSDP_SIZE=1 if the dense part fits " + "one GPU." + ) return plan, hf_config diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index c83a596a869..c322c86fdf4 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -194,6 +194,25 @@ def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monke hf_ptq._resolve_distribution(_args(torchrun(world_size=2))) +@pytest.mark.parametrize("ep_degree", [2, 4]) +def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused( + hf_ptq, monkeypatch, torchrun, ep_degree +): + """Below the world size each expert-parallel group is a replica, and FSDP2 shards the dense part + within each: a layout the export cannot write, so it must fail before the load and the + calibration rather than at the end. Expert parallelism over the whole world can still shard it.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", str(ep_degree)) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + args = _args(torchrun(world_size=4)) + + if ep_degree == 4: + assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(4, ep_size=4, dense="fsdp") + else: + with pytest.raises(NotImplementedError, match=r"MODELOPT_EP_DEGREE.*MODELOPT_FSDP_SIZE=1"): + hf_ptq._resolve_distribution(args) + + def test_a_dense_model_under_torchrun_runs_data_parallel(hf_ptq, monkeypatch, torchrun): """A dense model has no experts to shard, so it loads whole on every rank, through get_model.""" _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) From 58c8e6b06a9971f0c92900d80197926803ae46cc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 02:04:28 +0000 Subject: [PATCH 22/57] Answer the checkpoint-read gather in hf_ptq's simulated torchrun launches resolve_distribution_plan reads the checkpoint through the failure-carrying all_reduce, which gathers from every rank of a launch. The torchrun fixture answers that gather as it answers the plan's broadcast: every rank read the checkpoint as this one did. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index c322c86fdf4..de574acfb28 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -93,20 +93,27 @@ def _checkpoint_is_local(monkeypatch): ) +def _gathered_alike(gathered, mine, group=None): + """Answer a gather as if every rank handed it what this rank did.""" + gathered[:] = [mine] * len(gathered) + + @pytest.fixture def torchrun(monkeypatch): """Run as rank ``rank`` of a ``world_size``-rank torchrun launch; returns its ``dist_state``. - Every rank acts as its group's first, deciding for itself; how the ranks agree on what they - decide is tested with the collectives. + Every rank acts as its group's first, deciding for itself, and reads the checkpoint as this one + does; how the ranks agree on what they decide is tested with the collectives. """ monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) + monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike) def enter(world_size, rank=0): monkeypatch.setattr(dist, "is_initialized", lambda: world_size > 1) monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) monkeypatch.setattr(dist, "size", lambda group=None: world_size) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size) monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) monkeypatch.setattr(dist, "local_rank", lambda: rank) return _dist_state(world_size, rank) From c46981db084fe5aead6f3f4703b06db7d39fb703 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:09:20 +0000 Subject: [PATCH 23/57] Run hf_ptq expert-parallel for real on two GPUs Nothing ran examples/hf_ptq/hf_ptq.py expert-parallel: its wiring tests approximate a torchrun launch with stubbed collectives, loaders and forward loops. The new test launches it under torchrun on two GPUs with MODELOPT_EP_DEGREE=2 on a tiny Qwen3-MoE, calibrates FP8 on a local calibration file, and checks that every expert of every layer is exported FP8 with its scale: a rank whose share of the experts is lost, or written under the wrong expert index, fails it. It skips below two GPUs, so it runs on the multi-GPU runners. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_llm_ptq.py | 86 +++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index a820158fa9a..6c3c1bd77c4 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -13,10 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses +import json +import os import pytest +import torch +import transformers from _test_utils.examples.hf_ptq_utils import PTQCommand from _test_utils.examples.models import MIXTRAL_PATH, T5_PATH, TINY_LLAMA_PATH +from _test_utils.examples.run_command import run_example_command +from _test_utils.torch.distributed.utils import get_free_port +from _test_utils.torch.misc import minimum_gpu +from _test_utils.torch.transformers_models import create_tiny_qwen3_moe_dir +from packaging.version import Version +from safetensors.torch import load_file def _run_sharing_the_gpu(command: PTQCommand, model_path: str): @@ -85,3 +95,79 @@ def test_ptq_mixtral(command): @pytest.mark.usefixtures("shared_gpu") def test_ptq_llama(command): _run_sharing_the_gpu(command, TINY_LLAMA_PATH) + + +@minimum_gpu(2) +@pytest.mark.usefixtures("exclusive_gpu") +def test_ptq_qwen3_moe_expert_parallel(tmp_path): + """A torchrun launch shards a tiny Qwen3-MoE's experts over two ranks, calibrates and exports. + + Each rank holds and writes only its own experts, so every expert of every layer must come out + FP8-quantized: a rank whose share is lost, or written under the wrong expert index, fails it. + """ + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("Requires sm89 or higher") + if Version(transformers.__version__) < Version("5.7"): + pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7") + num_layers, num_experts = 2, 4 + model_dir = create_tiny_qwen3_moe_dir( + tmp_path, + with_tokenizer=True, + hidden_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + intermediate_size=128, + moe_intermediate_size=128, + num_experts=num_experts, + num_experts_per_tok=2, + num_hidden_layers=num_layers, + max_position_embeddings=128, + ) + # A local calibration set: no dataset download. + calib = tmp_path / "calib.jsonl" + calib.write_text( + "".join( + json.dumps({"text": f"Sample {i}: the quick brown fox jumps over the lazy dog."}) + "\n" + for i in range(8) + ) + ) + export_dir = tmp_path / "export" + run_example_command( + [ + "torchrun", + "--nproc_per_node=2", + f"--master_port={get_free_port()}", + "hf_ptq.py", + "--pyt_ckpt_path", + str(model_dir), + "--recipe", + "general/ptq/fp8_default-kv_fp8_cast", + "--dataset", + str(calib), + "--calib_size", + "8", + "--batch_size", + "2", + "--calib_seq", + "64", + "--skip_generate", + "--no-verbose", + "--export_path", + str(export_dir), + ], + "hf_ptq", + env={**os.environ, "MODELOPT_EP_DEGREE": "2"}, + ) + + config = json.loads((export_dir / "config.json").read_text()) + assert config["quantization_config"]["quant_algo"] == "FP8" + state_dict = {} + for shard in export_dir.glob("*.safetensors"): + state_dict.update(load_file(shard)) + for layer in range(num_layers): + for expert in range(num_experts): + for proj in ("gate_proj", "up_proj", "down_proj"): + name = f"model.layers.{layer}.mlp.experts.{expert}.{proj}" + assert state_dict[f"{name}.weight"].dtype == torch.float8_e4m3fn, name + assert f"{name}.weight_scale" in state_dict, name From a42aa25c8038651a5baf419f93f0fd57ea31b2e2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:28:24 +0000 Subject: [PATCH 24/57] Trim the hf_ptq expert-parallel tests: drop duplicates, merge tables - The single-process and dense-under-torchrun plan tests go: the library's single-process tests and the get_model load test fail on the same regressions. - The EP and --use_fsdp2 load_pretrained-call tests become one table; the get_model test stays, a different loader and device placement. - The two calibration-loop tests become one parametrized test, its helper inlined. - The memory-monitor pair becomes one parametrized test. - Kept: the Nemotron-VL preview test's full plan table, whose cases are the only guard of the VL preview's own lockstep decision (a VL preview deciding by expert parallelism alone fails only [fsdp2], one ignoring the world size only [one_rank_fsdp2]); and the main-owns-the-process- group test, the only guard of main's collective timeout above NCCL's default, which the GPU smoke test runs too briefly to reach. Signed-off-by: Shengliang Xu --- .../hf_ptq/test_hf_ptq_expert_parallel.py | 137 +++++++----------- tests/unit/torch/utils/test_memory_monitor.py | 22 ++- 2 files changed, 64 insertions(+), 95 deletions(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index de574acfb28..39ac8f2c66e 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -169,17 +169,6 @@ def _deepseek_v3(n_routed_experts): # -------------------------------------------------------------------------------- the decision -@pytest.mark.parametrize("one_rank_group", [False, True], ids=["no_process_group", "one_rank"]) -def test_a_single_process_run_is_never_expert_parallel( - hf_ptq, monkeypatch, torchrun, one_rank_group -): - """With no distributed run there is nothing to shard over, even for an EP-capable MoE.""" - dist_state = torchrun(world_size=1) if one_rank_group else _dist_state(world_size=1) - _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) - - assert hf_ptq._resolve_distribution(_args(dist_state))[0] == DistributionPlan() - - def test_use_fsdp2_plans_fsdp2_without_expert_parallelism(hf_ptq, monkeypatch, torchrun): """--use_fsdp2 names its load, so even an explicit EP degree for an EP-capable MoE must not divert the run to the expert-parallel one.""" @@ -220,13 +209,6 @@ def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused( hf_ptq._resolve_distribution(args) -def test_a_dense_model_under_torchrun_runs_data_parallel(hf_ptq, monkeypatch, torchrun): - """A dense model has no experts to shard, so it loads whole on every rank, through get_model.""" - _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) - - assert hf_ptq._resolve_distribution(_args(torchrun(world_size=4)))[0] == DistributionPlan(4) - - # ---------------------------------------------------------------------------------- load_model @@ -243,54 +225,54 @@ def load_pretrained_stub(*args, **kwargs): return calls -def test_load_model_loads_an_expert_parallel_plan_through_load_pretrained( - hf_ptq, monkeypatch, torchrun +@pytest.mark.parametrize( + ("make_config", "knobs", "world_size", "options", "plan", "passed"), + [ + pytest.param( + lambda: _deepseek_v3(n_routed_experts=8), + {"MODELOPT_FSDP_SIZE": "world"}, + 4, + {"trust_remote_code": True, "attn_implementation": "eager"}, + DistributionPlan(4, ep_size=4, dense="fsdp"), + {"trust_remote_code": True, "attn_implementation": "eager"}, + id="expert_parallel", + ), + pytest.param( + lambda: LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY), + {}, + 2, + { + "use_fsdp2": True, + "cpu_offload": True, + "specdec_offline_dataset": None, + "low_memory_mode": False, + }, + DistributionPlan(2, dense="fsdp"), + {"trust_remote_code": False, "cpu_offload": True, "attn_implementation": None}, + id="use_fsdp2", + ), + ], +) +def test_load_model_hands_load_pretrained_the_plan_and_the_config_it_was_decided_on( + hf_ptq, monkeypatch, torchrun, make_config, knobs, world_size, options, plan, passed ): - """With an EP plan, load_model hands load_pretrained the plan -- the decided degree and dense - sharding -- and the spec-prepared config it was decided on.""" - monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") - served = _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + """With an EP plan, or under --use_fsdp2, load_model hands load_pretrained the plan -- the + decided degree and dense sharding -- the config it was decided on, spec-prepared for EP, and + the caller's options, such as --use_fsdp2's CPU-offload choice.""" + for name, value in knobs.items(): + monkeypatch.setenv(name, value) + served = _serve(monkeypatch, hf_ptq, make_config()) calls = _record_load_pretrained(hf_ptq, monkeypatch) - args = _args(torchrun(world_size=4), trust_remote_code=True, attn_implementation="eager") with pytest.raises(_LoaderReachedError): - hf_ptq.load_model(args) + hf_ptq.load_model(_args(torchrun(world_size=world_size), **options)) (call,) = calls assert call["ckpt_path"] == _CKPT - assert call["distribution"] == DistributionPlan(4, ep_size=4, dense="fsdp") - assert call["kwargs"] == { - "hf_config": served, - "trust_remote_code": True, - "attn_implementation": "eager", - } - assert served.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan - - -def test_use_fsdp2_loads_through_load_pretrained(hf_ptq, monkeypatch, torchrun): - """--use_fsdp2 asks load_pretrained for its FSDP2 load, with the CPU-offload choice and the - config its checks already read.""" - served = _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY)) - calls = _record_load_pretrained(hf_ptq, monkeypatch) - args = _args( - torchrun(world_size=2), - use_fsdp2=True, - cpu_offload=True, - specdec_offline_dataset=None, - low_memory_mode=False, - ) - - with pytest.raises(_LoaderReachedError): - hf_ptq.load_model(args) - - (call,) = calls - assert call["distribution"] == DistributionPlan(2, dense="fsdp") - assert call["kwargs"] == { - "hf_config": served, - "trust_remote_code": False, - "cpu_offload": True, - "attn_implementation": None, - } + assert call["distribution"] == plan + assert call["kwargs"] == {"hf_config": served, **passed} + if plan.is_expert_parallel: + assert served.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan def test_a_whole_model_loads_by_the_plan_through_get_model(hf_ptq, monkeypatch, torchrun): @@ -430,9 +412,19 @@ def calib_loader(world_size, rank=0): assert calib_loader(world_size=1) is loader -def _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh=None): - """Run mono_quantize on a model distributed as ``plan``, its experts loaded over ``ep_mesh``, and - return which calibration loop it built.""" +@pytest.mark.parametrize( + ("plan", "loop"), + [(DistributionPlan(2, ep_size=2), "distributed"), (DistributionPlan(2, dense="fsdp"), "plain")], + ids=["expert_parallel", "fsdp2"], +) +def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_loop( + hf_ptq, monkeypatch, plan, loop +): + """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion + clobbers, so an expert-parallel run must calibrate through it, over the mesh load_and_prepare_ep + loaded the experts over. Without expert parallelism there is nothing to re-plant: FSDP2 keeps + the plain loop and its split-on-OOM retry, as on main.""" + ep_mesh = object() model = torch.nn.Linear(2, 2) set_distribution_plan(model, plan) chosen, quantized = [], [] @@ -470,26 +462,7 @@ def _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh=None): ) assert len(quantized) == 1 - return chosen - - -def test_expert_parallel_calibration_runs_through_the_distributed_forward_loop(hf_ptq, monkeypatch): - """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion - clobbers, so an expert-parallel run must calibrate through it rather than the plain loop.""" - ep_mesh = object() # the mesh load_and_prepare_ep loaded the experts over - plan = DistributionPlan(2, ep_size=2) - - assert _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh) == [ - ("distributed", ep_mesh) - ] - - -def test_fsdp2_calibration_keeps_the_plain_forward_loop(hf_ptq, monkeypatch): - """Without expert parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its - split-on-OOM retry, as on main.""" - plan = DistributionPlan(2, dense="fsdp") - - assert _calibration_loop_choice(hf_ptq, monkeypatch, plan, ep_mesh=object()) == ["plain"] + assert chosen == ([("distributed", ep_mesh)] if loop == "distributed" else ["plain"]) class _PreviewModel(torch.nn.Module): diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py index c90052a48fb..8bb3dcab5bf 100644 --- a/tests/unit/torch/utils/test_memory_monitor.py +++ b/tests/unit/torch/utils/test_memory_monitor.py @@ -70,20 +70,16 @@ def launched(): monitor.stop() -def test_the_monitor_starts_when_nvml_can_report_memory(monkeypatch, launched): +@pytest.mark.parametrize("supported", [True, False], ids=["reports_memory", "cannot"]) +def test_the_monitor_starts_only_where_nvml_can_report_memory(monkeypatch, launched, supported): """The probe used to run before NVML was initialized, so it always failed and the monitor - never started on any host.""" - nvml = _FakeNVML() + never started on any host. Hosts where NVML cannot report memory still skip the monitor, and + leave NVML shut down; elsewhere the probe's own init is balanced and the monitor holds one.""" + nvml = _FakeNVML(memory_info_supported=supported) nvml.install(monkeypatch) monitor = mm.launch_memory_monitor(monitor_interval=0.01) - assert monitor is not None - launched.append(monitor) - assert nvml.inits == 1 # the probe's own init is balanced; the monitor holds one + if monitor is not None: + launched.append(monitor) - -def test_the_monitor_is_skipped_where_nvml_cannot_report_memory(monkeypatch): - """Hosts where NVML cannot report memory still skip the monitor, and leave NVML shut down.""" - nvml = _FakeNVML(memory_info_supported=False) - nvml.install(monkeypatch) - assert mm.launch_memory_monitor() is None - assert nvml.inits == 0 + assert (monitor is not None) is supported + assert nvml.inits == (1 if supported else 0) From 55669db237b5d9b56960e49a5979423346ca8997 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 05:27:54 +0000 Subject: [PATCH 25/57] Refuse the TensorRT-LLM checkpoint export in a multi-process run That exporter reads the process count as Megatron's training tensor parallelism. Under torchrun with a replicated plan, rank 0 ran it alone and waited in its barrier for ranks that had already returned; run on every rank, it would merge whole replicas as if they were shards. hf_ptq now refuses an encoder-decoder, sparsity or int8_smoothquant export with more than one process before it loads the model. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 25 ++++++++++++----- .../hf_ptq/test_hf_ptq_expert_parallel.py | 27 +++++++++++++++++++ 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index d66f76495a1..fe941c64f86 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -347,9 +347,27 @@ def _resolve_distribution(args: argparse.Namespace): "size if it divides the expert count, or MODELOPT_FSDP_SIZE=1 if the dense part fits " "one GPU." ) + if plan.world_size > 1 and _exports_tensorrt_llm_checkpoint(args, hf_model_type(hf_config)): + # That exporter reads the process count as Megatron's training tensor parallelism. Rank 0 + # alone waits in its barrier for ranks that never arrive, and every rank would merge + # whole replicas as if they were shards. Refused before the load, like the layout above. + raise NotImplementedError( + f"{plan!r}: the TensorRT-LLM checkpoint export (encoder-decoder models, " + "--sparsity_fmt, int8_smoothquant) runs in a single process; launch hf_ptq without " + "torchrun for it." + ) return plan, hf_config +def _exports_tensorrt_llm_checkpoint(args: argparse.Namespace, model_type: str | None) -> bool: + """Whether the run exports through the deprecated TensorRT-LLM checkpoint exporter.""" + return ( + is_trtllm_enc_dec_export(model_type) + or args.sparsity_fmt != "dense" + or "int8_smoothquant" in args.qformat + ) + + def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False @@ -733,12 +751,7 @@ def export_quantized( save_source_config(args, export_path) start_time = time.time() - is_tensorrt_llm_export = ( - is_trtllm_enc_dec_export(model_type) - or args.sparsity_fmt != "dense" - or "int8_smoothquant" in args.qformat - ) - if is_tensorrt_llm_export: + if _exports_tensorrt_llm_checkpoint(args, model_type): if ( args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1 ) and args.qformat == "nvfp4_svdquant": diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 39ac8f2c66e..aa4b81ce170 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -137,6 +137,8 @@ def _args(dist_state, **overrides): local_checkpoint_path=_CKPT, trust_remote_code=False, attn_implementation=None, + sparsity_fmt="dense", + qformat="fp8", ) vars(args).update(overrides) return args @@ -209,6 +211,31 @@ def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused( hf_ptq._resolve_distribution(args) +@pytest.mark.parametrize( + ("world_size", "options"), + [ + (2, {"qformat": "int8_smoothquant"}), + (2, {"sparsity_fmt": "sparsegpt"}), + (1, {"sparsity_fmt": "sparsegpt"}), + ], + ids=["int8_smoothquant", "sparsity", "single_process"], +) +def test_a_tensorrt_llm_checkpoint_export_needs_a_single_process( + hf_ptq, monkeypatch, torchrun, world_size, options +): + """That exporter reads the process count as Megatron's training tensor parallelism: launched + with torchrun, it hangs in a barrier on rank 0 alone, and on every rank it merges replicas as + shards. So it fails before the load and the calibration; a single process exports as before.""" + _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) + args = _args(torchrun(world_size=world_size), **options) + + if world_size == 1: + assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1) + else: + with pytest.raises(NotImplementedError, match="single process"): + hf_ptq._resolve_distribution(args) + + # ---------------------------------------------------------------------------------- load_model From 5a0ee25b72b2dd0f43ce015109d19526bafa9ef9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 05:28:52 +0000 Subject: [PATCH 26/57] Calibrate image batches through the expert-parallel loop hf_ptq picked the VLM calibration loop for image calibration before it checked for expert parallelism, so an expert-parallel run with images never re-planted the all-to-all forwards that quantizer insertion drops. On transformers 5.16 and later, the first routed forward then raised for want of the EP process group. Nemotron-H Omni, which has an EP spec and calibrates on images by default, hit this with default flags. distributed_forward_loop now takes an optional per-batch forward, and the VLM loop's per-batch call is factored out as create_vlm_calibration_forward. An expert-parallel image run calibrates through distributed_forward_loop with that forward on the whole VLM. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 84 +++++++++++-------- examples/hf_ptq/hf_ptq.py | 16 ++-- modelopt/torch/utils/distributed.py | 15 ++-- .../hf_ptq/test_hf_ptq_expert_parallel.py | 35 +++++--- .../utils/test_distributed_calibration.py | 15 ++++ 5 files changed, 111 insertions(+), 54 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 2071c7db655..33cda9d3b8e 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -187,18 +187,18 @@ def is_nemotron_vl(model_or_config): return any("nemotron" in arch.lower() for arch in architectures) -def create_vlm_calibration_loop(full_model, calib_dataloader): - """Create a calibration loop for VLM models that handles multimodal inputs. +def create_vlm_calibration_forward(full_model): + """Return ``forward(batch)``, which runs one multimodal calibration batch through the VLM. - This function inspects the model's forward signature and filters batch kwargs - to only include supported parameters, then calls the appropriate forward method. + It filters the batch kwargs to the parameters the model's forward accepts, then calls the + appropriate forward method. :func:`create_vlm_calibration_loop` runs it over a dataloader; an + expert-parallel run hands it to ``distributed_forward_loop`` instead. Args: full_model: The full VLM model - calib_dataloader: DataLoader yielding multimodal batches Returns: - A calibration function that can be passed to mtq.quantize() + A function that runs one batch through ``full_model`` """ # Import here to avoid circular dependency from nemotron_vl_calib import safe_nemotron_vl_forward @@ -222,37 +222,55 @@ def create_vlm_calibration_loop(full_model, calib_dataloader): or (is_nemotron_vl(full_model) and not hasattr(full_model, "extract_feature")) ) - def calibrate_loop(_model): - # Inspect model's forward signature to determine what parameters it accepts - forward_params = inspect.signature(full_model.forward).parameters - accepts_kwargs = any( - p.kind == inspect.Parameter.VAR_KEYWORD for p in forward_params.values() - ) - allowed_keys = set(forward_params.keys()) + # Inspect model's forward signature to determine what parameters it accepts + forward_params = inspect.signature(full_model.forward).parameters + accepts_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in forward_params.values()) + allowed_keys = set(forward_params.keys()) + + def forward(batch): + # For encoder-decoder models, rename input_ids → decoder_input_ids + # and disable KV caching to avoid tuple index errors in decoder layers + if is_enc_dec and "input_ids" in batch and "pixel_values" in batch: + batch["decoder_input_ids"] = batch.pop("input_ids") + if "attention_mask" in batch: + batch["decoder_attention_mask"] = batch.pop("attention_mask") + batch["use_cache"] = False + + # Filter batch to only include parameters the model accepts + if accepts_kwargs: + call_kwargs = batch + else: + call_kwargs = {k: v for k, v in batch.items() if k in allowed_keys} + # Remove None values + call_kwargs = {k: v for k, v in call_kwargs.items() if v is not None} + + if use_safe_forward: + safe_nemotron_vl_forward(full_model, call_kwargs) + else: + full_model(**call_kwargs) + + return forward + + +def create_vlm_calibration_loop(full_model, calib_dataloader): + """Create a calibration loop for VLM models that handles multimodal inputs. + + It runs every batch through :func:`create_vlm_calibration_forward`. + + Args: + full_model: The full VLM model + calib_dataloader: DataLoader yielding multimodal batches + + Returns: + A calibration function that can be passed to mtq.quantize() + """ + def calibrate_loop(_model): + forward = create_vlm_calibration_forward(full_model) full_model.eval() with torch.no_grad(): for batch in calib_dataloader: - # For encoder-decoder models, rename input_ids → decoder_input_ids - # and disable KV caching to avoid tuple index errors in decoder layers - if is_enc_dec and "input_ids" in batch and "pixel_values" in batch: - batch["decoder_input_ids"] = batch.pop("input_ids") - if "attention_mask" in batch: - batch["decoder_attention_mask"] = batch.pop("attention_mask") - batch["use_cache"] = False - - # Filter batch to only include parameters the model accepts - if accepts_kwargs: - call_kwargs = batch - else: - call_kwargs = {k: v for k, v in batch.items() if k in allowed_keys} - # Remove None values - call_kwargs = {k: v for k, v in call_kwargs.items() if v is not None} - - if use_safe_forward: - safe_nemotron_vl_forward(full_model, call_kwargs) - else: - full_model(**call_kwargs) + forward(batch) return calibrate_loop diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index fe941c64f86..e2a7f3b1f21 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -36,6 +36,7 @@ HF_PTQ, _prepare_quant_cfg, build_quant_cfg, + create_vlm_calibration_forward, create_vlm_calibration_loop, generate_excludes_prompt, get_model, @@ -605,9 +606,7 @@ def mono_quantize( # Image calibration batches contain multimodal kwargs (for example pixel_values). # They must be consumed by the complete VLM even when only a nested component is the # quantization target; the full forward still exercises that component's quantizers. - if args.calib_with_images: - calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) - elif distribution_plan(full_model).is_expert_parallel: + if distribution_plan(full_model).is_expert_parallel: # An expert-parallel run has to calibrate through distributed_forward_loop, not the # plain one: it is the only caller of _reinstall_a2a_forwards, which re-plants the # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on @@ -617,13 +616,20 @@ def mono_quantize( # On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op -- # which is why this is invisible on older transformers and fatal on newer. # It also restores the use_cache guard the plain loop does not apply. Other - # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. + # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. Image + # batches go through the VLM's own forward, as they do without expert parallelism: + # Nemotron-H Omni, for one, calibrates on images by default. calibrate_loop = distributed_forward_loop( - language_model, + full_model if args.calib_with_images else language_model, calib_dataloader, device=args.dist_state.device, ep_mesh=expert_parallel_mesh(full_model), + forward=create_vlm_calibration_forward(full_model) + if args.calib_with_images + else None, ) + elif args.calib_with_images: + calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) else: calibrate_loop = create_forward_loop( dataloader=calib_dataloader, diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index b08d44172e8..ab155ed05c1 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -896,14 +896,16 @@ def distributed_forward_loop( dataloader: DataLoader, device: torch.device | str | None = None, ep_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, + forward: Callable[[dict], object] | 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 - 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 + all-reduce / all-to-all collectives fire inside that forward), or ``forward(batch)`` when given, + for batches that need their own call into ``model``, such as a VLM's image-text batches. 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. @@ -941,7 +943,10 @@ def calibrate(_inner_model): k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items() } - model(**batch) + if forward is None: + model(**batch) + else: + forward(batch) # Again for what runs next (amax sync, export collectives), past the forwards' pool. _release_cached_cuda_memory() diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index aa4b81ce170..cddf4d1bb71 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -439,21 +439,23 @@ def calib_loader(world_size, rank=0): assert calib_loader(world_size=1) is loader +@pytest.mark.parametrize("images", [False, True], ids=["text", "images"]) @pytest.mark.parametrize( ("plan", "loop"), [(DistributionPlan(2, ep_size=2), "distributed"), (DistributionPlan(2, dense="fsdp"), "plain")], ids=["expert_parallel", "fsdp2"], ) def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_loop( - hf_ptq, monkeypatch, plan, loop + hf_ptq, monkeypatch, plan, loop, images ): """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion clobbers, so an expert-parallel run must calibrate through it, over the mesh load_and_prepare_ep - loaded the experts over. Without expert parallelism there is nothing to re-plant: FSDP2 keeps - the plain loop and its split-on-OOM retry, as on main.""" - ep_mesh = object() - model = torch.nn.Linear(2, 2) - set_distribution_plan(model, plan) + loaded the experts over, image batches included: they run through the whole VLM's own forward + there (Nemotron-H Omni calibrates on images by default). Without expert parallelism there is + nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry, as on main.""" + ep_mesh, vlm_forward = object(), object() + full_model, language_model = torch.nn.Linear(2, 2), torch.nn.Linear(2, 2) + set_distribution_plan(full_model, plan) chosen, quantized = [], [] monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False) monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: True) @@ -461,7 +463,13 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ monkeypatch.setattr( hf_ptq, "distributed_forward_loop", - lambda *a, **k: chosen.append(("distributed", k["ep_mesh"])) or object(), + lambda model, loader, **k: chosen.append((model, k["ep_mesh"], k["forward"])) or object(), + ) + monkeypatch.setattr(hf_ptq, "create_vlm_calibration_forward", lambda model: vlm_forward) + monkeypatch.setattr( + hf_ptq, + "create_vlm_calibration_loop", + lambda model, loader: chosen.append("vlm") or object(), ) monkeypatch.setattr( hf_ptq, "create_forward_loop", lambda **k: chosen.append("plain") or object() @@ -471,7 +479,7 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ ) args = SimpleNamespace( qformat="fp8", - calib_with_images=False, + calib_with_images=images, dist_state=_dist_state(world_size=2), layerwise_export=False, specdec_offline_dataset=None, @@ -480,8 +488,8 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ hf_ptq.mono_quantize( args, {"algorithm": "max"}, - model, - model, + full_model, + language_model, model_type="qwen3_moe", calibration_only=False, calib_dataloader=object(), @@ -489,7 +497,12 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ ) assert len(quantized) == 1 - assert chosen == ([("distributed", ep_mesh)] if loop == "distributed" else ["plain"]) + if loop == "plain": + assert chosen == ["vlm" if images else "plain"] + elif images: + assert chosen == [(full_model, ep_mesh, vlm_forward)] + else: + assert chosen == [(language_model, ep_mesh, None)] class _PreviewModel(torch.nn.Module): diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index b4a3e47b2c7..af5cd27aefc 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -458,6 +458,21 @@ def test_forward_loop_moves_tensors_to_the_device_it_is_given(): assert [call["logits_to_keep"] for call in model.calls] == [1, 1] +def test_forward_loop_runs_each_batch_through_the_forward_it_is_given(): + """A batch that needs its own call into the model, as a VLM's image-text batch does, goes + through ``forward``, on the device and with the cache disabled like any other.""" + batches = [{"input_ids": torch.zeros(2, 3)}] * 2 + model = _Recorder() + + distributed_forward_loop( + model, batches, device="meta", forward=lambda batch: model(**batch, own_call=True) + )(model) + + assert [call["own_call"] for call in model.calls] == [True, True] + assert [call["input_ids"].device.type for call in model.calls] == ["meta"] * 2 + assert model.use_cache_seen == [False, False] + + 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() From 39c6e2b942be8e51cc022de9eac0d2c07a6adf14 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 01:48:16 +0000 Subject: [PATCH 27/57] Let the hf_ptq EP tests' checkpoint stub take the Hub options resolve_distribution_plan now hands ensure_local_checkpoint token, cache_dir and local_files_only. The stub these tests put in its place took only the path and the group, so every test that reached the plan's checkpoint read failed with an unexpected keyword. (Found by the epdp port.) Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index cddf4d1bb71..773468cfa78 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -89,7 +89,7 @@ def _no_distribution_knobs(monkeypatch): def _checkpoint_is_local(monkeypatch): """The checkpoint path names no files; resolving it must not reach the Hub.""" monkeypatch.setattr( - model_load_utils, "ensure_local_checkpoint", lambda path, group=None: (None, path) + model_load_utils, "ensure_local_checkpoint", lambda path, group=None, **_: (None, path) ) From fa320e69e983f3430a344f1b24fca639175a909d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 18:30:32 +0000 Subject: [PATCH 28/57] Keep generate from switching a MoE's experts to batched_mm in hf_ptq's preview On a GPU, Transformers' generate runs the experts of a grouped_mm MoE on batched_mm for the decoding stage and switches them back (_optimize_model_for_decode, Transformers 5.0 and later). Two things follow. An FSDP2-wrapped model, whose class is a subclass built outside its modeling module, is refused the switch. A Qwen3-30B-A3B run under --use_fsdp2 on 4 GPUs failed in the pre-quantization preview with "FSDPQwen3MoeForCausalLM does not support setting experts implementation" (Transformers 5.13). And ModelOpt's quantizers sit on the eager, grouped_mm and torch_a2a forwards, which batched_mm bypasses, so the preview after PTQ of any grouped_mm MoE on a GPU decoded on unquantized experts. That predates the FSDP2 and expert-parallel work. The pre- and post-quantization previews and the Nemotron VL preview now run inside a context manager that keeps generate's hook off the model and puts it back. A model without the hook is left alone. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 81 ++++++++++++------- examples/hf_ptq/hf_ptq.py | 27 ++++--- tests/examples/hf_ptq/test_example_utils.py | 42 +++++++++- .../hf_ptq/test_hf_ptq_expert_parallel.py | 20 ++++- 4 files changed, 122 insertions(+), 48 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 33cda9d3b8e..d8050a6a306 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -22,7 +22,7 @@ import os import warnings from collections.abc import Iterator -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from pathlib import Path from typing import Any @@ -96,6 +96,28 @@ def preview_generate_kwargs(model) -> dict[str, Any]: return {"synced_gpus": True} if sharded else {} +@contextmanager +def keep_the_experts_implementation(model) -> Iterator[None]: + """Keep ``generate`` from switching a MoE's experts to ``batched_mm`` while it decodes. + + On a GPU, ``generate`` runs the experts of a ``grouped_mm`` MoE on ``batched_mm`` for the + decoding stage and switches them back (``_optimize_model_for_decode``, Transformers 5.0 and + later). Two things follow. ModelOpt's quantizers sit on the ``eager``, ``grouped_mm`` and + ``torch_a2a`` forwards, which ``batched_mm`` bypasses, so the preview after PTQ would decode on + unquantized experts. And Transformers refuses the switch for an FSDP2-wrapped model, whose class + is a subclass built outside its modeling module: ``FSDPQwen3MoeForCausalLM does not support + setting experts implementation``. A model without the hook is left alone. + """ + if not hasattr(model, "_optimize_model_for_decode"): + yield + return + model._optimize_model_for_decode = nullcontext + try: + yield + finally: + del model._optimize_model_for_decode + + def run_nemotron_vl_preview( full_model, tokenizer, @@ -120,36 +142,37 @@ def run_nemotron_vl_preview( """ from vlm_utils import run_text_only_generation, run_vl_preview_generation - print(f"Running text-only preview generation for Nemotron VL model ({stage_name})...") - question = tokenizer.decode(input_ids[0], skip_special_tokens=True) - lockstep = preview_generate_kwargs(full_model) - generation_config = { - "max_new_tokens": 100, - "do_sample": False, - "eos_token_id": tokenizer.eos_token_id, - **lockstep, - } - - # Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse) - text_response = run_text_only_generation( - full_model, tokenizer, question, generation_config, pyt_ckpt_path, trust_remote_code - ) + with keep_the_experts_implementation(full_model): + print(f"Running text-only preview generation for Nemotron VL model ({stage_name})...") + question = tokenizer.decode(input_ids[0], skip_special_tokens=True) + lockstep = preview_generate_kwargs(full_model) + generation_config = { + "max_new_tokens": 100, + "do_sample": False, + "eos_token_id": tokenizer.eos_token_id, + **lockstep, + } + + # Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse) + text_response = run_text_only_generation( + full_model, tokenizer, question, generation_config, pyt_ckpt_path, trust_remote_code + ) - generated_ids = None - if text_response is not None: - print(f"✅ Text-only generation successful: {text_response[:100]}...") - generated_ids = text_response - elif allow_fallback: - print("Text-only generation failed, falling back to standard generate...") - generated_ids = full_model.generate(input_ids, max_new_tokens=100, **lockstep) - - # Run additional VL test with images - print(f"Running additional VL test with images ({stage_name})...") - run_vl_preview_generation( - full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code, lockstep - ) + generated_ids = None + if text_response is not None: + print(f"✅ Text-only generation successful: {text_response[:100]}...") + generated_ids = text_response + elif allow_fallback: + print("Text-only generation failed, falling back to standard generate...") + generated_ids = full_model.generate(input_ids, max_new_tokens=100, **lockstep) + + # Run additional VL test with images + print(f"Running additional VL test with images ({stage_name})...") + run_vl_preview_generation( + full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code, lockstep + ) - return generated_ids + return generated_ids def _is_multimodal_config(config): diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index e2a7f3b1f21..b65231cc18b 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -44,6 +44,7 @@ get_tokenizer, is_nemotron_vl, is_trtllm_enc_dec_export, + keep_the_experts_implementation, mlflow_run, preview_generate_kwargs, recipe_layerwise_blocks, @@ -876,12 +877,13 @@ def pre_quantize( trust_remote_code=args.trust_remote_code, ) else: - generated_ids_before_ptq = full_model.generate( - preview_input_ids, - attention_mask=preview_attention_mask, - max_new_tokens=100, - **preview_generate_kwargs(full_model), - ) + with keep_the_experts_implementation(full_model): + generated_ids_before_ptq = full_model.generate( + preview_input_ids, + attention_mask=preview_attention_mask, + max_new_tokens=100, + **preview_generate_kwargs(full_model), + ) return preview_input_ids, preview_attention_mask, generated_ids_before_ptq @@ -935,12 +937,13 @@ def post_quantize( # quantized layer run on CPU, which some kernels (e.g. NVFP4 dynamic block # quantization) don't support. Don't let that discard the completed calibration. try: - generated_ids_after_ptq = full_model.generate( - preview_input_ids, - attention_mask=preview_attention_mask, - max_new_tokens=100, - **preview_generate_kwargs(full_model), - ) + with keep_the_experts_implementation(full_model): + generated_ids_after_ptq = full_model.generate( + preview_input_ids, + attention_mask=preview_attention_mask, + max_new_tokens=100, + **preview_generate_kwargs(full_model), + ) except Exception as e: warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif is_nemotron_vl_model and tokenizer is not None: diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 7e294c0ed0b..808d8347f8f 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -21,16 +21,16 @@ """ import warnings -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from types import SimpleNamespace from unittest.mock import patch import pytest import torch from _test_utils.examples.hf_ptq_example_utils import example_utils -from _test_utils.torch.transformers_models import get_tiny_nemotron_h +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_nemotron_h from safetensors.torch import save_file -from transformers import AutoModelForCausalLM, LlamaConfig +from transformers import AutoModelForCausalLM, GenerationMixin, LlamaConfig from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -138,6 +138,42 @@ def test_build_quant_cfg_sets_moe_calib_experts_ratio_on_every_algorithm(algorit assert quant_cfg["algorithm"] == expected +def test_a_preview_keeps_generate_from_switching_the_experts_implementation(monkeypatch): + """On a GPU, ``generate`` switches a ``grouped_mm`` MoE's experts to ``batched_mm`` for decoding, + which ModelOpt's quantizers do not intercept (the preview after PTQ decodes unquantized experts) + and which Transformers refuses for an FSDP2-wrapped class (``FSDPQwen3MoeForCausalLM does not + support setting experts implementation``). The hook ``generate`` enters is out of the way for the + preview, and back after it.""" + entered = [] + + @contextmanager + def switch(self): + entered.append(type(self).__name__) + raise ValueError(f"{type(self).__name__} does not support setting experts implementation.") + yield + + monkeypatch.setattr(GenerationMixin, "_optimize_model_for_decode", switch) + model = get_tiny_llama() + input_ids = torch.tensor([[1, 2, 3]]) + + # Some releases enter the hook only for the decode steps after the first token. + with pytest.raises(ValueError, match="does not support setting experts implementation"): + model.generate(input_ids, max_new_tokens=3) + with example_utils.keep_the_experts_implementation(model): + model.generate(input_ids, max_new_tokens=3) + + assert entered == ["LlamaForCausalLM"] # only the first, unprotected call reached the hook + assert "_optimize_model_for_decode" not in vars(model) + + +def test_a_model_without_the_decode_hook_is_previewed_as_it_is(): + model = torch.nn.Linear(1, 1) + + with example_utils.keep_the_experts_implementation(model): + assert "_optimize_model_for_decode" not in vars(model) + assert not hasattr(model, "_optimize_model_for_decode") + + def test_get_original_hf_quant_method_mxfp4_dict(): # gpt-oss layout: quantization_config is a plain dict carrying quant_method. cfg = SimpleNamespace( diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 773468cfa78..a987431aade 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -506,18 +506,26 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ class _PreviewModel(torch.nn.Module): - """A model whose ``generate`` records the kwargs each preview passes it.""" + """A model whose ``generate`` records the kwargs each preview passes it, and how often it entered + the decode hook that Transformers' ``generate`` enters to switch a MoE's experts.""" def __init__(self, plan=None): super().__init__() self.config = SimpleNamespace(is_encoder_decoder=False, model_type="qwen3_moe") self.calls = [] + self.switches = 0 if plan is not None: set_distribution_plan(self, plan) + @contextlib.contextmanager + def _optimize_model_for_decode(self): + self.switches += 1 + yield + def generate(self, input_ids, **kwargs): self.calls.append(kwargs) - return torch.cat([input_ids, input_ids[:, :1]], dim=1) + with self._optimize_model_for_decode(): + return torch.cat([input_ids, input_ids[:, :1]], dim=1) _PREVIEW_PLANS = pytest.mark.parametrize( @@ -568,6 +576,7 @@ def test_a_sharded_model_previews_in_lockstep(hf_ptq, monkeypatch, plan, lockste synced = [call.get("synced_gpus", False) for call in model.calls] assert synced == [lockstep, lockstep] # before and after quantization + assert model.switches == 0 # neither preview lets generate move the experts off their backend @_PREVIEW_PLANS @@ -575,12 +584,13 @@ def test_a_sharded_nemotron_vl_model_previews_in_lockstep(monkeypatch, plan, loc """The Nemotron-VL preview's text, fallback and image generations keep a sharded model's ranks decoding together too.""" recorded = {} + keeps = lambda model: "_optimize_model_for_decode" in vars(model) # noqa: E731 vlm_utils = SimpleNamespace( run_text_only_generation=lambda model, tokenizer, question, config, *args: recorded.update( - text=config + text=config, text_keeps_the_experts=keeps(model) ), run_vl_preview_generation=lambda model, tokenizer, path, stage, trust, kwargs: ( - recorded.update(image=kwargs) + recorded.update(image=kwargs, image_keeps_the_experts=keeps(model)) ), ) monkeypatch.setitem(sys.modules, "vlm_utils", vlm_utils) @@ -593,6 +603,8 @@ def test_a_sharded_nemotron_vl_model_previews_in_lockstep(monkeypatch, plan, loc expected = {"synced_gpus": True} if lockstep else {} assert recorded["text"].get("synced_gpus", False) == lockstep assert recorded["image"] == expected + assert recorded["text_keeps_the_experts"] and recorded["image_keeps_the_experts"] + assert "_optimize_model_for_decode" not in vars(model) assert model.calls[0].get("synced_gpus", False) == lockstep # the fallback generate From bbeb61c8648343494458aceedf64d7d7414df00e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 19:00:48 +0000 Subject: [PATCH 29/57] Run hf_ptq's sharded previews for real, under expert parallelism and FSDP2 The two-GPU hf_ptq test skipped generation, so no test reached a sharded preview. A real Qwen3-30B-A3B run did: under FSDP2 Transformers refused the experts switch generate makes for an FSDP2-wrapped class, after the previews of an earlier round had hung ranks that finished decoding at different steps. The test now runs both plans, expert parallelism and FSDP2, with the previews on, on a tiny Qwen3-MoE. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_llm_ptq.py | 29 +++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index 6c3c1bd77c4..72c27541b98 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -99,16 +99,29 @@ def test_ptq_llama(command): @minimum_gpu(2) @pytest.mark.usefixtures("exclusive_gpu") -def test_ptq_qwen3_moe_expert_parallel(tmp_path): - """A torchrun launch shards a tiny Qwen3-MoE's experts over two ranks, calibrates and exports. +@pytest.mark.parametrize( + ("launch_args", "launch_env", "min_transformers"), + [ + # qwen3_moe's base_model_ep_plan ships in transformers 5.7. + ([], {"MODELOPT_EP_DEGREE": "2"}, "5.7"), + (["--use_fsdp2"], {}, None), + ], + ids=["expert_parallel", "fsdp2"], +) +def test_ptq_qwen3_moe_sharded(tmp_path, launch_args, launch_env, min_transformers): + """A torchrun launch shards a tiny Qwen3-MoE over two ranks, previews, calibrates and exports. - Each rank holds and writes only its own experts, so every expert of every layer must come out - FP8-quantized: a rank whose share is lost, or written under the wrong expert index, fails it. + The model's experts are sharded over the ranks, or the whole model is with FSDP2. Each rank holds + and writes only its own share, so every expert of every layer must come out FP8-quantized: a rank + whose share is lost, or written under the wrong expert index, fails it. The previews generate + before and after calibration: a sharded preview broke in three ways (ranks finishing at different + steps hung in a collective, Transformers refused an FSDP2-wrapped class's experts switch, and the + switch decoded unquantized experts), and only a real launch reaches it. """ if torch.cuda.get_device_capability() < (8, 9): pytest.skip("Requires sm89 or higher") - if Version(transformers.__version__) < Version("5.7"): - pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7") + if min_transformers and Version(transformers.__version__) < Version(min_transformers): + pytest.skip(f"Requires transformers {min_transformers}") num_layers, num_experts = 2, 4 model_dir = create_tiny_qwen3_moe_dir( tmp_path, @@ -151,13 +164,13 @@ def test_ptq_qwen3_moe_expert_parallel(tmp_path): "2", "--calib_seq", "64", - "--skip_generate", "--no-verbose", "--export_path", str(export_dir), + *launch_args, ], "hf_ptq", - env={**os.environ, "MODELOPT_EP_DEGREE": "2"}, + env={**os.environ, **launch_env}, ) config = json.loads((export_dir / "config.json").read_text()) From abbc48355ee6c89be6f9c868322bb756860cfc9b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 19:32:05 +0000 Subject: [PATCH 30/57] Refuse an AutoQuantize recipe on the expert-parallel and replicated runs hf_ptq refused the TensorRT-LLM exporter on a multi-rank plan but not AutoQuantize, whose search does not run through distributed_forward_loop and scores each rank's own calibration shard without reducing the scores across ranks. Ranks could choose different formats per layer, which under expert parallelism writes a checkpoint of mixed formats. A multi-rank run now refuses the recipe before the load; --use_fsdp2 keeps main's warning for it, and a single process runs it as before. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 9 +++++++ .../hf_ptq/test_hf_ptq_expert_parallel.py | 27 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index b65231cc18b..87b5232573d 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -358,6 +358,15 @@ def _resolve_distribution(args: argparse.Namespace): "--sparsity_fmt, int8_smoothquant) runs in a single process; launch hf_ptq without " "torchrun for it." ) + if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe): + # Its search scores each rank's own calibration shard and reduces no score across ranks, so + # the ranks can choose different formats per layer: under expert parallelism, a checkpoint + # of mixed formats. --use_fsdp2 keeps main's warning for it. + raise NotImplementedError( + f"{plan!r}: AutoQuantize scores each rank's own calibration data and reduces no score " + "across ranks, so ranks could choose different formats per layer; launch hf_ptq " + "without torchrun for an AutoQuantize recipe." + ) return plan, hf_config diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index a987431aade..4d436248bfe 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -236,6 +236,33 @@ def test_a_tensorrt_llm_checkpoint_export_needs_a_single_process( hf_ptq._resolve_distribution(args) +@pytest.mark.parametrize( + ("world_size", "env", "use_fsdp2"), + [(2, {"MODELOPT_EP_DEGREE": "2"}, False), (2, {}, False), (2, {}, True), (1, {}, False)], + ids=["expert_parallel", "replicated", "use_fsdp2", "single_process"], +) +def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds( + hf_ptq, monkeypatch, torchrun, world_size, env, use_fsdp2 +): + """AutoQuantize scores each rank's own calibration shard and reduces no score across ranks, so + the ranks could choose different formats per layer. An expert-parallel or replicated run refuses + the recipe before the load; --use_fsdp2 keeps main's warning, and a single process runs it.""" + for name, value in env.items(): + monkeypatch.setenv(name, value) + _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) + args = _args( + torchrun(world_size=world_size), + use_fsdp2=use_fsdp2, + recipe="general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits", + ) + + if world_size == 1 or use_fsdp2: + hf_ptq._resolve_distribution(args) + else: + with pytest.raises(NotImplementedError, match="AutoQuantize"): + hf_ptq._resolve_distribution(args) + + # ---------------------------------------------------------------------------------- load_model From ff00dbcd8297ca89fc14b103a73f7907c671b556 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 21:19:31 +0000 Subject: [PATCH 31/57] Keep each rank on its own GPU on every load path, and stop a sharded run whose preview fails Under torchrun each rank pins the whole model to its own GPU, and every load path now honours the pin: the pack-quantized and MXFP4 paths placed by maps of their own, and DiffusionGemma asked for a sequential map, so every rank's copy landed on the same GPUs. An explicit --use_seq_device_map, which spreads one model over every GPU, is refused where each rank holds its own copy. A preview that fails on a sharded model raises instead of being skipped: its ranks decode in lockstep, so a rank that walked on to the export left the others blocked in the next step's collective; raising aborts the run. The Nemotron-VL text and image previews do the same. The offline speculative-decoding loader is split across ranks like the others, and a loader of a kind that cannot be split is refused instead of fed whole to every rank. The load tests read the spec's plan off the decision's copy of the config, which now carries it. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 44 ++++--- examples/hf_ptq/hf_ptq.py | 27 +++- examples/hf_ptq/vlm_utils.py | 4 + tests/examples/hf_ptq/test_hf_ptq_args.py | 10 +- .../hf_ptq/test_hf_ptq_expert_parallel.py | 123 +++++++++++++++++- 5 files changed, 184 insertions(+), 24 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d8050a6a306..c1d2322bdab 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -605,19 +605,20 @@ def get_model( UserWarning, ) - device_map = "auto" - if device == "cpu": - device_map = "cpu" - elif ":" in str(device) and plan.world_size > 1: - # Under torchrun each rank asks for its own "cuda:". Do not fall through to - # "auto": accelerate runs its placement heuristic INDEPENDENTLY in every rank, and on a node - # whose GPUs are all idle every rank scores them identically and picks the SAME card. All ranks' - # weights, and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL - # collective binds the whole data-parallel group to that single device and aborts with - # "Duplicate GPU detected : rank 1 and rank 0 both on CUDA device ...". Honour the caller. - # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under - # data-parallel calibration each rank must own its own GPU. A single process keeps "auto". - device_map = {"": str(device)} + # Under torchrun each rank asks for its own "cuda:". Do not fall through to "auto": + # accelerate runs its placement heuristic INDEPENDENTLY in every rank, and on a node whose GPUs + # are all idle every rank scores them identically and picks the SAME card. All ranks' weights, + # and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL collective + # binds the whole data-parallel group to that single device and aborts with "Duplicate GPU + # detected : rank 1 and rank 0 both on CUDA device ...". Every load path below honours the pin. + # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under + # data-parallel calibration each rank must own its own GPU. A single process keeps "auto". + pinned = ( + {"": str(device)} + if device != "cpu" and ":" in str(device) and plan.world_size > 1 + else None + ) + device_map = "cpu" if device == "cpu" else pinned or "auto" # Prepare config kwargs for loading config_kwargs = {"trust_remote_code": trust_remote_code} if trust_remote_code else {} @@ -648,13 +649,24 @@ def get_model( # modules together for checkpoints that fit; larger ones can still spill and split a # tied pair, and need an explicit single-device map. Multi-GPU only: a single-GPU split # cannot separate a tied pair, and sequential would needlessly cap max_memory there. - if device != "cpu" and torch.cuda.device_count() > 1 and is_diffusion_gemma(hf_config): + if ( + not pinned + and device != "cpu" + and torch.cuda.device_count() > 1 + and is_diffusion_gemma(hf_config) + ): print( "Detected DiffusionGemma model. Using device_map='sequential'; the balanced " "'auto' mapping can split its tied encoder/decoder weights across GPUs." ) use_seq_device_map = True + if use_seq_device_map and pinned: + raise ValueError( + f"--use_seq_device_map spreads one model over every GPU, but each rank of this run " + f"holds its own copy on its own GPU ({device}). Drop it, or shard the model with " + "--use_fsdp2." + ) if use_seq_device_map: device_map = "sequential" # If we use sequential, set max_memory limit to ensure that the model does not occupy the full GPU @@ -719,7 +731,7 @@ def has_pack_quantized_config(config): AutoModelForCausalLM, model_type=hf_config.model_type, distribution=plan, - device_map="auto", + device_map=pinned or "auto", trust_remote_code=trust_remote_code, dtype="auto", ) @@ -746,7 +758,7 @@ def has_pack_quantized_config(config): AutoModelForCausalLM, model_type=hf_config.model_type, distribution=plan, - device_map="cpu" if device == "cpu" else "sequential", + device_map="cpu" if device == "cpu" else pinned or "sequential", **model_kwargs, ) else: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 87b5232573d..5365c142f32 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -192,6 +192,20 @@ def _to_device(value): return {k: _to_device(v) for k, v in batch.items()} +def _shard_calib_dataloader(calib_dataloader, rank: int, world_size: int): + """This rank's share of the calibration data, in the loader type it came in.""" + if isinstance(calib_dataloader, _DeviceDataLoader): + return _DeviceDataLoader( + shard_dataloader(calib_dataloader.dataloader, rank, world_size), calib_dataloader.device + ) + if isinstance(calib_dataloader, DataLoader): + return shard_dataloader(calib_dataloader, rank, world_size) + raise NotImplementedError( + f"A {type(calib_dataloader).__name__} calibration loader cannot be split across this " + f"run's {world_size} ranks." + ) + + def make_calib_dataloader( args: argparse.Namespace, language_model: torch.nn.Module, @@ -275,12 +289,14 @@ def make_calib_dataloader( include_labels=include_labels, ) world_size = distribution_plan(language_model).world_size - if world_size > 1 and calib_dataloader is not None and isinstance(calib_dataloader, DataLoader): + if world_size > 1 and calib_dataloader is not None: # Every distributed calibration path is data-parallel over the world, so each rank calibrates # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2 # shards the params but still feeds each rank its own data. - calib_dataloader = shard_dataloader(calib_dataloader, args.dist_state.rank, world_size) + calib_dataloader = _shard_calib_dataloader( + calib_dataloader, args.dist_state.rank, world_size + ) return calib_dataloader, first_text_speech_dataset @@ -945,15 +961,20 @@ def post_quantize( # part of the model to CPU (seen on unified-memory single-GPU hosts) can make a # quantized layer run on CPU, which some kernels (e.g. NVFP4 dynamic block # quantization) don't support. Don't let that discard the completed calibration. + lockstep = preview_generate_kwargs(full_model) try: with keep_the_experts_implementation(full_model): generated_ids_after_ptq = full_model.generate( preview_input_ids, attention_mask=preview_attention_mask, max_new_tokens=100, - **preview_generate_kwargs(full_model), + **lockstep, ) except Exception as e: + if lockstep: + # A sharded model's ranks decode in lockstep: a rank that skipped the preview would + # leave the others blocked in the next step's collective. Raising aborts the run. + raise warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif is_nemotron_vl_model and tokenizer is not None: generated_ids_after_ptq = run_nemotron_vl_preview( diff --git a/examples/hf_ptq/vlm_utils.py b/examples/hf_ptq/vlm_utils.py index 16c779c0d71..68884aa4c68 100644 --- a/examples/hf_ptq/vlm_utils.py +++ b/examples/hf_ptq/vlm_utils.py @@ -193,6 +193,8 @@ def run_vl_preview_generation( return response except Exception as e: + if (generation_kwargs or {}).get("synced_gpus"): + raise # the other ranks of a sharded model would wait for this one in lockstep print(f"❌ VL preview generation {stage_name} failed: {e}") print("This may indicate issues with the quantized model") return None @@ -270,5 +272,7 @@ def run_text_only_generation( return output_text[0] except Exception as e: + if generation_config.get("synced_gpus"): + raise # the other ranks of a sharded model would wait for this one in lockstep print(f"Text-only generation failed: {e}") return None diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index fa687017724..f7b1af23ac2 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -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, set_distribution_plan from modelopt.torch.utils.mlflow import describe_run, run_tags _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -956,9 +957,12 @@ def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch): NVBug 6752977.""" hf_ptq = _import_hf_ptq(monkeypatch) - full_model = SimpleNamespace( - generate=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) - ) + class _FailingModel(torch.nn.Module): + def generate(self, *args, **kwargs): + raise RuntimeError("boom") + + full_model = _FailingModel() + set_distribution_plan(full_model, DistributionPlan()) export_calls = [] monkeypatch.setattr( hf_ptq, diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 4d436248bfe..e4cf66e2382 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -324,9 +324,14 @@ def test_load_model_hands_load_pretrained_the_plan_and_the_config_it_was_decided (call,) = calls assert call["ckpt_path"] == _CKPT assert call["distribution"] == plan - assert call["kwargs"] == {"hf_config": served, **passed} + hf_config = call["kwargs"].pop("hf_config") + assert call["kwargs"] == passed if plan.is_expert_parallel: - assert served.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + # The spec's plan goes on the decision's own copy of the config. + assert hf_config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + assert served.base_model_ep_plan is None + else: + assert hf_config is served def test_a_whole_model_loads_by_the_plan_through_get_model(hf_ptq, monkeypatch, torchrun): @@ -392,6 +397,50 @@ def load_pretrained(ckpt_path, auto_class, **kwargs): assert placed == [device_map] +@pytest.mark.parametrize( + "checkpoint", + ["pack-quantized", "mxfp4", "diffusion-gemma"], +) +def test_every_load_path_keeps_the_rank_pin(monkeypatch, tmp_path, checkpoint): + """A pack-quantized or MXFP4 checkpoint loads through a path of its own, and DiffusionGemma asks + for a sequential map on a multi-GPU node; each would otherwise place every rank's copy over the + same GPUs.""" + quantization_config = { + "pack-quantized": {"quant_method": "compressed-tensors", "format": "pack-quantized"}, + "mxfp4": {"quant_method": "mxfp4"}, + }.get(checkpoint) + LlamaConfig( + architectures=["LlamaForCausalLM"], quantization_config=quantization_config, **_TINY + ).save_pretrained(tmp_path) + if checkpoint == "diffusion-gemma": + monkeypatch.setattr(example_utils, "is_diffusion_gemma", lambda config: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(example_utils, "_unpack_compressed_linear_weights", lambda *args: None) + placed = [] + + def load_pretrained(ckpt_path, auto_class, **kwargs): + placed.append(kwargs["device_map"]) + return torch.nn.Linear(1, 1) + + monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained) + + example_utils.get_model(str(tmp_path), device="cuda:1", distribution=DistributionPlan(4)) + + assert placed == [{"": "cuda:1"}] + + +def test_a_sequential_map_is_refused_where_each_rank_holds_its_own_copy(tmp_path): + LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) + + with pytest.raises(ValueError, match="--use_seq_device_map"): + example_utils.get_model( + str(tmp_path), + device="cuda:1", + distribution=DistributionPlan(4), + use_seq_device_map=True, + ) + + # ------------------------------------------------------------------ process group and dist_state @@ -635,6 +684,76 @@ def test_a_sharded_nemotron_vl_model_previews_in_lockstep(monkeypatch, plan, loc assert model.calls[0].get("synced_gpus", False) == lockstep # the fallback generate +def test_the_offline_specdec_loader_is_split_across_ranks_too(hf_ptq): + """It wraps its DataLoader to move each batch onto the device; each rank still calibrates on + its own share, as it would from the bare loader.""" + + def loader(): + return DataLoader(list(range(8)), batch_size=1, collate_fn=lambda b: {"x": torch.tensor(b)}) + + wrapped = hf_ptq._DeviceDataLoader(loader(), torch.device("cpu")) + + shard = hf_ptq._shard_calib_dataloader(wrapped, 1, 2) + + assert isinstance(shard, hf_ptq._DeviceDataLoader) + assert [b["x"].tolist() for b in shard] == [ + b["x"].tolist() for b in dist.shard_dataloader(loader(), 1, 2) + ] + with pytest.raises(NotImplementedError, match="cannot be split"): + hf_ptq._shard_calib_dataloader(iter(loader()), 1, 2) + + +@pytest.mark.parametrize( + ("plan", "raises"), + [(DistributionPlan(4), False), (DistributionPlan(4, ep_size=4), True)], + ids=["replicated", "expert_parallel"], +) +def test_a_failed_preview_of_a_sharded_model_stops_the_run(hf_ptq, monkeypatch, plan, raises): + """A rank that skipped a sharded model's preview would leave the others decoding in lockstep, + blocked in the next step's collective, so the failure is raised, which aborts the run. A model + every rank holds whole keeps the best-effort skip.""" + + class _FailingModel(torch.nn.Module): + def generate(self, *args, **kwargs): + raise RuntimeError("the preview failed on this rank") + + model = _FailingModel() + set_distribution_plan(model, plan) + monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None) + ids = torch.zeros(1, 4, dtype=torch.long) + args = SimpleNamespace(specdec_offline_dataset=None, verbose=False) + run = lambda: hf_ptq.post_quantize( # noqa: E731 + args, model, model, "llama", None, None, ids, None, ids, False, None, None + ) + + if raises: + with pytest.raises(RuntimeError, match="failed on this rank"): + run() + else: + with pytest.warns(UserWarning, match="sanity check failed"): + run() + + +@pytest.mark.parametrize(("lockstep", "raises"), [({}, False), ({"synced_gpus": True}, True)]) +def test_a_failed_text_only_vl_preview_stops_a_sharded_run(monkeypatch, lockstep, raises): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + vlm_utils = importlib.import_module("vlm_utils") + + class _ChatModel: + def chat(self, *args, **kwargs): + raise RuntimeError("the preview failed on this rank") + + run = lambda: vlm_utils.run_text_only_generation( # noqa: E731 + _ChatModel(), None, "question", dict(lockstep), _CKPT + ) + + if raises: + with pytest.raises(RuntimeError, match="failed on this rank"): + run() + else: + assert run() is None + + @pytest.mark.parametrize("generation_kwargs", [None, {"synced_gpus": True}]) def test_the_vl_image_preview_generates_with_the_kwargs_it_is_handed( monkeypatch, tmp_path, generation_kwargs From c3294aed9da3ae3c2d3348c78991b7d13c7c2222 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:29:39 +0000 Subject: [PATCH 32/57] Have a sharded model's ranks skip the image preview together Each node reads the VL preview's sample images from its own checkout, and the repository ships none, so the ranks could disagree on whether there is one. A rank without an image returned before generating while its peers decoded in lockstep, which left them blocked in the next step's collective until the process-group timeout. The ranks of a sharded model now decide together: every rank previews, or none does. Raising instead would stop every sharded Nemotron-VL run on a checkout without images. A failed preview that stops the run now says the export was not reached and that --skip_generate exports without the preview, with the preview's error as its cause. The GPU memory monitor reads every GPU on the node, so rank 0 starts it rather than every rank. _resolve_distribution no longer refuses FSDP2 within several expert-data-parallel replicas itself: no plan from resolve_distribution_plan reaches that check, since resolve_distribution_plan refuses the layout first, and load_and_prepare_ep keeps its own refusal for a plan handed to it. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 41 +++++------ examples/hf_ptq/vlm_utils.py | 68 ++++++++++-------- tests/examples/hf_ptq/test_hf_ptq_args.py | 25 +++++++ .../hf_ptq/test_hf_ptq_expert_parallel.py | 70 ++++++++++++++++--- 4 files changed, 142 insertions(+), 62 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 5365c142f32..8fbd7203174 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -120,6 +120,10 @@ RAND_SEED = 1234 # Most tokens per calibration batch when layerwise calibration picks the batch size (64 x 512). LAYERWISE_CALIB_TOKENS = 32768 +_PREVIEW_STOPS_THE_RUN = ( + "The preview after quantization failed, which stops the run before the export, so no " + "checkpoint is written; --skip_generate exports without the preview." +) def _layerwise_calib_batch_size(calib_seq: int, num_samples: int) -> int: @@ -357,14 +361,6 @@ def _resolve_distribution(args: argparse.Namespace): "or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU; an MoE can instead shard its " "experts (MODELOPT_EP_DEGREE above 1)." ) - if plan.is_expert_parallel and plan.is_fsdp_dense and plan.edp > 1: - # Refused before the load: the export, after the whole calibration, is what cannot do it. - raise NotImplementedError( - f"{plan!r} FSDP2-shards the dense part within each of {plan.edp} expert-data-parallel " - "replicas, a layout the export cannot write yet. Set MODELOPT_EP_DEGREE to the world " - "size if it divides the expert count, or MODELOPT_FSDP_SIZE=1 if the dense part fits " - "one GPU." - ) if plan.world_size > 1 and _exports_tensorrt_llm_checkpoint(args, hf_model_type(hf_config)): # That exporter reads the process count as Megatron's training tensor parallelism. Rank 0 # alone waits in its barrier for ranks that never arrive, and every rank would merge @@ -974,18 +970,21 @@ def post_quantize( if lockstep: # A sharded model's ranks decode in lockstep: a rank that skipped the preview would # leave the others blocked in the next step's collective. Raising aborts the run. - raise + raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif is_nemotron_vl_model and tokenizer is not None: - generated_ids_after_ptq = run_nemotron_vl_preview( - full_model, - tokenizer, - preview_input_ids, - args.local_checkpoint_path, - "after quantization", - allow_fallback=False, - trust_remote_code=args.trust_remote_code, - ) + try: + generated_ids_after_ptq = run_nemotron_vl_preview( + full_model, + tokenizer, + preview_input_ids, + args.local_checkpoint_path, + "after quantization", + allow_fallback=False, + trust_remote_code=args.trust_remote_code, + ) + except Exception as e: + raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e else: warnings.warn( "Llama4 Maverick generation after quantization has a bug. Skipping generation sample." @@ -1639,8 +1638,10 @@ def _main_inner(args: argparse.Namespace): # path. args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path) - # launch a memory monitor to read the currently used GPU memory. - launch_memory_monitor() + # launch a memory monitor to read the currently used GPU memory. It reads every GPU on the + # node, so a distributed run starts one, on rank 0, rather than one per rank. + if args.dist_state.is_main: + launch_memory_monitor() # Force eager execution for all model types. torch.compiler.set_stance("force_eager") diff --git a/examples/hf_ptq/vlm_utils.py b/examples/hf_ptq/vlm_utils.py index 68884aa4c68..87cd7e7ec3d 100644 --- a/examples/hf_ptq/vlm_utils.py +++ b/examples/hf_ptq/vlm_utils.py @@ -20,6 +20,36 @@ from PIL import Image from transformers import AutoImageProcessor, AutoProcessor +from modelopt.torch.utils.distributed import all_reduce + + +def _load_sample_image(images_dir): + """The first sample image in ``images_dir`` that opens, or None, saying why there is none.""" + if not os.path.exists(images_dir): + print(f"❌ Warning: Images directory not found at {images_dir}") + return None + + # Use single image for VL preview to avoid shape mismatch issues + image_files = ["example1a.jpeg", "example1b.jpeg", "example.jpg", "test.jpg", "sample.png"] + missing_files = [] + for img_file in image_files: + img_path = os.path.join(images_dir, img_file) + if os.path.exists(img_path): + try: + image = Image.open(img_path) + print(f" ✅ Successfully loaded: {img_file}") + return image # Use the first available image + except Exception as e: + print(f" ⚠️ Warning: Could not open {img_file}: {e}") + missing_files.append(f"{img_file} (corrupted)") + else: + missing_files.append(img_file) + + print(f"❌ Warning: No valid sample images found in {images_dir}") + print(f" Searched for: {', '.join(image_files)}") + print(f" Missing/invalid files: {', '.join(missing_files)}") + return None + def run_vl_preview_generation( model, tokenizer, model_path, stage_name, trust_remote_code=False, generation_kwargs=None @@ -41,37 +71,13 @@ def run_vl_preview_generation( # Load sample images from the images directory script_dir = os.path.dirname(os.path.abspath(__file__)) - images_dir = os.path.join(script_dir, "images") - - # Check if images directory exists - if not os.path.exists(images_dir): - print(f"❌ Warning: Images directory not found at {images_dir}") - print(" VL preview generation requires sample images to test vision capabilities.") - print(" Skipping VL preview generation.") - return None - - # Use single image for VL preview to avoid shape mismatch issues - image_files = ["example1a.jpeg", "example1b.jpeg", "example.jpg", "test.jpg", "sample.png"] - image = None - missing_files = [] - for img_file in image_files: - img_path = os.path.join(images_dir, img_file) - if os.path.exists(img_path): - try: - image = Image.open(img_path) - print(f" ✅ Successfully loaded: {img_file}") - break # Use the first available image - except Exception as e: - print(f" ⚠️ Warning: Could not open {img_file}: {e}") - missing_files.append(f"{img_file} (corrupted)") - else: - missing_files.append(img_file) - - if image is None: - print(f"❌ Warning: No valid sample images found in {images_dir}") - print(f" Searched for: {', '.join(image_files)}") - if missing_files: - print(f" Missing/invalid files: {', '.join(missing_files)}") + image = _load_sample_image(os.path.join(script_dir, "images")) + found = image is not None + if (generation_kwargs or {}).get("synced_gpus"): + # Each node reads the images from its own checkout, while the ranks of a sharded model + # decode in lockstep: every rank previews, or none does. + found = all_reduce("VL preview sample image", lambda: image is not None, all) + if not found: print(" VL preview generation requires sample images to test vision capabilities.") print(" Skipping VL preview generation.") return None diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index f7b1af23ac2..95884c0dbe1 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -1026,6 +1026,31 @@ def load_model(args): ] +@pytest.mark.parametrize("is_main", [True, False], ids=["rank_0", "other_rank"]) +def test_one_rank_runs_the_gpu_memory_monitor(monkeypatch, tmp_path, is_main): + """The monitor reads every GPU on the node: one per rank would poll them all once per rank and + print the same report once per rank.""" + hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path)) + args.dist_state = SimpleNamespace(is_main=is_main, world_size=2) + launched = [] + + class _LoadReachedError(Exception): + pass + + def load_model(args): + raise _LoadReachedError + + monkeypatch.setattr(torch.compiler, "set_stance", lambda *args, **kwargs: None) + monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: launched.append(True)) + monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", lambda path: (None, path)) + monkeypatch.setattr(hf_ptq, "load_model", load_model) + + with pytest.raises(_LoadReachedError): + hf_ptq._main_inner(args) + + assert len(launched) == is_main + + def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypatch, tmp_path): """The deprecated TensorRT-LLM exporter does not copy them itself; hf_ptq does, from disk.""" hf_ptq, args = _parse_hf_ptq_args( diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index e4cf66e2382..4a12c634e3c 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -704,31 +704,47 @@ def loader(): @pytest.mark.parametrize( - ("plan", "raises"), - [(DistributionPlan(4), False), (DistributionPlan(4, ep_size=4), True)], - ids=["replicated", "expert_parallel"], + ("plan", "nemotron_vl", "raises"), + [ + (DistributionPlan(4), False, False), + (DistributionPlan(4, ep_size=4), False, True), + (DistributionPlan(4, ep_size=4), True, True), + ], + ids=["replicated", "expert_parallel", "expert_parallel_nemotron_vl"], ) -def test_a_failed_preview_of_a_sharded_model_stops_the_run(hf_ptq, monkeypatch, plan, raises): +def test_a_failed_preview_of_a_sharded_model_stops_the_run( + hf_ptq, monkeypatch, plan, nemotron_vl, raises +): """A rank that skipped a sharded model's preview would leave the others decoding in lockstep, - blocked in the next step's collective, so the failure is raised, which aborts the run. A model - every rank holds whole keeps the best-effort skip.""" + blocked in the next step's collective, so the failure is raised, which aborts the run, with the + preview's error as its cause. A model every rank holds whole keeps the best-effort skip.""" + + def fail(*args, **kwargs): + raise RuntimeError("the preview failed on this rank") class _FailingModel(torch.nn.Module): - def generate(self, *args, **kwargs): - raise RuntimeError("the preview failed on this rank") + generate = fail model = _FailingModel() set_distribution_plan(model, plan) monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None) + monkeypatch.setattr(hf_ptq, "run_nemotron_vl_preview", fail) # as vlm_utils raises in lockstep ids = torch.zeros(1, 4, dtype=torch.long) - args = SimpleNamespace(specdec_offline_dataset=None, verbose=False) + args = SimpleNamespace( + specdec_offline_dataset=None, + verbose=False, + local_checkpoint_path=_CKPT, + trust_remote_code=False, + ) + tokenizer = SimpleNamespace() run = lambda: hf_ptq.post_quantize( # noqa: E731 - args, model, model, "llama", None, None, ids, None, ids, False, None, None + args, model, model, "llama", tokenizer, None, ids, None, ids, nemotron_vl, None, None ) if raises: - with pytest.raises(RuntimeError, match="failed on this rank"): + with pytest.raises(RuntimeError) as stopped: run() + assert "failed on this rank" in str(stopped.value.__cause__) else: with pytest.warns(UserWarning, match="sanity check failed"): run() @@ -784,6 +800,38 @@ def chat(self, tokenizer, question, generation_config, **image_features): assert received[0].get("synced_gpus", False) == (generation_kwargs is not None) +@pytest.mark.parametrize("has_image", [True, False], ids=["a_peer_has_none", "no_rank_has_one"]) +def test_the_ranks_of_a_sharded_model_skip_the_image_preview_together( + monkeypatch, tmp_path, has_image +): + """Each node reads the sample images from its own checkout, which may have none: the repository + ships none. A rank that previewed alone would wait for its peers in lockstep, so a rank with an + image skips when a peer has none, and when no rank has one, they all skip without stopping.""" + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + vlm_utils = importlib.import_module("vlm_utils") + if has_image: + (tmp_path / "images").mkdir() + Image.new("RGB", (4, 4)).save(tmp_path / "images" / "example1a.jpeg") + monkeypatch.setattr(vlm_utils, "__file__", str(tmp_path / "vlm_utils.py")) + monkeypatch.setattr( + vlm_utils.AutoImageProcessor, "from_pretrained", lambda *args, **kwargs: lambda images: {} + ) + # This rank's answer, and a peer's that found no image. + monkeypatch.setattr(vlm_utils, "all_reduce", lambda label, fn, combine: combine([fn(), False])) + + class _ChatModel: + device = torch.device("cpu") + + def chat(self, *args, **kwargs): + pytest.fail("this rank previewed without its peers") + + preview = vlm_utils.run_vl_preview_generation( + _ChatModel(), SimpleNamespace(eos_token_id=0), _CKPT, "stage", False, {"synced_gpus": True} + ) + + assert preview is None + + # ------------------------------------------------------------------------------------- export From 82b3e9b0944b73c14ccc9d775361d30a14227e42 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:05:13 +0000 Subject: [PATCH 33/57] Refuse single-process load options in hf_ptq's multi-rank runs An expert-parallel run never reaches get_model or the low-memory loader, so it accepted --low_memory_mode, --use_seq_device_map and --offload_folder and ignored them: torchrun with --low_memory_mode loaded an MoE at full precision without a word. A replicated run could not use them either. Its pin of each rank's copy to the rank's GPU leaves --offload_folder nothing to offload, and the low-memory loader cannot run under a process group: it all-reduces amax computed from meta weights, and hands accelerate its own device map next to the pin. Both routes now refuse the three options before loading, as they refuse the TensorRT-LLM exporter and AutoQuantize. A single process and --use_fsdp2 keep main's handling. get_model no longer sizes a pinned copy's memory over the node's GPUs: the pinned map ignores max_memory, so the limit it printed was never applied. The README no longer points GPT-OSS and Llama-4 at --use_fsdp2, which refuses their released checkpoints (an MXFP4 quantization_config, a multimodal config); it says to quantize them in a single process. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 +- examples/hf_ptq/example_utils.py | 4 +- examples/hf_ptq/hf_ptq.py | 17 ++++++++ .../hf_ptq/test_hf_ptq_expert_parallel.py | 43 +++++++++++++++++-- 4 files changed, 61 insertions(+), 5 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 7d28f3c8b72..34bc1cf54ad 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -678,7 +678,7 @@ Two environment variables override the automatic distribution plan: | `MODELOPT_EP_DEGREE` | `auto` | Expert-parallel degree. `auto` picks the largest degree that divides both the routed-expert count and the world size. Below the world size, each group of `MODELOPT_EP_DEGREE` ranks holds a full set of experts, and the groups calibrate on different data. | | `MODELOPT_FSDP_SIZE` | `auto` | Whether to FSDP2-shard the dense part. `auto` shards it when it would take more than 60% of one GPU's memory; `world` always shards it; `1` replicates it. With `MODELOPT_EP_DEGREE` below the world size the export cannot write a sharded dense part yet, so such a run fails before loading. | -A run that cannot shard the model the way the plan asks fails before loading, rather than loading the whole model on every rank. `--use_fsdp2` skips expert parallelism and uses the [FSDP2 path](#multi-node-post-training-quantization-with-fsdp2) instead. That is also the way to quantize GPT-OSS and Llama-4 across GPUs: their experts are stored transposed, which the expert-parallel export cannot write yet, so expert parallelism refuses them. +A run that cannot shard the model the way the plan asks fails before loading, rather than loading the whole model on every rank. `--use_fsdp2` skips expert parallelism and uses the [FSDP2 path](#multi-node-post-training-quantization-with-fsdp2) instead. Expert parallelism refuses GPT-OSS and Llama-4: their experts are stored transposed, which the expert-parallel export cannot write yet. `--use_fsdp2` refuses their released checkpoints too (GPT-OSS's carry an MXFP4 `quantization_config`, and Llama-4's are multimodal), so quantize those in a single process (without `torchrun`), which spreads the model over the node's GPUs. ## Evaluate Accuracy diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index c1d2322bdab..a15323c2840 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -844,7 +844,9 @@ def has_pack_quantized_config(config): f"Offload folder: {offload_folder}\n" "Weights exceeding GPU+CPU budgets will be streamed from disk." ) - else: + elif not pinned: + # A pinned copy goes onto its rank's GPU whole, by a map that ignores max_memory: sizing + # it over the node's GPUs would print a limit the load does not apply. inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) if "cpu" in inferred_device_map.values(): for _device in max_memory: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 8fbd7203174..4226f8fa930 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -379,6 +379,23 @@ def _resolve_distribution(args: argparse.Namespace): "across ranks, so ranks could choose different formats per layer; launch hf_ptq " "without torchrun for an AutoQuantize recipe." ) + if plan.world_size > 1 and not args.use_fsdp2: + # These configure loading one model in a single process: compressed as it loads, or + # spread over the node's GPUs, CPU and disk. An expert-parallel run never reaches those + # loaders, and a replicated one pins each rank's copy to the rank's GPU, which none of them + # works with. Refused before the load rather than ignored. --max_cpu_memory_gb and + # --max_gpu_memory_gb come only with --offload_folder. + for flag, given in ( + ("--low_memory_mode", args.low_memory_mode), + ("--use_seq_device_map", args.use_seq_device_map), + ("--offload_folder", args.offload_folder is not None), + ): + if given: + raise NotImplementedError( + f"{plan!r}: {flag} configures loading one model in a single process, but each " + "rank of this run loads its own share or copy of it onto its own GPU. Drop it, " + "or launch hf_ptq without torchrun to use it." + ) return plan, hf_config diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 4a12c634e3c..f9a66b3f46e 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -139,6 +139,9 @@ def _args(dist_state, **overrides): attn_implementation=None, sparsity_fmt="dense", qformat="fp8", + low_memory_mode=False, + use_seq_device_map=False, + offload_folder=None, ) vars(args).update(overrides) return args @@ -263,6 +266,37 @@ def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds( hf_ptq._resolve_distribution(args) +@pytest.mark.parametrize( + "option", + [{"low_memory_mode": True}, {"use_seq_device_map": True}, {"offload_folder": "/offload"}], + ids=["low_memory_mode", "use_seq_device_map", "offload_folder"], +) +@pytest.mark.parametrize( + ("make_config", "world_size"), + [ + (lambda: _deepseek_v3(n_routed_experts=8), 2), + (lambda: LlamaConfig(**_TINY), 2), + (lambda: LlamaConfig(**_TINY), 1), + ], + ids=["expert_parallel", "replicated", "single_process"], +) +def test_the_single_process_load_options_are_refused_on_the_routes_this_script_adds( + hf_ptq, monkeypatch, torchrun, make_config, world_size, option +): + """Each rank of an expert-parallel or replicated run loads its own share or copy of the model + onto its own GPU, where these options of a single process's load would be ignored or would break + the load. Both routes refuse them before the load; a single process loads by them as before.""" + _serve(monkeypatch, hf_ptq, make_config()) + args = _args(torchrun(world_size=world_size), **option) + + if world_size == 1: + assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1) + else: + (name,) = option + with pytest.raises(NotImplementedError, match=f"--{name}"): + hf_ptq._resolve_distribution(args) + + # ---------------------------------------------------------------------------------- load_model @@ -379,22 +413,25 @@ def test_a_rank_device_pins_the_whole_model_to_it( ): """Under torchrun each rank asks for its own cuda:; "auto" would let every rank's accelerate pick the same idle GPU. An unindexed "cuda" still spreads the model over the GPUs, - and so does a single process, whatever device it names.""" + and so does a single process, whatever device it names. Only a spread model is held to a memory + limit sized over the node's GPUs: the pinned map would ignore one.""" LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) + # No GPU on the node holds the model, so sizing it over the node limits its memory. + monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 0, "cpu": 1 << 40}) placed = [] plan = DistributionPlan(world_size) def load_pretrained(ckpt_path, auto_class, **kwargs): assert kwargs["distribution"] == plan - placed.append(kwargs["device_map"]) + placed.append((kwargs["device_map"], "max_memory" in kwargs)) return torch.nn.Linear(1, 1) monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained) example_utils.get_model(str(tmp_path), device=device, distribution=plan) - assert placed == [device_map] + assert placed == [(device_map, device_map == "auto")] @pytest.mark.parametrize( From 4971185c894c4406bda8b8042aa386c752256654 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 02:08:34 +0000 Subject: [PATCH 34/57] Refuse a sharded vLLM fake-quant export; re-plant a2a forwards after PTQ An expert-parallel run that quantized without running the calibration loop lost its experts' all-to-all forwards: quantizer insertion drops them, and only distributed_forward_loop re-planted them, before its first forward. A config without static activation scales (int8_weight_only, w4a16_nvfp4, ...) calibrates without the loop, and algorithm None (the mxfp4, mxfp6, mxfp8 and mxint8 presets) never calls the loop it is handed, so the preview after PTQ raised "torch_a2a experts forward requires the EP process_group" and stopped the run before the export. mono_quantize now re-plants them once the model is quantized; after a loop that ran, that is a no-op. An expert-parallel run refuses --vllm_fakequant_export before the load. That exporter writes the whole model with save_pretrained from each process it runs in, and every rank of a sharded model exports, so each rank would write its own share of the experts into the same directory. Rank 0 alone still writes a replicated model. keep_the_experts_implementation puts back an _optimize_model_for_decode the model held itself instead of deleting it, so a hook bound per instance survives the preview and the manager nests. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 7 +- examples/hf_ptq/hf_ptq.py | 16 +++- tests/examples/hf_ptq/test_example_utils.py | 14 ++++ .../hf_ptq/test_hf_ptq_expert_parallel.py | 81 +++++++++++++++++++ 4 files changed, 116 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index a15323c2840..bc1979e9694 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -111,11 +111,16 @@ def keep_the_experts_implementation(model) -> Iterator[None]: if not hasattr(model, "_optimize_model_for_decode"): yield return + # Put back what the model itself held, if anything: a nested use, or a hook bound per instance. + own = vars(model).get("_optimize_model_for_decode") model._optimize_model_for_decode = nullcontext try: yield finally: - del model._optimize_model_for_decode + if own is None: + del model._optimize_model_for_decode + else: + model._optimize_model_for_decode = own def run_nemotron_vl_preview( diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 4226f8fa930..6f04ae12a83 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -101,6 +101,7 @@ get_supported_datasets, ) from modelopt.torch.utils.distributed import ( + _reinstall_a2a_forwards, distributed_forward_loop, distributed_run, distribution_plan, @@ -370,6 +371,14 @@ def _resolve_distribution(args: argparse.Namespace): "--sparsity_fmt, int8_smoothquant) runs in a single process; launch hf_ptq without " "torchrun for it." ) + if (plan.is_expert_parallel or plan.is_fsdp_dense) and args.vllm_fakequant_export: + # That exporter writes the whole model with save_pretrained from each process it runs in, + # and every rank of a sharded model exports: each would write its own part into the same + # directory. Rank 0 alone writes a replicated model; --use_fsdp2 refuses it while parsing. + raise NotImplementedError( + f"{plan!r}: --vllm_fakequant_export writes the whole model from one process, which " + "holds only its own part of a sharded model; launch hf_ptq without torchrun for it." + ) if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe): # Its search scores each rank's own calibration shard and reduces no score across ranks, so # the ranks can choose different formats per layer: under expert parallelism, a checkpoint @@ -647,7 +656,7 @@ def mono_quantize( # quantization target; the full forward still exercises that component's quantizers. if distribution_plan(full_model).is_expert_parallel: # An expert-parallel run has to calibrate through distributed_forward_loop, not the - # plain one: it is the only caller of _reinstall_a2a_forwards, which re-plants the + # plain one: it calls _reinstall_a2a_forwards first, which re-plants the # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on # transformers >= 5.16 (where the styles install a forward rather than a pre-hook). # Without it the experts forward runs with process_group=None and raises @@ -686,6 +695,11 @@ def mono_quantize( ) else: language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop) + if distribution_plan(full_model).is_expert_parallel: + # The calibration loop re-plants the a2a forwards quantizer insertion drops, but a + # config without static activation scales calibrates without it, and algorithm None + # (the MX formats) never calls it; the preview and the export forward through them. + _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model)) # For VL models, update full_model to use the quantized language model if is_nemotron_vl_model and language_model is not full_model: diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 808d8347f8f..458c5aa1109 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -174,6 +174,20 @@ def test_a_model_without_the_decode_hook_is_previewed_as_it_is(): assert not hasattr(model, "_optimize_model_for_decode") +def test_a_decode_hook_on_the_model_itself_is_put_back_and_previews_nest(): + """A hook bound on the model instance is restored rather than dropped, and a preview inside + another leaves the outer one's stand-in in place.""" + model = torch.nn.Linear(1, 1) + own = model._optimize_model_for_decode = object() + + with example_utils.keep_the_experts_implementation(model): + with example_utils.keep_the_experts_implementation(model): + pass + assert model._optimize_model_for_decode is nullcontext + + assert model._optimize_model_for_decode is own + + def test_get_original_hf_quant_method_mxfp4_dict(): # gpt-oss layout: quantization_config is a plain dict carrying quant_method. cfg = SimpleNamespace( diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index f9a66b3f46e..fe80b2582cf 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -142,6 +142,7 @@ def _args(dist_state, **overrides): low_memory_mode=False, use_seq_device_map=False, offload_folder=None, + vllm_fakequant_export=False, ) vars(args).update(overrides) return args @@ -297,6 +298,32 @@ def test_the_single_process_load_options_are_refused_on_the_routes_this_script_a hf_ptq._resolve_distribution(args) +@pytest.mark.parametrize( + ("make_config", "world_size", "refused"), + [ + (lambda: _deepseek_v3(n_routed_experts=8), 2, True), + (lambda: LlamaConfig(**_TINY), 2, False), + (lambda: LlamaConfig(**_TINY), 1, False), + ], + ids=["expert_parallel", "replicated", "single_process"], +) +def test_the_vllm_fakequant_export_is_refused_for_a_sharded_model( + hf_ptq, monkeypatch, torchrun, make_config, world_size, refused +): + """That exporter writes the whole model with save_pretrained from each process it runs in, so + every rank of an expert-parallel run would write its own experts into the same directory. Such a + run refuses it before the load; rank 0 alone writes a replicated model, and a single process + exports as before.""" + _serve(monkeypatch, hf_ptq, make_config()) + args = _args(torchrun(world_size=world_size), vllm_fakequant_export=True) + + if refused: + with pytest.raises(NotImplementedError, match="--vllm_fakequant_export"): + hf_ptq._resolve_distribution(args) + else: + hf_ptq._resolve_distribution(args) + + # ---------------------------------------------------------------------------------- load_model @@ -618,6 +645,60 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ assert chosen == [(language_model, ep_mesh, None)] +@pytest.mark.parametrize( + "loop_handed", [False, True], ids=["calibration_free", "loop_never_called"] +) +@pytest.mark.parametrize( + ("plan", "replanted"), + [(DistributionPlan(2, ep_size=2), True), (DistributionPlan(2, dense="fsdp"), False)], + ids=["expert_parallel", "fsdp2"], +) +def test_an_expert_parallel_run_replants_the_all_to_all_forwards_after_quantizing( + hf_ptq, monkeypatch, plan, replanted, loop_handed +): + """The calibration loop re-plants the all-to-all forwards that quantizer insertion drops only + when the calibration runs it: a config without static activation scales calibrates without the + loop, and algorithm None (the MX formats) never calls the loop it is handed. The preview and the + export forward through the experts either way, so an expert-parallel run re-plants them once + quantized; without expert parallelism there is nothing to re-plant.""" + ep_mesh, full_model = object(), torch.nn.Linear(2, 2) + set_distribution_plan(full_model, plan) + events = [] + monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False) + monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: loop_handed) + monkeypatch.setattr(hf_ptq, "expert_parallel_mesh", lambda model: ep_mesh) + monkeypatch.setattr(hf_ptq, "distributed_forward_loop", lambda *args, **kwargs: object()) + monkeypatch.setattr(hf_ptq, "create_forward_loop", lambda **kwargs: object()) + monkeypatch.setattr( + hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: events.append("quantize") or m + ) + monkeypatch.setattr( + hf_ptq, + "_reinstall_a2a_forwards", + lambda model, mesh: events.append(("replant", model, mesh)), + ) + args = SimpleNamespace( + qformat="mxfp8", + calib_with_images=False, + dist_state=_dist_state(world_size=2), + layerwise_export=False, + specdec_offline_dataset=None, + ) + + hf_ptq.mono_quantize( + args, + {"algorithm": None}, + full_model, + full_model, + model_type="qwen3_moe", + calibration_only=False, + calib_dataloader=object(), + is_nemotron_vl_model=False, + ) + + assert events == ["quantize"] + ([("replant", full_model, ep_mesh)] if replanted else []) + + class _PreviewModel(torch.nn.Module): """A model whose ``generate`` records the kwargs each preview passes it, and how often it entered the decode hook that Transformers' ``generate`` enters to switch a MoE's experts.""" From 46a37d7dd50ac31e1bb7eaf8bc11abeee407bca7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:00:30 +0000 Subject: [PATCH 35/57] Describe auto's EP degree below a spec's floor, and who re-plants a2a The hf_ptq README said a model needing a newer transformers says so before loading, and described auto's EP degree without the floor; since [13/21], auto resolves to 1 there and the run is refused by the plan instead. _reinstall_a2a_forwards' docstring named the calibration loop as its only caller; hf_ptq now calls it after quantizing too. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 4 ++-- modelopt/torch/utils/distributed.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 34bc1cf54ad..964b958dc4f 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -658,7 +658,7 @@ See [Recipe-based Quantization](#recipe-based-quantization) for the recipe forma Under `torchrun`, `hf_ptq.py` shards an MoE model's routed experts across the ranks (expert parallelism) instead of loading the whole model on every GPU. Each rank holds `num_experts / ep_degree` experts and calibrates on its own shard of the calibration data. An all-to-all sends every token to the rank that owns its expert, so each expert's quantizers see the tokens routed to it from all ranks. The dense (non-expert) part is replicated on every rank, or FSDP2-sharded when it does not fit one GPU. Each rank then writes its share of the unified HF checkpoint, without gathering the model onto one GPU. -A model can be expert-parallel when its config declares a `base_model_ep_plan`. transformers ships one for several MoEs, and ModelOpt supplies the plan or an adapter for others (see the `ep_spec` entries in `modelopt/torch/models/*/specs.py`). Expert parallelism needs transformers 5.6 or newer; a model that needs a newer release says so before loading. +A model can be expert-parallel when its config declares a `base_model_ep_plan`. transformers ships one for several MoEs, and ModelOpt supplies the plan or an adapter for others (see the `ep_spec` entries in `modelopt/torch/models/*/specs.py`). Expert parallelism needs transformers 5.6 or newer, and some models a newer release; a run that cannot be expert-parallel is refused before loading. ```bash torchrun --nproc_per_node= \ @@ -675,7 +675,7 @@ Two environment variables override the automatic distribution plan: | Variable | Default | Effect | | :--- | :--- | :--- | -| `MODELOPT_EP_DEGREE` | `auto` | Expert-parallel degree. `auto` picks the largest degree that divides both the routed-expert count and the world size. Below the world size, each group of `MODELOPT_EP_DEGREE` ranks holds a full set of experts, and the groups calibrate on different data. | +| `MODELOPT_EP_DEGREE` | `auto` | Expert-parallel degree. `auto` picks the largest degree that divides both the routed-expert count and the world size, or 1 for a model whose spec refuses expert parallelism or needs a newer transformers. Below the world size, each group of `MODELOPT_EP_DEGREE` ranks holds a full set of experts, and the groups calibrate on different data. | | `MODELOPT_FSDP_SIZE` | `auto` | Whether to FSDP2-shard the dense part. `auto` shards it when it would take more than 60% of one GPU's memory; `world` always shards it; `1` replicates it. With `MODELOPT_EP_DEGREE` below the world size the export cannot write a sharded dense part yet, so such a run fails before loading. | A run that cannot shard the model the way the plan asks fails before loading, rather than loading the whole model on every rank. `--use_fsdp2` skips expert parallelism and uses the [FSDP2 path](#multi-node-post-training-quantization-with-fsdp2) instead. Expert parallelism refuses GPT-OSS and Llama-4: their experts are stored transposed, which the expert-parallel export cannot write yet. `--use_fsdp2` refuses their released checkpoints too (GPT-OSS's carry an MXFP4 `quantization_config`, and Llama-4's are multimodal), so quantize those in a single process (without `torchrun`), which spreads the model over the node's GPUs. diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index ab155ed05c1..591fe0c07b2 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -831,8 +831,9 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: 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. Only a forward quantizer insertion rebound is planted again: a + Called at the top of the calibration loop, and by ``hf_ptq`` right after quantizing (a recipe + with no calibration runs no loop): after conversion (so it wraps the QUANTIZED forward) 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 From 10deb6d823290b021effe9e08caa4291435932ba Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:51:50 +0000 Subject: [PATCH 36/57] Keep --use_fsdp2's best-effort preview; refuse a rank naming no GPU The preview after quantization raised on a failure whenever the plan sharded the model, which aborts the run before the export, for --use_fsdp2 as well as for expert parallelism. On main a failed --use_fsdp2 preview warned and the export ran, and a failure every rank meets at the same step (an unsupported kernel, the case that skip was written for) still exports that way. Only an expert-parallel preview now stops the run; --use_fsdp2 keeps main's skip. get_model pinned a rank's copy to its GPU only for a device with an index, and under a multi-rank plan fell back to "auto" for one without, the shared-GPU placement its comment warns about. It now refuses such a device; hf_ptq hands every rank of a launch an indexed cuda:. The re-plant after quantizing notes it matters from transformers 5.16: earlier releases install the TP transforms as forward hooks, which quantizer insertion leaves in place. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 11 ++--- examples/hf_ptq/hf_ptq.py | 14 +++--- .../hf_ptq/test_hf_ptq_expert_parallel.py | 45 ++++++++++++------- 3 files changed, 42 insertions(+), 28 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index bc1979e9694..9c85bb7551e 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -618,11 +618,12 @@ def get_model( # detected : rank 1 and rank 0 both on CUDA device ...". Every load path below honours the pin. # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under # data-parallel calibration each rank must own its own GPU. A single process keeps "auto". - pinned = ( - {"": str(device)} - if device != "cpu" and ":" in str(device) and plan.world_size > 1 - else None - ) + if plan.world_size > 1 and device != "cpu" and ":" not in str(device): + raise ValueError( + f"{plan!r} has each rank hold its own copy of the model on its own GPU, but " + f"device={device!r} names no GPU index; pass the rank's own GPU, cuda:." + ) + pinned = {"": str(device)} if device != "cpu" and plan.world_size > 1 else None device_map = "cpu" if device == "cpu" else pinned or "auto" # Prepare config kwargs for loading diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6f04ae12a83..b78328f68b0 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -696,9 +696,10 @@ def mono_quantize( else: language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop) if distribution_plan(full_model).is_expert_parallel: - # The calibration loop re-plants the a2a forwards quantizer insertion drops, but a - # config without static activation scales calibrates without it, and algorithm None - # (the MX formats) never calls it; the preview and the export forward through them. + # The calibration loop re-plants the a2a forwards quantizer insertion drops (on + # transformers >= 5.16; a no-op before), but a config without static activation scales + # calibrates without it, and algorithm None (the MX formats) never calls it; the + # preview and the export forward through them. _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model)) # For VL models, update full_model to use the quantized language model @@ -998,9 +999,10 @@ def post_quantize( **lockstep, ) except Exception as e: - if lockstep: - # A sharded model's ranks decode in lockstep: a rank that skipped the preview would - # leave the others blocked in the next step's collective. Raising aborts the run. + if distribution_plan(full_model).is_expert_parallel: + # The ranks decode in lockstep: a rank that skipped the preview would leave the + # others blocked in the next step's collective, so the run stops. --use_fsdp2 keeps + # main's skip, which still exports after a failure every rank meets at the same step. raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif is_nemotron_vl_model and tokenizer is not None: diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index fe80b2582cf..2ff38843b0d 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -430,18 +430,17 @@ def get_model(ckpt_path, device, **kwargs): [ ("cuda:1", 4, {"": "cuda:1"}), (torch.device("cuda:1"), 4, {"": "cuda:1"}), - ("cuda", 4, "auto"), ("cuda:1", 1, "auto"), ], - ids=["str", "torch_device", "unindexed", "single_process"], + ids=["str", "torch_device", "single_process"], ) def test_a_rank_device_pins_the_whole_model_to_it( monkeypatch, tmp_path, device, world_size, device_map ): """Under torchrun each rank asks for its own cuda:; "auto" would let every rank's - accelerate pick the same idle GPU. An unindexed "cuda" still spreads the model over the GPUs, - and so does a single process, whatever device it names. Only a spread model is held to a memory - limit sized over the node's GPUs: the pinned map would ignore one.""" + accelerate pick the same idle GPU. A single process spreads the model over the GPUs, whatever + device it names. Only a spread model is held to a memory limit sized over the node's GPUs: the + pinned map would ignore one.""" LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) # No GPU on the node holds the model, so sizing it over the node limits its memory. monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 0, "cpu": 1 << 40}) @@ -493,15 +492,23 @@ def load_pretrained(ckpt_path, auto_class, **kwargs): assert placed == [{"": "cuda:1"}] -def test_a_sequential_map_is_refused_where_each_rank_holds_its_own_copy(tmp_path): +@pytest.mark.parametrize( + ("device", "kwargs", "match"), + [ + ("cuda:1", {"use_seq_device_map": True}, "--use_seq_device_map"), + # "auto" would let every rank's accelerate pick the same idle GPU. + ("cuda", {}, "names no GPU"), + ], + ids=["sequential_map", "unindexed_device"], +) +def test_a_placement_over_shared_gpus_is_refused_where_each_rank_holds_its_own_copy( + tmp_path, device, kwargs, match +): LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) - with pytest.raises(ValueError, match="--use_seq_device_map"): + with pytest.raises(ValueError, match=match): example_utils.get_model( - str(tmp_path), - device="cuda:1", - distribution=DistributionPlan(4), - use_seq_device_map=True, + str(tmp_path), device=device, distribution=DistributionPlan(4), **kwargs ) @@ -825,17 +832,19 @@ def loader(): ("plan", "nemotron_vl", "raises"), [ (DistributionPlan(4), False, False), + (DistributionPlan(4, dense="fsdp"), False, False), (DistributionPlan(4, ep_size=4), False, True), (DistributionPlan(4, ep_size=4), True, True), ], - ids=["replicated", "expert_parallel", "expert_parallel_nemotron_vl"], + ids=["replicated", "fsdp2", "expert_parallel", "expert_parallel_nemotron_vl"], ) -def test_a_failed_preview_of_a_sharded_model_stops_the_run( +def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run( hf_ptq, monkeypatch, plan, nemotron_vl, raises ): - """A rank that skipped a sharded model's preview would leave the others decoding in lockstep, - blocked in the next step's collective, so the failure is raised, which aborts the run, with the - preview's error as its cause. A model every rank holds whole keeps the best-effort skip.""" + """A rank that skipped an expert-parallel model's preview would leave the others decoding in + lockstep, blocked in the next step's collective, so the failure is raised, which aborts the run, + with the preview's error as its cause. A model every rank holds whole keeps the best-effort + skip and exports, and so does --use_fsdp2, as on main.""" def fail(*args, **kwargs): raise RuntimeError("the preview failed on this rank") @@ -845,7 +854,8 @@ class _FailingModel(torch.nn.Module): model = _FailingModel() set_distribution_plan(model, plan) - monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None) + exported = [] + monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: exported.append(True)) monkeypatch.setattr(hf_ptq, "run_nemotron_vl_preview", fail) # as vlm_utils raises in lockstep ids = torch.zeros(1, 4, dtype=torch.long) args = SimpleNamespace( @@ -866,6 +876,7 @@ class _FailingModel(torch.nn.Module): else: with pytest.warns(UserWarning, match="sanity check failed"): run() + assert exported == ([] if raises else [True]) @pytest.mark.parametrize(("lockstep", "raises"), [({}, False), ({"synced_gpus": True}, True)]) From afd1b59ada2754ae5f9abbff2b4919996a299513 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 04:16:56 +0000 Subject: [PATCH 37/57] Re-plant a text-calibrated VLM's a2a forwards over the whole model An expert-parallel run that calibrated a VLM on text handed distributed_forward_loop the language model, and the loop re-plants the all-to-all forwards over the model it is handed. transformers keeps the expert-parallel plan on the model from_pretrained built: from 5.16, tp_plan returns it only for a module whose own config carries distributed_config, which from_pretrained sets on the top config alone. The language model's plan named no experts, so the forwards quantizer insertion drops stayed dropped, and the first batch raised "torch_a2a experts forward requires the EP process_group". The loop now takes the whole model, as it already did for image batches, and runs text batches through the language model; quantization still targets the language model alone. A CPU test builds a tiny Qwen3-VL-MoE as the expert-parallel load leaves it and calibrates it on text over a one-rank gloo group; without the change it fails on transformers 5.18. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 9 +- .../hf_ptq/test_hf_ptq_expert_parallel.py | 126 +++++++++++++++++- 2 files changed, 125 insertions(+), 10 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index b78328f68b0..6dc391151e6 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -666,15 +666,18 @@ def mono_quantize( # It also restores the use_cache guard the plain loop does not apply. Other # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. Image # batches go through the VLM's own forward, as they do without expert parallelism: - # Nemotron-H Omni, for one, calibrates on images by default. + # Nemotron-H Omni, for one, calibrates on images by default. Text batches go through + # the language model, but the loop re-plants over the model it is handed, so that is + # the full model: transformers keeps the expert-parallel plan on the model + # from_pretrained built, and a VLM's language model names no experts in its own. calibrate_loop = distributed_forward_loop( - full_model if args.calib_with_images else language_model, + full_model, calib_dataloader, device=args.dist_state.device, ep_mesh=expert_parallel_mesh(full_model), forward=create_vlm_calibration_forward(full_model) if args.calib_with_images - else None, + else lambda batch: language_model(**batch), ) elif args.calib_with_images: calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 2ff38843b0d..b4e4bb9a0f7 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -21,9 +21,10 @@ distributed plan calibrates on its own shard of the data, an expert-parallel one through the distributed forward loop, and a whole model is pinned to the rank's own GPU instead of ``"auto"``. -Nothing here needs a process group or a checkpoint: the process group's rank and size are answered -the way a torchrun launch would, the configs are tiny real transformers configs served from memory, -and the loaders are replaced by stand-ins that record how they were called. +Nothing here needs a checkpoint, and only the test of a VLM's all-to-all forwards needs a process +group, of one gloo rank: elsewhere the process group's rank and size are answered the way a torchrun +launch would, the configs are tiny real transformers configs served from memory, and the loaders are +replaced by stand-ins that record how they were called. """ import contextlib @@ -44,7 +45,7 @@ import modelopt.torch.utils.distributed as dist from modelopt.torch.models import get_spec -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan +from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -598,8 +599,10 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion clobbers, so an expert-parallel run must calibrate through it, over the mesh load_and_prepare_ep loaded the experts over, image batches included: they run through the whole VLM's own forward - there (Nemotron-H Omni calibrates on images by default). Without expert parallelism there is - nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry, as on main.""" + there (Nemotron-H Omni calibrates on images by default). The loop re-plants over the model it is + handed, the whole model, and text batches run through the language model. Without expert + parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry, + as on main.""" ep_mesh, vlm_forward = object(), object() full_model, language_model = torch.nn.Linear(2, 2), torch.nn.Linear(2, 2) set_distribution_plan(full_model, plan) @@ -649,7 +652,10 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ elif images: assert chosen == [(full_model, ep_mesh, vlm_forward)] else: - assert chosen == [(language_model, ep_mesh, None)] + ((model, mesh, forward),) = chosen + assert (model, mesh) == (full_model, ep_mesh) + x = torch.ones(1, 2) + assert torch.equal(forward({"input": x}), language_model(x)) @pytest.mark.parametrize( @@ -706,6 +712,112 @@ def test_an_expert_parallel_run_replants_the_all_to_all_forwards_after_quantizin assert events == ["quantize"] + ([("replant", full_model, ep_mesh)] if replanted else []) +@pytest.fixture +def gloo_world_of_one(tmp_path): + """A one-rank gloo process group, for the all-to-all forwards to run on.""" + torch.distributed.init_process_group( + "gloo", init_method=(tmp_path / "store").as_uri(), rank=0, world_size=1 + ) + yield torch.distributed.group.WORLD + torch.distributed.destroy_process_group() + + +class _EpMesh: + """The expert-parallel submesh of a two-rank plan, over this process's one-rank group.""" + + def size(self): + return 2 + + def get_group(self): + return torch.distributed.group.WORLD + + +def test_a_vlm_calibrated_on_text_gets_its_experts_all_to_all_forwards_back( + hf_ptq, gloo_world_of_one +): + """Text batches run through a VLM's language model, but transformers keeps the expert-parallel + plan on the model from_pretrained built, and the language model's own plan names no experts. + Re-planted over the language model, the experts' all-to-all forward that quantizer insertion + drops (transformers 5.16 and later) stayed dropped, and the first batch raised for the missing + EP process group.""" + pytest.importorskip("transformers", minversion="5.16") + from transformers import Qwen3VLMoeConfig, Qwen3VLMoeForConditionalGeneration + from transformers.distributed.configuration_utils import DistributedConfig + from transformers.distributed.tensor_parallel import ALL_PARALLEL_STYLES, _get_parameter_tp_plan + + from modelopt.torch.utils.moe import register_torch_a2a + + register_torch_a2a() + config = Qwen3VLMoeConfig( + text_config={ + **_TINY, + "head_dim": 8, + "num_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 8, + "decoder_sparse_step": 1, + }, + vision_config={ + "depth": 1, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + "patch_size": 2, + "spatial_merge_size": 1, + "temporal_patch_size": 1, + "deepstack_visual_indexes": [], + }, + ) + # The model as load_and_prepare_ep leaves it: its plan's all-to-all styles on the language + # model's config and expert parallelism on the top one, the forwards from_pretrained installs + # for them, the experts dispatching through the all-to-all, and their seeded parallel state. + config.text_config.base_model_ep_plan = { + "layers.*.mlp.gate": "torch_a2a_router", + "layers.*.mlp.experts": "torch_a2a_experts", + } + config.distributed_config = DistributedConfig(enable_expert_parallel=True) + model = Qwen3VLMoeForConditionalGeneration(config).eval() + mesh = _EpMesh() + for name, module in model.named_modules(): + style = _get_parameter_tp_plan(parameter_name=name, tp_plan=model.tp_plan, is_weight=False) + if style in ("torch_a2a_router", "torch_a2a_experts"): + ALL_PARALLEL_STYLES[style].install_forward(module, mesh) + model.config.text_config._experts_implementation = "torch_a2a" + model._device_mesh = mesh + set_distribution_plan(model, DistributionPlan(2, ep_size=2)) + model_load_utils._seed_parallel_states( + model, + [name for name, _ in model.named_modules() if name.endswith("mlp.experts")], + dense_state=ParallelState(data_parallel_group=-1), + expert_state=ParallelState( + data_parallel_group=-1, expert_model_parallel_group=gloo_world_of_one + ), + ) + language_model = model.model.language_model + args = SimpleNamespace( + qformat="int8", + calib_with_images=False, + dist_state=dist.DistributedState(rank=0, world_size=2, device=torch.device("cpu")), + layerwise_export=False, + specdec_offline_dataset=None, + ) + batch = {"input_ids": torch.tensor([[1, 2, 3, 4]]), "attention_mask": torch.ones(1, 4)} + + hf_ptq.mono_quantize( + args, + hf_ptq.mtq.INT8_DEFAULT_CFG, + model, + language_model, + model_type="qwen3_vl_moe", + calibration_only=False, + calib_dataloader=[batch], + is_nemotron_vl_model=False, + ) + + assert language_model.layers[0].mlp.experts.down_proj_input_quantizer.amax is not None + + class _PreviewModel(torch.nn.Module): """A model whose ``generate`` records the kwargs each preview passes it, and how often it entered the decode hook that Transformers' ``generate`` enters to switch a MoE's experts.""" From 3d70c1d2e01fc2452ce93c1215dd09f264d7c5f8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 05:34:12 +0000 Subject: [PATCH 38/57] Keep a preview's nested language model off batched_mm experts too keep_the_experts_implementation planted its stand-in for generate's _optimize_model_for_decode on the model it was handed only. A Nemotron-VL model's chat generates through its language model, whose generate then entered its own hook and switched the experts to batched_mm: under expert parallelism that forward is not the torch_a2a one and raises, which stops the run before the export; otherwise the preview after PTQ decoded on unquantized experts. The guard now covers every module of the model that has the hook, and puts back what each one held. The post-PTQ preview's comment also says what a failure on one --use_fsdp2 rank alone does: the other ranks stay blocked in generate's collectives until the collective timeout. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 24 +++++++++++---------- examples/hf_ptq/hf_ptq.py | 4 +++- tests/examples/hf_ptq/test_example_utils.py | 10 ++++++--- 3 files changed, 23 insertions(+), 15 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 9c85bb7551e..773a420a482 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -106,21 +106,23 @@ def keep_the_experts_implementation(model) -> Iterator[None]: ``torch_a2a`` forwards, which ``batched_mm`` bypasses, so the preview after PTQ would decode on unquantized experts. And Transformers refuses the switch for an FSDP2-wrapped model, whose class is a subclass built outside its modeling module: ``FSDPQwen3MoeForCausalLM does not support - setting experts implementation``. A model without the hook is left alone. + setting experts implementation``. Every module of the model that has the hook is kept from it: + a model can generate through a submodule, as a Nemotron-VL model's ``chat`` generates through + its language model. """ - if not hasattr(model, "_optimize_model_for_decode"): - yield - return - # Put back what the model itself held, if anything: a nested use, or a hook bound per instance. - own = vars(model).get("_optimize_model_for_decode") - model._optimize_model_for_decode = nullcontext + modules = [m for m in model.modules() if hasattr(m, "_optimize_model_for_decode")] + # Put back what each module itself held, if anything: a nested use, or a hook bound per instance. + own = [(module, vars(module).get("_optimize_model_for_decode")) for module in modules] + for module in modules: + module._optimize_model_for_decode = nullcontext try: yield finally: - if own is None: - del model._optimize_model_for_decode - else: - model._optimize_model_for_decode = own + for module, held in own: + if held is None: + del module._optimize_model_for_decode + else: + module._optimize_model_for_decode = held def run_nemotron_vl_preview( diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6dc391151e6..88238cb41a0 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -1005,7 +1005,9 @@ def post_quantize( if distribution_plan(full_model).is_expert_parallel: # The ranks decode in lockstep: a rank that skipped the preview would leave the # others blocked in the next step's collective, so the run stops. --use_fsdp2 keeps - # main's skip, which still exports after a failure every rank meets at the same step. + # main's skip, which still exports after a failure every rank meets at the same step; + # a failure on one rank alone, such as an out-of-memory error, leaves the other ranks + # blocked in generate's collectives until the collective timeout. raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif is_nemotron_vl_model and tokenizer is not None: diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 458c5aa1109..d39d2667fd4 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -138,12 +138,14 @@ def test_build_quant_cfg_sets_moe_calib_experts_ratio_on_every_algorithm(algorit assert quant_cfg["algorithm"] == expected -def test_a_preview_keeps_generate_from_switching_the_experts_implementation(monkeypatch): +@pytest.mark.parametrize("wrapped", [False, True], ids=["the-model", "a-model-around-it"]) +def test_a_preview_keeps_generate_from_switching_the_experts_implementation(monkeypatch, wrapped): """On a GPU, ``generate`` switches a ``grouped_mm`` MoE's experts to ``batched_mm`` for decoding, which ModelOpt's quantizers do not intercept (the preview after PTQ decodes unquantized experts) and which Transformers refuses for an FSDP2-wrapped class (``FSDPQwen3MoeForCausalLM does not support setting experts implementation``). The hook ``generate`` enters is out of the way for the - preview, and back after it.""" + preview, and back after it, also when the previewed model generates through a submodule, as a + Nemotron-VL model's ``chat`` generates through its language model.""" entered = [] @contextmanager @@ -154,12 +156,14 @@ def switch(self): monkeypatch.setattr(GenerationMixin, "_optimize_model_for_decode", switch) model = get_tiny_llama() + # The model around it has no hook of its own; the preview generates through the inner one. + previewed = torch.nn.ModuleDict({"language_model": model}) if wrapped else model input_ids = torch.tensor([[1, 2, 3]]) # Some releases enter the hook only for the decode steps after the first token. with pytest.raises(ValueError, match="does not support setting experts implementation"): model.generate(input_ids, max_new_tokens=3) - with example_utils.keep_the_experts_implementation(model): + with example_utils.keep_the_experts_implementation(previewed): model.generate(input_ids, max_new_tokens=3) assert entered == ["LlamaForCausalLM"] # only the first, unprotected call reached the hook From 9cf0049384295e0179f8af1fc2bab3ca1552d7c3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 06:26:01 +0000 Subject: [PATCH 39/57] Note hf_ptq's --device refusal under a multi-rank torchrun in the changelog hf_ptq now runs inside distributed_run, which refuses a device other than cuda under a launch; on main a multi-rank torchrun ignored --device. Record that in Backward Breaking Changes. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6753acb6ab5..888bbd8aead 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,6 +43,7 @@ Changelog **Backward Breaking Changes** +- ``examples/hf_ptq/hf_ptq.py`` launched by ``torchrun`` on more than one rank now refuses a ``--device`` other than ``cuda``, which it used to ignore: each rank runs on its own GPU, ``cuda:``. - ``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. From 50aed96a60e5557c1da27d35139552ed3a350021 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 06:44:15 +0000 Subject: [PATCH 40/57] Refuse a sharded model's speculative-decoding export on every rank export_speculative_decoding writes the whole checkpoint from one process and knows no plan, so under an expert-parallel or FSDP2 plan every rank reached it and wrote the same files over each other. export_quantized now refuses that export for a sharded model, on every rank alike, before writing anything; a replicated model still exports it from rank 0. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 7 +++++++ .../hf_ptq/test_hf_ptq_expert_parallel.py | 20 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 88238cb41a0..1cd149d725a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -801,6 +801,13 @@ def export_quantized( # Early exit for speculative decoding checkpoints # No tokenizer saving needed for spec ckpts if has_spec_opt(full_model): + if plan.is_fsdp_dense or plan.is_expert_parallel: + # export_speculative_decoding writes the whole checkpoint from one process and + # knows no plan: each rank of a sharded model would write it over the others. + raise NotImplementedError( + f"Exporting a speculative-decoding model sharded as {plan!r} is not " + "supported; quantize it in a single process." + ) export_speculative_decoding(full_model, export_dir=export_path) args.checkpoint_exported = True print(f"Quantized speculative decoding checkpoint exported to: {export_path}") diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index b4e4bb9a0f7..ee88ba5da56 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -1113,3 +1113,23 @@ def test_only_rank_0_writes_a_replicated_model(hf_ptq, monkeypatch, plan, rank, hf_ptq.export_quantized(args, model, model, "qwen3_moe") assert written == ["config"] * (rank == 0) + ["checkpoint"] * exports + + +@pytest.mark.parametrize( + "plan", [DistributionPlan(4, dense="fsdp"), DistributionPlan(4, ep_size=4)], ids=["fsdp2", "ep"] +) +def test_a_sharded_speculative_decoding_model_is_refused_at_export(hf_ptq, monkeypatch, plan): + """The speculative-decoding exporter writes the whole checkpoint from one process and has no + rank gate, so every rank of a sharded model would write it over the others: it is refused on + every rank instead.""" + written = [] + monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: True) + monkeypatch.setattr( + hf_ptq, "export_speculative_decoding", lambda model, export_dir: written.append(export_dir) + ) + args = SimpleNamespace(dist_state=_dist_state(plan.world_size, 1), export_path="/export") + + model = _PreviewModel(plan) + with pytest.raises(NotImplementedError, match="speculative-decoding model sharded"): + hf_ptq.export_quantized(args, model, model, "qwen3_moe") + assert written == [] From 3c3f2d31e13f63e5102ace280b543d600f992683 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 07:27:14 +0000 Subject: [PATCH 41/57] Name transformers 5.7 as the expert-parallel floor in the hf_ptq README The expert-parallel load now requires 5.7, the first release that builds a model with torch_a2a. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 964b958dc4f..8780f55a3df 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -658,7 +658,7 @@ See [Recipe-based Quantization](#recipe-based-quantization) for the recipe forma Under `torchrun`, `hf_ptq.py` shards an MoE model's routed experts across the ranks (expert parallelism) instead of loading the whole model on every GPU. Each rank holds `num_experts / ep_degree` experts and calibrates on its own shard of the calibration data. An all-to-all sends every token to the rank that owns its expert, so each expert's quantizers see the tokens routed to it from all ranks. The dense (non-expert) part is replicated on every rank, or FSDP2-sharded when it does not fit one GPU. Each rank then writes its share of the unified HF checkpoint, without gathering the model onto one GPU. -A model can be expert-parallel when its config declares a `base_model_ep_plan`. transformers ships one for several MoEs, and ModelOpt supplies the plan or an adapter for others (see the `ep_spec` entries in `modelopt/torch/models/*/specs.py`). Expert parallelism needs transformers 5.6 or newer, and some models a newer release; a run that cannot be expert-parallel is refused before loading. +A model can be expert-parallel when its config declares a `base_model_ep_plan`. transformers ships one for several MoEs, and ModelOpt supplies the plan or an adapter for others (see the `ep_spec` entries in `modelopt/torch/models/*/specs.py`). Expert parallelism needs transformers 5.7 or newer, and some models a newer release; a run that cannot be expert-parallel is refused before loading. ```bash torchrun --nproc_per_node= \ From ddc9f7e318c905c97d4e2f4237b6b5eeb077f216 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 07:45:46 +0000 Subject: [PATCH 42/57] Skip hf_ptq's preview of a model that has no generate pre_quantize called full_model.generate for every model it did not skip by type, and a remote-code class that does not inherit GenerationMixin has no generate: Kimi-K3's KimiK3ForConditionalGeneration stopped the run there with an AttributeError, before quantization, unless --skip_generate was passed. Such a model now skips the preview, before quantization and so after it too. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 3 +++ tests/examples/hf_ptq/test_hf_ptq_model_type.py | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 1cd149d725a..5cb9fd1ad2c 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -939,6 +939,9 @@ def pre_quantize( allow_fallback=False, trust_remote_code=args.trust_remote_code, ) + elif not hasattr(full_model, "generate"): + # A remote-code class that does not inherit GenerationMixin has none (Kimi-K3's). + generated_ids_before_ptq = None else: with keep_the_experts_implementation(full_model): generated_ids_before_ptq = full_model.generate( diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py index 34b5194e1c2..beaf5639618 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -63,6 +63,16 @@ def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: assert len(generate_calls) == (0 if skips_generation else 1) +def test_pre_quantize_skips_the_preview_of_a_model_without_generate(hf_ptq): + """A remote-code class that does not inherit GenerationMixin (Kimi-K3's) has no generate.""" + args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) + batch = {"input_ids": torch.ones(2, 4, dtype=torch.long)} + + _, _, generated = hf_ptq.pre_quantize(args, torch.nn.Module(), "kimi_k3", None, [batch], False) + + assert generated is None + + @pytest.mark.parametrize( ("ids", "mask"), [([0, 0, 1, 2], [0, 0, 1, 1]), ([1, 2, 0, 0], [1, 1, 0, 0])] ) From 7e5c4d79ddeb4602c0355b6ffe2203b738fef77f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 08:57:56 +0000 Subject: [PATCH 43/57] Drop hf_ptq's preview skip for a model without generate The skip was for Kimi-K3, whose model had no generate, but only rank 0's lacked it: the plan decision's meta build had registered the remote class with the auto class, so rank 0's load built it as registered, without the GenerationMixin every other rank's first load adds. [5/21] now undoes that registration, and every rank's Kimi-K3 has generate. The skip, decided on each rank alone, instead let rank 0 move on to calibration while the others decoded in lockstep, a deadlock where the missing method would have raised. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 3 --- tests/examples/hf_ptq/test_hf_ptq_model_type.py | 10 ---------- 2 files changed, 13 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 5cb9fd1ad2c..1cd149d725a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -939,9 +939,6 @@ def pre_quantize( allow_fallback=False, trust_remote_code=args.trust_remote_code, ) - elif not hasattr(full_model, "generate"): - # A remote-code class that does not inherit GenerationMixin has none (Kimi-K3's). - generated_ids_before_ptq = None else: with keep_the_experts_implementation(full_model): generated_ids_before_ptq = full_model.generate( diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py index beaf5639618..34b5194e1c2 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -63,16 +63,6 @@ def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: assert len(generate_calls) == (0 if skips_generation else 1) -def test_pre_quantize_skips_the_preview_of_a_model_without_generate(hf_ptq): - """A remote-code class that does not inherit GenerationMixin (Kimi-K3's) has no generate.""" - args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) - batch = {"input_ids": torch.ones(2, 4, dtype=torch.long)} - - _, _, generated = hf_ptq.pre_quantize(args, torch.nn.Module(), "kimi_k3", None, [batch], False) - - assert generated is None - - @pytest.mark.parametrize( ("ids", "mask"), [([0, 0, 1, 2], [0, 0, 1, 1]), ([1, 2, 0, 0], [1, 1, 0, 0])] ) From 58465b76127cd3803b2824fa8fcdcdb62e43d39c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 18:35:42 +0000 Subject: [PATCH 44/57] Skip hf_ptq's DeepSeek-V3 expert-parallel routes below transformers 5.7 [13/21] holds every expert-parallel spec to the expert-parallel load's transformers floor, so on transformers 5.5, the CI floor, a multi-rank DeepSeek-V3 run is refused before the hf_ptq route these cases exercise. They skip below the floor; their single-process and --use_fsdp2 cases still run there. Signed-off-by: Shengliang Xu --- .../hf_ptq/test_hf_ptq_expert_parallel.py | 26 ++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index ee88ba5da56..376bb739ebf 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -37,8 +37,10 @@ import pytest import torch +import transformers from _test_utils.examples.hf_ptq_example_utils import example_utils from _test_utils.torch.transformers_models import get_tiny_tokenizer +from packaging.version import Version from PIL import Image from torch.utils.data import DataLoader from transformers import AutoConfig, DeepseekV3Config, LlamaConfig @@ -46,6 +48,7 @@ import modelopt.torch.utils.distributed as dist from modelopt.torch.models import get_spec from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan +from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -155,6 +158,14 @@ def _serve(monkeypatch, hf_ptq, config): return config +# DeepSeek-V3 is expert-parallel only from the transformers release the expert-parallel load needs; +# below it, a multi-rank run of it is refused before the route the test exercises. +needs_ep_transformers = pytest.mark.skipif( + Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS), + reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}", +) + + def _deepseek_v3(n_routed_experts): """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable. @@ -197,6 +208,7 @@ def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monke hf_ptq._resolve_distribution(_args(torchrun(world_size=2))) +@needs_ep_transformers @pytest.mark.parametrize("ep_degree", [2, 4]) def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused( hf_ptq, monkeypatch, torchrun, ep_degree @@ -243,7 +255,12 @@ def test_a_tensorrt_llm_checkpoint_export_needs_a_single_process( @pytest.mark.parametrize( ("world_size", "env", "use_fsdp2"), - [(2, {"MODELOPT_EP_DEGREE": "2"}, False), (2, {}, False), (2, {}, True), (1, {}, False)], + [ + pytest.param(2, {"MODELOPT_EP_DEGREE": "2"}, False, marks=needs_ep_transformers), + pytest.param(2, {}, False, marks=needs_ep_transformers), + (2, {}, True), + (1, {}, False), + ], ids=["expert_parallel", "replicated", "use_fsdp2", "single_process"], ) def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds( @@ -276,7 +293,7 @@ def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds( @pytest.mark.parametrize( ("make_config", "world_size"), [ - (lambda: _deepseek_v3(n_routed_experts=8), 2), + pytest.param(lambda: _deepseek_v3(n_routed_experts=8), 2, marks=needs_ep_transformers), (lambda: LlamaConfig(**_TINY), 2), (lambda: LlamaConfig(**_TINY), 1), ], @@ -302,7 +319,9 @@ def test_the_single_process_load_options_are_refused_on_the_routes_this_script_a @pytest.mark.parametrize( ("make_config", "world_size", "refused"), [ - (lambda: _deepseek_v3(n_routed_experts=8), 2, True), + pytest.param( + lambda: _deepseek_v3(n_routed_experts=8), 2, True, marks=needs_ep_transformers + ), (lambda: LlamaConfig(**_TINY), 2, False), (lambda: LlamaConfig(**_TINY), 1, False), ], @@ -352,6 +371,7 @@ def load_pretrained_stub(*args, **kwargs): DistributionPlan(4, ep_size=4, dense="fsdp"), {"trust_remote_code": True, "attn_implementation": "eager"}, id="expert_parallel", + marks=needs_ep_transformers, ), pytest.param( lambda: LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY), From 1e6b74ce973fccd69a7e4b4d929d1ade1dda3224 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:25:47 +0000 Subject: [PATCH 45/57] Keep hf_ptq's Whisper calibration loader whole on every rank The preview prints the transcript of the set's first sample next to each rank's first output, so a sharded loader paired every rank but rank 0 with another sample's transcript. With the amax synced across the ranks, the calibration comes out the same as from shards. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 5 +++- .../hf_ptq/test_hf_ptq_expert_parallel.py | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 1cd149d725a..1256002c6df 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -294,7 +294,10 @@ def make_calib_dataloader( include_labels=include_labels, ) world_size = distribution_plan(language_model).world_size - if world_size > 1 and calib_dataloader is not None: + # Whisper's loader stays whole on every rank: the preview prints the transcript of the set's first + # sample next to each rank's first output, and with the amax synced across the ranks the + # calibration comes out the same as from shards. + if world_size > 1 and calib_dataloader is not None and model_type != "whisper": # Every distributed calibration path is data-parallel over the world, so each rank calibrates # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2 diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index 376bb739ebf..e0da9413c17 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -607,6 +607,34 @@ def calib_loader(world_size, rank=0): assert calib_loader(world_size=1) is loader +def test_a_distributed_whisper_run_calibrates_every_rank_on_the_whole_set(hf_ptq, monkeypatch): + """The preview prints the transcript of the set's first sample next to each rank's first + output, which a shard would mismatch on every rank but rank 0.""" + loader = DataLoader(list(range(4)), batch_size=1) + monkeypatch.setattr(hf_ptq, "WhisperProcessor", SimpleNamespace) + monkeypatch.setattr( + hf_ptq, "get_speech_dataset_dataloader", lambda **kwargs: (loader, "the first transcript") + ) + args = SimpleNamespace( + specdec_offline_dataset=None, + calib_with_images=False, + dataset=None, + batch_size=1, + calib_size=[4], + dist_state=_dist_state(2, rank=1), + ) + model = torch.nn.Linear(2, 2) + model.dtype = torch.float32 + set_distribution_plan(model, DistributionPlan(2)) + + built, first = hf_ptq.make_calib_dataloader( + args, model, SimpleNamespace(), None, torch.device("cpu"), "whisper" + ) + + assert built is loader + assert first == "the first transcript" + + @pytest.mark.parametrize("images", [False, True], ids=["text", "images"]) @pytest.mark.parametrize( ("plan", "loop"), From 31e6a3cfb9089770fff6d85bd1c33a82b240bb6b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 09:44:15 +0000 Subject: [PATCH 46/57] Append get_model's distribution parameter after the existing ones It sat third, so a caller passing gpu_mem_percentage or a later option by position handed it to distribution instead. Every caller in the tree passes it by keyword, as they keep doing. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 773a420a482..b408abd0d4c 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -590,7 +590,6 @@ def _fmt_max_memory(max_memory: dict) -> str: def get_model( ckpt_path, device="cuda", - distribution: DistributionPlan | None = None, gpu_mem_percentage=0.8, trust_remote_code=False, use_seq_device_map=False, @@ -598,6 +597,7 @@ def get_model( offload_folder=None, max_cpu_memory_gb=None, max_gpu_memory_gb=None, + distribution: DistributionPlan | None = None, ): print(f"Initializing model from {ckpt_path}") # The plan to load the whole model by on every rank; a single process's by default. From 959b25a1429407636a39b0ce5db7232eda909195 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 09:45:11 +0000 Subject: [PATCH 47/57] Drop hf_ptq's claim that the plain calibration loop skips the use_cache guard It does not: the plain loop runs under the same guard, so the distributed loop restores nothing there. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 1256002c6df..2e7782519c4 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -665,8 +665,7 @@ def mono_quantize( # Without it the experts forward runs with process_group=None and raises # "torch_a2a experts forward requires the EP `process_group`" at calibration step 0. # On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op -- - # which is why this is invisible on older transformers and fatal on newer. - # It also restores the use_cache guard the plain loop does not apply. Other + # which is why this is invisible on older transformers and fatal on newer. Other # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. Image # batches go through the VLM's own forward, as they do without expert parallelism: # Nemotron-H Omni, for one, calibrates on images by default. Text batches go through From b3c2c28c6573ab05a3ffbc1eeb123c5793276656 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 17:23:47 +0000 Subject: [PATCH 48/57] Drop the transformers import the rebase left unused in hf_ptq's argument tests Main replaced the FSDP2 KV-AutoQuantize refusal test this PR had pointed at transformers.AutoConfig, so nothing in the file uses the import any more. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_args.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 95884c0dbe1..1f2444cdf55 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -23,7 +23,6 @@ import pytest import torch -import transformers import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 From 37bdee1a48f2537b9ba970108525457e288e084c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 17:24:57 +0000 Subject: [PATCH 49/57] Point main's FSDP2 KV-AutoQuantize load test at load_pretrained Main's new test checks that an FSDP2 KV-AutoQuantize run reaches the distributed load by stubbing parallel_load_and_prepare_fsdp2, which hf_ptq no longer calls: it loads every distributed plan through load_pretrained. The test stubs that instead, as its MTP neighbour does. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_args.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 1f2444cdf55..e24d21ffa83 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -409,7 +409,8 @@ def distributed_load(checkpoint, *_args, **_kwargs): assert checkpoint == str(tmp_path) raise RuntimeError("distributed loader reached") - monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load) + # hf_ptq loads an FSDP2 plan through load_pretrained. + monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load) with pytest.raises(RuntimeError, match="distributed loader reached"): hf_ptq.load_model(args) From bc74f743215ad0f3b47e75eca145b5544679746b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 17:58:10 +0000 Subject: [PATCH 50/57] Run --use_fsdp2's checks before deciding the plan On a launch of several ranks, deciding the plan builds the model on the meta device, also when --use_fsdp2 pins both knobs. Some checkpoints validate_fsdp2_supported refuses (VILA, multimodal, pack-quantized) cannot be built there, and that build's error names neither --use_fsdp2 nor the way out. The checks read the config alone, so they now run first, as they did before the run decided its plan. The DeepSeek-V3 test config names its architecture, as a checkpoint's does, since the checks read it. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 12 +++++++++-- .../hf_ptq/test_hf_ptq_expert_parallel.py | 20 +++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 2e7782519c4..dfd5038cd33 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -54,6 +54,7 @@ ) from torch.utils.data import DataLoader from transformers import ( + AutoConfig, AutoModelForCausalLM, AutoProcessor, PreTrainedTokenizer, @@ -353,6 +354,15 @@ def _resolve_distribution(args: argparse.Namespace): parallelism. Without --use_fsdp2 a plan that FSDP2-shards the whole model is refused: the FSDP2 run is the flag's, which checks what this script cannot do with it. """ + if args.use_fsdp2: + # Its checks read the config alone, and come first: deciding builds the model on the meta + # device, which some checkpoints they refuse (VILA, multimodal, pack-quantized) cannot do. + validate_fsdp2_supported( + args, + AutoConfig.from_pretrained( + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code + ), + ) plan, hf_config = resolve_distribution_plan( args.local_checkpoint_path, trust_remote_code=args.trust_remote_code, @@ -424,8 +434,6 @@ def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False plan, hf_config = _resolve_distribution(args) - if plan.is_fsdp_dense and not plan.is_expert_parallel: - validate_fsdp2_supported(args, hf_config) if plan.is_expert_parallel: full_model = load_pretrained( args.local_checkpoint_path, diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index e0da9413c17..e783494ae00 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -172,6 +172,7 @@ def _deepseek_v3(n_routed_experts): Its one layer is an MoE layer: DeepSeek-V3's first layers are dense by default. """ config = DeepseekV3Config( + architectures=["DeepseekV3ForCausalLM"], n_routed_experts=n_routed_experts, first_k_dense_replace=0, moe_intermediate_size=8, @@ -197,6 +198,25 @@ def test_use_fsdp2_plans_fsdp2_without_expert_parallelism(hf_ptq, monkeypatch, t assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(2, dense="fsdp") +def test_use_fsdp2_refuses_a_checkpoint_before_building_it_to_decide(hf_ptq, monkeypatch, torchrun): + """--use_fsdp2's checks read the config alone. They come before the decision, which builds the + model on the meta device: some checkpoints they refuse cannot be built, and that build's error + would name neither --use_fsdp2 nor the way out.""" + _serve( + monkeypatch, + hf_ptq, + LlamaConfig( + architectures=["LlamaForCausalLM"], quantization_config={"quant_method": "fp8"}, **_TINY + ), + ) + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built the model") + ) + + with pytest.raises(NotImplementedError, match="--use_fsdp2 does not support"): + hf_ptq._resolve_distribution(_args(torchrun(world_size=2), use_fsdp2=True)) + + def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monkeypatch, torchrun): """Without expert parallelism a plan that shards the dense part is an FSDP2 run, which this script runs only under --use_fsdp2, the flag its FSDP2 checks hang off; so it must fail and point From 8c9e88cf71e2011692e0b9246bdee2e82a6e5520 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:00:12 +0000 Subject: [PATCH 51/57] Warn about an off-GPU model for a rank's own cuda device too The check compared the device with the string "cuda", so a replicated run, whose ranks pass torch.device("cuda:"), never printed the warning, though a pinned copy that accelerate spilled to the CPU is where it matters most. It now compares the device type. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index b408abd0d4c..6538d321847 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -889,7 +889,7 @@ def has_pack_quantized_config(config): print(f"Moving model to {device} device...") model = model.to(device) - if device == "cuda" and not is_model_on_gpu(model): + if torch.device(device).type == "cuda" and not is_model_on_gpu(model): print("Warning: Some parameters are not on a GPU. Calibration can be slow or hit OOM") return model From 72171e7f58b6e54db71278b88d6fdc5d28c79b74 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:00:39 +0000 Subject: [PATCH 52/57] Drop the Whisper exemption from sharding the calibration set A run of several ranks never calibrates Whisper: _resolve_distribution refuses every encoder-decoder model the TensorRT-LLM exporter takes, Whisper included, before the load. The exemption and its test covered a configuration no launch reaches. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 5 +--- .../hf_ptq/test_hf_ptq_expert_parallel.py | 28 ------------------- 2 files changed, 1 insertion(+), 32 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index dfd5038cd33..ea32fdd28f1 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -295,10 +295,7 @@ def make_calib_dataloader( include_labels=include_labels, ) world_size = distribution_plan(language_model).world_size - # Whisper's loader stays whole on every rank: the preview prints the transcript of the set's first - # sample next to each rank's first output, and with the amax synced across the ranks the - # calibration comes out the same as from shards. - if world_size > 1 and calib_dataloader is not None and model_type != "whisper": + if world_size > 1 and calib_dataloader is not None: # Every distributed calibration path is data-parallel over the world, so each rank calibrates # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2 diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index e783494ae00..c2c85aed714 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -627,34 +627,6 @@ def calib_loader(world_size, rank=0): assert calib_loader(world_size=1) is loader -def test_a_distributed_whisper_run_calibrates_every_rank_on_the_whole_set(hf_ptq, monkeypatch): - """The preview prints the transcript of the set's first sample next to each rank's first - output, which a shard would mismatch on every rank but rank 0.""" - loader = DataLoader(list(range(4)), batch_size=1) - monkeypatch.setattr(hf_ptq, "WhisperProcessor", SimpleNamespace) - monkeypatch.setattr( - hf_ptq, "get_speech_dataset_dataloader", lambda **kwargs: (loader, "the first transcript") - ) - args = SimpleNamespace( - specdec_offline_dataset=None, - calib_with_images=False, - dataset=None, - batch_size=1, - calib_size=[4], - dist_state=_dist_state(2, rank=1), - ) - model = torch.nn.Linear(2, 2) - model.dtype = torch.float32 - set_distribution_plan(model, DistributionPlan(2)) - - built, first = hf_ptq.make_calib_dataloader( - args, model, SimpleNamespace(), None, torch.device("cpu"), "whisper" - ) - - assert built is loader - assert first == "the first transcript" - - @pytest.mark.parametrize("images", [False, True], ids=["text", "images"]) @pytest.mark.parametrize( ("plan", "loop"), From e114765a4d14849458115721e12ea0aa5a85b47d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:01:39 +0000 Subject: [PATCH 53/57] Start the GPU memory monitor once per node It reads every GPU on its node, so one per rank polls them all once per rank, but gating it on global rank 0 left every node but the first unreported on a multi-node run, where expert-parallel ranks hold different experts. Each node's local rank 0 starts it now. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 6 ++++-- tests/examples/hf_ptq/test_hf_ptq_args.py | 17 ++++++++++++----- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index ea32fdd28f1..1c136679817 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -106,6 +106,7 @@ distributed_forward_loop, distributed_run, distribution_plan, + local_rank, shard_dataloader, ) from modelopt.torch.utils.memory_monitor import launch_memory_monitor @@ -1691,8 +1692,9 @@ def _main_inner(args: argparse.Namespace): args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path) # launch a memory monitor to read the currently used GPU memory. It reads every GPU on the - # node, so a distributed run starts one, on rank 0, rather than one per rank. - if args.dist_state.is_main: + # node, so a distributed run starts one per node, on its local rank 0, rather than one per + # rank. + if args.dist_state.world_size == 1 or local_rank() == 0: launch_memory_monitor() # Force eager execution for all model types. diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index e24d21ffa83..0250444c4ad 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -1026,12 +1026,19 @@ def load_model(args): ] -@pytest.mark.parametrize("is_main", [True, False], ids=["rank_0", "other_rank"]) -def test_one_rank_runs_the_gpu_memory_monitor(monkeypatch, tmp_path, is_main): +@pytest.mark.parametrize( + ("rank", "node_rank", "launches"), + [(0, 0, True), (1, 1, False), (4, 0, True)], + ids=["rank_0", "other_rank", "first_rank_of_another_node"], +) +def test_one_rank_per_node_runs_the_gpu_memory_monitor( + monkeypatch, tmp_path, rank, node_rank, launches +): """The monitor reads every GPU on the node: one per rank would poll them all once per rank and - print the same report once per rank.""" + print the same report once per rank, and one on rank 0 alone would report node 0 alone.""" hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path)) - args.dist_state = SimpleNamespace(is_main=is_main, world_size=2) + args.dist_state = SimpleNamespace(is_main=rank == 0, rank=rank, world_size=8) + monkeypatch.setattr(hf_ptq, "local_rank", lambda: node_rank) launched = [] class _LoadReachedError(Exception): @@ -1048,7 +1055,7 @@ def load_model(args): with pytest.raises(_LoadReachedError): hf_ptq._main_inner(args) - assert len(launched) == is_main + assert len(launched) == launches def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypatch, tmp_path): From fbff5cc3bcf78ebd354bfc28a7d9d5988450f6eb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:02:48 +0000 Subject: [PATCH 54/57] Refuse a layerwise-export recipe before a distributed load Per-layer export writes the same shard files from every rank. assert_layerwise_export_supported refuses that, but only once calibration binds the exporter, after the expert-parallel load placed the model on every rank. _resolve_distribution refuses it before the load, as it does the other single-process-only options. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 14 +++++++++++++ .../hf_ptq/test_hf_ptq_expert_parallel.py | 20 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 1c136679817..2c5aff786c9 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -399,6 +399,20 @@ def _resolve_distribution(args: argparse.Namespace): "across ranks, so ranks could choose different formats per layer; launch hf_ptq " "without torchrun for an AutoQuantize recipe." ) + if ( + plan.world_size > 1 + and args.recipe is not None + and any( + block.get("export_dir") is not None and block.get("enable", False) + for block in recipe_layerwise_blocks(load_recipe(args.recipe)) + ) + ): + # Every rank would write the same per-layer shard files. The exporter refuses that too, but + # only once calibration binds it, after the load. + raise NotImplementedError( + f"{plan!r}: the recipe's layerwise export writes its checkpoint shards from one " + "process; launch hf_ptq without torchrun for it." + ) if plan.world_size > 1 and not args.use_fsdp2: # These configure loading one model in a single process: compressed as it loads, or # spread over the node's GPUs, CPU and disk. An expert-parallel run never reaches those diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py index c2c85aed714..abfa49b6471 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -305,6 +305,26 @@ def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds( hf_ptq._resolve_distribution(args) +@pytest.mark.parametrize("world_size", [2, 1], ids=["torchrun", "single_process"]) +def test_a_layerwise_export_recipe_is_refused_before_a_distributed_load( + hf_ptq, monkeypatch, torchrun, world_size +): + """Every rank would write the same per-layer shard files, which the exporter refuses only once + calibration binds it, after the load; a launch of several ranks refuses the recipe before it. + A single process exports as before.""" + _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY)) + args = _args( + torchrun(world_size=world_size), + recipe="general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export", + ) + + if world_size == 1: + assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1) + else: + with pytest.raises(NotImplementedError, match="layerwise export"): + hf_ptq._resolve_distribution(args) + + @pytest.mark.parametrize( "option", [{"low_memory_mode": True}, {"use_seq_device_map": True}, {"offload_folder": "/offload"}], From 57d104bf0c47e44b8b8f9614e2d235d2402ac9f0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:03:02 +0000 Subject: [PATCH 55/57] Say which hf_ptq load paths skip the rank's GPU pin The Nemotron VL load replaces the device map with None and moves the model to the rank's GPU only at the end, through host memory, so not every path below honours the pin as the comment said. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 6538d321847..18bbb0d096a 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -617,7 +617,9 @@ def get_model( # are all idle every rank scores them identically and picks the SAME card. All ranks' weights, # and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL collective # binds the whole data-parallel group to that single device and aborts with "Duplicate GPU - # detected : rank 1 and rank 0 both on CUDA device ...". Every load path below honours the pin. + # detected : rank 1 and rank 0 both on CUDA device ...". The load paths below honour the pin, + # but for the ones that load with no device map (Nemotron VL), which move the model to the + # rank's GPU at the end, through host memory. # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under # data-parallel calibration each rank must own its own GPU. A single process keeps "auto". if plan.world_size > 1 and device != "cpu" and ":" not in str(device): From 86f2b9174ee401c89e65b923d91bf1c21c383404 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:52:24 +0000 Subject: [PATCH 56/57] Say plainly which loads skip the rank's GPU pin, and what it costs The note read as the opposite of what it meant. The loads with no device map (Nemotron VL) skip the pin: each rank builds a whole copy in host memory before moving it to its GPU, so a node needs one host copy per rank. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 18bbb0d096a..2deaf3283b5 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -617,9 +617,9 @@ def get_model( # are all idle every rank scores them identically and picks the SAME card. All ranks' weights, # and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL collective # binds the whole data-parallel group to that single device and aborts with "Duplicate GPU - # detected : rank 1 and rank 0 both on CUDA device ...". The load paths below honour the pin, - # but for the ones that load with no device map (Nemotron VL), which move the model to the - # rank's GPU at the end, through host memory. + # detected : rank 1 and rank 0 both on CUDA device ...". Every load path below honours the pin + # except those that load with no device map (Nemotron VL): each rank builds a whole copy in host + # memory and only then moves it to its own GPU, so the node needs one host copy per rank. # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under # data-parallel calibration each rank must own its own GPU. A single process keeps "auto". if plan.world_size > 1 and device != "cpu" and ":" not in str(device): From 34485bb134b9b8b249b4ca3601000252aaae2356 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:54:37 +0000 Subject: [PATCH 57/57] Decide a layerwise export in one place The refusal before a distributed load and quantize_main's export switch each spelled out the test for a layerwise block that exports, and only the two staying identical kept the refusal covering the runs it is for. Both now ask layerwise_export_requested. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 7 +++++++ examples/hf_ptq/hf_ptq.py | 10 +++------- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 2deaf3283b5..c97656fb87e 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -1054,6 +1054,13 @@ def recipe_layerwise_blocks(recipe) -> list[dict]: return blocks +def layerwise_export_requested(blocks) -> bool: + """Whether a recipe's ``layerwise`` blocks (:func:`recipe_layerwise_blocks`) export layerwise.""" + return any( + block.get("export_dir") is not None and block.get("enable", False) for block in blocks + ) + + def _layerwise_checkpoint_dir(algorithm) -> str | None: """First ``layerwise.checkpoint_dir`` across the algorithm entries, or None.""" return next( diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 2c5aff786c9..6697627f140 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -45,6 +45,7 @@ is_nemotron_vl, is_trtllm_enc_dec_export, keep_the_experts_implementation, + layerwise_export_requested, mlflow_run, preview_generate_kwargs, recipe_layerwise_blocks, @@ -402,10 +403,7 @@ def _resolve_distribution(args: argparse.Namespace): if ( plan.world_size > 1 and args.recipe is not None - and any( - block.get("export_dir") is not None and block.get("enable", False) - for block in recipe_layerwise_blocks(load_recipe(args.recipe)) - ) + and layerwise_export_requested(recipe_layerwise_blocks(load_recipe(args.recipe))) ): # Every rank would write the same per-layer shard files. The exporter refuses that too, but # only once calibration binds it, after the load. @@ -1138,9 +1136,7 @@ def quantize_main( is_layerwise = any(cfg.get("enable", False) for cfg in layerwise_cfgs) # The value is a placeholder, replaced with --export_path below; presence is the switch. - args.layerwise_export = any( - cfg.get("export_dir") is not None and cfg.get("enable", False) for cfg in layerwise_cfgs - ) + args.layerwise_export = layerwise_export_requested(layerwise_cfgs) if not args.layerwise_export and any( cfg.get("export_dir") is not None for cfg in layerwise_cfgs ):