From b29a214b825c7e07c1ed037e375302f3813a58af Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:08:33 +0000 Subject: [PATCH 01/10] Run expert-parallel PTQ from hf_ptq.py hf_ptq.py decides once how to distribute the model, with resolve_distribution_plan, and loads it by that plan through load_pretrained: expert-parallel, FSDP2 under --use_fsdp2, or whole on every rank. Everything after reads the plan back from the model. - main() runs inside distributed_run, which owns the process group and yields this process's rank, size and device; the example's DistributedState and its process-group helpers go. - Refused before the load: an FSDP2 plan for the whole model without --use_fsdp2, the TensorRT-LLM checkpoint export, layerwise-export and AutoQuantize recipes on the routes this script adds, the single-process load options there, and --vllm_fakequant_export and --experts_implementation under expert parallelism. - Every rank of a distributed plan calibrates on its own shard of the data. An expert-parallel run calibrates through distributed_forward_loop, image batches through the VLM's own forward, and re-plants the all-to-all forwards once quantized. - A replicated model is exported by rank 0 alone, a sharded speculative-decoding export is refused, and a failed preview of an expert-parallel model stops the run. - The GPU memory monitor runs on each node's first rank. [3/3] of the split of this PR: [1/3] adds the library pieces and [2/3] the lockstep previews and the per-rank placement this builds on. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 + examples/hf_ptq/README.md | 28 +- examples/hf_ptq/example_utils.py | 126 ++- examples/hf_ptq/hf_ptq.py | 352 +++++-- modelopt/torch/utils/distributed.py | 5 +- tests/examples/hf_ptq/test_hf_ptq_args.py | 85 +- .../hf_ptq/test_hf_ptq_expert_parallel.py | 963 ++++++++++++++++++ .../hf_ptq/test_hf_ptq_vision_quantization.py | 6 + tests/examples/hf_ptq/test_llm_ptq.py | 99 ++ .../quantization/test_fsdp2_autoquant.py | 3 + 10 files changed, 1445 insertions(+), 224 deletions(-) create mode 100644 tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6ae2032ad55..6fe276643cb 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,6 +25,7 @@ Changelog - ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration. - Speed up Hugging Face QAT and QAD with ``QATTrainer``, ``QADTrainer`` or ``--manual_gc``: the ``gc.collect()`` run before each training step now skips the model and optimizer state, saving about 1 s per step (30%) on Qwen3-30B-A3B. - 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* @@ -48,6 +49,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, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). It also raises when a calibration batch runs out of GPU memory, instead of retrying it in smaller pieces: set ``--batch_size`` to one that fits. - ``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. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 71be2080d53..83de1a3be49 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -620,7 +620,7 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *`nvfp4_awq_clip` still works with `--use_fsdp2` on more than one GPU: it calibrates with AWQ-Clip alone, which folds no smoothing scale into the weights, and its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* -> *With `--use_fsdp2` on more than one GPU, a calibration batch that runs out of memory is not retried in smaller pieces, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives. With `--batch_size 0`, each rank estimates the batch size from one forward and the ranks use the smallest, which is not probed for out-of-memory either: set `--batch_size` if calibration runs out of memory.* +> *On more than one GPU, with `--use_fsdp2` or in any other `torchrun` launch of several ranks, a calibration batch that runs out of memory is not retried in smaller pieces, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives. With `--batch_size 0`, each rank estimates the batch size from one forward and the ranks use the smallest, which is not probed for out-of-memory either: set `--batch_size` if calibration runs out of memory.* ### Usage @@ -660,6 +660,32 @@ Qwen checkpoints with `model_type` `qwen3_5` or `qwen3_5_moe` support this path > *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. It does not overwrite a checkpoint already in the export directory.* +## 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.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= \ + 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, 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. + ## Evaluate Accuracy ### TensorRT-LLM Validation diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 54f805faabd..e67466fd166 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -23,8 +23,6 @@ import warnings from collections.abc import Iterator from contextlib import contextmanager, nullcontext -from dataclasses import dataclass -from datetime import timedelta from pathlib import Path from typing import Any @@ -65,39 +63,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 FSDP2 off).""" - if getattr(args, "use_fsdp2", False): - # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can - # exceed it, and barrier() has no per-call timeout (it 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): - dist_utils.cleanup() - - def _awq_calibrations(args) -> set[str]: """The AWQ calibrations (``awq_lite``, ``awq_full``) the run's --recipe or --qformat uses.""" if getattr(args, "recipe", None) is not None: @@ -289,18 +254,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 @@ -324,37 +289,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 @@ -1125,6 +1108,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 753f4be5c5d..04a9d02cf4b 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -18,6 +18,7 @@ import random import time import warnings +from datetime import timedelta from pathlib import Path from typing import Any @@ -35,7 +36,7 @@ HF_PTQ, _prepare_quant_cfg, build_quant_cfg, - cleanup_distributed, + create_vlm_calibration_forward, create_vlm_calibration_loop, generate_excludes_prompt, get_model, @@ -44,12 +45,12 @@ is_nemotron_vl, is_trtllm_enc_dec_export, keep_the_experts_implementation, + layerwise_export_requested, mlflow_run, preview_generate_kwargs, recipe_layerwise_blocks, run_nemotron_vl_preview, save_source_config, - setup_distributed_args, validate_fsdp2_supported, ) from torch.utils.data import DataLoader @@ -101,17 +102,22 @@ get_max_batch_size, get_supported_datasets, ) -from modelopt.torch.utils.distributed import DistributionPlan, abort, set_distribution_plan +from modelopt.torch.utils.distributed import ( + _reinstall_a2a_forwards, + distributed_forward_loop, + distributed_run, + distribution_plan, + local_rank, + 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 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 @@ -194,6 +200,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 +295,15 @@ 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 - ), + ) + world_size = distribution_plan(language_model).world_size + 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_calib_dataloader( + calib_dataloader, args.dist_state.rank, world_size ) return calib_dataloader, first_text_speech_dataset @@ -323,37 +346,140 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: ) +def _resolve_distribution(args: argparse.Namespace): + """Decide how this run distributes the model: ``(plan, hf_config)``. + + ``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. + """ + 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, + ep_size=1 if args.use_fsdp2 else None, + dense="fsdp" if args.use_fsdp2 else None, + ) + if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2: + raise RuntimeError( + 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)." + ) + 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." + ) + 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.is_expert_parallel and args.experts_implementation is not None: + # The expert-parallel load runs the routed experts through torch_a2a, whatever backend the + # config names, so the flag would be dropped without a word. + raise NotImplementedError( + f"{plan!r}: --experts_implementation picks the experts backend a model is loaded with, " + "but the expert-parallel load runs the routed experts through torch_a2a whatever it " + "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a " + "single process." + ) + 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." + ) + if ( + plan.world_size > 1 + and args.recipe is not None + 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. + 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 + # 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 + + +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 - if args.use_fsdp2: - hf_config = AutoConfig.from_pretrained( - args.local_checkpoint_path, trust_remote_code=args.trust_remote_code + plan, hf_config = _resolve_distribution(args) + if plan.is_expert_parallel: + full_model = load_pretrained( + args.local_checkpoint_path, + distribution=plan, + hf_config=hf_config, + trust_remote_code=args.trust_remote_code, + attn_implementation=args.attn_implementation, ) + elif plan.is_fsdp_dense: if args.experts_implementation is not None: # Set on the config, as build_meta_causal_lm sets a pre-fetched config's # attn_implementation: the meta model is built from it, with no kwarg of its own. The # setter carries it to the sub-configs, such as the text config a VLM's experts read. hf_config._experts_implementation = args.experts_implementation - 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=plan, + 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( - 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, @@ -378,9 +504,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 @@ -390,6 +513,7 @@ def load_model(args: argparse.Namespace): args.local_checkpoint_path, AutoModelForCausalLM, model_type=hf_config.model_type, + distribution=plan, **model_kwargs, ) calibration_only = True @@ -401,7 +525,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 @@ -568,7 +692,33 @@ 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: + if distribution_plan(full_model).is_expert_parallel: + # An expert-parallel run has to calibrate through distributed_forward_loop, not 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 + # "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. Other + # distributed runs (FSDP2, or a model replicated on every rank) keep the plain loop, + # which retries no batch that runs out of memory on a plan of several ranks: the + # error reaches distributed_run, which aborts the rank and with it the run. 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 + # 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, + 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 lambda batch: language_model(**batch), + ) + elif args.calib_with_images: calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader) else: calibrate_loop = create_forward_loop( @@ -587,6 +737,12 @@ 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 (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 if is_nemotron_vl_model and language_model is not full_model: @@ -671,6 +827,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(): @@ -679,6 +840,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}") @@ -686,17 +854,13 @@ 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() - 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": @@ -890,15 +1054,18 @@ def post_quantize( 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." @@ -985,9 +1152,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 ): @@ -1547,27 +1712,48 @@ def main(args: argparse.Namespace): random.seed(RAND_SEED) np.random.seed(RAND_SEED) - 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 - # 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 - ) - - # launch a memory monitor to read the currently used GPU memory. + # 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): + 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. It reads every GPU on the + # 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. - 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, @@ -1575,33 +1761,7 @@ def main(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, - ) - except BaseException: - # A rank that fails alone, such as in a sharded model's lockstep preview, leaves the others - # in a collective it never reaches, and cleanup's barrier and destroy_process_group would - # wait on them with the traceback unprinted. abort() prints it and exits, so the launcher - # stops the other ranks. - if args.dist_state.world_size > 1: - abort() - raise - finally: - cleanup_distributed(args) + ) if __name__ == "__main__": diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 75d4eb4a7f5..4bd97c5ac95 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -860,8 +860,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 diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index e4300892572..9181ea33ddf 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +26,7 @@ import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 -from transformers import Qwen3_5MoeConfig, Qwen3Config, Qwen3MoeConfig +from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -38,11 +38,7 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.utils import mlflow as mlflow_lib -from modelopt.torch.utils.distributed import ( - DistributionPlan, - distribution_plan, - set_distribution_plan, -) +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" @@ -414,7 +410,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) @@ -449,7 +446,7 @@ def distributed_load(*_args, hf_config, **_kwargs): assert hf_config.get_text_config()._experts_implementation == "eager" raise RuntimeError("distributed loader reached") - monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load) + monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load) with pytest.raises(RuntimeError, match="distributed loader reached"): hf_ptq.load_model(args) @@ -475,7 +472,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 @@ -485,27 +482,6 @@ def distributed_load(*args, **kwargs): hf_ptq.load_model(args) -def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path): - """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an - FSDP2-sharded model without one is refused rather than taken for a single process's.""" - Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path) - monkeypatch.setenv("RANK", "0") - hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2") - args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=2) - args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path) # as main() sets them - model = torch.nn.Linear(1, 1) - monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *args, **kwargs: model) - - def loaded(model): - raise RuntimeError("loaded") - - monkeypatch.setattr(hf_ptq, "prepare_model_for_calibration", loaded) - - with pytest.raises(RuntimeError, match="loaded"): - hf_ptq.load_model(args) - assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") - - def test_mlflow_flag_defaults_the_experiment_name(monkeypatch): monkeypatch.setattr(getpass, "getuser", lambda: "tester") hf_ptq, args = _parse_hf_ptq_args( @@ -1111,41 +1087,36 @@ def load_model(args): ] -@pytest.mark.parametrize("world_size", [1, 2], ids=["single_process", "two_ranks"]) -def test_a_rank_that_fails_aborts_a_run_of_several_ranks(monkeypatch, world_size): - """The other ranks may wait in a collective the failed rank never reaches, where tearing the - process group down would wait on them with the error unprinted: a run of several ranks aborts - instead, which exits, so the launcher stops the others. A single process raises as before.""" - hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "org/model") - events = [] - - class _AbortedError(BaseException): - """abort() exits the process; this stand-in raises instead, so the test sees it.""" - - def setup_distributed_args(args): - args.dist_state = SimpleNamespace(rank=0, world_size=world_size, is_main=True) +@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, 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=rank == 0, rank=rank, world_size=8) + monkeypatch.setattr(hf_ptq, "local_rank", lambda: node_rank) + launched = [] - def abort(): - events.append("abort") - raise _AbortedError + class _LoadReachedError(Exception): + pass def load_model(args): - raise RuntimeError("this rank's load failed") + raise _LoadReachedError - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.compiler, "set_stance", lambda *args, **kwargs: None) - monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: None) - monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", lambda path: (path, path)) - monkeypatch.setattr(hf_ptq, "setup_distributed_args", setup_distributed_args) - monkeypatch.setattr(hf_ptq, "cleanup_distributed", lambda args: events.append("cleanup")) - monkeypatch.setattr(hf_ptq, "abort", abort) + 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(_AbortedError if world_size > 1 else RuntimeError): - hf_ptq.main(args) + with pytest.raises(_LoadReachedError): + hf_ptq._main_inner(args) - # The real abort() exits before the teardown runs. - assert events == (["abort", "cleanup"] if world_size > 1 else ["cleanup"]) + assert len(launched) == launches def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypatch, tmp_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 new file mode 100644 index 00000000000..62cc6ca6b2d --- /dev/null +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -0,0 +1,963 @@ +# 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``. + +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 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 +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.torch.transformers_models import get_tiny_tokenizer +from packaging.version import Version +from torch.utils.data import DataLoader +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, 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 + +_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) + + +@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) + ) + + +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, 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) + + return enter + + +def _dist_state(world_size, 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}") + ) + + +def _args(dist_state, **overrides): + """What load_model reads before it picks a loader, as main() leaves it.""" + args = SimpleNamespace( + dist_state=dist_state, + use_fsdp2=False, + recipe=None, + local_checkpoint_path=_CKPT, + trust_remote_code=False, + attn_implementation=None, + experts_implementation=None, + sparsity_fmt="dense", + qformat="fp8", + low_memory_mode=False, + use_seq_device_map=False, + offload_folder=None, + vllm_fakequant_export=False, + ) + vars(args).update(overrides) + return args + + +def _serve(monkeypatch, hf_ptq, config): + """Answer every config read with ``config`` rather than a checkpoint on disk.""" + monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *args, **kwargs: 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. + + 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, + 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 + + +# -------------------------------------------------------------------------------- the decision + + +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._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 + there rather than load the model whole on every rank.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) + + with pytest.raises(RuntimeError, match="--use_fsdp2"): + 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 +): + """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) + + +@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) + + +@pytest.mark.parametrize( + ("world_size", "env", "use_fsdp2"), + [ + 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( + 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) + + +@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"}], + ids=["low_memory_mode", "use_seq_device_map", "offload_folder"], +) +@pytest.mark.parametrize( + ("make_config", "world_size"), + [ + pytest.param(lambda: _deepseek_v3(n_routed_experts=8), 2, marks=needs_ep_transformers), + (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) + + +@pytest.mark.parametrize( + ("make_config", "world_size", "refused"), + [ + 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), + ], + 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) + + +@pytest.mark.parametrize( + ("make_config", "world_size", "use_fsdp2", "refused"), + [ + pytest.param( + lambda: _deepseek_v3(n_routed_experts=8), 2, False, True, marks=needs_ep_transformers + ), + (lambda: _deepseek_v3(n_routed_experts=8), 2, True, False), + (lambda: LlamaConfig(**_TINY), 2, False, False), + (lambda: LlamaConfig(**_TINY), 1, False, False), + ], + ids=["expert_parallel", "use_fsdp2", "replicated", "single_process"], +) +def test_an_experts_implementation_is_refused_where_the_load_picks_the_backend( + hf_ptq, monkeypatch, torchrun, make_config, world_size, use_fsdp2, refused +): + """The expert-parallel load runs the routed experts through torch_a2a whatever backend is asked + for, so an expert-parallel run refuses --experts_implementation before the load rather than drop + it. --use_fsdp2, a replicated run and a single process load by it.""" + _serve(monkeypatch, hf_ptq, make_config()) + args = _args( + torchrun(world_size=world_size), use_fsdp2=use_fsdp2, experts_implementation="eager" + ) + + if refused: + with pytest.raises(NotImplementedError, match="--experts_implementation"): + hf_ptq._resolve_distribution(args) + else: + hf_ptq._resolve_distribution(args) + + +# ---------------------------------------------------------------------------------- load_model + + +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 + + 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 + + +@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", + marks=needs_ep_transformers, + ), + 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, 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) + + with pytest.raises(_LoaderReachedError): + hf_ptq.load_model(_args(torchrun(world_size=world_size), **options)) + + (call,) = calls + assert call["ckpt_path"] == _CKPT + assert call["distribution"] == plan + hf_config = call["kwargs"].pop("hf_config") + assert call["kwargs"] == passed + if plan.is_expert_parallel: + # 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): + """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))] + + +# ------------------------------------------------------------------ process group and dist_state + + +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. 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=torch.device("cpu")) + + @contextlib.contextmanager + def distributed_run(timeout=None, device="cuda"): + calls.append((timeout, device)) + events.append("enter") + try: + 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) + 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=False, device="cpu")) + + assert events == ["enter", "run", "exit"] + timeout, device = calls[0] + assert timeout > timedelta(minutes=30) + assert device == "cpu" + assert len(seen) == 1 + assert seen[0] is state + + +# --------------------------------------------------------------------------------- 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(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, model, None, tokenizer, torch.device("cpu"), "llama" + ) + return built + + 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(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, 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, image batches included: they run through the whole VLM's own forward + 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) + 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 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() + ) + monkeypatch.setattr( + hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: quantized.append(forward_loop) or m + ) + args = SimpleNamespace( + qformat="fp8", + calib_with_images=images, + dist_state=_dist_state(world_size=2), + layerwise_export=False, + specdec_offline_dataset=None, + ) + + hf_ptq.mono_quantize( + args, + {"algorithm": "max"}, + full_model, + language_model, + model_type="qwen3_moe", + calibration_only=False, + calib_dataloader=object(), + is_nemotron_vl_model=False, + ) + + assert len(quantized) == 1 + if loop == "plain": + assert chosen == ["vlm" if images else "plain"] + elif images: + assert chosen == [(full_model, ep_mesh, vlm_forward)] + else: + ((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( + "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 []) + + +@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, monkeypatch +): + """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)) + # One process cannot build the plan's two-rank expert-parallel row, which calibration checks the + # experts' group against, so the one-rank group answers as that row. + group_ranks = torch.distributed.get_process_group_ranks + monkeypatch.setattr( + torch.distributed, + "get_process_group_ranks", + lambda group: [0, 1] if group is gloo_world_of_one else group_ranks(group), + ) + 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 + + +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", "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", "fsdp2", "expert_parallel", "expert_parallel_nemotron_vl"], +) +def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run( + hf_ptq, monkeypatch, plan, nemotron_vl, raises +): + """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") + + class _FailingModel(torch.nn.Module): + generate = fail + + model = _FailingModel() + set_distribution_plan(model, plan) + 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( + 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", tokenizer, None, ids, None, ids, nemotron_vl, None, None + ) + + if raises: + 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() + assert exported == ([] if raises else [True]) + + +# ------------------------------------------------------------------------------------- export + + +def _planned(plan): + """A module carrying ``plan``, as the load records it.""" + model = torch.nn.Linear(1, 1) + set_distribution_plan(model, plan) + return model + + +@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 = _planned(plan) + 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 = _planned(plan) + with pytest.raises(NotImplementedError, match="speculative-decoding model sharded"): + hf_ptq.export_quantized(args, model, model, "qwen3_moe") + assert written == [] 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 48165090687..20c86f2f1a8 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" @@ -65,6 +67,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, diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index a820158fa9a..72c27541b98 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,92 @@ 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") +@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. + + 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 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, + 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", + "--no-verbose", + "--export_path", + str(export_dir), + *launch_args, + ], + "hf_ptq", + env={**os.environ, **launch_env}, + ) + + 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 diff --git a/tests/gpu/torch/quantization/test_fsdp2_autoquant.py b/tests/gpu/torch/quantization/test_fsdp2_autoquant.py index e2f45ee9507..3a259465a57 100644 --- a/tests/gpu/torch/quantization/test_fsdp2_autoquant.py +++ b/tests/gpu/torch/quantization/test_fsdp2_autoquant.py @@ -113,6 +113,9 @@ def _run_kv_autoquant_example(rank, size, checkpoint_dir, search_dir, composed): calib_with_images=False, specdec_offline_dataset=None, low_memory_mode=False, + sparsity_fmt="dense", + qformat="fp8", + vllm_fakequant_export=False, dataset=["local-test-data"], calib_size=[2 * size], batch_size=1, From 849c119ec4109836f375c64c9ddfe55c9f0dcea4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:08:40 +0000 Subject: [PATCH 02/10] Say in the changelog what a multi-rank torchrun launch of hf_ptq now does The Backward Breaking entry named only the --device refusal. A torchrun launch of several ranks without --use_fsdp2 used to run one independent copy of hf_ptq per process; it now runs one distributed job, expert-parallel for an MoE that can shard its experts and replicated otherwise, and refuses the options its routes cannot honor. The entry now says so. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6fe276643cb..13b9a1b9188 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -49,7 +49,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`` launched by ``torchrun`` on more than one rank without ``--use_fsdp2`` now runs one distributed job instead of one independent run per process: an MoE that can shard its experts runs expert-parallel, any other model is replicated with each rank calibrating its own share of the data, and only rank 0 exports a replicated model. Such a launch refuses ``--low_memory_mode``, ``--use_seq_device_map``, ``--offload_folder``, AutoQuantize recipes and, under expert parallelism, ``--vllm_fakequant_export`` and ``--experts_implementation``, and any multi-rank launch now refuses layerwise-export recipes, the TensorRT-LLM checkpoint export and a ``--device`` other than ``cuda``: launch without ``torchrun`` for them. - ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). It also raises when a calibration batch runs out of GPU memory, instead of retrying it in smaller pieces: set ``--batch_size`` to one that fits. - ``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 c3ffbf0e949cb79cc9a9d7081d6dd1608a1e12f7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 06:00:37 +0000 Subject: [PATCH 03/10] Drop the calibration test's claim that FSDP2 keeps the out-of-memory retry #2752 dropped the plain loop's split-on-OOM retry for a plan of several ranks, so the docstring now says only that FSDP2 keeps the plain loop, as there is nothing to re-plant without expert parallelism. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 3 +-- 1 file changed, 1 insertion(+), 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 62cc6ca6b2d..f5ade5dde92 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -599,8 +599,7 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_ 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). 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.""" + parallelism there is nothing to re-plant, so FSDP2 keeps the plain loop.""" 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) From f927c2aa3b0b5e465676f3f2252b98a504d4eb85 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:12:00 +0000 Subject: [PATCH 04/10] Drop the failed-preview test's --use_fsdp2 row, which [2/3] now tests [2/3] stops the run when a sharded model's preview fails, --use_fsdp2 included, so this PR's expert-parallel-only raise and its message went down into it. The row that expected --use_fsdp2 to skip the failure and export no longer holds, and [2/3]'s test covers the case. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 5 ++--- 1 file changed, 2 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 f5ade5dde92..729cfc37849 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -846,11 +846,10 @@ 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", "fsdp2", "expert_parallel", "expert_parallel_nemotron_vl"], + ids=["replicated", "expert_parallel", "expert_parallel_nemotron_vl"], ) def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run( hf_ptq, monkeypatch, plan, nemotron_vl, raises @@ -858,7 +857,7 @@ def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run( """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.""" + skip and exports.""" def fail(*args, **kwargs): raise RuntimeError("the preview failed on this rank") From d2b9dd96e943117d3c78bd1a8a046d661482155e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:56:33 +0000 Subject: [PATCH 05/10] Move hf_ptq's distribution decision into example_utils resolve_distribution, next to validate_fsdp2_supported, holds the example-specific checks of a distributed hf_ptq run around the library's resolve_distribution_plan, with shorter comments; hf_ptq.py keeps the PTQ workflow. exports_tensorrt_llm_checkpoint moves with it. The tests call and patch the names where they now live. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 99 ++++++++++++++- examples/hf_ptq/hf_ptq.py | 114 +----------------- tests/examples/hf_ptq/test_hf_ptq_args.py | 1 - .../hf_ptq/test_hf_ptq_expert_parallel.py | 34 +++--- .../examples/hf_ptq/test_hf_ptq_model_type.py | 3 +- .../hf_ptq/test_hf_ptq_vision_quantization.py | 4 +- 6 files changed, 123 insertions(+), 132 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index e67466fd166..a76611e8735 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -30,6 +30,7 @@ import transformers from accelerate import infer_auto_device_map, init_empty_weights from accelerate.utils import get_max_memory +from autoquant_utils import _recipe_is_auto_quantize from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static from safetensors import safe_open from transformers import ( @@ -56,7 +57,11 @@ from modelopt.torch.utils import distributed as dist_utils 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 is_fsdp2_text_only_vlm, load_pretrained +from modelopt.torch.utils.plugins.model_load_utils import ( + is_fsdp2_text_only_vlm, + load_pretrained, + resolve_distribution_plan, +) logger = logging.getLogger(__name__) @@ -120,6 +125,98 @@ def validate_fsdp2_supported(args, config): ) +def exports_tensorrt_llm_checkpoint(args, 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 resolve_distribution(args) -> tuple[DistributionPlan, Any]: + """Decide how hf_ptq distributes the model, refusing what that run cannot do: ``(plan, config)``. + + :func:`resolve_distribution_plan` decides; ``--use_fsdp2`` asks it for FSDP2 without expert + parallelism. + """ + if args.use_fsdp2: + # Before the decision, which builds the model on the meta device: some checkpoints these + # checks refuse (VILA, multimodal, pack-quantized) cannot be built there. + 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, + ep_size=1 if args.use_fsdp2 else None, + dense="fsdp" if args.use_fsdp2 else None, + ) + if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2: + raise RuntimeError( + 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)." + ) + if plan.world_size > 1 and exports_tensorrt_llm_checkpoint(args, hf_model_type(hf_config)): + # That exporter takes the process count for tensor parallelism. + 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." + ) + if (plan.is_expert_parallel or plan.is_fsdp_dense) and args.vllm_fakequant_export: + # It writes the whole model from every rank that exports, and every rank of a sharded + # model exports. + 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.is_expert_parallel and args.experts_implementation is not None: + # The expert-parallel load runs the experts through torch_a2a whatever the flag names. + raise NotImplementedError( + f"{plan!r}: --experts_implementation picks the experts backend a model is loaded with, " + "but the expert-parallel load runs the routed experts through torch_a2a whatever it " + "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a " + "single process." + ) + if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe): + # Each rank would score its own data, and could pick other formats per layer. + 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." + ) + if ( + plan.world_size > 1 + and args.recipe is not None + and layerwise_export_requested(recipe_layerwise_blocks(load_recipe(args.recipe))) + ): + # Every rank would write the same shard files. + 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: + # Options for loading one model in a single process, refused 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 + + def preview_generate_kwargs(model) -> dict[str, Any]: """Extra ``generate`` kwargs that keep every rank of a sharded model decoding a preview. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 04a9d02cf4b..11bf68d6bef 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -38,24 +38,23 @@ build_quant_cfg, create_vlm_calibration_forward, create_vlm_calibration_loop, + exports_tensorrt_llm_checkpoint, generate_excludes_prompt, get_model, get_processor, get_tokenizer, is_nemotron_vl, - is_trtllm_enc_dec_export, keep_the_experts_implementation, layerwise_export_requested, mlflow_run, preview_generate_kwargs, recipe_layerwise_blocks, + resolve_distribution, run_nemotron_vl_preview, save_source_config, - validate_fsdp2_supported, ) from torch.utils.data import DataLoader from transformers import ( - AutoConfig, AutoModelForCausalLM, AutoProcessor, PreTrainedTokenizer, @@ -117,7 +116,7 @@ copy_non_model_files, ensure_local_checkpoint, ) -from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan +from modelopt.torch.utils.plugins.model_load_utils import load_pretrained from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader @@ -346,113 +345,10 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: ) -def _resolve_distribution(args: argparse.Namespace): - """Decide how this run distributes the model: ``(plan, hf_config)``. - - ``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. - """ - 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, - ep_size=1 if args.use_fsdp2 else None, - dense="fsdp" if args.use_fsdp2 else None, - ) - if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2: - raise RuntimeError( - 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)." - ) - 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." - ) - 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.is_expert_parallel and args.experts_implementation is not None: - # The expert-parallel load runs the routed experts through torch_a2a, whatever backend the - # config names, so the flag would be dropped without a word. - raise NotImplementedError( - f"{plan!r}: --experts_implementation picks the experts backend a model is loaded with, " - "but the expert-parallel load runs the routed experts through torch_a2a whatever it " - "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a " - "single process." - ) - 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." - ) - if ( - plan.world_size > 1 - and args.recipe is not None - 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. - 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 - # 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 - - -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 - plan, hf_config = _resolve_distribution(args) + plan, hf_config = resolve_distribution(args) if plan.is_expert_parallel: full_model = load_pretrained( args.local_checkpoint_path, @@ -860,7 +756,7 @@ def export_quantized( save_source_config(args, export_path) start_time = time.time() - if _exports_tensorrt_llm_checkpoint(args, model_type): + 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_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 9181ea33ddf..0ea0dab0264 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -1135,7 +1135,6 @@ def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypa copies = [] monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: False) monkeypatch.setattr(hf_ptq, "is_multimodal_model", lambda model: False) - monkeypatch.setattr(hf_ptq, "is_trtllm_enc_dec_export", lambda model_type: False) monkeypatch.setattr(hf_ptq, "remove_hook_from_module", lambda *args, **kwargs: None) monkeypatch.setattr(hf_ptq, "export_tensorrt_llm_checkpoint", lambda *args, **kwargs: None) monkeypatch.setattr(hf_ptq, "copy_non_model_files", lambda *args: copies.append(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 729cfc37849..c668b63d240 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -193,7 +193,7 @@ def test_use_fsdp2_plans_fsdp2_without_expert_parallelism(hf_ptq, monkeypatch, t _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8)) args = _args(torchrun(world_size=2), use_fsdp2=True) - assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(2, dense="fsdp") + 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): @@ -212,7 +212,7 @@ def test_use_fsdp2_refuses_a_checkpoint_before_building_it_to_decide(hf_ptq, mon ) with pytest.raises(NotImplementedError, match="--use_fsdp2 does not support"): - hf_ptq._resolve_distribution(_args(torchrun(world_size=2), use_fsdp2=True)) + 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): @@ -223,7 +223,7 @@ def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monke _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY)) with pytest.raises(RuntimeError, match="--use_fsdp2"): - hf_ptq._resolve_distribution(_args(torchrun(world_size=2))) + hf_ptq.resolve_distribution(_args(torchrun(world_size=2))) @needs_ep_transformers @@ -240,10 +240,10 @@ def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused( args = _args(torchrun(world_size=4)) if ep_degree == 4: - assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(4, ep_size=4, dense="fsdp") + 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) + hf_ptq.resolve_distribution(args) @pytest.mark.parametrize( @@ -265,10 +265,10 @@ def test_a_tensorrt_llm_checkpoint_export_needs_a_single_process( args = _args(torchrun(world_size=world_size), **options) if world_size == 1: - assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1) + assert hf_ptq.resolve_distribution(args)[0] == DistributionPlan(1) else: with pytest.raises(NotImplementedError, match="single process"): - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) @pytest.mark.parametrize( @@ -297,10 +297,10 @@ def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds( ) if world_size == 1 or use_fsdp2: - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) else: with pytest.raises(NotImplementedError, match="AutoQuantize"): - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) @pytest.mark.parametrize("world_size", [2, 1], ids=["torchrun", "single_process"]) @@ -317,10 +317,10 @@ def test_a_layerwise_export_recipe_is_refused_before_a_distributed_load( ) if world_size == 1: - assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1) + assert hf_ptq.resolve_distribution(args)[0] == DistributionPlan(1) else: with pytest.raises(NotImplementedError, match="layerwise export"): - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) @pytest.mark.parametrize( @@ -347,11 +347,11 @@ def test_the_single_process_load_options_are_refused_on_the_routes_this_script_a args = _args(torchrun(world_size=world_size), **option) if world_size == 1: - assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(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) + hf_ptq.resolve_distribution(args) @pytest.mark.parametrize( @@ -377,9 +377,9 @@ def test_the_vllm_fakequant_export_is_refused_for_a_sharded_model( if refused: with pytest.raises(NotImplementedError, match="--vllm_fakequant_export"): - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) else: - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) @pytest.mark.parametrize( @@ -407,9 +407,9 @@ def test_an_experts_implementation_is_refused_where_the_load_picks_the_backend( if refused: with pytest.raises(NotImplementedError, match="--experts_implementation"): - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) else: - hf_ptq._resolve_distribution(args) + hf_ptq.resolve_distribution(args) # ---------------------------------------------------------------------------------- load_model 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..2539311313e 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -93,7 +93,8 @@ def test_pre_quantize_preview_removes_padding(hf_ptq, ids, mask): ], ) def test_is_trtllm_enc_dec_export(hf_ptq, model_type: str, expected: bool): - assert hf_ptq.is_trtllm_enc_dec_export(model_type) is expected + example_utils = importlib.import_module("example_utils") # imported with hf_ptq + assert example_utils.is_trtllm_enc_dec_export(model_type) is expected def _config(model_type: str, is_encoder_decoder: bool) -> SimpleNamespace: 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 20c86f2f1a8..dbf194c2b3f 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py +++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py @@ -68,9 +68,7 @@ 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) - ) + monkeypatch.setattr(hf_ptq, "resolve_distribution", lambda args: (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 39a9ad49de4a1f0848e10c7973a2ce1bf22f2bb4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:56:42 +0000 Subject: [PATCH 06/10] Trim the comments hf_ptq's expert-parallel run adds Each now says why in a line or two. The calibration loop's comment no longer gives a second, contradicting version boundary for the a2a forwards' re-plant: transformers 5.16, as everywhere else. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 42 ++++++++++----------------------------- 1 file changed, 11 insertions(+), 31 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 11bf68d6bef..1b2c07a8fc5 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -297,10 +297,8 @@ def make_calib_dataloader( ) world_size = distribution_plan(language_model).world_size 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. + # Every distributed plan is data-parallel over the world: each rank calibrates on its own + # shard, and the quantizers' amax is synced across the ranks. calib_dataloader = _shard_calib_dataloader( calib_dataloader, args.dist_state.rank, world_size ) @@ -589,22 +587,10 @@ def mono_quantize( # 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 distribution_plan(full_model).is_expert_parallel: - # An expert-parallel run has to calibrate through distributed_forward_loop, not 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 - # "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. Other - # distributed runs (FSDP2, or a model replicated on every rank) keep the plain loop, - # which retries no batch that runs out of memory on a plan of several ranks: the - # error reaches distributed_run, which aborts the rank and with it the run. 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 - # 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. + # The distributed loop first re-plants the torch_a2a forwards that quantizer + # insertion clobbers on transformers >= 5.16. It is handed the full model, whose + # expert-parallel plan names the experts; text batches still run through the + # language model. calibrate_loop = distributed_forward_loop( full_model, calib_dataloader, @@ -634,10 +620,8 @@ 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 (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. + # A config that runs no calibration loop never re-planted the a2a forwards, which the + # preview and the export need. _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model)) # For VL models, update full_model to use the quantized language model @@ -1608,10 +1592,7 @@ 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, 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. + # The run owns a distributed launch's process group, and ends it even if loading raises. with distributed_run(timeout=timedelta(hours=3), device=args.device) as dist_state: args.dist_state = dist_state _main_inner(args) @@ -1624,9 +1605,8 @@ 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. It reads every GPU on the - # node, so a distributed run starts one per node, on its local rank 0, rather than one per - # rank. + # launch a memory monitor to read the currently used GPU memory, one per node: it reads all + # of the node's GPUs. if args.dist_state.world_size == 1 or local_rank() == 0: launch_memory_monitor() From 9988f1deb20a648eaa701933fddb4137c31c1148 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:57:43 +0000 Subject: [PATCH 07/10] Add a Slurm example of expert-parallel PTQ across nodes multinode_ep_ptq.slurm follows multinode_fsdp2_ptq.slurm: one torchrun per node, launched without --use_fsdp2, so hf_ptq.py shards the MoE's experts over every rank. It quantizes Qwen3-235B-A22B to NVFP4 by default, and names the two knobs that override the plan. The README's expert-parallel section points to it. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 6 ++ examples/hf_ptq/slurm/multinode_ep_ptq.slurm | 93 ++++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 examples/hf_ptq/slurm/multinode_ep_ptq.slurm diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 83de1a3be49..24a17e754ac 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -677,6 +677,12 @@ torchrun --nproc_per_node= \ 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. +Under Slurm, [`slurm/multinode_ep_ptq.slurm`](./slurm/multinode_ep_ptq.slurm) is a ready-to-run example that quantizes Qwen3-235B-A22B to NVFP4 across nodes. Edit its `CONFIG` block and submit: + +```bash +sbatch --nodes=2 slurm/multinode_ep_ptq.slurm +``` + Two environment variables override the automatic distribution plan: | Variable | Default | Effect | diff --git a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm new file mode 100644 index 00000000000..758ee429843 --- /dev/null +++ b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm @@ -0,0 +1,93 @@ +#!/bin/bash + +# 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. +# +# Multi-node expert-parallel post-training quantization of an MoE model, ready to run under Slurm. +# +# Slurm allocates the nodes and launches one task per node; each task starts a `torchrun` that +# spawns one process per GPU. Without `--use_fsdp2`, `hf_ptq.py` shards the MoE's routed experts +# across all ranks (world_size = num_nodes * gpus_per_node), calibrates each rank on its own share +# of the data, and has every rank write its share of the checkpoint. Expert parallelism needs +# transformers 5.7 or newer; see the README's expert-parallel section for the supported models. +# +# The defaults below quantize Qwen/Qwen3-235B-A22B to NVFP4 with an FP8 KV cache. Edit the CONFIG +# block for your cluster/model, then submit with e.g.: +# +# sbatch --nodes=2 multinode_ep_ptq.slurm +# +# For a single-node run, `--nodes=1` works unchanged. + +#SBATCH --job-name=ep-ptq +#SBATCH --account={account} +#SBATCH --partition={partition} +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 # one torchrun launcher per node; it fans out to the GPUs +#SBATCH --gpus-per-node=8 +#SBATCH --exclusive +#SBATCH --time=04:00:00 +#SBATCH --output=%x_%j.log + +set -euo pipefail + +# --------------------------------------------------------------------------- +# CONFIG — edit these for your cluster and model. They are exported so `srun` +# propagates them into the container task below. +# --------------------------------------------------------------------------- +export CONTAINER_IMAGE={container_image} # e.g. nvcr.io#nvidia/pytorch:25.10-py3 +export MODELOPT_PATH={path_to_modelopt_repo} # host clone of TensorRT-Model-Optimizer +export MODEL_PATH=Qwen/Qwen3-235B-A22B # HF repo id (auto-downloaded) or a local dir +export EXPORT_PATH={path_to_export_dir} # where the quantized checkpoint is written +export HF_HOME={path_to_hf_cache} # HF cache; persist so a repo id isn't re-downloaded each run +# export HF_TOKEN={hf_token} # required for gated repos (or `huggingface-cli login` on host) +export RECIPE=general/ptq/nvfp4_default-kv_fp8 # built-in recipe name or /path/to/recipe.yaml +export CALIB_SIZE=512 # total over all ranks +export BATCH_SIZE=4 +# export MODELOPT_EP_DEGREE=16 # expert-parallel degree (default: auto) +# export MODELOPT_FSDP_SIZE=1 # 1 replicates the dense part, world shards it (default: auto) + +# Rendezvous: node 0 is the master; all ranks meet at MASTER_ADDR:MASTER_PORT. +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) +export MASTER_PORT=29531 + +# Mount the repo, export dir, and HF cache; add the model dir only when MODEL_PATH is a local path +# (a repo id is downloaded into HF_HOME instead). Run from the hf_ptq example dir. +CONTAINER_MOUNTS="${MODELOPT_PATH}:/modelopt,${EXPORT_PATH}:${EXPORT_PATH},${HF_HOME}:${HF_HOME}" +if [ -d "${MODEL_PATH}" ]; then + CONTAINER_MOUNTS="${CONTAINER_MOUNTS},${MODEL_PATH}:${MODEL_PATH}" +fi +srun --container-image="${CONTAINER_IMAGE}" \ + --container-mounts="${CONTAINER_MOUNTS}" \ + --container-workdir=/modelopt/examples/hf_ptq \ + bash -c ' + set -euo pipefail + export PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + # ModelOpt + its HF deps (transformers/accelerate/datasets/...) from the mounted source, then example extras. + pip install -q -e "/modelopt[hf]" --no-build-isolation + pip install -q -r requirements.txt + + torchrun \ + --nnodes="${SLURM_NNODES}" \ + --node_rank="${SLURM_NODEID}" \ + --nproc_per_node="$(nvidia-smi -L | wc -l)" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \ + hf_ptq.py \ + --pyt_ckpt_path "${MODEL_PATH}" \ + --recipe "${RECIPE}" \ + --calib_size "${CALIB_SIZE}" \ + --batch_size "${BATCH_SIZE}" \ + --export_path "${EXPORT_PATH}" + ' From 8d5185ed5dafc88214ae6139cd01c2973f6b9c61 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:14:22 +0000 Subject: [PATCH 08/10] Install the transformers expert parallelism needs in the EP Slurm example The header says expert parallelism needs transformers 5.7 or newer, but the install kept whatever the image ships when it meets ModelOpt's own floor, such as 5.5, where the run would not shard the experts. Ask pip for 5.7 or newer together with ModelOpt, and print the version into the job log. Signed-off-by: Shengliang Xu --- examples/hf_ptq/slurm/multinode_ep_ptq.slurm | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm index 758ee429843..1d0516bf35a 100644 --- a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm +++ b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm @@ -75,8 +75,10 @@ srun --container-image="${CONTAINER_IMAGE}" \ set -euo pipefail export PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True # ModelOpt + its HF deps (transformers/accelerate/datasets/...) from the mounted source, then example extras. - pip install -q -e "/modelopt[hf]" --no-build-isolation + # Expert parallelism needs transformers 5.7 or newer, which the image may not ship. + pip install -q -e "/modelopt[hf]" "transformers>=5.7" --no-build-isolation pip install -q -r requirements.txt + python -c "import transformers as t; print(t.__name__, t.__version__)" torchrun \ --nnodes="${SLURM_NNODES}" \ From aa6d756578d91e1a4e48c65c220af6f7ddfa4084 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:39:22 +0000 Subject: [PATCH 09/10] Keep hf_ptq's expert-parallel run off private helpers hf_ptq.py re-planted the experts' all-to-all forwards through the private _reinstall_a2a_forwards. modelopt.torch.utils.distributed now exports reinstall_a2a_forwards, a thin public wrapper for a caller that quantizes without distributed_forward_loop, and hf_ptq calls it. example_utils.py checks for an AutoQuantize recipe with the public recipe API instead of importing autoquant_utils' private helper. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 10 +++++++--- examples/hf_ptq/hf_ptq.py | 4 ++-- modelopt/torch/utils/distributed.py | 15 +++++++++++++-- .../hf_ptq/test_hf_ptq_expert_parallel.py | 2 +- 4 files changed, 23 insertions(+), 8 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index a76611e8735..b9cc0158d36 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -30,7 +30,6 @@ import transformers from accelerate import infer_auto_device_map, init_empty_weights from accelerate.utils import get_max_memory -from autoquant_utils import _recipe_is_auto_quantize from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static from safetensors import safe_open from transformers import ( @@ -43,7 +42,7 @@ ProcessorMixin, ) -from modelopt.recipe import load_recipe +from modelopt.recipe import ModelOptAutoQuantizeRecipe, load_recipe from modelopt.recipe.presets import QUANT_CFG_CHOICES from modelopt.torch.export import has_spec_opt from modelopt.torch.export.model_utils import is_multimodal_model @@ -183,7 +182,12 @@ def resolve_distribution(args) -> tuple[DistributionPlan, Any]: "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a " "single process." ) - if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe): + if ( + plan.world_size > 1 + and not args.use_fsdp2 + and args.recipe is not None + and isinstance(load_recipe(args.recipe), ModelOptAutoQuantizeRecipe) + ): # Each rank would score its own data, and could pick other formats per layer. raise NotImplementedError( f"{plan!r}: AutoQuantize scores each rank's own calibration data and reduces no score " diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 1b2c07a8fc5..18ff07274d1 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -102,11 +102,11 @@ get_supported_datasets, ) from modelopt.torch.utils.distributed import ( - _reinstall_a2a_forwards, distributed_forward_loop, distributed_run, distribution_plan, local_rank, + reinstall_a2a_forwards, shard_dataloader, ) from modelopt.torch.utils.memory_monitor import launch_memory_monitor @@ -622,7 +622,7 @@ def mono_quantize( if distribution_plan(full_model).is_expert_parallel: # A config that runs no calibration loop never re-planted the a2a forwards, which the # preview and the export need. - _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model)) + 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/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 4bd97c5ac95..cd3b2c2d4cd 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -69,6 +69,7 @@ "is_master", "rank", "recorded_distribution_plan", + "reinstall_a2a_forwards", "set_distribution_plan", "shard_dataloader", "size", @@ -860,8 +861,8 @@ 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, and by ``hf_ptq`` right after quantizing (a recipe - with no calibration runs no loop): after conversion (so it wraps the QUANTIZED forward) and + Called at the top of the calibration loop, and through :func:`reinstall_a2a_forwards` right + after quantizing without one: 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 @@ -909,6 +910,16 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int: return n +def reinstall_a2a_forwards(model, ep_mesh) -> int: + """Re-plant the routed experts' all-to-all forwards that quantizer insertion dropped. + + For a caller that quantizes an expert-parallel model without :func:`distributed_forward_loop`, + which does it itself: a recipe that runs no calibration runs no loop. ``ep_mesh`` is the mesh the + experts are sharded over. Returns how many forwards it planted again. + """ + return _reinstall_a2a_forwards(model, ep_mesh) + + def _release_cached_cuda_memory(): """Hand the caching allocator's reserved-but-unused pool back to the driver. 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 c668b63d240..b382e0fe054 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -684,7 +684,7 @@ def test_an_expert_parallel_run_replants_the_all_to_all_forwards_after_quantizin ) monkeypatch.setattr( hf_ptq, - "_reinstall_a2a_forwards", + "reinstall_a2a_forwards", lambda model, mesh: events.append(("replant", model, mesh)), ) args = SimpleNamespace( From 1a8e0072462873e63f2715b2901e4d7f87ac3cb0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:41:40 +0000 Subject: [PATCH 10/10] Warn about LSQ on a multi-rank hf_ptq run, and say which calibrations were tested LSQ's first scale calibration is synced or refused as the algorithm it runs is, but nothing has checked its learnable scales across ranks, so an expert-parallel or replicated run now warns before the load. The calibration methods a run uses come from one helper, which the AWQ check reads too. The README's expert-parallel section says which calibrations have run on GPUs on more than one rank, which are supported but not yet run there, and which are refused. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 ++ examples/hf_ptq/example_utils.py | 20 ++++++++++++--- .../hf_ptq/test_hf_ptq_expert_parallel.py | 25 +++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 24a17e754ac..80efe6785dc 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -690,6 +690,8 @@ 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, 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. | +On more than one rank, `max` and `mse` calibration have run on GPUs; `gptq`, `smoothquant`, `awq_clip` and `nvfp4_act_headroom` are supported but not yet run on GPUs; `awq_lite`, `awq_full`, `svdquant`, `local_hessian` and `imatrix`, and `max` or `mse` with `distributed_sync` off, are refused before quantizing; and an LSQ recipe runs with a warning, since nothing has checked its learnable scales across ranks yet. + 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 b9cc0158d36..f23d173d44a 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -55,6 +55,7 @@ from modelopt.torch.quantization.mode import get_modelike_from_algo_cfg from modelopt.torch.utils import distributed as dist_utils from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan +from modelopt.torch.utils.logging import warn_rank_0 from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run from modelopt.torch.utils.plugins.model_load_utils import ( is_fsdp2_text_only_vlm, @@ -67,8 +68,8 @@ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] -def _awq_calibrations(args) -> set[str]: - """The AWQ calibrations (``awq_lite``, ``awq_full``) the run's --recipe or --qformat uses.""" +def _calibration_methods(args) -> set[str]: + """The calibration methods (``max``, ``awq_lite``, ...) the run's --recipe or --qformat uses.""" if getattr(args, "recipe", None) is not None: recipe = load_recipe(args.recipe) configs = [getattr(recipe, "quantize", None)] @@ -80,12 +81,16 @@ def _awq_calibrations(args) -> set[str]: else: preset = QUANT_CFG_CHOICES.get(getattr(args, "qformat", None) or "") algorithms = [preset.get("algorithm")] if preset is not None else [] - methods = { + return { mode.removesuffix("_calibrate") for algorithm in algorithms for mode, _ in get_modelike_from_algo_cfg(algorithm) } - return methods & {"awq_lite", "awq_full"} + + +def _awq_calibrations(args) -> set[str]: + """The AWQ calibrations (``awq_lite``, ``awq_full``) the run's --recipe or --qformat uses.""" + return _calibration_methods(args) & {"awq_lite", "awq_full"} def validate_fsdp2_supported(args, config): @@ -194,6 +199,13 @@ def resolve_distribution(args) -> tuple[DistributionPlan, Any]: "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 and "lsq" in _calibration_methods(args): + # Its scale calibration is the algorithm it runs first, synced or refused as that one is, + # but nothing has checked its learnable scales across ranks yet. + warn_rank_0( + f"{plan!r}: LSQ has not been validated on a run of several ranks: nothing checks that " + "the ranks' learnable scales match. Launch hf_ptq without torchrun for the tested path." + ) if ( plan.world_size > 1 and args.recipe 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 b382e0fe054..176d416f034 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -30,6 +30,7 @@ import contextlib import importlib import inspect +import warnings from datetime import timedelta from pathlib import Path from types import SimpleNamespace @@ -303,6 +304,30 @@ def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds( hf_ptq.resolve_distribution(args) +@pytest.mark.parametrize( + ("world_size", "use_fsdp2", "warns"), + [(2, False, True), (2, True, False), (1, False, False)], + ids=["replicated", "use_fsdp2", "single_process"], +) +def test_an_lsq_recipe_warns_on_the_routes_this_script_adds( + hf_ptq, monkeypatch, torchrun, world_size, use_fsdp2, warns +): + """Nothing has checked LSQ's learnable scales across ranks: a replicated or expert-parallel run + warns before the load; --use_fsdp2 and a single process run it as before.""" + _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY)) + args = _args( + torchrun(world_size=world_size), + use_fsdp2=use_fsdp2, + recipe="general/qad/nvfp4_lsq-mse_init-fp8_kv", + ) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + hf_ptq.resolve_distribution(args) + + assert any("LSQ" in str(warning.message) for warning in caught) is warns + + @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