diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 1b9783d8dec..888bbd8aead 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -20,6 +20,7 @@ Changelog
 - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype.
 - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2.
 - Add an opt-in no-gather HuggingFace export for FSDP2-sharded models: with ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` sends each tensor to a single owner rank and every rank writes its own shards in parallel, instead of gathering the full state dict on one host. MoE experts stay sharded through the export, so each rank only materializes its own experts.
+- Add expert-parallel PTQ for HuggingFace MoE models to ``examples/hf_ptq/hf_ptq.py``: under ``torchrun``, each rank holds only its share of the routed experts, tokens reach them by all-to-all, and every rank calibrates on its own slice of the data. ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` override the automatic choice of expert-parallel degree and of FSDP2-sharding the dense layers, and the quantized checkpoint is written by every rank in parallel without gathering the model.
 
 *Speculative Decoding*
 
@@ -42,6 +43,7 @@ Changelog
 
 **Backward Breaking Changes**
 
+- ``examples/hf_ptq/hf_ptq.py`` launched by ``torchrun`` on more than one rank now refuses a ``--device`` other than ``cuda``, which it used to ignore: each rank runs on its own GPU, ``cuda:<local_rank>``.
 - ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank: AWQ, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate. Each rank calibrates on its own share of the data, so these calibrated the ranks' shards of a weight differently; run them without ``--use_fsdp2``.
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later.
diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 4fe27160db7..8780f55a3df 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -654,6 +654,32 @@ See [Recipe-based Quantization](#recipe-based-quantization) for the recipe forma
 
 > *Set `MODELOPT_FSDP2_EXPORT=nogather` to have every rank write its own shards of the checkpoint in parallel instead of gathering the full state dict on one host.*
 
+## Expert-Parallel Post-Training Quantization for MoE Models
+
+Under `torchrun`, `hf_ptq.py` shards an MoE model's routed experts across the ranks (expert parallelism) instead of loading the whole model on every GPU. Each rank holds `num_experts / ep_degree` experts and calibrates on its own shard of the calibration data. An all-to-all sends every token to the rank that owns its expert, so each expert's quantizers see the tokens routed to it from all ranks. The dense (non-expert) part is replicated on every rank, or FSDP2-sharded when it does not fit one GPU. Each rank then writes its share of the unified HF checkpoint, without gathering the model onto one GPU.
+
+A model can be expert-parallel when its config declares a `base_model_ep_plan`. transformers ships one for several MoEs, and ModelOpt supplies the plan or an adapter for others (see the `ep_spec` entries in `modelopt/torch/models/*/specs.py`). Expert parallelism needs transformers 5.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=<num_gpus_per_node> \
+    hf_ptq.py \
+    --pyt_ckpt_path <path_to_moe_model> \
+    --recipe general/ptq/nvfp4_default-kv_fp8 \
+    --calib_size <num_calib_samples> \
+    --export_path <export_path>
+```
+
+To run across nodes, add the `--nnodes`, `--node_rank`, `--master_addr` and `--master_port` flags shown in [Manual (run on each node)](#manual-run-on-each-node). `--calib_size` is the total over all ranks.
+
+Two environment variables override the automatic distribution plan:
+
+| Variable | Default | Effect |
+| :--- | :--- | :--- |
+| `MODELOPT_EP_DEGREE` | `auto` | Expert-parallel degree. `auto` picks the largest degree that divides both the routed-expert count and the world size, or 1 for a model whose spec refuses expert parallelism or needs a newer transformers. Below the world size, each group of `MODELOPT_EP_DEGREE` ranks holds a full set of experts, and the groups calibrate on different data. |
+| `MODELOPT_FSDP_SIZE` | `auto` | Whether to FSDP2-shard the dense part. `auto` shards it when it would take more than 60% of one GPU's memory; `world` always shards it; `1` replicates it. With `MODELOPT_EP_DEGREE` below the world size the export cannot write a sharded dense part yet, so such a run fails before loading. |
+
+A run that cannot shard the model the way the plan asks fails before loading, rather than loading the whole model on every rank. `--use_fsdp2` skips expert parallelism and uses the [FSDP2 path](#multi-node-post-training-quantization-with-fsdp2) instead. Expert parallelism refuses GPT-OSS and Llama-4: their experts are stored transposed, which the expert-parallel export cannot write yet. `--use_fsdp2` refuses their released checkpoints too (GPT-OSS's carry an MXFP4 `quantization_config`, and Llama-4's are multimodal), so quantize those in a single process (without `torchrun`), which spreads the model over the node's GPUs.
+
 ## Evaluate Accuracy
 
 ### TensorRT-LLM Validation
diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index e90b034960c..c97656fb87e 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
 
@@ -52,7 +50,7 @@
     prefers_bundled_code,
     prepare_model_for_loading,
 )
-from modelopt.torch.utils import distributed as dist_utils
+from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
 from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run
 from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
 
@@ -61,39 +59,6 @@
 SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"]
 
 
-@dataclass
-class DistributedState:
-    """Example-local distributed state for model loading, dataloader sharding, and rank-0 output."""
-
-    rank: int
-    world_size: int
-    device: torch.device | str
-    is_main: bool
-
-
-def setup_distributed_args(args):
-    """Initialize and attach ``args.dist_state`` (single-process if FSDP2 off)."""
-    if getattr(args, "use_fsdp2", False):
-        # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can
-        # exceed it, and PyTorch 2.8 has no per-call barrier() timeout (must be set at PG creation).
-        dist_utils.setup(timeout=timedelta(hours=2))
-        rank = dist_utils.rank()
-        args.dist_state = DistributedState(
-            rank=rank,
-            world_size=dist_utils.size(),
-            device=torch.device(f"cuda:{dist_utils.local_rank()}"),
-            is_main=rank == 0,
-        )
-    else:
-        args.dist_state = DistributedState(rank=0, world_size=1, device=args.device, is_main=True)
-
-
-def cleanup_distributed(args):
-    """Destroy the process group if ``--use_fsdp2`` set it up."""
-    if getattr(args, "use_fsdp2", False):
-        dist_utils.cleanup()
-
-
 def validate_fsdp2_supported(args, config):
     """Raise ``NotImplementedError`` for model/CLI combos the FSDP2 path doesn't support yet."""
     issues = []
@@ -118,6 +83,48 @@ def validate_fsdp2_supported(args, config):
         )
 
 
+def preview_generate_kwargs(model) -> dict[str, Any]:
+    """Extra ``generate`` kwargs that keep every rank of a sharded model decoding a preview.
+
+    Each rank previews the first sample of its own calibration shard, so the ranks finish at
+    different steps. A model its plan shards runs a collective in every forward (the routed experts'
+    all-to-all, FSDP2's all-gathers): a rank that stopped decoding would leave the others blocked in
+    it. ``synced_gpus`` keeps every rank stepping until all of them have finished.
+    """
+    plan = distribution_plan(model)
+    sharded = plan.world_size > 1 and (plan.is_expert_parallel or plan.is_fsdp_dense)
+    return {"synced_gpus": True} if sharded else {}
+
+
+@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,
@@ -142,32 +149,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):
@@ -205,18 +217,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
@@ -240,37 +252,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
 
@@ -567,8 +597,11 @@ 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.
+    plan = distribution or DistributionPlan()
 
     _disk_offload = offload_folder is not None
     if _disk_offload and max_cpu_memory_gb is None:
@@ -579,9 +612,23 @@ def get_model(
             UserWarning,
         )
 
-    device_map = "auto"
-    if device == "cpu":
-        device_map = "cpu"
+    # Under torchrun each rank asks for its own "cuda:<local_rank>". 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".
+    if plan.world_size > 1 and device != "cpu" and ":" not in str(device):
+        raise ValueError(
+            f"{plan!r} has each rank hold its own copy of the model on its own GPU, but "
+            f"device={device!r} names no GPU index; pass the rank's own GPU, cuda:<local_rank>."
+        )
+    pinned = {"": str(device)} if device != "cpu" and plan.world_size > 1 else None
+    device_map = "cpu" if device == "cpu" else pinned or "auto"
 
     # Prepare config kwargs for loading
     config_kwargs = {"trust_remote_code": trust_remote_code} if trust_remote_code else {}
@@ -612,13 +659,24 @@ def get_model(
     # modules together for checkpoints that fit; larger ones can still spill and split a
     # tied pair, and need an explicit single-device map. Multi-GPU only: a single-GPU split
     # cannot separate a tied pair, and sequential would needlessly cap max_memory there.
-    if device != "cpu" and torch.cuda.device_count() > 1 and is_diffusion_gemma(hf_config):
+    if (
+        not pinned
+        and device != "cpu"
+        and torch.cuda.device_count() > 1
+        and is_diffusion_gemma(hf_config)
+    ):
         print(
             "Detected DiffusionGemma model. Using device_map='sequential'; the balanced "
             "'auto' mapping can split its tied encoder/decoder weights across GPUs."
         )
         use_seq_device_map = True
 
+    if use_seq_device_map and pinned:
+        raise ValueError(
+            f"--use_seq_device_map spreads one model over every GPU, but each rank of this run "
+            f"holds its own copy on its own GPU ({device}). Drop it, or shard the model with "
+            "--use_fsdp2."
+        )
     if use_seq_device_map:
         device_map = "sequential"
         # If we use sequential, set max_memory limit to ensure that the model does not occupy the full GPU
@@ -670,6 +728,7 @@ def has_pack_quantized_config(config):
             ckpt_path,
             AutoModelForCausalLM,
             model_type=hf_config.model_type,
+            distribution=plan,
             device_map=device_map,
             **model_kwargs,
         )
@@ -681,7 +740,8 @@ def has_pack_quantized_config(config):
                 ckpt_path,
                 AutoModelForCausalLM,
                 model_type=hf_config.model_type,
-                device_map="auto",
+                distribution=plan,
+                device_map=pinned or "auto",
                 trust_remote_code=trust_remote_code,
                 dtype="auto",
             )
@@ -707,7 +767,8 @@ def has_pack_quantized_config(config):
             ckpt_path,
             AutoModelForCausalLM,
             model_type=hf_config.model_type,
-            device_map="cpu" if device == "cpu" else "sequential",
+            distribution=plan,
+            device_map="cpu" if device == "cpu" else pinned or "sequential",
             **model_kwargs,
         )
     else:
@@ -793,7 +854,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:
@@ -815,6 +878,7 @@ def has_pack_quantized_config(config):
             ckpt_path,
             auto_model_module,
             model_type=hf_config.model_type,
+            distribution=plan,
             device_map=device_map,
             **model_kwargs2,
         )
@@ -827,7 +891,7 @@ def has_pack_quantized_config(config):
         print(f"Moving model to {device} device...")
         model = model.to(device)
 
-    if device == "cuda" and not is_model_on_gpu(model):
+    if torch.device(device).type == "cuda" and not is_model_on_gpu(model):
         print("Warning: Some parameters are not on a GPU. Calibration can be slow or hit OOM")
 
     return model
@@ -990,6 +1054,13 @@ def recipe_layerwise_blocks(recipe) -> list[dict]:
     return blocks
 
 
+def layerwise_export_requested(blocks) -> bool:
+    """Whether a recipe's ``layerwise`` blocks (:func:`recipe_layerwise_blocks`) export layerwise."""
+    return any(
+        block.get("export_dir") is not None and block.get("enable", False) for block in blocks
+    )
+
+
 def _layerwise_checkpoint_dir(algorithm) -> str | None:
     """First ``layerwise.checkpoint_dir`` across the algorithm entries, or None."""
     return next(
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 78e175b95a3..6697627f140 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -18,6 +18,7 @@
 import random
 import time
 import warnings
+from datetime import timedelta
 from pathlib import Path
 from typing import Any
 
@@ -35,7 +36,7 @@
     HF_PTQ,
     _prepare_quant_cfg,
     build_quant_cfg,
-    cleanup_distributed,
+    create_vlm_calibration_forward,
     create_vlm_calibration_loop,
     generate_excludes_prompt,
     get_model,
@@ -43,11 +44,13 @@
     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,
     run_nemotron_vl_preview,
     save_source_config,
-    setup_distributed_args,
     validate_fsdp2_supported,
 )
 from torch.utils.data import DataLoader
@@ -99,23 +102,32 @@
     get_max_batch_size,
     get_supported_datasets,
 )
-from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+from modelopt.torch.utils.distributed import (
+    _reinstall_a2a_forwards,
+    distributed_forward_loop,
+    distributed_run,
+    distribution_plan,
+    local_rank,
+    shard_dataloader,
+)
 from modelopt.torch.utils.memory_monitor import launch_memory_monitor
 from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args
+from modelopt.torch.utils.moe import expert_parallel_mesh
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
     copy_non_model_files,
     ensure_local_checkpoint,
 )
-from modelopt.torch.utils.plugins.model_load_utils import (
-    load_pretrained,
-    parallel_load_and_prepare_fsdp2,
-)
+from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan
 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
 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 +200,20 @@ def _to_device(value):
     return {k: _to_device(v) for k, v in batch.items()}
 
 
+def _shard_calib_dataloader(calib_dataloader, rank: int, world_size: int):
+    """This rank's share of the calibration data, in the loader type it came in."""
+    if isinstance(calib_dataloader, _DeviceDataLoader):
+        return _DeviceDataLoader(
+            shard_dataloader(calib_dataloader.dataloader, rank, world_size), calib_dataloader.device
+        )
+    if isinstance(calib_dataloader, DataLoader):
+        return shard_dataloader(calib_dataloader, rank, world_size)
+    raise NotImplementedError(
+        f"A {type(calib_dataloader).__name__} calibration loader cannot be split across this "
+        f"run's {world_size} ranks."
+    )
+
+
 def make_calib_dataloader(
     args: argparse.Namespace,
     language_model: torch.nn.Module,
@@ -269,12 +295,15 @@ def make_calib_dataloader(
             max_sample_length=args.calib_seq,
             device=device,
             include_labels=include_labels,
-            distributed=args.use_fsdp2,
-            sampler_kwargs=(
-                {"num_replicas": args.dist_state.world_size, "rank": args.dist_state.rank}
-                if args.use_fsdp2
-                else None
-            ),
+        )
+    world_size = distribution_plan(language_model).world_size
+    if world_size > 1 and calib_dataloader is not None:
+        # Every distributed calibration path is data-parallel over the world, so each rank calibrates
+        # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel
+        # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2
+        # shards the params but still feeds each rank its own data.
+        calib_dataloader = _shard_calib_dataloader(
+            calib_dataloader, args.dist_state.rank, world_size
         )
     return calib_dataloader, first_text_speech_dataset
 
@@ -317,32 +346,126 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None:
         )
 
 
+def _resolve_distribution(args: argparse.Namespace):
+    """Decide how this run distributes the model: ``(plan, hf_config)``.
+
+    ``resolve_distribution_plan`` decides, and --use_fsdp2 asks it for FSDP2 and no expert
+    parallelism. Without --use_fsdp2 a plan that FSDP2-shards the whole model is refused: the
+    FSDP2 run is the flag's, which checks what this script cannot do with it.
+    """
+    if args.use_fsdp2:
+        # Its checks read the config alone, and come first: deciding builds the model on the meta
+        # device, which some checkpoints they refuse (VILA, multimodal, pack-quantized) cannot do.
+        validate_fsdp2_supported(
+            args,
+            AutoConfig.from_pretrained(
+                args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
+            ),
+        )
+    plan, hf_config = resolve_distribution_plan(
+        args.local_checkpoint_path,
+        trust_remote_code=args.trust_remote_code,
+        ep_size=1 if args.use_fsdp2 else None,
+        dense="fsdp" if args.use_fsdp2 else None,
+    )
+    if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2:
+        raise RuntimeError(
+            f"{plan!r} FSDP2-shards the model, which this script does under --use_fsdp2. Pass it, "
+            "or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU; an MoE can instead shard its "
+            "experts (MODELOPT_EP_DEGREE above 1)."
+        )
+    if plan.world_size > 1 and _exports_tensorrt_llm_checkpoint(args, hf_model_type(hf_config)):
+        # That exporter reads the process count as Megatron's training tensor parallelism. Rank 0
+        # alone waits in its barrier for ranks that never arrive, and every rank would merge
+        # whole replicas as if they were shards. Refused before the load, like the layout above.
+        raise NotImplementedError(
+            f"{plan!r}: the TensorRT-LLM checkpoint export (encoder-decoder models, "
+            "--sparsity_fmt, int8_smoothquant) runs in a single process; launch hf_ptq without "
+            "torchrun for it."
+        )
+    if (plan.is_expert_parallel or plan.is_fsdp_dense) and args.vllm_fakequant_export:
+        # That exporter writes the whole model with save_pretrained from each process it runs in,
+        # and every rank of a sharded model exports: each would write its own part into the same
+        # directory. Rank 0 alone writes a replicated model; --use_fsdp2 refuses it while parsing.
+        raise NotImplementedError(
+            f"{plan!r}: --vllm_fakequant_export writes the whole model from one process, which "
+            "holds only its own part of a sharded model; launch hf_ptq without torchrun for it."
+        )
+    if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
+        # Its search scores each rank's own calibration shard and reduces no score across ranks, so
+        # the ranks can choose different formats per layer: under expert parallelism, a checkpoint
+        # of mixed formats. --use_fsdp2 keeps main's warning for it.
+        raise NotImplementedError(
+            f"{plan!r}: AutoQuantize scores each rank's own calibration data and reduces no score "
+            "across ranks, so ranks could choose different formats per layer; launch hf_ptq "
+            "without torchrun for an AutoQuantize recipe."
+        )
+    if (
+        plan.world_size > 1
+        and args.recipe is not None
+        and layerwise_export_requested(recipe_layerwise_blocks(load_recipe(args.recipe)))
+    ):
+        # Every rank would write the same per-layer shard files. The exporter refuses that too, but
+        # only once calibration binds it, after the load.
+        raise NotImplementedError(
+            f"{plan!r}: the recipe's layerwise export writes its checkpoint shards from one "
+            "process; launch hf_ptq without torchrun for it."
+        )
+    if plan.world_size > 1 and not args.use_fsdp2:
+        # These configure loading one model in a single process: compressed as it loads, or
+        # spread over the node's GPUs, CPU and disk. An expert-parallel run never reaches those
+        # loaders, and a replicated one pins each rank's copy to the rank's GPU, which none of them
+        # works with. Refused before the load rather than ignored. --max_cpu_memory_gb and
+        # --max_gpu_memory_gb come only with --offload_folder.
+        for flag, given in (
+            ("--low_memory_mode", args.low_memory_mode),
+            ("--use_seq_device_map", args.use_seq_device_map),
+            ("--offload_folder", args.offload_folder is not None),
+        ):
+            if given:
+                raise NotImplementedError(
+                    f"{plan!r}: {flag} configures loading one model in a single process, but each "
+                    "rank of this run loads its own share or copy of it onto its own GPU. Drop it, "
+                    "or launch hf_ptq without torchrun to use it."
+                )
+    return plan, hf_config
+
+
+def _exports_tensorrt_llm_checkpoint(args: argparse.Namespace, model_type: str | None) -> bool:
+    """Whether the run exports through the deprecated TensorRT-LLM checkpoint exporter."""
+    return (
+        is_trtllm_enc_dec_export(model_type)
+        or args.sparsity_fmt != "dense"
+        or "int8_smoothquant" in args.qformat
+    )
+
+
 def load_model(args: argparse.Namespace):
     # If low memory mode is enabled, we compress the model while loading the HF checkpoint.
     calibration_only = False
-    if args.use_fsdp2:
-        hf_config = AutoConfig.from_pretrained(
-            args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
+    plan, hf_config = _resolve_distribution(args)
+    if plan.is_expert_parallel:
+        full_model = load_pretrained(
+            args.local_checkpoint_path,
+            distribution=plan,
+            hf_config=hf_config,
+            trust_remote_code=args.trust_remote_code,
+            attn_implementation=args.attn_implementation,
         )
-        validate_fsdp2_supported(args, hf_config)
-        full_model = parallel_load_and_prepare_fsdp2(
+    elif plan.is_fsdp_dense:
+        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,
@@ -366,9 +489,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
@@ -376,6 +496,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
@@ -387,7 +508,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
@@ -554,7 +675,31 @@ def mono_quantize(
             # Image calibration batches contain multimodal kwargs (for example pixel_values).
             # They must be consumed by the complete VLM even when only a nested component is the
             # quantization target; the full forward still exercises that component's quantizers.
-            if args.calib_with_images:
+            if distribution_plan(full_model).is_expert_parallel:
+                # An expert-parallel run has to calibrate through distributed_forward_loop, not the
+                # plain one: it calls _reinstall_a2a_forwards first, which re-plants the
+                # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on
+                # transformers >= 5.16 (where the styles install a forward rather than a pre-hook).
+                # Without it the experts forward runs with process_group=None and raises
+                # "torch_a2a experts forward requires the EP `process_group`" at calibration step 0.
+                # On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op --
+                # which is why this is invisible on older transformers and fatal on newer. Other
+                # distributed runs (FSDP2) keep the plain loop and its split-on-OOM retry. Image
+                # batches go through the VLM's own forward, as they do without expert parallelism:
+                # Nemotron-H Omni, for one, calibrates on images by default. Text batches go through
+                # the language model, but the loop re-plants over the model it is handed, so that is
+                # the full model: transformers keeps the expert-parallel plan on the model
+                # from_pretrained built, and a VLM's language model names no experts in its own.
+                calibrate_loop = distributed_forward_loop(
+                    full_model,
+                    calib_dataloader,
+                    device=args.dist_state.device,
+                    ep_mesh=expert_parallel_mesh(full_model),
+                    forward=create_vlm_calibration_forward(full_model)
+                    if args.calib_with_images
+                    else lambda batch: language_model(**batch),
+                )
+            elif args.calib_with_images:
                 calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader)
             else:
                 calibrate_loop = create_forward_loop(
@@ -573,6 +718,12 @@ def mono_quantize(
             )
         else:
             language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop)
+        if distribution_plan(full_model).is_expert_parallel:
+            # The calibration loop re-plants the a2a forwards quantizer insertion drops (on
+            # transformers >= 5.16; a no-op before), but a config without static activation scales
+            # calibrates without it, and algorithm None (the MX formats) never calls it; the
+            # preview and the export forward through them.
+            _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model))
 
         # For VL models, update full_model to use the quantized language model
         if is_nemotron_vl_model and language_model is not full_model:
@@ -657,6 +808,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():
@@ -665,6 +821,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}")
@@ -672,17 +835,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":
@@ -801,11 +960,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
 
@@ -858,24 +1019,37 @@ 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 distribution_plan(full_model).is_expert_parallel:
+                # The ranks decode in lockstep: a rank that skipped the preview would leave the
+                # others blocked in the next step's collective, so the run stops. --use_fsdp2 keeps
+                # main's skip, which still exports after a failure every rank meets at the same step;
+                # a failure on one rank alone, such as an out-of-memory error, leaves the other ranks
+                # blocked in generate's collectives until the collective timeout.
+                raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e
             warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
     elif is_nemotron_vl_model and tokenizer is not None:
-        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."
@@ -962,9 +1136,7 @@ def quantize_main(
     is_layerwise = any(cfg.get("enable", False) for cfg in layerwise_cfgs)
 
     # The value is a placeholder, replaced with --export_path below; presence is the switch.
-    args.layerwise_export = any(
-        cfg.get("export_dir") is not None and cfg.get("enable", False) for cfg in layerwise_cfgs
-    )
+    args.layerwise_export = layerwise_export_requested(layerwise_cfgs)
     if not args.layerwise_export and any(
         cfg.get("export_dir") is not None for cfg in layerwise_cfgs
     ):
@@ -1513,27 +1685,48 @@ def main(args: argparse.Namespace):
     random.seed(RAND_SEED)
     np.random.seed(RAND_SEED)
 
-    setup_distributed_args(args)
-
-    try:
-        # Entered inside the try: opening the run is fatal by design, and skipping
-        # cleanup_distributed would leave the other ranks blocked on the first collective
-        # until the NCCL timeout.
-        with mlflow_run(args):
-            # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole
-            # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local
-            # path.
-            args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(
-                args.pyt_ckpt_path
-            )
-
-            # launch a memory monitor to read the currently used GPU memory.
+    # A distributed launch owns its process group for the whole run, and the run takes this
+    # process's rank, size and device from it. Entered here rather than inside load_model so
+    # cleanup still runs if loading raises -- otherwise the surviving ranks block on the first
+    # collective until the NCCL timeout.
+    with distributed_run(timeout=timedelta(hours=3), device=args.device) as dist_state:
+        args.dist_state = dist_state
+        _main_inner(args)
+
+
+def _main_inner(args: argparse.Namespace):
+    with mlflow_run(args):
+        # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole
+        # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local
+        # path.
+        args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path)
+
+        # launch a memory monitor to read the currently used GPU memory. It reads every GPU on the
+        # node, so a distributed run starts one per node, on its local rank 0, rather than one per
+        # rank.
+        if args.dist_state.world_size == 1 or local_rank() == 0:
             launch_memory_monitor()
 
-            # Force eager execution for all model types.
-            torch.compiler.set_stance("force_eager")
+        # Force eager execution for all model types.
+        torch.compiler.set_stance("force_eager")
 
-            (
+        (
+            full_model,
+            language_model,
+            model_type,
+            calibration_only,
+            processor,
+            tokenizer,
+            device,
+        ) = load_model(args)
+
+        if args.sparsity_fmt != "dense":
+            # Sparse
+            sparsity_main(args, full_model, tokenizer, device)
+        else:
+            # Quantize
+            quantize_main(
+                args,
                 full_model,
                 language_model,
                 model_type,
@@ -1541,25 +1734,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/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 b08d44172e8..591fe0c07b2 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -831,8 +831,9 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int:
     rebinds ``module.forward`` on the same instance, dropping the wrapper -- after which
     ``torch_a2a_experts_forward`` receives ``process_group=None`` and raises.
 
-    Called at the top of the calibration loop: after conversion (so it wraps the QUANTIZED forward)
-    and before the first forward. Only a forward quantizer insertion rebound is planted again: a
+    Called at the top of the calibration loop, and by ``hf_ptq`` right after quantizing (a recipe
+    with no calibration runs no loop): after conversion (so it wraps the QUANTIZED forward) and
+    before the first forward. Only a forward quantizer insertion rebound is planted again: a
     wrapper still in place -- transformers' from the load, on a module quantization left alone, or
     one planted here on an earlier pass -- is kept, since wrapping it again would run its transforms
     twice. ``ep_mesh`` is the mesh the routed experts are sharded over; without one there is nothing
@@ -896,14 +897,16 @@ def distributed_forward_loop(
     dataloader: DataLoader,
     device: torch.device | str | None = None,
     ep_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None,
+    forward: Callable[[dict], object] | None = None,
 ) -> Callable[[torch.nn.Module], None]:
     """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model.
 
     Moves each batch onto ``device``, when given, and runs the captured ``model`` forward (the EP
-    all-reduce / all-to-all collectives fire inside that forward). A model whose plan
-    (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its experts are sharded
-    over, for the all-to-all forwards that quantizer insertion drops to be re-installed before the
-    first forward; any other model passes none. ``mtq.quantize`` hands the
+    all-reduce / all-to-all collectives fire inside that forward), or ``forward(batch)`` when given,
+    for batches that need their own call into ``model``, such as a VLM's image-text batches. A
+    model whose plan (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its
+    experts are sharded over, for the all-to-all forwards that quantizer insertion drops to be
+    re-installed before the first forward; any other model passes none. ``mtq.quantize`` hands the
     ``forward_loop`` callback an inner module that this closure ignores, calling the captured
     top-level ``model`` instead so the full distributed forward (and its collectives) runs.
 
@@ -941,7 +944,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..0b1fbbb1fd1 100644
--- a/modelopt/torch/utils/memory_monitor.py
+++ b/modelopt/torch/utils/memory_monitor.py
@@ -140,8 +140,15 @@ def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor | N
     Returns:
         GPUMemoryMonitor: The monitor instance that was launched
     """
+    # Probe that NVML can report memory on this host before starting the monitor. NVML must be
+    # initialized for the probe itself; without it every call raises NVMLError_Uninitialized, and
+    # the monitor would never start anywhere.
     try:
-        nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0))
+        nvmlInit()
+        try:
+            nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0))
+        finally:
+            nvmlShutdown()
     except Exception as e:
         print(f"Failed to get GPU memory info: {e}. Stopping GPU memory monitor.")
         return None
diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py
index 7e294c0ed0b..d39d2667fd4 100644
--- a/tests/examples/hf_ptq/test_example_utils.py
+++ b/tests/examples/hf_ptq/test_example_utils.py
@@ -21,16 +21,16 @@
 """
 
 import warnings
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
 from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 import torch
 from _test_utils.examples.hf_ptq_example_utils import example_utils
-from _test_utils.torch.transformers_models import get_tiny_nemotron_h
+from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_nemotron_h
 from safetensors.torch import save_file
-from transformers import AutoModelForCausalLM, LlamaConfig
+from transformers import AutoModelForCausalLM, GenerationMixin, LlamaConfig
 
 from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
 
@@ -138,6 +138,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(
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 608a03e1beb..0250444c4ad 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -37,6 +37,7 @@
 from modelopt.torch.quantization import tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.utils import mlflow as mlflow_lib
+from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
 from modelopt.torch.utils.mlflow import describe_run, run_tags
 
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
@@ -408,7 +409,8 @@ def distributed_load(checkpoint, *_args, **_kwargs):
         assert checkpoint == str(tmp_path)
         raise RuntimeError("distributed loader reached")
 
-    monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load)
+    # hf_ptq loads an FSDP2 plan through load_pretrained.
+    monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load)
 
     with pytest.raises(RuntimeError, match="distributed loader reached"):
         hf_ptq.load_model(args)
@@ -435,7 +437,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
@@ -445,31 +447,6 @@ def distributed_load(*args, **kwargs):
         hf_ptq.load_model(args)
 
 
-def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path):
-    """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an
-    FSDP2-sharded model without one is refused rather than taken for a single process's."""
-    from transformers import Qwen3Config
-
-    from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
-
-    Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path)
-    monkeypatch.setenv("RANK", "0")
-    hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2")
-    args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=2)
-    args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path)  # as main() sets them
-    model = torch.nn.Linear(1, 1)
-    monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *args, **kwargs: model)
-
-    def loaded(model):
-        raise RuntimeError("loaded")
-
-    monkeypatch.setattr(hf_ptq, "prepare_model_for_calibration", loaded)
-
-    with pytest.raises(RuntimeError, match="loaded"):
-        hf_ptq.load_model(args)
-    assert distribution_plan(model) == DistributionPlan(2, dense="fsdp")
-
-
 def test_mlflow_flag_defaults_the_experiment_name(monkeypatch):
     monkeypatch.setattr(getpass, "getuser", lambda: "tester")
     hf_ptq, args = _parse_hf_ptq_args(
@@ -980,9 +957,12 @@ def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch):
     NVBug 6752977."""
     hf_ptq = _import_hf_ptq(monkeypatch)
 
-    full_model = SimpleNamespace(
-        generate=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))
-    )
+    class _FailingModel(torch.nn.Module):
+        def generate(self, *args, **kwargs):
+            raise RuntimeError("boom")
+
+    full_model = _FailingModel()
+    set_distribution_plan(full_model, DistributionPlan())
     export_calls = []
     monkeypatch.setattr(
         hf_ptq,
@@ -1046,6 +1026,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(
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..abfa49b6471
--- /dev/null
+++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
@@ -0,0 +1,1195 @@
+# 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 sys
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+import transformers
+from _test_utils.examples.hf_ptq_example_utils import example_utils
+from _test_utils.torch.transformers_models import get_tiny_tokenizer
+from packaging.version import Version
+from PIL import Image
+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,
+        sparsity_fmt="dense",
+        qformat="fp8",
+        low_memory_mode=False,
+        use_seq_device_map=False,
+        offload_folder=None,
+        vllm_fakequant_export=False,
+    )
+    vars(args).update(overrides)
+    return args
+
+
+def _serve(monkeypatch, hf_ptq, config):
+    """Answer every config read with ``config`` rather than a checkpoint on disk."""
+    monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *args, **kwargs: config)
+    return config
+
+
+# DeepSeek-V3 is expert-parallel only from the transformers release the expert-parallel load needs;
+# below it, a multi-rank run of it is refused before the route the test exercises.
+needs_ep_transformers = pytest.mark.skipif(
+    Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS),
+    reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}",
+)
+
+
+def _deepseek_v3(n_routed_experts):
+    """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable.
+
+    Its one layer is an MoE layer: DeepSeek-V3's first layers are dense by default.
+    """
+    config = DeepseekV3Config(
+        architectures=["DeepseekV3ForCausalLM"],
+        n_routed_experts=n_routed_experts,
+        first_k_dense_replace=0,
+        moe_intermediate_size=8,
+        n_group=1,
+        topk_group=1,
+        num_experts_per_tok=1,
+        **_TINY,
+    )
+    config.base_model_ep_plan = None  # holds even if a later transformers ships a plan
+    return config
+
+
+# -------------------------------------------------------------------------------- the decision
+
+
+def test_use_fsdp2_plans_fsdp2_without_expert_parallelism(hf_ptq, monkeypatch, torchrun):
+    """--use_fsdp2 names its load, so even an explicit EP degree for an EP-capable MoE must not
+    divert the run to the expert-parallel one."""
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "2")
+    _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8))
+    args = _args(torchrun(world_size=2), use_fsdp2=True)
+
+    assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(2, dense="fsdp")
+
+
+def test_use_fsdp2_refuses_a_checkpoint_before_building_it_to_decide(hf_ptq, monkeypatch, torchrun):
+    """--use_fsdp2's checks read the config alone. They come before the decision, which builds the
+    model on the meta device: some checkpoints they refuse cannot be built, and that build's error
+    would name neither --use_fsdp2 nor the way out."""
+    _serve(
+        monkeypatch,
+        hf_ptq,
+        LlamaConfig(
+            architectures=["LlamaForCausalLM"], quantization_config={"quant_method": "fp8"}, **_TINY
+        ),
+    )
+    monkeypatch.setattr(
+        model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built the model")
+    )
+
+    with pytest.raises(NotImplementedError, match="--use_fsdp2 does not support"):
+        hf_ptq._resolve_distribution(_args(torchrun(world_size=2), use_fsdp2=True))
+
+
+def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monkeypatch, torchrun):
+    """Without expert parallelism a plan that shards the dense part is an FSDP2 run, which this
+    script runs only under --use_fsdp2, the flag its FSDP2 checks hang off; so it must fail and point
+    there rather than load the model whole on every rank."""
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY))
+
+    with pytest.raises(RuntimeError, match="--use_fsdp2"):
+        hf_ptq._resolve_distribution(_args(torchrun(world_size=2)))
+
+
+@needs_ep_transformers
+@pytest.mark.parametrize("ep_degree", [2, 4])
+def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused(
+    hf_ptq, monkeypatch, torchrun, ep_degree
+):
+    """Below the world size each expert-parallel group is a replica, and FSDP2 shards the dense part
+    within each: a layout the export cannot write, so it must fail before the load and the
+    calibration rather than at the end. Expert parallelism over the whole world can still shard it."""
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", str(ep_degree))
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8))
+    args = _args(torchrun(world_size=4))
+
+    if ep_degree == 4:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(4, ep_size=4, dense="fsdp")
+    else:
+        with pytest.raises(NotImplementedError, match=r"MODELOPT_EP_DEGREE.*MODELOPT_FSDP_SIZE=1"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    ("world_size", "options"),
+    [
+        (2, {"qformat": "int8_smoothquant"}),
+        (2, {"sparsity_fmt": "sparsegpt"}),
+        (1, {"sparsity_fmt": "sparsegpt"}),
+    ],
+    ids=["int8_smoothquant", "sparsity", "single_process"],
+)
+def test_a_tensorrt_llm_checkpoint_export_needs_a_single_process(
+    hf_ptq, monkeypatch, torchrun, world_size, options
+):
+    """That exporter reads the process count as Megatron's training tensor parallelism: launched
+    with torchrun, it hangs in a barrier on rank 0 alone, and on every rank it merges replicas as
+    shards. So it fails before the load and the calibration; a single process exports as before."""
+    _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY))
+    args = _args(torchrun(world_size=world_size), **options)
+
+    if world_size == 1:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+    else:
+        with pytest.raises(NotImplementedError, match="single process"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    ("world_size", "env", "use_fsdp2"),
+    [
+        pytest.param(2, {"MODELOPT_EP_DEGREE": "2"}, False, marks=needs_ep_transformers),
+        pytest.param(2, {}, False, marks=needs_ep_transformers),
+        (2, {}, True),
+        (1, {}, False),
+    ],
+    ids=["expert_parallel", "replicated", "use_fsdp2", "single_process"],
+)
+def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds(
+    hf_ptq, monkeypatch, torchrun, world_size, env, use_fsdp2
+):
+    """AutoQuantize scores each rank's own calibration shard and reduces no score across ranks, so
+    the ranks could choose different formats per layer. An expert-parallel or replicated run refuses
+    the recipe before the load; --use_fsdp2 keeps main's warning, and a single process runs it."""
+    for name, value in env.items():
+        monkeypatch.setenv(name, value)
+    _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8))
+    args = _args(
+        torchrun(world_size=world_size),
+        use_fsdp2=use_fsdp2,
+        recipe="general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits",
+    )
+
+    if world_size == 1 or use_fsdp2:
+        hf_ptq._resolve_distribution(args)
+    else:
+        with pytest.raises(NotImplementedError, match="AutoQuantize"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize("world_size", [2, 1], ids=["torchrun", "single_process"])
+def test_a_layerwise_export_recipe_is_refused_before_a_distributed_load(
+    hf_ptq, monkeypatch, torchrun, world_size
+):
+    """Every rank would write the same per-layer shard files, which the exporter refuses only once
+    calibration binds it, after the load; a launch of several ranks refuses the recipe before it.
+    A single process exports as before."""
+    _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY))
+    args = _args(
+        torchrun(world_size=world_size),
+        recipe="general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export",
+    )
+
+    if world_size == 1:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+    else:
+        with pytest.raises(NotImplementedError, match="layerwise export"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    "option",
+    [{"low_memory_mode": True}, {"use_seq_device_map": True}, {"offload_folder": "/offload"}],
+    ids=["low_memory_mode", "use_seq_device_map", "offload_folder"],
+)
+@pytest.mark.parametrize(
+    ("make_config", "world_size"),
+    [
+        pytest.param(lambda: _deepseek_v3(n_routed_experts=8), 2, marks=needs_ep_transformers),
+        (lambda: LlamaConfig(**_TINY), 2),
+        (lambda: LlamaConfig(**_TINY), 1),
+    ],
+    ids=["expert_parallel", "replicated", "single_process"],
+)
+def test_the_single_process_load_options_are_refused_on_the_routes_this_script_adds(
+    hf_ptq, monkeypatch, torchrun, make_config, world_size, option
+):
+    """Each rank of an expert-parallel or replicated run loads its own share or copy of the model
+    onto its own GPU, where these options of a single process's load would be ignored or would break
+    the load. Both routes refuse them before the load; a single process loads by them as before."""
+    _serve(monkeypatch, hf_ptq, make_config())
+    args = _args(torchrun(world_size=world_size), **option)
+
+    if world_size == 1:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+    else:
+        (name,) = option
+        with pytest.raises(NotImplementedError, match=f"--{name}"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    ("make_config", "world_size", "refused"),
+    [
+        pytest.param(
+            lambda: _deepseek_v3(n_routed_experts=8), 2, True, marks=needs_ep_transformers
+        ),
+        (lambda: LlamaConfig(**_TINY), 2, False),
+        (lambda: LlamaConfig(**_TINY), 1, False),
+    ],
+    ids=["expert_parallel", "replicated", "single_process"],
+)
+def test_the_vllm_fakequant_export_is_refused_for_a_sharded_model(
+    hf_ptq, monkeypatch, torchrun, make_config, world_size, refused
+):
+    """That exporter writes the whole model with save_pretrained from each process it runs in, so
+    every rank of an expert-parallel run would write its own experts into the same directory. Such a
+    run refuses it before the load; rank 0 alone writes a replicated model, and a single process
+    exports as before."""
+    _serve(monkeypatch, hf_ptq, make_config())
+    args = _args(torchrun(world_size=world_size), vllm_fakequant_export=True)
+
+    if refused:
+        with pytest.raises(NotImplementedError, match="--vllm_fakequant_export"):
+            hf_ptq._resolve_distribution(args)
+    else:
+        hf_ptq._resolve_distribution(args)
+
+
+# ---------------------------------------------------------------------------------- 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))]
+
+
+# ------------------------------------------------------------------------- per-rank placement
+
+
+@pytest.mark.parametrize(
+    ("device", "world_size", "device_map"),
+    [
+        ("cuda:1", 4, {"": "cuda:1"}),
+        (torch.device("cuda:1"), 4, {"": "cuda:1"}),
+        ("cuda:1", 1, "auto"),
+    ],
+    ids=["str", "torch_device", "single_process"],
+)
+def test_a_rank_device_pins_the_whole_model_to_it(
+    monkeypatch, tmp_path, device, world_size, device_map
+):
+    """Under torchrun each rank asks for its own cuda:<local_rank>; "auto" would let every rank's
+    accelerate pick the same idle GPU. A single process spreads the model over the GPUs, whatever
+    device it names. Only a spread model is held to a memory limit sized over the node's GPUs: the
+    pinned map would ignore one."""
+    LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path)
+    # No GPU on the node holds the model, so sizing it over the node limits its memory.
+    monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 0, "cpu": 1 << 40})
+    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, device_map == "auto")]
+
+
+@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
+        )
+
+
+# ------------------------------------------------------------------ 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="cpu")
+
+    @contextlib.contextmanager
+    def distributed_run(timeout=None, device="cuda"):
+        calls.append((timeout, device))
+        events.append("enter")
+        try:
+            yield state
+        finally:
+            events.append("exit")
+
+    def main_inner(args):
+        events.append("run")
+        seen.append(args.dist_state)
+        raise _RunFailedError
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(hf_ptq, "distributed_run", distributed_run)
+    monkeypatch.setattr(hf_ptq, "_main_inner", main_inner)
+
+    with pytest.raises(_RunFailedError):
+        hf_ptq.main(SimpleNamespace(use_fsdp2=False, device="cpu"))
+
+    assert events == ["enter", "run", "exit"]
+    timeout, device = calls[0]
+    assert timeout > timedelta(minutes=30)
+    assert device == "cpu"
+    assert len(seen) == 1
+    assert seen[0] is state
+
+
+# --------------------------------------------------------------------------------- calibration
+
+
+def test_distributed_calibration_gives_each_rank_its_own_shard(hf_ptq, monkeypatch):
+    """Every distributed path calibrates data-parallel: the ranks get equal, disjoint shards that
+    together cover the calibration set. A single process keeps the loader as built."""
+    loader = DataLoader(list(range(8)), batch_size=1)
+    monkeypatch.setattr(hf_ptq, "get_dataset_dataloader", lambda **kwargs: loader)
+    args = SimpleNamespace(
+        specdec_offline_dataset=None,
+        calib_with_images=False,
+        dataset=["cnn_dailymail"],
+        batch_size=1,
+        calib_size=[8],
+        calib_seq=16,
+    )
+    tokenizer = get_tiny_tokenizer()
+
+    def calib_loader(world_size, rank=0):
+        args.dist_state = _dist_state(world_size, rank)
+        model = torch.nn.Linear(2, 2)
+        set_distribution_plan(model, DistributionPlan(world_size))
+        built, _ = hf_ptq.make_calib_dataloader(
+            args, model, None, tokenizer, torch.device("cpu"), "llama"
+        )
+        return built
+
+    shards = [[int(b) for b in calib_loader(world_size=2, rank=r)] for r in range(2)]
+
+    assert len(shards[0]) == len(shards[1]) == 4
+    assert sorted(shards[0] + shards[1]) == list(range(8))
+    assert calib_loader(world_size=1) is loader
+
+
+@pytest.mark.parametrize("images", [False, True], ids=["text", "images"])
+@pytest.mark.parametrize(
+    ("plan", "loop"),
+    [(DistributionPlan(2, ep_size=2), "distributed"), (DistributionPlan(2, dense="fsdp"), "plain")],
+    ids=["expert_parallel", "fsdp2"],
+)
+def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_loop(
+    hf_ptq, monkeypatch, plan, loop, images
+):
+    """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion
+    clobbers, so an expert-parallel run must calibrate through it, over the mesh load_and_prepare_ep
+    loaded the experts over, image batches included: they run through the whole VLM's own forward
+    there (Nemotron-H Omni calibrates on images by default). The loop re-plants over the model it is
+    handed, the whole model, and text batches run through the language model. Without expert
+    parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry,
+    as on main."""
+    ep_mesh, vlm_forward = object(), object()
+    full_model, language_model = torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)
+    set_distribution_plan(full_model, plan)
+    chosen, quantized = [], []
+    monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False)
+    monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: True)
+    monkeypatch.setattr(hf_ptq, "expert_parallel_mesh", lambda full_model: ep_mesh)
+    monkeypatch.setattr(
+        hf_ptq,
+        "distributed_forward_loop",
+        lambda model, loader, **k: chosen.append((model, k["ep_mesh"], k["forward"])) or object(),
+    )
+    monkeypatch.setattr(hf_ptq, "create_vlm_calibration_forward", lambda model: vlm_forward)
+    monkeypatch.setattr(
+        hf_ptq,
+        "create_vlm_calibration_loop",
+        lambda model, loader: chosen.append("vlm") or object(),
+    )
+    monkeypatch.setattr(
+        hf_ptq, "create_forward_loop", lambda **k: chosen.append("plain") or object()
+    )
+    monkeypatch.setattr(
+        hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: quantized.append(forward_loop) or m
+    )
+    args = SimpleNamespace(
+        qformat="fp8",
+        calib_with_images=images,
+        dist_state=_dist_state(world_size=2),
+        layerwise_export=False,
+        specdec_offline_dataset=None,
+    )
+
+    hf_ptq.mono_quantize(
+        args,
+        {"algorithm": "max"},
+        full_model,
+        language_model,
+        model_type="qwen3_moe",
+        calibration_only=False,
+        calib_dataloader=object(),
+        is_nemotron_vl_model=False,
+    )
+
+    assert len(quantized) == 1
+    if loop == "plain":
+        assert chosen == ["vlm" if images else "plain"]
+    elif images:
+        assert chosen == [(full_model, ep_mesh, vlm_forward)]
+    else:
+        ((model, mesh, forward),) = chosen
+        assert (model, mesh) == (full_model, ep_mesh)
+        x = torch.ones(1, 2)
+        assert torch.equal(forward({"input": x}), language_model(x))
+
+
+@pytest.mark.parametrize(
+    "loop_handed", [False, True], ids=["calibration_free", "loop_never_called"]
+)
+@pytest.mark.parametrize(
+    ("plan", "replanted"),
+    [(DistributionPlan(2, ep_size=2), True), (DistributionPlan(2, dense="fsdp"), False)],
+    ids=["expert_parallel", "fsdp2"],
+)
+def test_an_expert_parallel_run_replants_the_all_to_all_forwards_after_quantizing(
+    hf_ptq, monkeypatch, plan, replanted, loop_handed
+):
+    """The calibration loop re-plants the all-to-all forwards that quantizer insertion drops only
+    when the calibration runs it: a config without static activation scales calibrates without the
+    loop, and algorithm None (the MX formats) never calls the loop it is handed. The preview and the
+    export forward through the experts either way, so an expert-parallel run re-plants them once
+    quantized; without expert parallelism there is nothing to re-plant."""
+    ep_mesh, full_model = object(), torch.nn.Linear(2, 2)
+    set_distribution_plan(full_model, plan)
+    events = []
+    monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False)
+    monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: loop_handed)
+    monkeypatch.setattr(hf_ptq, "expert_parallel_mesh", lambda model: ep_mesh)
+    monkeypatch.setattr(hf_ptq, "distributed_forward_loop", lambda *args, **kwargs: object())
+    monkeypatch.setattr(hf_ptq, "create_forward_loop", lambda **kwargs: object())
+    monkeypatch.setattr(
+        hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: events.append("quantize") or m
+    )
+    monkeypatch.setattr(
+        hf_ptq,
+        "_reinstall_a2a_forwards",
+        lambda model, mesh: events.append(("replant", model, mesh)),
+    )
+    args = SimpleNamespace(
+        qformat="mxfp8",
+        calib_with_images=False,
+        dist_state=_dist_state(world_size=2),
+        layerwise_export=False,
+        specdec_offline_dataset=None,
+    )
+
+    hf_ptq.mono_quantize(
+        args,
+        {"algorithm": None},
+        full_model,
+        full_model,
+        model_type="qwen3_moe",
+        calibration_only=False,
+        calib_dataloader=object(),
+        is_nemotron_vl_model=False,
+    )
+
+    assert events == ["quantize"] + ([("replant", full_model, ep_mesh)] if replanted else [])
+
+
+@pytest.fixture
+def gloo_world_of_one(tmp_path):
+    """A one-rank gloo process group, for the all-to-all forwards to run on."""
+    torch.distributed.init_process_group(
+        "gloo", init_method=(tmp_path / "store").as_uri(), rank=0, world_size=1
+    )
+    yield torch.distributed.group.WORLD
+    torch.distributed.destroy_process_group()
+
+
+class _EpMesh:
+    """The expert-parallel submesh of a two-rank plan, over this process's one-rank group."""
+
+    def size(self):
+        return 2
+
+    def get_group(self):
+        return torch.distributed.group.WORLD
+
+
+def test_a_vlm_calibrated_on_text_gets_its_experts_all_to_all_forwards_back(
+    hf_ptq, gloo_world_of_one
+):
+    """Text batches run through a VLM's language model, but transformers keeps the expert-parallel
+    plan on the model from_pretrained built, and the language model's own plan names no experts.
+    Re-planted over the language model, the experts' all-to-all forward that quantizer insertion
+    drops (transformers 5.16 and later) stayed dropped, and the first batch raised for the missing
+    EP process group."""
+    pytest.importorskip("transformers", minversion="5.16")
+    from transformers import Qwen3VLMoeConfig, Qwen3VLMoeForConditionalGeneration
+    from transformers.distributed.configuration_utils import DistributedConfig
+    from transformers.distributed.tensor_parallel import ALL_PARALLEL_STYLES, _get_parameter_tp_plan
+
+    from modelopt.torch.utils.moe import register_torch_a2a
+
+    register_torch_a2a()
+    config = Qwen3VLMoeConfig(
+        text_config={
+            **_TINY,
+            "head_dim": 8,
+            "num_experts": 4,
+            "num_experts_per_tok": 2,
+            "moe_intermediate_size": 8,
+            "decoder_sparse_step": 1,
+        },
+        vision_config={
+            "depth": 1,
+            "hidden_size": 16,
+            "intermediate_size": 16,
+            "num_heads": 2,
+            "out_hidden_size": 16,
+            "patch_size": 2,
+            "spatial_merge_size": 1,
+            "temporal_patch_size": 1,
+            "deepstack_visual_indexes": [],
+        },
+    )
+    # The model as load_and_prepare_ep leaves it: its plan's all-to-all styles on the language
+    # model's config and expert parallelism on the top one, the forwards from_pretrained installs
+    # for them, the experts dispatching through the all-to-all, and their seeded parallel state.
+    config.text_config.base_model_ep_plan = {
+        "layers.*.mlp.gate": "torch_a2a_router",
+        "layers.*.mlp.experts": "torch_a2a_experts",
+    }
+    config.distributed_config = DistributedConfig(enable_expert_parallel=True)
+    model = Qwen3VLMoeForConditionalGeneration(config).eval()
+    mesh = _EpMesh()
+    for name, module in model.named_modules():
+        style = _get_parameter_tp_plan(parameter_name=name, tp_plan=model.tp_plan, is_weight=False)
+        if style in ("torch_a2a_router", "torch_a2a_experts"):
+            ALL_PARALLEL_STYLES[style].install_forward(module, mesh)
+    model.config.text_config._experts_implementation = "torch_a2a"
+    model._device_mesh = mesh
+    set_distribution_plan(model, DistributionPlan(2, ep_size=2))
+    model_load_utils._seed_parallel_states(
+        model,
+        [name for name, _ in model.named_modules() if name.endswith("mlp.experts")],
+        dense_state=ParallelState(data_parallel_group=-1),
+        expert_state=ParallelState(
+            data_parallel_group=-1, expert_model_parallel_group=gloo_world_of_one
+        ),
+    )
+    language_model = model.model.language_model
+    args = SimpleNamespace(
+        qformat="int8",
+        calib_with_images=False,
+        dist_state=dist.DistributedState(rank=0, world_size=2, device=torch.device("cpu")),
+        layerwise_export=False,
+        specdec_offline_dataset=None,
+    )
+    batch = {"input_ids": torch.tensor([[1, 2, 3, 4]]), "attention_mask": torch.ones(1, 4)}
+
+    hf_ptq.mono_quantize(
+        args,
+        hf_ptq.mtq.INT8_DEFAULT_CFG,
+        model,
+        language_model,
+        model_type="qwen3_vl_moe",
+        calibration_only=False,
+        calib_dataloader=[batch],
+        is_nemotron_vl_model=False,
+    )
+
+    assert language_model.layers[0].mlp.experts.down_proj_input_quantizer.amax is not None
+
+
+class _PreviewModel(torch.nn.Module):
+    """A model whose ``generate`` records the kwargs each preview passes it, and how often it entered
+    the decode hook that Transformers' ``generate`` enters to switch a MoE's experts."""
+
+    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
+
+
+def test_the_offline_specdec_loader_is_split_across_ranks_too(hf_ptq):
+    """It wraps its DataLoader to move each batch onto the device; each rank still calibrates on
+    its own share, as it would from the bare loader."""
+
+    def loader():
+        return DataLoader(list(range(8)), batch_size=1, collate_fn=lambda b: {"x": torch.tensor(b)})
+
+    wrapped = hf_ptq._DeviceDataLoader(loader(), torch.device("cpu"))
+
+    shard = hf_ptq._shard_calib_dataloader(wrapped, 1, 2)
+
+    assert isinstance(shard, hf_ptq._DeviceDataLoader)
+    assert [b["x"].tolist() for b in shard] == [
+        b["x"].tolist() for b in dist.shard_dataloader(loader(), 1, 2)
+    ]
+    with pytest.raises(NotImplementedError, match="cannot be split"):
+        hf_ptq._shard_calib_dataloader(iter(loader()), 1, 2)
+
+
+@pytest.mark.parametrize(
+    ("plan", "nemotron_vl", "raises"),
+    [
+        (DistributionPlan(4), False, False),
+        (DistributionPlan(4, dense="fsdp"), False, False),
+        (DistributionPlan(4, ep_size=4), False, True),
+        (DistributionPlan(4, ep_size=4), True, True),
+    ],
+    ids=["replicated", "fsdp2", "expert_parallel", "expert_parallel_nemotron_vl"],
+)
+def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run(
+    hf_ptq, monkeypatch, plan, nemotron_vl, raises
+):
+    """A rank that skipped an expert-parallel model's preview would leave the others decoding in
+    lockstep, blocked in the next step's collective, so the failure is raised, which aborts the run,
+    with the preview's error as its cause. A model every rank holds whole keeps the best-effort
+    skip and exports, and so does --use_fsdp2, as on main."""
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("the preview failed on this rank")
+
+    class _FailingModel(torch.nn.Module):
+        generate = fail
+
+    model = _FailingModel()
+    set_distribution_plan(model, plan)
+    exported = []
+    monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: exported.append(True))
+    monkeypatch.setattr(hf_ptq, "run_nemotron_vl_preview", fail)  # as vlm_utils raises in lockstep
+    ids = torch.zeros(1, 4, dtype=torch.long)
+    args = SimpleNamespace(
+        specdec_offline_dataset=None,
+        verbose=False,
+        local_checkpoint_path=_CKPT,
+        trust_remote_code=False,
+    )
+    tokenizer = SimpleNamespace()
+    run = lambda: hf_ptq.post_quantize(  # noqa: E731
+        args, model, model, "llama", tokenizer, None, ids, None, ids, nemotron_vl, None, None
+    )
+
+    if raises:
+        with pytest.raises(RuntimeError) as stopped:
+            run()
+        assert "failed on this rank" in str(stopped.value.__cause__)
+    else:
+        with pytest.warns(UserWarning, match="sanity check failed"):
+            run()
+    assert exported == ([] if raises else [True])
+
+
+@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
+
+
+# ------------------------------------------------------------------------------------- export
+
+
+@pytest.mark.parametrize(
+    ("plan", "rank", "exports"),
+    [
+        (DistributionPlan(4), 0, True),
+        (DistributionPlan(4), 1, False),
+        (DistributionPlan(4, dense="fsdp"), 1, True),
+        (DistributionPlan(4, ep_size=4), 1, True),
+    ],
+    ids=["replicated_rank0", "replicated_rank1", "fsdp2_rank1", "expert_parallel_rank1"],
+)
+def test_only_rank_0_writes_a_replicated_model(hf_ptq, monkeypatch, plan, rank, exports):
+    """A replicated model is whole on every rank, whose exports would write one checkpoint into the
+    same directory over each other; a sharded one needs every rank for its part. The source config
+    is rank 0's either way: another rank's, written after rank 0's export, would drop the
+    quantization config."""
+    written = []
+    monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: False)
+    monkeypatch.setattr(hf_ptq, "is_multimodal_model", lambda model: True)
+    monkeypatch.setattr(hf_ptq, "save_source_config", lambda args, path: written.append("config"))
+    monkeypatch.setattr(
+        hf_ptq, "export_hf_checkpoint", lambda model, export_dir: written.append("checkpoint")
+    )
+    args = SimpleNamespace(
+        dist_state=_dist_state(plan.world_size, rank),
+        export_path="/export",
+        layerwise_export=False,
+        sparsity_fmt="dense",
+        qformat="fp8",
+        inference_tensor_parallel=1,
+        inference_pipeline_parallel=1,
+        vllm_fakequant_export=False,
+    )
+
+    model = _PreviewModel(plan)
+    hf_ptq.export_quantized(args, model, model, "qwen3_moe")
+
+    assert written == ["config"] * (rank == 0) + ["checkpoint"] * exports
+
+
+@pytest.mark.parametrize(
+    "plan", [DistributionPlan(4, dense="fsdp"), DistributionPlan(4, ep_size=4)], ids=["fsdp2", "ep"]
+)
+def test_a_sharded_speculative_decoding_model_is_refused_at_export(hf_ptq, monkeypatch, plan):
+    """The speculative-decoding exporter writes the whole checkpoint from one process and has no
+    rank gate, so every rank of a sharded model would write it over the others: it is refused on
+    every rank instead."""
+    written = []
+    monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: True)
+    monkeypatch.setattr(
+        hf_ptq, "export_speculative_decoding", lambda model, export_dir: written.append(export_dir)
+    )
+    args = SimpleNamespace(dist_state=_dist_state(plan.world_size, 1), export_path="/export")
+
+    model = _PreviewModel(plan)
+    with pytest.raises(NotImplementedError, match="speculative-decoding model sharded"):
+        hf_ptq.export_quantized(args, model, model, "qwen3_moe")
+    assert written == []
diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py
index 328dcd5e564..34b5194e1c2 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py
@@ -24,6 +24,16 @@
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
 
 
+class _Model(torch.nn.Module):
+    """A model whose ``generate`` is ``generate``: the previews read the module's distribution plan."""
+
+    def __init__(self, generate, config=None):
+        super().__init__()
+        self.generate = generate
+        if config is not None:
+            self.config = config
+
+
 @pytest.fixture
 def hf_ptq(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.syspath_prepend(str(_EXAMPLES_DIR))
@@ -43,7 +53,7 @@ def hf_ptq(monkeypatch: pytest.MonkeyPatch):
 )
 def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: bool):
     generate_calls = []
-    full_model = SimpleNamespace(generate=lambda *args, **kwargs: generate_calls.append(1) or "out")
+    full_model = _Model(lambda *args, **kwargs: generate_calls.append(1) or "out")
     args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False)
     batch = {"input_ids": torch.ones(2, 4, dtype=torch.long)}
 
@@ -58,7 +68,7 @@ def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation:
 )
 def test_pre_quantize_preview_removes_padding(hf_ptq, ids, mask):
     calls = []
-    model = SimpleNamespace(generate=lambda ids, **kwargs: calls.append((ids, kwargs)) or ids)
+    model = _Model(lambda ids, **kwargs: calls.append((ids, kwargs)) or ids)
     args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False)
     batch = {"input_ids": torch.tensor([ids]), "attention_mask": torch.tensor([mask])}
     prompt, attention_mask, _ = hf_ptq.pre_quantize(args, model, "llama", None, [batch], False)
@@ -122,9 +132,7 @@ def test_post_quantize_preview_keeps_whole_answer(
 ):
     """The preview must decode the answer, not slice a prompt-length prefix off it."""
     monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None)
-    full_model = SimpleNamespace(
-        config=config, generate=lambda *args, **kwargs: torch.tensor(generated)
-    )
+    full_model = _Model(lambda *args, **kwargs: torch.tensor(generated), config)
     tokenizer = SimpleNamespace(
         batch_decode=lambda ids, **kwargs: [" ".join(map(str, row.tolist())) for row in ids]
     )
diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
index d37d3aba3b1..a0a56fbd4b4 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
@@ -20,6 +20,8 @@
 import pytest
 import torch
 
+from modelopt.torch.utils.distributed import DistributionPlan
+
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
 
 
@@ -64,6 +66,10 @@ def test_image_calibration_model_target_follows_recipe(
     )
 
     monkeypatch.setattr(hf_ptq, "get_model", lambda *args, **kwargs: full_model)
+    # A single process's plan, decided without reading the "dummy" checkpoint.
+    monkeypatch.setattr(
+        hf_ptq, "resolve_distribution_plan", lambda *args, **kwargs: (DistributionPlan(), None)
+    )
     # A VLM's extracted language model reports its own sub-config type; hf_ptq must keep the root's.
     monkeypatch.setattr(
         hf_ptq,
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/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py
index b4a3e47b2c7..af5cd27aefc 100644
--- a/tests/unit/torch/utils/test_distributed_calibration.py
+++ b/tests/unit/torch/utils/test_distributed_calibration.py
@@ -458,6 +458,21 @@ def test_forward_loop_moves_tensors_to_the_device_it_is_given():
     assert [call["logits_to_keep"] for call in model.calls] == [1, 1]
 
 
+def test_forward_loop_runs_each_batch_through_the_forward_it_is_given():
+    """A batch that needs its own call into the model, as a VLM's image-text batch does, goes
+    through ``forward``, on the device and with the cache disabled like any other."""
+    batches = [{"input_ids": torch.zeros(2, 3)}] * 2
+    model = _Recorder()
+
+    distributed_forward_loop(
+        model, batches, device="meta", forward=lambda batch: model(**batch, own_call=True)
+    )(model)
+
+    assert [call["own_call"] for call in model.calls] == [True, True]
+    assert [call["input_ids"].device.type for call in model.calls] == ["meta"] * 2
+    assert model.use_cache_seen == [False, False]
+
+
 def test_forward_loop_disables_use_cache_only_while_calibrating():
     """KV/state caching is off for the calibration forwards (it breaks hybrid Mamba models)."""
     model = _Recorder()
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..8bb3dcab5bf
--- /dev/null
+++ b/tests/unit/torch/utils/test_memory_monitor.py
@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for launch_memory_monitor's NVML availability probe, against a fake NVML."""
+
+import atexit
+from types import SimpleNamespace
+
+import pytest
+
+import modelopt.torch.utils.memory_monitor as mm
+
+
+class _FakeNVML:
+    """Reference-counted like NVML: every call but init raises while uninitialized."""
+
+    def __init__(self, memory_info_supported: bool = True):
+        self.inits = 0
+        self.memory_info_supported = memory_info_supported
+
+    def install(self, monkeypatch):
+        for name in ("Init", "Shutdown", "DeviceGetCount", "DeviceGetHandleByIndex"):
+            monkeypatch.setattr(mm, f"nvml{name}", getattr(self, name))
+        monkeypatch.setattr(mm, "nvmlDeviceGetMemoryInfo", self.DeviceGetMemoryInfo)
+
+    def _require_init(self):
+        if self.inits <= 0:
+            raise RuntimeError("Uninitialized")
+
+    def Init(self):  # noqa: N802
+        self.inits += 1
+
+    def Shutdown(self):  # noqa: N802
+        self._require_init()
+        self.inits -= 1
+
+    def DeviceGetCount(self):  # noqa: N802
+        self._require_init()
+        return 1
+
+    def DeviceGetHandleByIndex(self, index):  # noqa: N802
+        self._require_init()
+        return index
+
+    def DeviceGetMemoryInfo(self, handle):  # noqa: N802
+        self._require_init()
+        if not self.memory_info_supported:
+            raise RuntimeError("Not Supported")
+        return SimpleNamespace(used=2 * 1024**3)
+
+
+@pytest.fixture
+def launched():
+    monitors = []
+    yield monitors
+    for monitor in monitors:
+        atexit.unregister(monitor.stop)
+        monitor.stop()
+
+
+@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)