diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a05d4351bab..1c068e7e34b 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`` moves each tensor once, to the one rank that writes it, instead of all-gathering each FSDP2 unit on every rank as the default export does. Fused MoE experts stay sharded, each rank writing its own, and a conversion mapping that splits or merges tensors, which the default export refuses, is reversed as a single-process export reverses it. +- 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* @@ -51,6 +52,7 @@ Changelog **Backward Breaking Changes** - ``QDQAutotunerBase.export_onnx(output_path=...)`` now returns ``None`` instead of serialized bytes. Omit ``output_path`` to keep the bytes-returning behavior, which remains subject to protobuf size limits. +- ``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. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index b95e61719fa..42307b8126a 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,40 @@ Qwen checkpoints with `model_type` `qwen3_5` or `qwen3_5_moe` support this path > *Set `MODELOPT_FSDP2_EXPORT=nogather` to move each tensor once, to the one rank that writes it, instead of all-gathering each FSDP2 unit on every rank as the default export does. Fused MoE experts then stay sharded, and a conversion mapping that splits or merges tensors, which the default export refuses, is reversed as a single-process export reverses it. 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. The export directory and the Hugging Face cache (`HF_HOME`) must be on storage every node can see: every rank writes its share of the checkpoint, and rank 0 downloads a Hub model for the others to read. + +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 | +| :--- | :--- | :--- | +| `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 ### TensorRT-LLM Validation diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 0ef0c4fa389..f23d173d44a 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -22,9 +22,7 @@ import os import warnings from collections.abc import Iterator -from contextlib import contextmanager -from dataclasses import dataclass -from datetime import timedelta +from contextlib import contextmanager, nullcontext from pathlib import Path from typing import Any @@ -44,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 @@ -56,49 +54,22 @@ ) 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, 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__) 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.""" +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)] @@ -110,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): @@ -154,6 +129,152 @@ 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 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 " + "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 + 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. + + Each rank previews the first sample of its own calibration shard, so the ranks finish at + different steps. A model its plan shards runs a collective in every forward (the routed experts' + all-to-all, FSDP2's all-gathers): a rank that stopped decoding would leave the others blocked in + it. ``synced_gpus`` keeps every rank stepping until all of them have finished. + """ + plan = distribution_plan(model) + sharded = plan.world_size > 1 and (plan.is_expert_parallel or plan.is_fsdp_dense) + return {"synced_gpus": True} if sharded else {} + + +@contextmanager +def keep_the_experts_implementation(model) -> Iterator[None]: + """Keep ``generate`` from switching a MoE's experts to ``batched_mm`` while it decodes. + + On a GPU, ``generate`` runs the experts of a ``grouped_mm`` MoE on ``batched_mm`` for the + decoding stage and switches them back (``_optimize_model_for_decode``, Transformers 5.0 and + later). Two things follow. ModelOpt's quantizers sit on the ``eager``, ``grouped_mm`` and + ``torch_a2a`` forwards, which ``batched_mm`` bypasses, so the preview after PTQ would decode on + unquantized experts. And Transformers refuses the switch for an FSDP2-wrapped model, whose class + is a subclass built outside its modeling module: ``FSDPQwen3MoeForCausalLM does not support + setting experts implementation``. Every module of the model that has the hook is kept from it: + a model can generate through a submodule, as a Nemotron-VL model's ``chat`` generates through + its language model. + """ + modules = [m for m in model.modules() if hasattr(m, "_optimize_model_for_decode")] + # Put back what each module itself held, if anything: a nested use, or a hook bound per instance. + own = [(module, vars(module).get("_optimize_model_for_decode")) for module in modules] + for module in modules: + module._optimize_model_for_decode = nullcontext + try: + yield + finally: + for module, held in own: + if held is None: + del module._optimize_model_for_decode + else: + module._optimize_model_for_decode = held + + def run_nemotron_vl_preview( full_model, tokenizer, @@ -178,32 +299,37 @@ def run_nemotron_vl_preview( """ from vlm_utils import run_text_only_generation, run_vl_preview_generation - print(f"Running text-only preview generation for Nemotron VL model ({stage_name})...") - question = tokenizer.decode(input_ids[0], skip_special_tokens=True) - generation_config = { - "max_new_tokens": 100, - "do_sample": False, - "eos_token_id": tokenizer.eos_token_id, - } - - # Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse) - text_response = run_text_only_generation( - full_model, tokenizer, question, generation_config, pyt_ckpt_path, trust_remote_code - ) - - generated_ids = None - if text_response is not None: - print(f"✅ Text-only generation successful: {text_response[:100]}...") - generated_ids = text_response - elif allow_fallback: - print("Text-only generation failed, falling back to standard generate...") - generated_ids = full_model.generate(input_ids, max_new_tokens=100) + with keep_the_experts_implementation(full_model): + print(f"Running text-only preview generation for Nemotron VL model ({stage_name})...") + question = tokenizer.decode(input_ids[0], skip_special_tokens=True) + lockstep = preview_generate_kwargs(full_model) + generation_config = { + "max_new_tokens": 100, + "do_sample": False, + "eos_token_id": tokenizer.eos_token_id, + **lockstep, + } + + # Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse) + text_response = run_text_only_generation( + full_model, tokenizer, question, generation_config, pyt_ckpt_path, trust_remote_code + ) - # Run additional VL test with images - print(f"Running additional VL test with images ({stage_name})...") - run_vl_preview_generation(full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code) + generated_ids = None + if text_response is not None: + print(f"✅ Text-only generation successful: {text_response[:100]}...") + generated_ids = text_response + elif allow_fallback: + print("Text-only generation failed, falling back to standard generate...") + generated_ids = full_model.generate(input_ids, max_new_tokens=100, **lockstep) + + # Run additional VL test with images + print(f"Running additional VL test with images ({stage_name})...") + run_vl_preview_generation( + full_model, tokenizer, pyt_ckpt_path, stage_name, trust_remote_code, lockstep + ) - return generated_ids + return generated_ids def _is_multimodal_config(config): @@ -241,18 +367,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 @@ -276,37 +402,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 @@ -604,8 +748,19 @@ def get_model( offload_folder=None, max_cpu_memory_gb=None, max_gpu_memory_gb=None, + distribution: DistributionPlan | None = None, ): print(f"Initializing model from {ckpt_path}") + # The plan to load the whole model by on every rank; a single process's by default. Only a plan + # the caller gave reaches load_pretrained: without one it decides itself as before, and refuses + # MODELOPT_* knobs asking to shard over a launch whose process group is not set up. + plan = distribution or DistributionPlan() + if plan.is_expert_parallel or plan.is_fsdp_dense: + raise ValueError( + f"{plan!r} shards the model over the ranks, but get_model loads the whole model on " + "every rank: load it with load_pretrained." + ) + routing = {} if distribution is None else {"distribution": distribution} _disk_offload = offload_folder is not None if _disk_offload and max_cpu_memory_gb is None: @@ -616,9 +771,25 @@ def get_model( UserWarning, ) - device_map = "auto" - if device == "cpu": - device_map = "cpu" + # Under torchrun each rank asks for its own "cuda:". Do not fall through to "auto": + # accelerate runs its placement heuristic INDEPENDENTLY in every rank, and on a node whose GPUs + # are all idle every rank scores them identically and picks the SAME card. All ranks' weights, + # and with them every quantizer's _amax buffer, then live on one GPU; the first NCCL collective + # binds the whole data-parallel group to that single device and aborts with "Duplicate GPU + # detected : rank 1 and rank 0 both on CUDA device ...". Every load path below honours the pin + # except those that load with no device map (Nemotron VL): each rank builds a whole copy in host + # memory and only then moves it to its own GPU, so the node needs one host copy per rank. + # A dense model too large to replicate per rank wants --use_fsdp2, not "auto" sharding: under + # data-parallel calibration each rank must own its own GPU. A single process keeps "auto". + # hf_ptq hands in its run's torch.device, which never equals a string such as "cpu". + on_cpu = torch.device(device).type == "cpu" + if plan.world_size > 1 and not on_cpu and ":" not in str(device): + raise ValueError( + f"{plan!r} has each rank hold its own copy of the model on its own GPU, but " + f"device={device!r} names no GPU index; pass the rank's own GPU, cuda:." + ) + pinned = {"": str(device)} if not on_cpu and plan.world_size > 1 else None + device_map = "cpu" if on_cpu else pinned or "auto" # Prepare config kwargs for loading config_kwargs = {"trust_remote_code": trust_remote_code} if trust_remote_code else {} @@ -654,13 +825,24 @@ def get_model( # modules together for checkpoints that fit; larger ones can still spill and split a # tied pair, and need an explicit single-device map. Multi-GPU only: a single-GPU split # cannot separate a tied pair, and sequential would needlessly cap max_memory there. - if device != "cpu" and torch.cuda.device_count() > 1 and is_diffusion_gemma(hf_config): + if ( + not pinned + and not on_cpu + and torch.cuda.device_count() > 1 + and is_diffusion_gemma(hf_config) + ): print( "Detected DiffusionGemma model. Using device_map='sequential'; the balanced " "'auto' mapping can split its tied encoder/decoder weights across GPUs." ) use_seq_device_map = True + if use_seq_device_map and pinned: + raise ValueError( + f"--use_seq_device_map spreads one model over every GPU, but each rank of this run " + f"holds its own copy on its own GPU ({device}). Drop it, or shard the model with " + "--use_fsdp2." + ) if use_seq_device_map: device_map = "sequential" # If we use sequential, set max_memory limit to ensure that the model does not occupy the full GPU @@ -712,6 +894,7 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, + **routing, device_map=device_map, **model_kwargs, ) @@ -723,7 +906,8 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, - device_map="auto", + **routing, + device_map=pinned or "auto", trust_remote_code=trust_remote_code, dtype="auto", **experts_kwargs, @@ -750,7 +934,8 @@ def has_pack_quantized_config(config): ckpt_path, AutoModelForCausalLM, model_type=hf_config.model_type, - device_map="cpu" if device == "cpu" else "sequential", + **routing, + device_map="cpu" if on_cpu else pinned or "sequential", **model_kwargs, ) else: @@ -836,7 +1021,9 @@ def has_pack_quantized_config(config): f"Offload folder: {offload_folder}\n" "Weights exceeding GPU+CPU budgets will be streamed from disk." ) - else: + elif not pinned: + # A pinned copy goes onto its rank's GPU whole, by a map that ignores max_memory: sizing + # it over the node's GPUs would print a limit the load does not apply. inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) if "cpu" in inferred_device_map.values(): for _device in max_memory: @@ -858,6 +1045,7 @@ def has_pack_quantized_config(config): ckpt_path, auto_model_module, model_type=hf_config.model_type, + **routing, device_map=device_map, **model_kwargs2, ) @@ -866,11 +1054,11 @@ def has_pack_quantized_config(config): _unpack_compressed_linear_weights(model, ckpt_path) # If device_map was disabled (None), manually move model to target device - if device_map is None and device != "cpu": + if device_map is None and not on_cpu: print(f"Moving model to {device} device...") model = model.to(device) - if device == "cuda" and not is_model_on_gpu(model): + if torch.device(device).type == "cuda" and not is_model_on_gpu(model): print("Warning: Some parameters are not on a GPU. Calibration can be slow or hit OOM") return model @@ -1033,6 +1221,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 c79c25ccf98..18ff07274d1 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,24 +36,25 @@ HF_PTQ, _prepare_quant_cfg, build_quant_cfg, - cleanup_distributed, + 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, - setup_distributed_args, - validate_fsdp2_supported, ) from torch.utils.data import DataLoader from transformers import ( - AutoConfig, AutoModelForCausalLM, AutoProcessor, PreTrainedTokenizer, @@ -99,23 +101,32 @@ get_max_batch_size, get_supported_datasets, ) -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan +from modelopt.torch.utils.distributed import ( + distributed_forward_loop, + distributed_run, + distribution_plan, + local_rank, + reinstall_a2a_forwards, + 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 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader RAND_SEED = 1234 # Most tokens per calibration batch when layerwise calibration picks the batch size (64 x 512). LAYERWISE_CALIB_TOKENS = 32768 +_PREVIEW_STOPS_THE_RUN = ( + "The preview after quantization failed, which stops the run before the export, so no " + "checkpoint is written; --skip_generate exports without the preview." +) def _layerwise_calib_batch_size(calib_seq: int, num_samples: int) -> int: @@ -188,6 +199,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, @@ -269,12 +294,13 @@ 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 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 ) return calib_dataloader, first_text_speech_dataset @@ -320,34 +346,34 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False - if args.use_fsdp2: - 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, @@ -372,9 +398,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 @@ -384,6 +407,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 @@ -395,7 +419,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 @@ -562,7 +586,21 @@ 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: + # 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, + 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( @@ -581,6 +619,10 @@ def mono_quantize( ) else: language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop) + if distribution_plan(full_model).is_expert_parallel: + # 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 if is_nemotron_vl_model and language_model is not full_model: @@ -665,6 +707,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(): @@ -673,6 +720,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}") @@ -680,17 +734,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": @@ -809,11 +859,13 @@ def pre_quantize( trust_remote_code=args.trust_remote_code, ) else: - generated_ids_before_ptq = full_model.generate( - preview_input_ids, - attention_mask=preview_attention_mask, - max_new_tokens=100, - ) + with keep_the_experts_implementation(full_model): + generated_ids_before_ptq = full_model.generate( + preview_input_ids, + attention_mask=preview_attention_mask, + max_new_tokens=100, + **preview_generate_kwargs(full_model), + ) return preview_input_ids, preview_attention_mask, generated_ids_before_ptq @@ -866,24 +918,34 @@ def post_quantize( # part of the model to CPU (seen on unified-memory single-GPU hosts) can make a # quantized layer run on CPU, which some kernels (e.g. NVFP4 dynamic block # quantization) don't support. Don't let that discard the completed calibration. + lockstep = preview_generate_kwargs(full_model) try: - generated_ids_after_ptq = full_model.generate( - preview_input_ids, - attention_mask=preview_attention_mask, - max_new_tokens=100, - ) + with keep_the_experts_implementation(full_model): + generated_ids_after_ptq = full_model.generate( + preview_input_ids, + attention_mask=preview_attention_mask, + max_new_tokens=100, + **lockstep, + ) except Exception as e: + if lockstep: + # The ranks decode in lockstep: a rank that skipped the preview would go on to the + # export's collectives while the others still wait in generate's, so the run stops. + 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." @@ -970,9 +1032,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 ): @@ -1532,27 +1592,44 @@ 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 - ) + # 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) - # launch a memory monitor to read the currently used GPU memory. + +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, 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() - # 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, @@ -1560,25 +1637,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, - ) - finally: - cleanup_distributed(args) + ) if __name__ == "__main__": 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..04fb0b61a9d --- /dev/null +++ b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm @@ -0,0 +1,102 @@ +#!/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_PATH and HF_HOME must be on storage every node can see: every rank writes its share of +# the checkpoint to EXPORT_PATH, and rank 0 downloads a repo id into HF_HOME for the others to read. +# --------------------------------------------------------------------------- +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. + # Expert parallelism needs transformers 5.7 or newer, which the image may not ship. The editable + # install writes into the one clone every node mounts, so the nodes take turns where the + # filesystem offers flock. + mkdir -p /modelopt/build + ( + flock 9 || true + pip install -q -e "/modelopt[hf]" "transformers>=5.7" --no-build-isolation + ) 9>/modelopt/build/.pip-install.lock + pip install -q -r requirements.txt + python -c "import transformers as t; print(t.__name__, t.__version__)" + + torchrun \ + --nnodes="${SLURM_NNODES}" \ + --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}" + ' diff --git a/examples/hf_ptq/vlm_utils.py b/examples/hf_ptq/vlm_utils.py index abfebbd4f01..87cd7e7ec3d 100644 --- a/examples/hf_ptq/vlm_utils.py +++ b/examples/hf_ptq/vlm_utils.py @@ -20,8 +20,40 @@ from PIL import Image from transformers import AutoImageProcessor, AutoProcessor +from modelopt.torch.utils.distributed import all_reduce -def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_remote_code=False): + +def _load_sample_image(images_dir): + """The first sample image in ``images_dir`` that opens, or None, saying why there is none.""" + if not os.path.exists(images_dir): + print(f"❌ Warning: Images directory not found at {images_dir}") + return None + + # Use single image for VL preview to avoid shape mismatch issues + image_files = ["example1a.jpeg", "example1b.jpeg", "example.jpg", "test.jpg", "sample.png"] + missing_files = [] + for img_file in image_files: + img_path = os.path.join(images_dir, img_file) + if os.path.exists(img_path): + try: + image = Image.open(img_path) + print(f" ✅ Successfully loaded: {img_file}") + return image # Use the first available image + except Exception as e: + print(f" ⚠️ Warning: Could not open {img_file}: {e}") + missing_files.append(f"{img_file} (corrupted)") + else: + missing_files.append(img_file) + + print(f"❌ Warning: No valid sample images found in {images_dir}") + print(f" Searched for: {', '.join(image_files)}") + print(f" Missing/invalid files: {', '.join(missing_files)}") + return None + + +def run_vl_preview_generation( + model, tokenizer, model_path, stage_name, trust_remote_code=False, generation_kwargs=None +): """Run preview generation for VL models using sample images. Args: @@ -30,6 +62,7 @@ def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_re model_path: Path to the model (for loading image processor) stage_name: Description of the stage (e.g., "before quantization") trust_remote_code: Whether to trust remote code for Huggingface models and tokenizers + generation_kwargs: Extra ``generate`` kwargs, such as ``synced_gpus`` for a sharded model Returns: Generated response text for logging/comparison """ @@ -38,37 +71,13 @@ def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_re # Load sample images from the images directory script_dir = os.path.dirname(os.path.abspath(__file__)) - images_dir = os.path.join(script_dir, "images") - - # Check if images directory exists - if not os.path.exists(images_dir): - print(f"❌ Warning: Images directory not found at {images_dir}") - print(" VL preview generation requires sample images to test vision capabilities.") - print(" Skipping VL preview generation.") - return None - - # Use single image for VL preview to avoid shape mismatch issues - image_files = ["example1a.jpeg", "example1b.jpeg", "example.jpg", "test.jpg", "sample.png"] - image = None - missing_files = [] - for img_file in image_files: - img_path = os.path.join(images_dir, img_file) - if os.path.exists(img_path): - try: - image = Image.open(img_path) - print(f" ✅ Successfully loaded: {img_file}") - break # Use the first available image - except Exception as e: - print(f" ⚠️ Warning: Could not open {img_file}: {e}") - missing_files.append(f"{img_file} (corrupted)") - else: - missing_files.append(img_file) - - if image is None: - print(f"❌ Warning: No valid sample images found in {images_dir}") - print(f" Searched for: {', '.join(image_files)}") - if missing_files: - print(f" Missing/invalid files: {', '.join(missing_files)}") + image = _load_sample_image(os.path.join(script_dir, "images")) + found = image is not None + if (generation_kwargs or {}).get("synced_gpus"): + # Each node reads the images from its own checkout, while the ranks of a sharded model + # decode in lockstep: every rank previews, or none does. + found = all_reduce("VL preview sample image", lambda: image is not None, all) + if not found: print(" VL preview generation requires sample images to test vision capabilities.") print(" Skipping VL preview generation.") return None @@ -79,6 +88,7 @@ def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_re "max_new_tokens": 50, "do_sample": False, "eos_token_id": tokenizer.eos_token_id, + **(generation_kwargs or {}), } print(f"Generating VL response ({stage_name})...") @@ -189,6 +199,8 @@ def run_vl_preview_generation(model, tokenizer, model_path, stage_name, trust_re return response except Exception as e: + if (generation_kwargs or {}).get("synced_gpus"): + raise # the other ranks of a sharded model would wait for this one in lockstep print(f"❌ VL preview generation {stage_name} failed: {e}") print("This may indicate issues with the quantized model") return None @@ -266,5 +278,7 @@ def run_text_only_generation( return output_text[0] except Exception as e: + if generation_config.get("synced_gpus"): + raise # the other ranks of a sharded model would wait for this one in lockstep print(f"Text-only generation failed: {e}") return None diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 2e5855b1384..8f2f99158b2 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", @@ -843,15 +844,18 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad ) -def _reinstall_a2a_forwards(model, ep_mesh) -> int: - """Re-plant the torch_a2a TP forward wrappers that quantizer insertion clobbered. +def reinstall_a2a_forwards(model, ep_mesh) -> int: + """Re-plant the routed experts' all-to-all forwards that quantizer insertion dropped. transformers installs its TP transforms by replacing ``module.forward`` on the instance (``install_forward``); for torch_a2a_experts that wrapper passes the EP ``process_group`` to the experts forward. Quantizer insertion rebinds ``module.forward``, dropping it, and the experts - forward then raises. Called after conversion, before the first forward. Only a forward quantizer - insertion rebound is planted again: a wrapper still in place is kept, since wrapping it again - would run its transforms twice. Without an ``ep_mesh`` there is nothing to re-plant. + forward then raises. Call this after conversion, before the first forward: + :func:`distributed_forward_loop` does, and so must a caller that quantizes an expert-parallel + model without that loop (a recipe that runs no calibration runs no loop). Only a forward + quantizer insertion rebound is planted again: a wrapper still in place is kept, since wrapping + it again would run its transforms twice. ``ep_mesh`` is the mesh the experts are sharded over; + without one there is nothing to re-plant. Returns how many forwards it planted again. """ if ep_mesh is None: return 0 @@ -908,14 +912,16 @@ def distributed_forward_loop( dataloader: DataLoader, device: torch.device | str | None = None, ep_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, + forward: Callable[[Mapping], object] | None = None, ) -> Callable[[torch.nn.Module], None]: """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. Moves each batch onto ``device``, when given, and runs the captured ``model`` forward (the EP - all-reduce / all-to-all collectives fire inside that forward). A model whose plan - (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its experts are sharded - over, for the all-to-all forwards that quantizer insertion drops to be re-installed before the - first forward; any other model passes none. ``mtq.quantize`` hands the + all-reduce / all-to-all collectives fire inside that forward), or ``forward(batch)`` when given, + for batches that need their own call into ``model``, such as a VLM's image-text batches. A + model whose plan (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its + experts are sharded over, for the all-to-all forwards that quantizer insertion drops to be + re-installed before the first forward; any other model passes none. ``mtq.quantize`` hands the ``forward_loop`` callback an inner module that this closure ignores, calling the captured top-level ``model`` instead so the full distributed forward (and its collectives) runs. @@ -923,7 +929,8 @@ def distributed_forward_loop( the single-node ``_forward_loop`` runs them: a hybrid Mamba/attention model mutates its cache in place. Each batch is a mapping of the forward's keyword arguments, as ``model(**batch)`` takes; only its tensor values move to ``device``, so tensors nested in a value (a list of images, say) - must already be there. + must already be there. ``forward`` gets a new dict when ``device`` is given, and the loader's own + mapping otherwise: copy that before changing it. """ # Local import to avoid any import cycle within modelopt.torch.utils. from modelopt.torch.utils.dataset_utils import _disable_use_cache @@ -938,7 +945,7 @@ def distributed_forward_loop( raise ValueError(f"{plan!r}: the ep_mesh spans {ep_mesh.size()} ranks, not {plan.ep_size}.") def calibrate(_inner_model): - _reinstall_a2a_forwards(model, ep_mesh) + reinstall_a2a_forwards(model, ep_mesh) with _disable_use_cache(model), torch.no_grad(): # Hand the caching allocator's unused pool back: NCCL allocates its buffers outside it. _release_cached_cuda_memory() @@ -953,7 +960,10 @@ def calibrate(_inner_model): k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items() } - model(**batch) + if forward is None: + model(**batch) + else: + forward(batch) # Again for what runs next (amax sync, export collectives), past the forwards' pool. _release_cached_cuda_memory() diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py index 94ed1cab850..8001ebf4076 100644 --- a/modelopt/torch/utils/memory_monitor.py +++ b/modelopt/torch/utils/memory_monitor.py @@ -85,6 +85,7 @@ def __init__(self, monitor_interval: float = 10.0): self.peak_memory = {} # GPU index to peak memory mapping self.is_running = False self.monitor_thread = None + self._stopped = False nvmlInit() self.device_count = nvmlDeviceGetCount() @@ -116,9 +117,17 @@ def stop(self): GPU that was monitored, and properly shuts down the NVML interface. It will wait for the monitoring thread to complete before returning. - The peak memory usage is displayed in GB for each GPU index. + The peak memory usage is displayed in GB for each GPU index. Stopping a stopped monitor does + nothing, as when a caller stops it and the exit hook :func:`launch_memory_monitor` registered + runs too. """ + if self._stopped: + return + self._stopped = True self.is_running = False + # The thread can be inside a sampling pass, still adding to peak_memory: let it finish first. + if self.monitor_thread: + self.monitor_thread.join() # Print peak memory usage print("########") for gpu_idx, peak_mem in self.peak_memory.items(): @@ -126,8 +135,6 @@ def stop(self): f"GPU {gpu_idx}: Peak memory usage = {peak_mem:.2f} GB for all processes on the GPU" ) print("########") - if self.monitor_thread: - self.monitor_thread.join() nvmlShutdown() @@ -140,8 +147,15 @@ def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor | N Returns: GPUMemoryMonitor: The monitor instance that was launched """ + # Probe that NVML can report memory on this host before starting the monitor. NVML must be + # initialized for the probe itself; without it every call raises NVMLError_Uninitialized, and + # the monitor would never start anywhere. try: - nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) + nvmlInit() + try: + nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) + finally: + nvmlShutdown() except Exception as e: print(f"Failed to get GPU memory info: {e}. Stopping GPU memory monitor.") return None diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 01e0634c93a..37580707ea9 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -21,17 +21,22 @@ """ import warnings -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from types import SimpleNamespace from unittest.mock import patch import pytest import torch from _test_utils.examples.hf_ptq_example_utils import example_utils -from _test_utils.torch.transformers_models import create_tiny_qwen3_moe_dir, get_tiny_nemotron_h +from _test_utils.torch.transformers_models import ( + create_tiny_qwen3_moe_dir, + get_tiny_llama, + get_tiny_nemotron_h, +) from safetensors.torch import save_file -from transformers import AutoModelForCausalLM, LlamaConfig +from transformers import AutoModelForCausalLM, GenerationMixin, LlamaConfig +from modelopt.torch.utils.distributed import DistributionPlan from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -192,6 +197,60 @@ def test_build_quant_cfg_sets_moe_calib_experts_ratio_on_every_algorithm(algorit assert quant_cfg["algorithm"] == expected +@pytest.mark.parametrize("wrapped", [False, True], ids=["the-model", "a-model-around-it"]) +def test_a_preview_keeps_generate_from_switching_the_experts_implementation(monkeypatch, wrapped): + """On a GPU, ``generate`` switches a ``grouped_mm`` MoE's experts to ``batched_mm`` for decoding, + which ModelOpt's quantizers do not intercept (the preview after PTQ decodes unquantized experts) + and which Transformers refuses for an FSDP2-wrapped class (``FSDPQwen3MoeForCausalLM does not + support setting experts implementation``). The hook ``generate`` enters is out of the way for the + preview, and back after it, also when the previewed model generates through a submodule, as a + Nemotron-VL model's ``chat`` generates through its language model.""" + entered = [] + + @contextmanager + def switch(self): + entered.append(type(self).__name__) + raise ValueError(f"{type(self).__name__} does not support setting experts implementation.") + yield + + monkeypatch.setattr(GenerationMixin, "_optimize_model_for_decode", switch) + model = get_tiny_llama() + # The model around it has no hook of its own; the preview generates through the inner one. + previewed = torch.nn.ModuleDict({"language_model": model}) if wrapped else model + input_ids = torch.tensor([[1, 2, 3]]) + + # Some releases enter the hook only for the decode steps after the first token. + with pytest.raises(ValueError, match="does not support setting experts implementation"): + model.generate(input_ids, max_new_tokens=3) + with example_utils.keep_the_experts_implementation(previewed): + model.generate(input_ids, max_new_tokens=3) + + assert entered == ["LlamaForCausalLM"] # only the first, unprotected call reached the hook + assert "_optimize_model_for_decode" not in vars(model) + + +def test_a_model_without_the_decode_hook_is_previewed_as_it_is(): + model = torch.nn.Linear(1, 1) + + with example_utils.keep_the_experts_implementation(model): + assert "_optimize_model_for_decode" not in vars(model) + assert not hasattr(model, "_optimize_model_for_decode") + + +def test_a_decode_hook_on_the_model_itself_is_put_back_and_previews_nest(): + """A hook bound on the model instance is restored rather than dropped, and a preview inside + another leaves the outer one's stand-in in place.""" + model = torch.nn.Linear(1, 1) + own = model._optimize_model_for_decode = object() + + with example_utils.keep_the_experts_implementation(model): + with example_utils.keep_the_experts_implementation(model): + pass + assert model._optimize_model_for_decode is nullcontext + + assert model._optimize_model_for_decode is own + + def test_get_original_hf_quant_method_mxfp4_dict(): # gpt-oss layout: quantization_config is a plain dict carrying quant_method. cfg = SimpleNamespace( @@ -716,3 +775,131 @@ def __init__(self, algorithm): def test_recipe_layerwise_blocks(recipe, expected): """Both recipe shapes normalize to dicts, so callers need no shape-aware access.""" assert example_utils.recipe_layerwise_blocks(recipe) == expected + + +# Small enough that the dense part fits any GPU. +_TINY = { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "vocab_size": 32, +} + + +@pytest.mark.parametrize( + ("device", "world_size", "device_map"), + [ + ("cuda:1", 4, {"": "cuda:1"}), + (torch.device("cuda:1"), 4, {"": "cuda:1"}), + ("cuda:1", 1, "auto"), + (torch.device("cpu"), 1, "cpu"), + ], + ids=["str", "torch_device", "single_process", "cpu"], +) +def test_a_rank_device_pins_the_whole_model_to_it( + monkeypatch, tmp_path, device, world_size, device_map +): + """Under torchrun each rank asks for its own cuda:; "auto" would let every rank's + accelerate pick the same idle GPU. A single process spreads the model over the GPUs, whatever + GPU it names, and keeps it on a CPU it names, as the torch.device hf_ptq's run hands in. Only a + single process's load is held to a memory limit sized over the node's GPUs: the pinned map would + ignore one.""" + LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) + # No GPU on the node holds the model, so sizing it over the node limits its memory. + monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 0, "cpu": 1 << 40}) + placed = [] + + plan = DistributionPlan(world_size) + + def load_pretrained(ckpt_path, auto_class, **kwargs): + assert kwargs["distribution"] == plan + placed.append((kwargs["device_map"], "max_memory" in kwargs)) + return torch.nn.Linear(1, 1) + + monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained) + + example_utils.get_model(str(tmp_path), device=device, distribution=plan) + + assert placed == [(device_map, world_size == 1)] + + +@pytest.mark.parametrize( + "checkpoint", + ["pack-quantized", "mxfp4", "diffusion-gemma"], +) +def test_every_load_path_keeps_the_rank_pin(monkeypatch, tmp_path, checkpoint): + """A pack-quantized or MXFP4 checkpoint loads through a path of its own, and DiffusionGemma asks + for a sequential map on a multi-GPU node; each would otherwise place every rank's copy over the + same GPUs.""" + quantization_config = { + "pack-quantized": {"quant_method": "compressed-tensors", "format": "pack-quantized"}, + "mxfp4": {"quant_method": "mxfp4"}, + }.get(checkpoint) + LlamaConfig( + architectures=["LlamaForCausalLM"], quantization_config=quantization_config, **_TINY + ).save_pretrained(tmp_path) + if checkpoint == "diffusion-gemma": + monkeypatch.setattr(example_utils, "is_diffusion_gemma", lambda config: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(example_utils, "_unpack_compressed_linear_weights", lambda *args: None) + placed = [] + + def load_pretrained(ckpt_path, auto_class, **kwargs): + placed.append(kwargs["device_map"]) + return torch.nn.Linear(1, 1) + + monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained) + + example_utils.get_model(str(tmp_path), device="cuda:1", distribution=DistributionPlan(4)) + + assert placed == [{"": "cuda:1"}] + + +@pytest.mark.parametrize( + ("device", "kwargs", "match"), + [ + ("cuda:1", {"use_seq_device_map": True}, "--use_seq_device_map"), + # "auto" would let every rank's accelerate pick the same idle GPU. + ("cuda", {}, "names no GPU"), + ], + ids=["sequential_map", "unindexed_device"], +) +def test_a_placement_over_shared_gpus_is_refused_where_each_rank_holds_its_own_copy( + tmp_path, device, kwargs, match +): + LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) + + with pytest.raises(ValueError, match=match): + example_utils.get_model( + str(tmp_path), device=device, distribution=DistributionPlan(4), **kwargs + ) + + +@pytest.mark.parametrize( + "plan", + [DistributionPlan(4, ep_size=4), DistributionPlan(4, dense="fsdp")], + ids=["expert_parallel", "fsdp"], +) +def test_a_plan_that_shards_the_model_is_refused(monkeypatch, tmp_path, plan): + """get_model loads a whole model on every rank; a plan that shards it goes to load_pretrained.""" + LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) + monkeypatch.setattr(example_utils, "load_pretrained", lambda *a, **k: pytest.fail("loaded")) + + with pytest.raises(ValueError, match="shards the model over the ranks"): + example_utils.get_model(str(tmp_path), device="cuda:1", distribution=plan) + + +def test_without_a_plan_the_load_decides_and_refuses_sharding_knobs_it_cannot_honor( + monkeypatch, tmp_path +): + """get_model hands load_pretrained only a plan its caller gave. Without one the load decides + as before, so a MODELOPT_* knob asking to shard over a launch whose process group is not set up + is refused, rather than ignored for a one-rank plan.""" + LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path) + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + + with pytest.raises(RuntimeError, match="MODELOPT_FSDP_SIZE=world ask to shard"): + example_utils.get_model(str(tmp_path), device="cuda:0") diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 5543e9de902..0ea0dab0264 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +26,7 @@ import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 -from transformers import Qwen3_5MoeConfig, Qwen3Config, Qwen3MoeConfig +from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -38,7 +38,7 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.utils import mlflow as mlflow_lib -from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan +from modelopt.torch.utils.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" @@ -410,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) @@ -445,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) @@ -471,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 @@ -481,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( @@ -1012,9 +992,12 @@ def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch): NVBug 6752977.""" hf_ptq = _import_hf_ptq(monkeypatch) - full_model = SimpleNamespace( - generate=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) - ) + class _FailingModel(torch.nn.Module): + def generate(self, *args, **kwargs): + raise RuntimeError("boom") + + full_model = _FailingModel() + set_distribution_plan(full_model, DistributionPlan()) export_calls = [] monkeypatch.setattr( hf_ptq, @@ -1044,6 +1027,32 @@ def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch): assert len(export_calls) == 1 +def test_post_quantize_stops_the_run_when_a_sharded_models_preview_fails(monkeypatch): + """A sharded model's ranks decode the preview in lockstep: a rank that skipped its failure would + go on to the export's collectives while the others still wait in generate's.""" + hf_ptq = _import_hf_ptq(monkeypatch) + + class _FailingModel(torch.nn.Module): + def generate(self, *args, **kwargs): + raise RuntimeError("boom") + + full_model = _FailingModel() + set_distribution_plan(full_model, DistributionPlan(2, dense="fsdp")) + export_calls = [] + monkeypatch.setattr(hf_ptq, "export_quantized", lambda *a, **k: export_calls.append((a, k))) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + ids = torch.zeros(1, 4, dtype=torch.long) + args = SimpleNamespace(specdec_offline_dataset=None, verbose=False) + + with pytest.raises(RuntimeError, match="--skip_generate") as stopped: + hf_ptq.post_quantize( + args, full_model, full_model, "llama", None, None, ids, None, ids, False, None, None + ) + + assert str(stopped.value.__cause__) == "boom" + assert export_calls == [] + + @pytest.mark.usefixtures("clean_env") def test_main_makes_a_hub_checkpoint_local_before_loading(monkeypatch, tmp_path): """--pyt_ckpt_path keeps the Hub ID; the model loads from the local copy resolved first.""" @@ -1078,6 +1087,38 @@ def load_model(args): ] +@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 = [] + + class _LoadReachedError(Exception): + pass + + def load_model(args): + raise _LoadReachedError + + monkeypatch.setattr(torch.compiler, "set_stance", lambda *args, **kwargs: None) + monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: launched.append(True)) + monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", lambda path: (None, path)) + monkeypatch.setattr(hf_ptq, "load_model", load_model) + + with pytest.raises(_LoadReachedError): + hf_ptq._main_inner(args) + + assert len(launched) == launches + + def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypatch, tmp_path): """The deprecated TensorRT-LLM exporter does not copy them itself; hf_ptq does, from disk.""" hf_ptq, args = _parse_hf_ptq_args( @@ -1094,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 new file mode 100644 index 00000000000..176d416f034 --- /dev/null +++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py @@ -0,0 +1,986 @@ +# 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 +import warnings +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", "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 +): + """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, 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) + 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, ep_size=4), False, True), + (DistributionPlan(4, ep_size=4), True, True), + ], + ids=["replicated", "expert_parallel", "expert_parallel_nemotron_vl"], +) +def test_a_failed_preview_of_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.""" + + 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_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py index 328dcd5e564..2539311313e 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -24,6 +24,16 @@ _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" +class _Model(torch.nn.Module): + """A model whose ``generate`` is ``generate``: the previews read the module's distribution plan.""" + + def __init__(self, generate, config=None): + super().__init__() + self.generate = generate + if config is not None: + self.config = config + + @pytest.fixture def hf_ptq(monkeypatch: pytest.MonkeyPatch): monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) @@ -43,7 +53,7 @@ def hf_ptq(monkeypatch: pytest.MonkeyPatch): ) def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: bool): generate_calls = [] - full_model = SimpleNamespace(generate=lambda *args, **kwargs: generate_calls.append(1) or "out") + full_model = _Model(lambda *args, **kwargs: generate_calls.append(1) or "out") args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) batch = {"input_ids": torch.ones(2, 4, dtype=torch.long)} @@ -58,7 +68,7 @@ def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: ) def test_pre_quantize_preview_removes_padding(hf_ptq, ids, mask): calls = [] - model = SimpleNamespace(generate=lambda ids, **kwargs: calls.append((ids, kwargs)) or ids) + model = _Model(lambda ids, **kwargs: calls.append((ids, kwargs)) or ids) args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) batch = {"input_ids": torch.tensor([ids]), "attention_mask": torch.tensor([mask])} prompt, attention_mask, _ = hf_ptq.pre_quantize(args, model, "llama", None, [batch], False) @@ -83,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: @@ -122,9 +133,7 @@ def test_post_quantize_preview_keeps_whole_answer( ): """The preview must decode the answer, not slice a prompt-length prefix off it.""" monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None) - full_model = SimpleNamespace( - config=config, generate=lambda *args, **kwargs: torch.tensor(generated) - ) + full_model = _Model(lambda *args, **kwargs: torch.tensor(generated), config) tokenizer = SimpleNamespace( batch_decode=lambda ids, **kwargs: [" ".join(map(str, row.tolist())) for row in ids] ) diff --git a/tests/examples/hf_ptq/test_hf_ptq_previews.py b/tests/examples/hf_ptq/test_hf_ptq_previews.py new file mode 100644 index 00000000000..b339722c050 --- /dev/null +++ b/tests/examples/hf_ptq/test_hf_ptq_previews.py @@ -0,0 +1,229 @@ +# 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. + +"""hf_ptq's generation previews of a model its distribution plan shards. + +Each rank previews the first sample of its own calibration shard, so the ranks finish decoding at +different steps; a sharded model's previews keep every rank decoding until all have finished, and +keep ``generate`` from moving a MoE's experts off the backend quantization hooks. +""" + +import contextlib +import importlib +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from _test_utils.examples.hf_ptq_example_utils import example_utils +from _test_utils.torch.transformers_models import get_tiny_tokenizer +from PIL import Image + +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" +_CKPT = "/ckpt/model" + + +@pytest.fixture +def hf_ptq(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + return importlib.import_module("hf_ptq") + + +class _PreviewModel(torch.nn.Module): + """A model whose ``generate`` records the kwargs each preview passes it, and how often it entered + the decode hook that Transformers' ``generate`` enters to switch a MoE's experts.""" + + def __init__(self, plan=None): + super().__init__() + self.config = SimpleNamespace(is_encoder_decoder=False, model_type="qwen3_moe") + self.calls = [] + self.switches = 0 + if plan is not None: + set_distribution_plan(self, plan) + + @contextlib.contextmanager + def _optimize_model_for_decode(self): + self.switches += 1 + yield + + def generate(self, input_ids, **kwargs): + self.calls.append(kwargs) + with self._optimize_model_for_decode(): + return torch.cat([input_ids, input_ids[:, :1]], dim=1) + + +_PREVIEW_PLANS = pytest.mark.parametrize( + ("plan", "lockstep"), + [ + (None, False), + (DistributionPlan(4), False), + (DistributionPlan(4, ep_size=4), True), + (DistributionPlan(4, dense="fsdp"), True), + (DistributionPlan(1, dense="fsdp"), False), + ], + ids=["single_process", "replicated", "expert_parallel", "fsdp2", "one_rank_fsdp2"], +) + + +@_PREVIEW_PLANS +def test_a_sharded_model_previews_in_lockstep(hf_ptq, monkeypatch, plan, lockstep): + """Each rank previews its own calibration shard's first sample, so the ranks finish decoding at + different steps. A sharded model runs a collective in every forward, and a rank that stopped + early left the others blocked in it: Qwen3-30B-A3B NVFP4 at EP4 hung in the post-PTQ preview. + Both previews keep every rank of a sharded model decoding until all have finished.""" + model = _PreviewModel(plan) + loader = [{"input_ids": torch.tensor([[1, 2, 3]]), "attention_mask": torch.ones(1, 3)}] + args = SimpleNamespace( + specdec_offline_dataset=None, + skip_generate=False, + verbose=False, + local_checkpoint_path=_CKPT, + trust_remote_code=False, + ) + monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None) + + ids, mask, before = hf_ptq.pre_quantize(args, model, "qwen3_moe", None, loader, False) + hf_ptq.post_quantize( + args, + model, + model, + "qwen3_moe", + get_tiny_tokenizer(), + None, + ids, + mask, + before, + False, + None, + loader, + ) + + synced = [call.get("synced_gpus", False) for call in model.calls] + assert synced == [lockstep, lockstep] # before and after quantization + assert model.switches == 0 # neither preview lets generate move the experts off their backend + + +@_PREVIEW_PLANS +def test_a_sharded_nemotron_vl_model_previews_in_lockstep(monkeypatch, plan, lockstep): + """The Nemotron-VL preview's text, fallback and image generations keep a sharded model's ranks + decoding together too.""" + recorded = {} + keeps = lambda model: "_optimize_model_for_decode" in vars(model) # noqa: E731 + vlm_utils = SimpleNamespace( + run_text_only_generation=lambda model, tokenizer, question, config, *args: recorded.update( + text=config, text_keeps_the_experts=keeps(model) + ), + run_vl_preview_generation=lambda model, tokenizer, path, stage, trust, kwargs: ( + recorded.update(image=kwargs, image_keeps_the_experts=keeps(model)) + ), + ) + monkeypatch.setitem(sys.modules, "vlm_utils", vlm_utils) + model = _PreviewModel(plan) + + example_utils.run_nemotron_vl_preview( + model, get_tiny_tokenizer(), torch.tensor([[1, 2, 3]]), _CKPT, "stage", allow_fallback=True + ) + + expected = {"synced_gpus": True} if lockstep else {} + assert recorded["text"].get("synced_gpus", False) == lockstep + assert recorded["image"] == expected + assert recorded["text_keeps_the_experts"] and recorded["image_keeps_the_experts"] + assert "_optimize_model_for_decode" not in vars(model) + assert model.calls[0].get("synced_gpus", False) == lockstep # the fallback generate + + +@pytest.mark.parametrize(("lockstep", "raises"), [({}, False), ({"synced_gpus": True}, True)]) +def test_a_failed_text_only_vl_preview_stops_a_sharded_run(monkeypatch, lockstep, raises): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + vlm_utils = importlib.import_module("vlm_utils") + + class _ChatModel: + def chat(self, *args, **kwargs): + raise RuntimeError("the preview failed on this rank") + + run = lambda: vlm_utils.run_text_only_generation( # noqa: E731 + _ChatModel(), None, "question", dict(lockstep), _CKPT + ) + + if raises: + with pytest.raises(RuntimeError, match="failed on this rank"): + run() + else: + assert run() is None + + +@pytest.mark.parametrize("generation_kwargs", [None, {"synced_gpus": True}]) +def test_the_vl_image_preview_generates_with_the_kwargs_it_is_handed( + monkeypatch, tmp_path, generation_kwargs +): + """run_vl_preview_generation hands the kwargs it is given to the model's chat(), in its config.""" + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + vlm_utils = importlib.import_module("vlm_utils") + (tmp_path / "images").mkdir() + Image.new("RGB", (4, 4)).save(tmp_path / "images" / "example1a.jpeg") + monkeypatch.setattr(vlm_utils, "__file__", str(tmp_path / "vlm_utils.py")) + monkeypatch.setattr( + vlm_utils.AutoImageProcessor, "from_pretrained", lambda *args, **kwargs: lambda images: {} + ) + received = [] + + class _ChatModel: + device = torch.device("cpu") + + def chat(self, tokenizer, question, generation_config, **image_features): + received.append(generation_config) + return "a picture" + + vlm_utils.run_vl_preview_generation( + _ChatModel(), SimpleNamespace(eos_token_id=0), _CKPT, "stage", False, generation_kwargs + ) + + assert len(received) == 1 + assert received[0].get("synced_gpus", False) == (generation_kwargs is not None) + + +@pytest.mark.parametrize("has_image", [True, False], ids=["a_peer_has_none", "no_rank_has_one"]) +def test_the_ranks_of_a_sharded_model_skip_the_image_preview_together( + monkeypatch, tmp_path, has_image +): + """Each node reads the sample images from its own checkout, which may have none: the repository + ships none. A rank that previewed alone would wait for its peers in lockstep, so a rank with an + image skips when a peer has none, and when no rank has one, they all skip without stopping.""" + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + vlm_utils = importlib.import_module("vlm_utils") + if has_image: + (tmp_path / "images").mkdir() + Image.new("RGB", (4, 4)).save(tmp_path / "images" / "example1a.jpeg") + monkeypatch.setattr(vlm_utils, "__file__", str(tmp_path / "vlm_utils.py")) + monkeypatch.setattr( + vlm_utils.AutoImageProcessor, "from_pretrained", lambda *args, **kwargs: lambda images: {} + ) + # This rank's answer, and a peer's that found no image. + monkeypatch.setattr(vlm_utils, "all_reduce", lambda label, fn, combine: combine([fn(), False])) + + class _ChatModel: + device = torch.device("cpu") + + def chat(self, *args, **kwargs): + pytest.fail("this rank previewed without its peers") + + preview = vlm_utils.run_vl_preview_generation( + _ChatModel(), SimpleNamespace(eos_token_id=0), _CKPT, "stage", False, {"synced_gpus": True} + ) + + assert preview is None 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..dbf194c2b3f 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,8 @@ 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", 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, 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, diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index f5a21d23dc3..a0f750c7312 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -51,8 +51,8 @@ from modelopt.torch.utils.distributed import ( DistributionPlan, ParallelState, - _reinstall_a2a_forwards, distributed_forward_loop, + reinstall_a2a_forwards, set_distribution_plan, shard_dataloader, ) @@ -262,7 +262,7 @@ def test_sharded_loader_keeps_the_original_loader_settings(dataset): assert getattr(shard, setting) == getattr(loader, setting), setting -# -------------------------------------------------------------------------- _reinstall_a2a_forwards +# -------------------------------------------------------------------------- reinstall_a2a_forwards _TP_MODULE = "transformers.distributed.tensor_parallel" _LEGACY_TP_MODULE = "transformers.integrations.tensor_parallel" @@ -398,7 +398,7 @@ def test_reinstall_is_a_noop_without_an_ep_mesh(monkeypatch): """A single-process or non-EP model is left untouched, even with the torch_a2a styles around.""" styles = _install_fake_tp(monkeypatch) - assert _reinstall_a2a_forwards(_MoEModel(), None) == 0 + assert reinstall_a2a_forwards(_MoEModel(), None) == 0 assert not any(style.installed for style in styles.values()) @@ -423,7 +423,7 @@ def test_reinstall_skips_styles_whose_transforms_are_pre_hooks(monkeypatch, expe for layer in model.layers: # quantizer insertion rebinds the planned modules' forwards _drop_tp_wrapper(layer.gate) _drop_tp_wrapper(layer.experts) - assert _reinstall_a2a_forwards(model, _FakeMesh()) == planted + assert reinstall_a2a_forwards(model, _FakeMesh()) == planted def test_reinstall_leaves_a_module_the_load_did_not_wrap(monkeypatch): @@ -431,7 +431,7 @@ def test_reinstall_leaves_a_module_the_load_did_not_wrap(monkeypatch): the load did not wrap keeps running as it did, without the style's transforms.""" styles = _install_fake_tp(monkeypatch) - assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0 + assert reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0 assert not any(style.installed for style in styles.values()) @@ -484,6 +484,21 @@ def test_forward_loop_moves_tensors_to_the_device_it_is_given(): assert [call["logits_to_keep"] for call in model.calls] == [1, 1] +def test_forward_loop_runs_each_batch_through_the_forward_it_is_given(): + """A batch that needs its own call into the model, as a VLM's image-text batch does, goes + through ``forward``, on the device and with the cache disabled like any other.""" + batches = [{"input_ids": torch.zeros(2, 3)}] * 2 + model = _Recorder() + + distributed_forward_loop( + model, batches, device="meta", forward=lambda batch: model(**batch, own_call=True) + )(model) + + assert [call["own_call"] for call in model.calls] == [True, True] + assert [call["input_ids"].device.type for call in model.calls] == ["meta"] * 2 + assert model.use_cache_seen == [False, False] + + def test_forward_loop_disables_use_cache_only_while_calibrating(): """KV/state caching is off for the calibration forwards (it breaks hybrid Mamba models).""" model = _Recorder() @@ -504,7 +519,7 @@ def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch, """ events = [] monkeypatch.setattr( - mto_dist, "_reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args)) + mto_dist, "reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args)) ) monkeypatch.setattr(torch.cuda, "is_initialized", lambda: cuda_started) monkeypatch.setattr(mto_dist.gc, "collect", lambda: events.append("collect")) diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py new file mode 100644 index 00000000000..eac9aa8d9c2 --- /dev/null +++ b/tests/unit/torch/utils/test_memory_monitor.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for launch_memory_monitor's NVML availability probe, against a fake NVML.""" + +import threading +from types import SimpleNamespace + +import pytest + +import modelopt.torch.utils.memory_monitor as mm + + +class _FakeNVML: + """Reference-counted like NVML: every call but init raises while uninitialized.""" + + def __init__(self, memory_info_supported: bool = True, devices: int = 1, on_sample=None): + self.inits = 0 + self.memory_info_supported = memory_info_supported + self.devices = devices + self.on_sample = on_sample # called with a device's handle before its memory is read + + def install(self, monkeypatch): + for name in ("Init", "Shutdown", "DeviceGetCount", "DeviceGetHandleByIndex"): + monkeypatch.setattr(mm, f"nvml{name}", getattr(self, name)) + monkeypatch.setattr(mm, "nvmlDeviceGetMemoryInfo", self.DeviceGetMemoryInfo) + + def _require_init(self): + if self.inits <= 0: + raise RuntimeError("Uninitialized") + + def Init(self): # noqa: N802 + self.inits += 1 + + def Shutdown(self): # noqa: N802 + self._require_init() + self.inits -= 1 + + def DeviceGetCount(self): # noqa: N802 + self._require_init() + return self.devices + + def DeviceGetHandleByIndex(self, index): # noqa: N802 + self._require_init() + return index + + def DeviceGetMemoryInfo(self, handle): # noqa: N802 + self._require_init() + if not self.memory_info_supported: + raise RuntimeError("Not Supported") + if self.on_sample is not None: + self.on_sample(handle) + return SimpleNamespace(used=2 * 1024**3) + + +@pytest.fixture +def launched(): + monitors = [] + yield monitors + for monitor in monitors: + monitor.stop() + + +@pytest.mark.parametrize("supported", [True, False], ids=["reports_memory", "cannot"]) +def test_the_monitor_starts_only_where_nvml_can_report_memory(monkeypatch, launched, supported): + """The probe used to run before NVML was initialized, so it always failed and the monitor + never started on any host. Hosts where NVML cannot report memory still skip the monitor, and + leave NVML shut down; elsewhere the probe's own init is balanced and the monitor holds one.""" + nvml = _FakeNVML(memory_info_supported=supported) + nvml.install(monkeypatch) + monitor = mm.launch_memory_monitor(monitor_interval=0.01) + if monitor is not None: + launched.append(monitor) + + assert (monitor is not None) is supported + assert nvml.inits == (1 if supported else 0) + + +def test_stopping_a_stopped_monitor_does_nothing(monkeypatch, capsys): + """launch_memory_monitor registers stop() to run at exit, so a monitor its caller stopped is + stopped again then: the second stop neither reports again nor shuts NVML down a second time.""" + nvml = _FakeNVML() + nvml.install(monkeypatch) + monitor = mm.launch_memory_monitor(monitor_interval=0.01) + + monitor.stop() + # The report's frame, which prints whether or not the monitor thread sampled a GPU yet. + assert "########" in capsys.readouterr().out + monitor.stop() + + assert capsys.readouterr().out == "" + assert nvml.inits == 0 + + +def test_the_report_waits_for_a_sampling_pass_in_flight(monkeypatch, capsys): + """stop() can come while the thread is between two GPUs of a sampling pass: the report waits for + the pass to finish rather than miss the last GPU, or iterate the samples while they grow.""" + in_flight, joining = threading.Event(), threading.Event() + + def hold_the_second_gpu(handle): + if handle == 1 and threading.current_thread() is not threading.main_thread(): + # The first GPU is sampled; hold the second until stop() waits for the thread (5 s max). + in_flight.set() + joining.wait(timeout=5) + + nvml = _FakeNVML(devices=2, on_sample=hold_the_second_gpu) + nvml.install(monkeypatch) + monitor = mm.launch_memory_monitor(monitor_interval=0.01) + join = monitor.monitor_thread.join + + def join_and_release(*args, **kwargs): + joining.set() + return join(*args, **kwargs) + + monkeypatch.setattr(monitor.monitor_thread, "join", join_and_release) + assert in_flight.wait(timeout=5) + + monitor.stop() + + report = capsys.readouterr().out + assert "GPU 0: Peak memory usage" in report + assert "GPU 1: Peak memory usage" in report + assert nvml.inits == 0