From d4c8fb694e5e7a75428a39d3ec9c56ff46536bf2 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:08:02 +0000
Subject: [PATCH 01/24] Keep hf_ptq's previews of a sharded model in lockstep,
 on the quantized experts

Each rank previews the first sample of its own calibration shard, so the
ranks finish decoding at different steps. A model its plan shards runs a
collective in every forward (FSDP2's all-gathers, the routed experts'
all-to-all), and a rank that stopped decoding left the others blocked in it.
The previews now generate with synced_gpus when the recorded plan shards the
model (preview_generate_kwargs), Nemotron-VL's text and image previews too.
The ranks decide together whether to run the image preview, since each node
reads the sample images from its own checkout, and a preview that fails
under synced_gpus raises instead of leaving its rank behind.

While a preview decodes, generate is kept from switching a MoE's experts to
batched_mm (keep_the_experts_implementation), on every module that has the
decode hook: batched_mm bypasses the quantizers, and transformers refuses the
switch on an FSDP2-wrapped model. A hook a module holds itself is put back.

Tests: the previews' kwargs by plan, Nemotron-VL's included, the image-preview
vote, the decode hook kept and put back, nested previews. The previews read
the model's plan, so the stand-in models of the existing preview tests are
modules now.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/example_utils.py              |  98 ++++++--
 examples/hf_ptq/hf_ptq.py                     |  27 ++-
 examples/hf_ptq/vlm_utils.py                  |  78 +++---
 tests/examples/hf_ptq/test_example_utils.py   |  64 ++++-
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  15 +-
 .../examples/hf_ptq/test_hf_ptq_model_type.py |  18 +-
 tests/examples/hf_ptq/test_hf_ptq_previews.py | 229 ++++++++++++++++++
 7 files changed, 450 insertions(+), 79 deletions(-)
 create mode 100644 tests/examples/hf_ptq/test_hf_ptq_previews.py

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index 0ef0c4fa389..ebdb974806f 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -22,7 +22,7 @@
 import os
 import warnings
 from collections.abc import Iterator
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from datetime import timedelta
 from pathlib import Path
@@ -56,6 +56,7 @@
 )
 from modelopt.torch.quantization.mode import get_modelike_from_algo_cfg
 from modelopt.torch.utils import distributed as dist_utils
+from modelopt.torch.utils.distributed import distribution_plan
 from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run
 from modelopt.torch.utils.plugins.model_load_utils import is_fsdp2_text_only_vlm, load_pretrained
 
@@ -154,6 +155,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,
@@ -178,32 +221,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):
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index c79c25ccf98..ec63f2af9b4 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -43,7 +43,9 @@
     get_tokenizer,
     is_nemotron_vl,
     is_trtllm_enc_dec_export,
+    keep_the_experts_implementation,
     mlflow_run,
+    preview_generate_kwargs,
     recipe_layerwise_blocks,
     run_nemotron_vl_preview,
     save_source_config,
@@ -809,11 +811,13 @@ def pre_quantize(
             trust_remote_code=args.trust_remote_code,
         )
     else:
-        generated_ids_before_ptq = full_model.generate(
-            preview_input_ids,
-            attention_mask=preview_attention_mask,
-            max_new_tokens=100,
-        )
+        with keep_the_experts_implementation(full_model):
+            generated_ids_before_ptq = full_model.generate(
+                preview_input_ids,
+                attention_mask=preview_attention_mask,
+                max_new_tokens=100,
+                **preview_generate_kwargs(full_model),
+            )
 
     return preview_input_ids, preview_attention_mask, generated_ids_before_ptq
 
@@ -866,12 +870,15 @@ 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:
             warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
     elif is_nemotron_vl_model and tokenizer is not None:
diff --git a/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/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py
index 01e0634c93a..5437e51617b 100644
--- a/tests/examples/hf_ptq/test_example_utils.py
+++ b/tests/examples/hf_ptq/test_example_utils.py
@@ -21,16 +21,20 @@
 """
 
 import warnings
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
 from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 import torch
 from _test_utils.examples.hf_ptq_example_utils import example_utils
-from _test_utils.torch.transformers_models import create_tiny_qwen3_moe_dir, get_tiny_nemotron_h
+from _test_utils.torch.transformers_models import (
+    create_tiny_qwen3_moe_dir,
+    get_tiny_llama,
+    get_tiny_nemotron_h,
+)
 from safetensors.torch import save_file
-from transformers import AutoModelForCausalLM, LlamaConfig
+from transformers import AutoModelForCausalLM, GenerationMixin, LlamaConfig
 
 from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
 
@@ -192,6 +196,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 5543e9de902..c36896fe8ab 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -38,7 +38,11 @@
 from modelopt.torch.quantization import tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.utils import mlflow as mlflow_lib
-from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
+from modelopt.torch.utils.distributed import (
+    DistributionPlan,
+    distribution_plan,
+    set_distribution_plan,
+)
 from modelopt.torch.utils.mlflow import describe_run, run_tags
 
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
@@ -1012,9 +1016,12 @@ def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch):
     NVBug 6752977."""
     hf_ptq = _import_hf_ptq(monkeypatch)
 
-    full_model = SimpleNamespace(
-        generate=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))
-    )
+    class _FailingModel(torch.nn.Module):
+        def generate(self, *args, **kwargs):
+            raise RuntimeError("boom")
+
+    full_model = _FailingModel()
+    set_distribution_plan(full_model, DistributionPlan())
     export_calls = []
     monkeypatch.setattr(
         hf_ptq,
diff --git a/tests/examples/hf_ptq/test_hf_ptq_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_previews.py b/tests/examples/hf_ptq/test_hf_ptq_previews.py
new file mode 100644
index 00000000000..b339722c050
--- /dev/null
+++ b/tests/examples/hf_ptq/test_hf_ptq_previews.py
@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""hf_ptq's generation previews of a model its distribution plan shards.
+
+Each rank previews the first sample of its own calibration shard, so the ranks finish decoding at
+different steps; a sharded model's previews keep every rank decoding until all have finished, and
+keep ``generate`` from moving a MoE's experts off the backend quantization hooks.
+"""
+
+import contextlib
+import importlib
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from _test_utils.examples.hf_ptq_example_utils import example_utils
+from _test_utils.torch.transformers_models import get_tiny_tokenizer
+from PIL import Image
+
+from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+
+_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
+_CKPT = "/ckpt/model"
+
+
+@pytest.fixture
+def hf_ptq(monkeypatch):
+    monkeypatch.syspath_prepend(str(_EXAMPLES_DIR))
+    return importlib.import_module("hf_ptq")
+
+
+class _PreviewModel(torch.nn.Module):
+    """A model whose ``generate`` records the kwargs each preview passes it, and how often it entered
+    the decode hook that Transformers' ``generate`` enters to switch a MoE's experts."""
+
+    def __init__(self, plan=None):
+        super().__init__()
+        self.config = SimpleNamespace(is_encoder_decoder=False, model_type="qwen3_moe")
+        self.calls = []
+        self.switches = 0
+        if plan is not None:
+            set_distribution_plan(self, plan)
+
+    @contextlib.contextmanager
+    def _optimize_model_for_decode(self):
+        self.switches += 1
+        yield
+
+    def generate(self, input_ids, **kwargs):
+        self.calls.append(kwargs)
+        with self._optimize_model_for_decode():
+            return torch.cat([input_ids, input_ids[:, :1]], dim=1)
+
+
+_PREVIEW_PLANS = pytest.mark.parametrize(
+    ("plan", "lockstep"),
+    [
+        (None, False),
+        (DistributionPlan(4), False),
+        (DistributionPlan(4, ep_size=4), True),
+        (DistributionPlan(4, dense="fsdp"), True),
+        (DistributionPlan(1, dense="fsdp"), False),
+    ],
+    ids=["single_process", "replicated", "expert_parallel", "fsdp2", "one_rank_fsdp2"],
+)
+
+
+@_PREVIEW_PLANS
+def test_a_sharded_model_previews_in_lockstep(hf_ptq, monkeypatch, plan, lockstep):
+    """Each rank previews its own calibration shard's first sample, so the ranks finish decoding at
+    different steps. A sharded model runs a collective in every forward, and a rank that stopped
+    early left the others blocked in it: Qwen3-30B-A3B NVFP4 at EP4 hung in the post-PTQ preview.
+    Both previews keep every rank of a sharded model decoding until all have finished."""
+    model = _PreviewModel(plan)
+    loader = [{"input_ids": torch.tensor([[1, 2, 3]]), "attention_mask": torch.ones(1, 3)}]
+    args = SimpleNamespace(
+        specdec_offline_dataset=None,
+        skip_generate=False,
+        verbose=False,
+        local_checkpoint_path=_CKPT,
+        trust_remote_code=False,
+    )
+    monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: None)
+
+    ids, mask, before = hf_ptq.pre_quantize(args, model, "qwen3_moe", None, loader, False)
+    hf_ptq.post_quantize(
+        args,
+        model,
+        model,
+        "qwen3_moe",
+        get_tiny_tokenizer(),
+        None,
+        ids,
+        mask,
+        before,
+        False,
+        None,
+        loader,
+    )
+
+    synced = [call.get("synced_gpus", False) for call in model.calls]
+    assert synced == [lockstep, lockstep]  # before and after quantization
+    assert model.switches == 0  # neither preview lets generate move the experts off their backend
+
+
+@_PREVIEW_PLANS
+def test_a_sharded_nemotron_vl_model_previews_in_lockstep(monkeypatch, plan, lockstep):
+    """The Nemotron-VL preview's text, fallback and image generations keep a sharded model's ranks
+    decoding together too."""
+    recorded = {}
+    keeps = lambda model: "_optimize_model_for_decode" in vars(model)  # noqa: E731
+    vlm_utils = SimpleNamespace(
+        run_text_only_generation=lambda model, tokenizer, question, config, *args: recorded.update(
+            text=config, text_keeps_the_experts=keeps(model)
+        ),
+        run_vl_preview_generation=lambda model, tokenizer, path, stage, trust, kwargs: (
+            recorded.update(image=kwargs, image_keeps_the_experts=keeps(model))
+        ),
+    )
+    monkeypatch.setitem(sys.modules, "vlm_utils", vlm_utils)
+    model = _PreviewModel(plan)
+
+    example_utils.run_nemotron_vl_preview(
+        model, get_tiny_tokenizer(), torch.tensor([[1, 2, 3]]), _CKPT, "stage", allow_fallback=True
+    )
+
+    expected = {"synced_gpus": True} if lockstep else {}
+    assert recorded["text"].get("synced_gpus", False) == lockstep
+    assert recorded["image"] == expected
+    assert recorded["text_keeps_the_experts"] and recorded["image_keeps_the_experts"]
+    assert "_optimize_model_for_decode" not in vars(model)
+    assert model.calls[0].get("synced_gpus", False) == lockstep  # the fallback generate
+
+
+@pytest.mark.parametrize(("lockstep", "raises"), [({}, False), ({"synced_gpus": True}, True)])
+def test_a_failed_text_only_vl_preview_stops_a_sharded_run(monkeypatch, lockstep, raises):
+    monkeypatch.syspath_prepend(str(_EXAMPLES_DIR))
+    vlm_utils = importlib.import_module("vlm_utils")
+
+    class _ChatModel:
+        def chat(self, *args, **kwargs):
+            raise RuntimeError("the preview failed on this rank")
+
+    run = lambda: vlm_utils.run_text_only_generation(  # noqa: E731
+        _ChatModel(), None, "question", dict(lockstep), _CKPT
+    )
+
+    if raises:
+        with pytest.raises(RuntimeError, match="failed on this rank"):
+            run()
+    else:
+        assert run() is None
+
+
+@pytest.mark.parametrize("generation_kwargs", [None, {"synced_gpus": True}])
+def test_the_vl_image_preview_generates_with_the_kwargs_it_is_handed(
+    monkeypatch, tmp_path, generation_kwargs
+):
+    """run_vl_preview_generation hands the kwargs it is given to the model's chat(), in its config."""
+    monkeypatch.syspath_prepend(str(_EXAMPLES_DIR))
+    vlm_utils = importlib.import_module("vlm_utils")
+    (tmp_path / "images").mkdir()
+    Image.new("RGB", (4, 4)).save(tmp_path / "images" / "example1a.jpeg")
+    monkeypatch.setattr(vlm_utils, "__file__", str(tmp_path / "vlm_utils.py"))
+    monkeypatch.setattr(
+        vlm_utils.AutoImageProcessor, "from_pretrained", lambda *args, **kwargs: lambda images: {}
+    )
+    received = []
+
+    class _ChatModel:
+        device = torch.device("cpu")
+
+        def chat(self, tokenizer, question, generation_config, **image_features):
+            received.append(generation_config)
+            return "a picture"
+
+    vlm_utils.run_vl_preview_generation(
+        _ChatModel(), SimpleNamespace(eos_token_id=0), _CKPT, "stage", False, generation_kwargs
+    )
+
+    assert len(received) == 1
+    assert received[0].get("synced_gpus", False) == (generation_kwargs is not None)
+
+
+@pytest.mark.parametrize("has_image", [True, False], ids=["a_peer_has_none", "no_rank_has_one"])
+def test_the_ranks_of_a_sharded_model_skip_the_image_preview_together(
+    monkeypatch, tmp_path, has_image
+):
+    """Each node reads the sample images from its own checkout, which may have none: the repository
+    ships none. A rank that previewed alone would wait for its peers in lockstep, so a rank with an
+    image skips when a peer has none, and when no rank has one, they all skip without stopping."""
+    monkeypatch.syspath_prepend(str(_EXAMPLES_DIR))
+    vlm_utils = importlib.import_module("vlm_utils")
+    if has_image:
+        (tmp_path / "images").mkdir()
+        Image.new("RGB", (4, 4)).save(tmp_path / "images" / "example1a.jpeg")
+    monkeypatch.setattr(vlm_utils, "__file__", str(tmp_path / "vlm_utils.py"))
+    monkeypatch.setattr(
+        vlm_utils.AutoImageProcessor, "from_pretrained", lambda *args, **kwargs: lambda images: {}
+    )
+    # This rank's answer, and a peer's that found no image.
+    monkeypatch.setattr(vlm_utils, "all_reduce", lambda label, fn, combine: combine([fn(), False]))
+
+    class _ChatModel:
+        device = torch.device("cpu")
+
+        def chat(self, *args, **kwargs):
+            pytest.fail("this rank previewed without its peers")
+
+    preview = vlm_utils.run_vl_preview_generation(
+        _ChatModel(), SimpleNamespace(eos_token_id=0), _CKPT, "stage", False, {"synced_gpus": True}
+    )
+
+    assert preview is None

From 7b6586ff946cb3617eecb2df6ccda07d0f6a66f7 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:08:08 +0000
Subject: [PATCH 02/24] Pin a whole-model load to the rank's own GPU by its
 distribution plan

get_model takes the distribution plan it loads the model by. When the plan
gives each of several ranks the whole model, every load path pins the rank's
copy to the device it was handed, cuda:<local_rank>, the pack-quantized,
MXFP4 and DiffusionGemma loads included. Left to accelerate's "auto", every
rank resolves the same placement, and all their copies land on one GPU. A
device that names no GPU index, and --use_seq_device_map, which spreads one
model over every GPU, are refused there. A pinned copy is not sized against
the node's GPUs, a limit its one-device map would ignore. The device is
compared by its type, as a run's state carries a torch.device, which never
equals a string such as "cpu". A single process keeps its placement.

hf_ptq passes the plan in the next PR of this split.

Tests: the pin on every load path, by a string and by a torch.device, a CPU
device, and both refusals.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/example_utils.py            |  56 +++++++++--
 tests/examples/hf_ptq/test_example_utils.py | 101 ++++++++++++++++++++
 2 files changed, 147 insertions(+), 10 deletions(-)

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index ebdb974806f..c3cc4b74d9b 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -56,7 +56,7 @@
 )
 from modelopt.torch.quantization.mode import get_modelike_from_algo_cfg
 from modelopt.torch.utils import distributed as dist_utils
-from modelopt.torch.utils.distributed import distribution_plan
+from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
 from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run
 from modelopt.torch.utils.plugins.model_load_utils import is_fsdp2_text_only_vlm, load_pretrained
 
@@ -652,8 +652,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:
@@ -664,9 +667,25 @@ 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".
+    # hf_ptq hands in its run's torch.device, which never equals a string such as "cpu".
+    on_cpu = torch.device(device).type == "cpu"
+    if plan.world_size > 1 and not on_cpu and ":" not in str(device):
+        raise ValueError(
+            f"{plan!r} has each rank hold its own copy of the model on its own GPU, but "
+            f"device={device!r} names no GPU index; pass the rank's own GPU, cuda:<local_rank>."
+        )
+    pinned = {"": str(device)} if not on_cpu and plan.world_size > 1 else None
+    device_map = "cpu" if on_cpu else pinned or "auto"
 
     # Prepare config kwargs for loading
     config_kwargs = {"trust_remote_code": trust_remote_code} if trust_remote_code else {}
@@ -702,13 +721,24 @@ def get_model(
     # modules together for checkpoints that fit; larger ones can still spill and split a
     # tied pair, and need an explicit single-device map. Multi-GPU only: a single-GPU split
     # cannot separate a tied pair, and sequential would needlessly cap max_memory there.
-    if device != "cpu" and torch.cuda.device_count() > 1 and is_diffusion_gemma(hf_config):
+    if (
+        not pinned
+        and not on_cpu
+        and torch.cuda.device_count() > 1
+        and is_diffusion_gemma(hf_config)
+    ):
         print(
             "Detected DiffusionGemma model. Using device_map='sequential'; the balanced "
             "'auto' mapping can split its tied encoder/decoder weights across GPUs."
         )
         use_seq_device_map = True
 
+    if use_seq_device_map and pinned:
+        raise ValueError(
+            f"--use_seq_device_map spreads one model over every GPU, but each rank of this run "
+            f"holds its own copy on its own GPU ({device}). Drop it, or shard the model with "
+            "--use_fsdp2."
+        )
     if use_seq_device_map:
         device_map = "sequential"
         # If we use sequential, set max_memory limit to ensure that the model does not occupy the full GPU
@@ -760,6 +790,7 @@ def has_pack_quantized_config(config):
             ckpt_path,
             AutoModelForCausalLM,
             model_type=hf_config.model_type,
+            distribution=plan,
             device_map=device_map,
             **model_kwargs,
         )
@@ -771,7 +802,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",
                 **experts_kwargs,
@@ -798,7 +830,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 on_cpu else pinned or "sequential",
             **model_kwargs,
         )
     else:
@@ -884,7 +917,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:
@@ -906,6 +941,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,
         )
@@ -914,11 +950,11 @@ def has_pack_quantized_config(config):
         _unpack_compressed_linear_weights(model, ckpt_path)
 
     # If device_map was disabled (None), manually move model to target device
-    if device_map is None and device != "cpu":
+    if device_map is None and not on_cpu:
         print(f"Moving model to {device} device...")
         model = model.to(device)
 
-    if device == "cuda" and not is_model_on_gpu(model):
+    if torch.device(device).type == "cuda" and not is_model_on_gpu(model):
         print("Warning: Some parameters are not on a GPU. Calibration can be slow or hit OOM")
 
     return model
diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py
index 5437e51617b..0fd24ee36d3 100644
--- a/tests/examples/hf_ptq/test_example_utils.py
+++ b/tests/examples/hf_ptq/test_example_utils.py
@@ -36,6 +36,7 @@
 from safetensors.torch import save_file
 from transformers import AutoModelForCausalLM, GenerationMixin, LlamaConfig
 
+from modelopt.torch.utils.distributed import DistributionPlan
 from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
 
 
@@ -774,3 +775,103 @@ def __init__(self, algorithm):
 def test_recipe_layerwise_blocks(recipe, expected):
     """Both recipe shapes normalize to dicts, so callers need no shape-aware access."""
     assert example_utils.recipe_layerwise_blocks(recipe) == expected
+
+
+# Small enough that the dense part fits any GPU.
+_TINY = {
+    "hidden_size": 16,
+    "intermediate_size": 32,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "vocab_size": 32,
+}
+
+
+@pytest.mark.parametrize(
+    ("device", "world_size", "device_map"),
+    [
+        ("cuda:1", 4, {"": "cuda:1"}),
+        (torch.device("cuda:1"), 4, {"": "cuda:1"}),
+        ("cuda:1", 1, "auto"),
+        (torch.device("cpu"), 1, "cpu"),
+    ],
+    ids=["str", "torch_device", "single_process", "cpu"],
+)
+def test_a_rank_device_pins_the_whole_model_to_it(
+    monkeypatch, tmp_path, device, world_size, device_map
+):
+    """Under torchrun each rank asks for its own cuda:<local_rank>; "auto" would let every rank's
+    accelerate pick the same idle GPU. A single process spreads the model over the GPUs, whatever
+    GPU it names, and keeps it on a CPU it names, as the torch.device hf_ptq's run hands in. Only a
+    single process's load is held to a memory limit sized over the node's GPUs: the pinned map would
+    ignore one."""
+    LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path)
+    # No GPU on the node holds the model, so sizing it over the node limits its memory.
+    monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 0, "cpu": 1 << 40})
+    placed = []
+
+    plan = DistributionPlan(world_size)
+
+    def load_pretrained(ckpt_path, auto_class, **kwargs):
+        assert kwargs["distribution"] == plan
+        placed.append((kwargs["device_map"], "max_memory" in kwargs))
+        return torch.nn.Linear(1, 1)
+
+    monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained)
+
+    example_utils.get_model(str(tmp_path), device=device, distribution=plan)
+
+    assert placed == [(device_map, world_size == 1)]
+
+
+@pytest.mark.parametrize(
+    "checkpoint",
+    ["pack-quantized", "mxfp4", "diffusion-gemma"],
+)
+def test_every_load_path_keeps_the_rank_pin(monkeypatch, tmp_path, checkpoint):
+    """A pack-quantized or MXFP4 checkpoint loads through a path of its own, and DiffusionGemma asks
+    for a sequential map on a multi-GPU node; each would otherwise place every rank's copy over the
+    same GPUs."""
+    quantization_config = {
+        "pack-quantized": {"quant_method": "compressed-tensors", "format": "pack-quantized"},
+        "mxfp4": {"quant_method": "mxfp4"},
+    }.get(checkpoint)
+    LlamaConfig(
+        architectures=["LlamaForCausalLM"], quantization_config=quantization_config, **_TINY
+    ).save_pretrained(tmp_path)
+    if checkpoint == "diffusion-gemma":
+        monkeypatch.setattr(example_utils, "is_diffusion_gemma", lambda config: True)
+        monkeypatch.setattr(torch.cuda, "device_count", lambda: 2)
+    monkeypatch.setattr(example_utils, "_unpack_compressed_linear_weights", lambda *args: None)
+    placed = []
+
+    def load_pretrained(ckpt_path, auto_class, **kwargs):
+        placed.append(kwargs["device_map"])
+        return torch.nn.Linear(1, 1)
+
+    monkeypatch.setattr(example_utils, "load_pretrained", load_pretrained)
+
+    example_utils.get_model(str(tmp_path), device="cuda:1", distribution=DistributionPlan(4))
+
+    assert placed == [{"": "cuda:1"}]
+
+
+@pytest.mark.parametrize(
+    ("device", "kwargs", "match"),
+    [
+        ("cuda:1", {"use_seq_device_map": True}, "--use_seq_device_map"),
+        # "auto" would let every rank's accelerate pick the same idle GPU.
+        ("cuda", {}, "names no GPU"),
+    ],
+    ids=["sequential_map", "unindexed_device"],
+)
+def test_a_placement_over_shared_gpus_is_refused_where_each_rank_holds_its_own_copy(
+    tmp_path, device, kwargs, match
+):
+    LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path)
+
+    with pytest.raises(ValueError, match=match):
+        example_utils.get_model(
+            str(tmp_path), device=device, distribution=DistributionPlan(4), **kwargs
+        )

From 6761d722255cc0ec121f792599e4fc7f2ab87a6e Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:57:00 +0000
Subject: [PATCH 03/24] Hand load_pretrained only a distribution plan
 get_model's caller gave

get_model forwarded a single process's plan when its caller gave none, so
load_pretrained no longer decided the layout itself: a MODELOPT_EP_DEGREE or
MODELOPT_FSDP_SIZE asking to shard over a launch whose process group is not
set up was ignored instead of refused, and every rank fell back to "auto"
placement. get_model now forwards only a plan it was given, and pins by plan
as before. The new test fails without it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/example_utils.py            | 13 ++++++++-----
 tests/examples/hf_ptq/test_example_utils.py | 14 ++++++++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index c3cc4b74d9b..b68b3f90bf4 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -655,8 +655,11 @@ def get_model(
     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.
+    # The plan to load the whole model by on every rank; a single process's by default. Only a plan
+    # the caller gave reaches load_pretrained: without one it decides itself as before, and refuses
+    # MODELOPT_* knobs asking to shard over a launch whose process group is not set up.
     plan = distribution or DistributionPlan()
+    routing = {} if distribution is None else {"distribution": distribution}
 
     _disk_offload = offload_folder is not None
     if _disk_offload and max_cpu_memory_gb is None:
@@ -790,7 +793,7 @@ def has_pack_quantized_config(config):
             ckpt_path,
             AutoModelForCausalLM,
             model_type=hf_config.model_type,
-            distribution=plan,
+            **routing,
             device_map=device_map,
             **model_kwargs,
         )
@@ -802,7 +805,7 @@ def has_pack_quantized_config(config):
                 ckpt_path,
                 AutoModelForCausalLM,
                 model_type=hf_config.model_type,
-                distribution=plan,
+                **routing,
                 device_map=pinned or "auto",
                 trust_remote_code=trust_remote_code,
                 dtype="auto",
@@ -830,7 +833,7 @@ def has_pack_quantized_config(config):
             ckpt_path,
             AutoModelForCausalLM,
             model_type=hf_config.model_type,
-            distribution=plan,
+            **routing,
             device_map="cpu" if on_cpu else pinned or "sequential",
             **model_kwargs,
         )
@@ -941,7 +944,7 @@ def has_pack_quantized_config(config):
             ckpt_path,
             auto_model_module,
             model_type=hf_config.model_type,
-            distribution=plan,
+            **routing,
             device_map=device_map,
             **model_kwargs2,
         )
diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py
index 0fd24ee36d3..76765e6b6df 100644
--- a/tests/examples/hf_ptq/test_example_utils.py
+++ b/tests/examples/hf_ptq/test_example_utils.py
@@ -875,3 +875,17 @@ def test_a_placement_over_shared_gpus_is_refused_where_each_rank_holds_its_own_c
         example_utils.get_model(
             str(tmp_path), device=device, distribution=DistributionPlan(4), **kwargs
         )
+
+
+def test_without_a_plan_the_load_decides_and_refuses_sharding_knobs_it_cannot_honor(
+    monkeypatch, tmp_path
+):
+    """get_model hands load_pretrained only a plan its caller gave. Without one the load decides
+    as before, so a MODELOPT_* knob asking to shard over a launch whose process group is not set up
+    is refused, rather than ignored for a one-rank plan."""
+    LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path)
+    monkeypatch.setenv("WORLD_SIZE", "2")
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+
+    with pytest.raises(RuntimeError, match="MODELOPT_FSDP_SIZE=world ask to shard"):
+        example_utils.get_model(str(tmp_path), device="cuda:0")

From 9bd65912db17b5bc11d06243f218d815b6a220af Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 18:59:39 +0000
Subject: [PATCH 04/24] Refuse a plan that shards the model in get_model, which
 loads it whole

get_model pins a whole copy of the model to each rank's GPU for a plan of
several ranks. An expert-parallel or FSDP2 plan shards the model instead, and
hf_ptq hands those to load_pretrained directly; one passed to get_model would
get a pinned whole-model load, and advice that does not fit. Refuse it up
front.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/example_utils.py            |  5 +++++
 tests/examples/hf_ptq/test_example_utils.py | 14 ++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index b68b3f90bf4..54f805faabd 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -659,6 +659,11 @@ def get_model(
     # the caller gave reaches load_pretrained: without one it decides itself as before, and refuses
     # MODELOPT_* knobs asking to shard over a launch whose process group is not set up.
     plan = distribution or DistributionPlan()
+    if plan.is_expert_parallel or plan.is_fsdp_dense:
+        raise ValueError(
+            f"{plan!r} shards the model over the ranks, but get_model loads the whole model on "
+            "every rank: load it with load_pretrained."
+        )
     routing = {} if distribution is None else {"distribution": distribution}
 
     _disk_offload = offload_folder is not None
diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py
index 76765e6b6df..37580707ea9 100644
--- a/tests/examples/hf_ptq/test_example_utils.py
+++ b/tests/examples/hf_ptq/test_example_utils.py
@@ -877,6 +877,20 @@ def test_a_placement_over_shared_gpus_is_refused_where_each_rank_holds_its_own_c
         )
 
 
+@pytest.mark.parametrize(
+    "plan",
+    [DistributionPlan(4, ep_size=4), DistributionPlan(4, dense="fsdp")],
+    ids=["expert_parallel", "fsdp"],
+)
+def test_a_plan_that_shards_the_model_is_refused(monkeypatch, tmp_path, plan):
+    """get_model loads a whole model on every rank; a plan that shards it goes to load_pretrained."""
+    LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY).save_pretrained(tmp_path)
+    monkeypatch.setattr(example_utils, "load_pretrained", lambda *a, **k: pytest.fail("loaded"))
+
+    with pytest.raises(ValueError, match="shards the model over the ranks"):
+        example_utils.get_model(str(tmp_path), device="cuda:1", distribution=plan)
+
+
 def test_without_a_plan_the_load_decides_and_refuses_sharding_knobs_it_cannot_honor(
     monkeypatch, tmp_path
 ):

From d5556a2777af2ad87a34c8c8af0fc709e2e8cef2 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 19:10:08 +0000
Subject: [PATCH 05/24] Stop the run when a sharded model's preview after
 quantization fails

The ranks of a sharded model decode the preview in lockstep (synced_gpus). A
rank that warned and skipped a failed preview would go on to the export's
collectives while the others still wait in generate's. Raise instead, with a
message that names --skip_generate. A single process keeps the best-effort
skip, as in NVBug 6752977's case.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/hf_ptq.py                 |  8 +++++++
 tests/examples/hf_ptq/test_hf_ptq_args.py | 26 +++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index ec63f2af9b4..d0ff1c61c74 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -118,6 +118,10 @@
 RAND_SEED = 1234
 # Most tokens per calibration batch when layerwise calibration picks the batch size (64 x 512).
 LAYERWISE_CALIB_TOKENS = 32768
+_PREVIEW_STOPS_THE_RUN = (
+    "The preview after quantization failed, which stops the run before the export, so no "
+    "checkpoint is written; --skip_generate exports without the preview."
+)
 
 
 def _layerwise_calib_batch_size(calib_seq: int, num_samples: int) -> int:
@@ -880,6 +884,10 @@ def post_quantize(
                     **lockstep,
                 )
         except Exception as e:
+            if lockstep:
+                # The ranks decode in lockstep: a rank that skipped the preview would go on to the
+                # export's collectives while the others still wait in generate's, so the run stops.
+                raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e
             warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
     elif is_nemotron_vl_model and tokenizer is not None:
         generated_ids_after_ptq = run_nemotron_vl_preview(
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index c36896fe8ab..80fa5a0b51b 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -1051,6 +1051,32 @@ def generate(self, *args, **kwargs):
     assert len(export_calls) == 1
 
 
+def test_post_quantize_stops_the_run_when_a_sharded_models_preview_fails(monkeypatch):
+    """A sharded model's ranks decode the preview in lockstep: a rank that skipped its failure would
+    go on to the export's collectives while the others still wait in generate's."""
+    hf_ptq = _import_hf_ptq(monkeypatch)
+
+    class _FailingModel(torch.nn.Module):
+        def generate(self, *args, **kwargs):
+            raise RuntimeError("boom")
+
+    full_model = _FailingModel()
+    set_distribution_plan(full_model, DistributionPlan(2, dense="fsdp"))
+    export_calls = []
+    monkeypatch.setattr(hf_ptq, "export_quantized", lambda *a, **k: export_calls.append((a, k)))
+    monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
+    ids = torch.zeros(1, 4, dtype=torch.long)
+    args = SimpleNamespace(specdec_offline_dataset=None, verbose=False)
+
+    with pytest.raises(RuntimeError, match="--skip_generate") as stopped:
+        hf_ptq.post_quantize(
+            args, full_model, full_model, "llama", None, None, ids, None, ids, False, None, None
+        )
+
+    assert str(stopped.value.__cause__) == "boom"
+    assert export_calls == []
+
+
 @pytest.mark.usefixtures("clean_env")
 def test_main_makes_a_hub_checkpoint_local_before_loading(monkeypatch, tmp_path):
     """--pyt_ckpt_path keeps the Hub ID; the model loads from the local copy resolved first."""

From 1af92ade04f3278e75bc201443d04e97a174b3f4 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:50:51 +0000
Subject: [PATCH 06/24] Abort a failed rank of a multi-rank hf_ptq run instead
 of tearing the group down

A rank that fails alone, such as in a sharded model's lockstep preview,
leaves the others in a collective it never reaches. main() then went
through cleanup_distributed, whose barrier and destroy_process_group wait
on those ranks, with the traceback still unprinted: the run hung until the
collective timeout instead of failing. A run of several ranks now exits
through abort(), which prints the error and exits, so the launcher stops
the others; a single process raises as before.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/hf_ptq.py                 | 10 +++++-
 tests/examples/hf_ptq/test_hf_ptq_args.py | 37 +++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index d0ff1c61c74..753f4be5c5d 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -101,7 +101,7 @@
     get_max_batch_size,
     get_supported_datasets,
 )
-from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+from modelopt.torch.utils.distributed import DistributionPlan, abort, set_distribution_plan
 from modelopt.torch.utils.memory_monitor import launch_memory_monitor
 from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
@@ -1592,6 +1592,14 @@ def main(args: argparse.Namespace):
                     tokenizer,
                     device,
                 )
+    except BaseException:
+        # A rank that fails alone, such as in a sharded model's lockstep preview, leaves the others
+        # in a collective it never reaches, and cleanup's barrier and destroy_process_group would
+        # wait on them with the traceback unprinted. abort() prints it and exits, so the launcher
+        # stops the other ranks.
+        if args.dist_state.world_size > 1:
+            abort()
+        raise
     finally:
         cleanup_distributed(args)
 
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 80fa5a0b51b..e4300892572 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -1111,6 +1111,43 @@ def load_model(args):
     ]
 
 
+@pytest.mark.parametrize("world_size", [1, 2], ids=["single_process", "two_ranks"])
+def test_a_rank_that_fails_aborts_a_run_of_several_ranks(monkeypatch, world_size):
+    """The other ranks may wait in a collective the failed rank never reaches, where tearing the
+    process group down would wait on them with the error unprinted: a run of several ranks aborts
+    instead, which exits, so the launcher stops the others. A single process raises as before."""
+    hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "org/model")
+    events = []
+
+    class _AbortedError(BaseException):
+        """abort() exits the process; this stand-in raises instead, so the test sees it."""
+
+    def setup_distributed_args(args):
+        args.dist_state = SimpleNamespace(rank=0, world_size=world_size, is_main=True)
+
+    def abort():
+        events.append("abort")
+        raise _AbortedError
+
+    def load_model(args):
+        raise RuntimeError("this rank's load failed")
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.compiler, "set_stance", lambda *args, **kwargs: None)
+    monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: None)
+    monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", lambda path: (path, path))
+    monkeypatch.setattr(hf_ptq, "setup_distributed_args", setup_distributed_args)
+    monkeypatch.setattr(hf_ptq, "cleanup_distributed", lambda args: events.append("cleanup"))
+    monkeypatch.setattr(hf_ptq, "abort", abort)
+    monkeypatch.setattr(hf_ptq, "load_model", load_model)
+
+    with pytest.raises(_AbortedError if world_size > 1 else RuntimeError):
+        hf_ptq.main(args)
+
+    # The real abort() exits before the teardown runs.
+    assert events == (["abort", "cleanup"] if world_size > 1 else ["cleanup"])
+
+
 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(

From 223a3ee422323ec86fa29227784c6ecde45c6615 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:03:29 +0000
Subject: [PATCH 07/24] Let distributed_forward_loop run each batch through a
 forward it is given

distributed_forward_loop calibrates by calling model(**batch). A batch that
needs its own call into the model cannot go through it: a VLM's image-text
batch, say, which runs through the VLM's own forward with the keyword
arguments that forward accepts. The loop now takes an optional
forward(batch), and keeps moving each batch to the device and calibrating
with the cache disabled. hf_ptq's expert-parallel image calibration uses it.

The new test runs two batches through a given forward and checks the device
and the cache.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py               | 15 ++++++++++-----
 .../torch/utils/test_distributed_calibration.py   | 15 +++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 2e5855b1384..6e2b3c6d50b 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -908,14 +908,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.
 
@@ -953,7 +955,10 @@ def calibrate(_inner_model):
                         k: (v.to(device) if isinstance(v, torch.Tensor) else v)
                         for k, v in batch.items()
                     }
-                model(**batch)
+                if forward is None:
+                    model(**batch)
+                else:
+                    forward(batch)
             # Again for what runs next (amax sync, export collectives), past the forwards' pool.
             _release_cached_cuda_memory()
 
diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py
index f5a21d23dc3..95a954b9272 100644
--- a/tests/unit/torch/utils/test_distributed_calibration.py
+++ b/tests/unit/torch/utils/test_distributed_calibration.py
@@ -484,6 +484,21 @@ def test_forward_loop_moves_tensors_to_the_device_it_is_given():
     assert [call["logits_to_keep"] for call in model.calls] == [1, 1]
 
 
+def test_forward_loop_runs_each_batch_through_the_forward_it_is_given():
+    """A batch that needs its own call into the model, as a VLM's image-text batch does, goes
+    through ``forward``, on the device and with the cache disabled like any other."""
+    batches = [{"input_ids": torch.zeros(2, 3)}] * 2
+    model = _Recorder()
+
+    distributed_forward_loop(
+        model, batches, device="meta", forward=lambda batch: model(**batch, own_call=True)
+    )(model)
+
+    assert [call["own_call"] for call in model.calls] == [True, True]
+    assert [call["input_ids"].device.type for call in model.calls] == ["meta"] * 2
+    assert model.use_cache_seen == [False, False]
+
+
 def test_forward_loop_disables_use_cache_only_while_calibrating():
     """KV/state caching is off for the calibration forwards (it breaks hybrid Mamba models)."""
     model = _Recorder()

From 2650415080a1e00bc5134b45bc7dc099b1a65421 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:03:44 +0000
Subject: [PATCH 08/24] Initialize NVML before probing it in
 launch_memory_monitor

launch_memory_monitor checks that NVML can report memory before starting
the monitor, so hosts where it cannot (#374) skip monitoring instead of
crashing. But the probe ran before nvmlInit, which only GPUMemoryMonitor's
constructor calls, so every NVML call in it raised NVMLError_Uninitialized:
the monitor printed "Failed to get GPU memory info: Uninitialized" and never
started on any host. hf_ptq.py launches it on every run, so no run has
reported its per-GPU peak memory since.

The probe now initializes NVML for itself and shuts it down afterwards, so
the init count stays balanced with the monitor's own. Hosts where memory
info is unsupported still skip the monitor.

Unit tests run launch_memory_monitor against a fake, reference-counted
NVML: the monitor starts when memory info is available (this fails on the
old code), and is skipped with NVML left shut down when it is not.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/memory_monitor.py        |  9 +-
 tests/unit/torch/utils/test_memory_monitor.py | 85 +++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/torch/utils/test_memory_monitor.py

diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py
index 94ed1cab850..0b1fbbb1fd1 100644
--- a/modelopt/torch/utils/memory_monitor.py
+++ b/modelopt/torch/utils/memory_monitor.py
@@ -140,8 +140,15 @@ def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor | N
     Returns:
         GPUMemoryMonitor: The monitor instance that was launched
     """
+    # Probe that NVML can report memory on this host before starting the monitor. NVML must be
+    # initialized for the probe itself; without it every call raises NVMLError_Uninitialized, and
+    # the monitor would never start anywhere.
     try:
-        nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0))
+        nvmlInit()
+        try:
+            nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0))
+        finally:
+            nvmlShutdown()
     except Exception as e:
         print(f"Failed to get GPU memory info: {e}. Stopping GPU memory monitor.")
         return None
diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py
new file mode 100644
index 00000000000..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)

From a7691fede118de0f2b4050a0aad74cf5b9cc7854 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:31:50 +0000
Subject: [PATCH 09/24] Stop a GPU memory monitor once, however often stop() is
 called

launch_memory_monitor registers the monitor's stop() to run at exit. Now that
the monitor starts, a caller that stops it, as the module's own example does,
had it stopped again at exit: the peak-memory report printed twice, and the
second NVML shutdown raised inside the exit hook. A stopped monitor's stop()
now does nothing, and the test fixture no longer unregisters the hook to work
around it.

The new test stops a monitor twice; it fails without the change.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/memory_monitor.py        |  8 +++++++-
 tests/unit/torch/utils/test_memory_monitor.py | 17 +++++++++++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py
index 0b1fbbb1fd1..918b161f2a8 100644
--- a/modelopt/torch/utils/memory_monitor.py
+++ b/modelopt/torch/utils/memory_monitor.py
@@ -85,6 +85,7 @@ def __init__(self, monitor_interval: float = 10.0):
         self.peak_memory = {}  # GPU index to peak memory mapping
         self.is_running = False
         self.monitor_thread = None
+        self._stopped = False
         nvmlInit()
         self.device_count = nvmlDeviceGetCount()
 
@@ -116,8 +117,13 @@ def stop(self):
         GPU that was monitored, and properly shuts down the NVML interface. It will
         wait for the monitoring thread to complete before returning.
 
-        The peak memory usage is displayed in GB for each GPU index.
+        The peak memory usage is displayed in GB for each GPU index. Stopping a stopped monitor does
+        nothing, as when a caller stops it and the exit hook :func:`launch_memory_monitor` registered
+        runs too.
         """
+        if self._stopped:
+            return
+        self._stopped = True
         self.is_running = False
         # Print peak memory usage
         print("########")
diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py
index 8bb3dcab5bf..df146afcad5 100644
--- a/tests/unit/torch/utils/test_memory_monitor.py
+++ b/tests/unit/torch/utils/test_memory_monitor.py
@@ -15,7 +15,6 @@
 
 """Unit tests for launch_memory_monitor's NVML availability probe, against a fake NVML."""
 
-import atexit
 from types import SimpleNamespace
 
 import pytest
@@ -66,7 +65,6 @@ def launched():
     monitors = []
     yield monitors
     for monitor in monitors:
-        atexit.unregister(monitor.stop)
         monitor.stop()
 
 
@@ -83,3 +81,18 @@ def test_the_monitor_starts_only_where_nvml_can_report_memory(monkeypatch, launc
 
     assert (monitor is not None) is supported
     assert nvml.inits == (1 if supported else 0)
+
+
+def test_stopping_a_stopped_monitor_does_nothing(monkeypatch, capsys):
+    """launch_memory_monitor registers stop() to run at exit, so a monitor its caller stopped is
+    stopped again then: the second stop neither reports again nor shuts NVML down a second time."""
+    nvml = _FakeNVML()
+    nvml.install(monkeypatch)
+    monitor = mm.launch_memory_monitor(monitor_interval=0.01)
+
+    monitor.stop()
+    assert "Peak memory usage" in capsys.readouterr().out
+    monitor.stop()
+
+    assert capsys.readouterr().out == ""
+    assert nvml.inits == 0

From a3b9b22d95949e365c0fbf0085dc5901d39b9362 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 01:56:15 +0000
Subject: [PATCH 10/24] Check the first stop's report without waiting for a GPU
 sample

The idempotence test asserted that the first stop() printed a peak-memory
line, which exists only once the monitor thread has sampled a GPU. stop()
right after launch can win that race, so the test could fail on a slow runner.
It now checks the report's frame, which stop() always prints, and keeps the
checks that the second stop prints nothing and leaves NVML's count at 0. It
still fails when stop() is not idempotent.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/utils/test_memory_monitor.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py
index df146afcad5..136ee8a2b08 100644
--- a/tests/unit/torch/utils/test_memory_monitor.py
+++ b/tests/unit/torch/utils/test_memory_monitor.py
@@ -91,7 +91,8 @@ def test_stopping_a_stopped_monitor_does_nothing(monkeypatch, capsys):
     monitor = mm.launch_memory_monitor(monitor_interval=0.01)
 
     monitor.stop()
-    assert "Peak memory usage" in capsys.readouterr().out
+    # The report's frame, which prints whether or not the monitor thread sampled a GPU yet.
+    assert "########" in capsys.readouterr().out
     monitor.stop()
 
     assert capsys.readouterr().out == ""

From b2c562d4e12782637210204c267d539ed94891dc Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:14:42 +0000
Subject: [PATCH 11/24] Wait for the memory monitor's sampling pass before
 reporting its peaks

stop() cleared is_running and printed peak_memory before joining the thread,
but the thread could be between two GPUs of a sampling pass. The report then
missed the last GPU's sample, or, on a host with several GPUs, iterated
peak_memory while the thread added to it and raised before the join and the
NVML shutdown. stop() now joins the thread first.

The new test holds the second of two fake GPUs' samples until stop() waits for
the thread, and checks that the report has both GPUs; it fails without the
change, in each of 10 runs, and takes milliseconds.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/memory_monitor.py        |  5 ++-
 tests/unit/torch/utils/test_memory_monitor.py | 40 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py
index 918b161f2a8..8001ebf4076 100644
--- a/modelopt/torch/utils/memory_monitor.py
+++ b/modelopt/torch/utils/memory_monitor.py
@@ -125,6 +125,9 @@ def stop(self):
             return
         self._stopped = True
         self.is_running = False
+        # The thread can be inside a sampling pass, still adding to peak_memory: let it finish first.
+        if self.monitor_thread:
+            self.monitor_thread.join()
         # Print peak memory usage
         print("########")
         for gpu_idx, peak_mem in self.peak_memory.items():
@@ -132,8 +135,6 @@ def stop(self):
                 f"GPU {gpu_idx}: Peak memory usage = {peak_mem:.2f} GB for all processes on the GPU"
             )
         print("########")
-        if self.monitor_thread:
-            self.monitor_thread.join()
         nvmlShutdown()
 
 
diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py
index 136ee8a2b08..eac9aa8d9c2 100644
--- a/tests/unit/torch/utils/test_memory_monitor.py
+++ b/tests/unit/torch/utils/test_memory_monitor.py
@@ -15,6 +15,7 @@
 
 """Unit tests for launch_memory_monitor's NVML availability probe, against a fake NVML."""
 
+import threading
 from types import SimpleNamespace
 
 import pytest
@@ -25,9 +26,11 @@
 class _FakeNVML:
     """Reference-counted like NVML: every call but init raises while uninitialized."""
 
-    def __init__(self, memory_info_supported: bool = True):
+    def __init__(self, memory_info_supported: bool = True, devices: int = 1, on_sample=None):
         self.inits = 0
         self.memory_info_supported = memory_info_supported
+        self.devices = devices
+        self.on_sample = on_sample  # called with a device's handle before its memory is read
 
     def install(self, monkeypatch):
         for name in ("Init", "Shutdown", "DeviceGetCount", "DeviceGetHandleByIndex"):
@@ -47,7 +50,7 @@ def Shutdown(self):  # noqa: N802
 
     def DeviceGetCount(self):  # noqa: N802
         self._require_init()
-        return 1
+        return self.devices
 
     def DeviceGetHandleByIndex(self, index):  # noqa: N802
         self._require_init()
@@ -57,6 +60,8 @@ def DeviceGetMemoryInfo(self, handle):  # noqa: N802
         self._require_init()
         if not self.memory_info_supported:
             raise RuntimeError("Not Supported")
+        if self.on_sample is not None:
+            self.on_sample(handle)
         return SimpleNamespace(used=2 * 1024**3)
 
 
@@ -97,3 +102,34 @@ def test_stopping_a_stopped_monitor_does_nothing(monkeypatch, capsys):
 
     assert capsys.readouterr().out == ""
     assert nvml.inits == 0
+
+
+def test_the_report_waits_for_a_sampling_pass_in_flight(monkeypatch, capsys):
+    """stop() can come while the thread is between two GPUs of a sampling pass: the report waits for
+    the pass to finish rather than miss the last GPU, or iterate the samples while they grow."""
+    in_flight, joining = threading.Event(), threading.Event()
+
+    def hold_the_second_gpu(handle):
+        if handle == 1 and threading.current_thread() is not threading.main_thread():
+            # The first GPU is sampled; hold the second until stop() waits for the thread (5 s max).
+            in_flight.set()
+            joining.wait(timeout=5)
+
+    nvml = _FakeNVML(devices=2, on_sample=hold_the_second_gpu)
+    nvml.install(monkeypatch)
+    monitor = mm.launch_memory_monitor(monitor_interval=0.01)
+    join = monitor.monitor_thread.join
+
+    def join_and_release(*args, **kwargs):
+        joining.set()
+        return join(*args, **kwargs)
+
+    monkeypatch.setattr(monitor.monitor_thread, "join", join_and_release)
+    assert in_flight.wait(timeout=5)
+
+    monitor.stop()
+
+    report = capsys.readouterr().out
+    assert "GPU 0: Peak memory usage" in report
+    assert "GPU 1: Peak memory usage" in report
+    assert nvml.inits == 0

From 653a6bc9631e0f4b72ddc338efb96d2338c2c5d6 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Mon, 2 Nov 2026 21:47:05 +0000
Subject: [PATCH 12/24] Type the forward distributed_forward_loop calls as
 taking a mapping

Without a device, the callback gets the loader's own batch, which can be
any mapping rather than a dict. Annotate it as a Mapping, and say that it
gets a dict of its own to change only when a device is given.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 6e2b3c6d50b..8a41448af03 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -908,7 +908,7 @@ 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,
+    forward: Callable[[Mapping], object] | None = None,
 ) -> Callable[[torch.nn.Module], None]:
     """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model.
 
@@ -925,7 +925,8 @@ def distributed_forward_loop(
     the single-node ``_forward_loop`` runs them: a hybrid Mamba/attention model mutates its cache in
     place. Each batch is a mapping of the forward's keyword arguments, as ``model(**batch)`` takes;
     only its tensor values move to ``device``, so tensors nested in a value (a list of images, say)
-    must already be there.
+    must already be there. ``forward`` gets a new dict when ``device`` is given, and the loader's own
+    mapping otherwise: copy that before changing it.
     """
     # Local import to avoid any import cycle within modelopt.torch.utils.
     from modelopt.torch.utils.dataset_utils import _disable_use_cache

From 26f07714e1f443d464e44e0dd55a4de946b75516 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:08:33 +0000
Subject: [PATCH 13/24] Run expert-parallel PTQ from hf_ptq.py

hf_ptq.py decides once how to distribute the model, with
resolve_distribution_plan, and loads it by that plan through load_pretrained:
expert-parallel, FSDP2 under --use_fsdp2, or whole on every rank. Everything
after reads the plan back from the model.

- main() runs inside distributed_run, which owns the process group and yields
  this process's rank, size and device; the example's DistributedState and its
  process-group helpers go.
- Refused before the load: an FSDP2 plan for the whole model without
  --use_fsdp2, the TensorRT-LLM checkpoint export, layerwise-export and
  AutoQuantize recipes on the routes this script adds, the single-process load
  options there, and --vllm_fakequant_export and --experts_implementation under
  expert parallelism.
- Every rank of a distributed plan calibrates on its own shard of the data. An
  expert-parallel run calibrates through distributed_forward_loop, image
  batches through the VLM's own forward, and re-plants the all-to-all forwards
  once quantized.
- A replicated model is exported by rank 0 alone, a sharded
  speculative-decoding export is refused, and a failed preview of an
  expert-parallel model stops the run.
- The GPU memory monitor runs on each node's first rank.

[3/3] of the split of this PR: [1/3] adds the library pieces and [2/3] the
lockstep previews and the per-rank placement this builds on.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst                                 |   2 +
 examples/hf_ptq/README.md                     |  28 +-
 examples/hf_ptq/example_utils.py              | 126 ++-
 examples/hf_ptq/hf_ptq.py                     | 352 +++++--
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  85 +-
 .../hf_ptq/test_hf_ptq_expert_parallel.py     | 963 ++++++++++++++++++
 .../hf_ptq/test_hf_ptq_vision_quantization.py |   6 +
 tests/examples/hf_ptq/test_llm_ptq.py         |  99 ++
 .../quantization/test_fsdp2_autoquant.py      |   3 +
 9 files changed, 1442 insertions(+), 222 deletions(-)
 create mode 100644 tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index a05d4351bab..e9b5aa0d28d 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -25,6 +25,7 @@ Changelog
 - ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration.
 - Speed up Hugging Face QAT and QAD with ``QATTrainer``, ``QADTrainer`` or ``--manual_gc``: the ``gc.collect()`` run before each training step now skips the model and optimizer state, saving about 1 s per step (30%) on Qwen3-30B-A3B.
 - Add an opt-in no-gather HuggingFace export for FSDP2-sharded models: with ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` moves each tensor once, to the one rank that writes it, instead of all-gathering each FSDP2 unit on every rank as the default export does. Fused MoE experts stay sharded, each rank writing its own, and a conversion mapping that splits or merges tensors, which the default export refuses, is reversed as a single-process export reverses it.
+- Add expert-parallel PTQ for HuggingFace MoE models to ``examples/hf_ptq/hf_ptq.py``: under ``torchrun``, each rank holds only its share of the routed experts, tokens reach them by all-to-all, and every rank calibrates on its own slice of the data. ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` override the automatic choice of expert-parallel degree and of FSDP2-sharding the dense layers, and the quantized checkpoint is written by every rank in parallel without gathering the model.
 
 *Speculative Decoding*
 
@@ -51,6 +52,7 @@ Changelog
 **Backward Breaking Changes**
 
 - ``QDQAutotunerBase.export_onnx(output_path=...)`` now returns ``None`` instead of serialized bytes. Omit ``output_path`` to keep the bytes-returning behavior, which remains subject to protobuf size limits.
+- ``examples/hf_ptq/hf_ptq.py`` launched by ``torchrun`` on more than one rank 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, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). It also raises when a calibration batch runs out of GPU memory, instead of retrying it in smaller pieces: set ``--batch_size`` to one that fits.
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later.
diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index b95e61719fa..7db3517ae49 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -620,7 +620,7 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for
 
 > *`nvfp4_awq_clip` still works with `--use_fsdp2` on more than one GPU: it calibrates with AWQ-Clip alone, which folds no smoothing scale into the weights, and its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.*
 
-> *With `--use_fsdp2` on more than one GPU, a calibration batch that runs out of memory is not retried in smaller pieces, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives. With `--batch_size 0`, each rank estimates the batch size from one forward and the ranks use the smallest, which is not probed for out-of-memory either: set `--batch_size` if calibration runs out of memory.*
+> *On more than one GPU, with `--use_fsdp2` or in any other `torchrun` launch of several ranks, a calibration batch that runs out of memory is not retried in smaller pieces, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives. With `--batch_size 0`, each rank estimates the batch size from one forward and the ranks use the smallest, which is not probed for out-of-memory either: set `--batch_size` if calibration runs out of memory.*
 
 ### Usage
 
@@ -660,6 +660,32 @@ Qwen checkpoints with `model_type` `qwen3_5` or `qwen3_5_moe` support this path
 
 > *Set `MODELOPT_FSDP2_EXPORT=nogather` to move each tensor once, to the one rank that writes it, instead of all-gathering each FSDP2 unit on every rank as the default export does. Fused MoE experts then stay sharded, and a conversion mapping that splits or merges tensors, which the default export refuses, is reversed as a single-process export reverses it. It does not overwrite a checkpoint already in the export directory.*
 
+## Expert-Parallel Post-Training Quantization for MoE Models
+
+Under `torchrun`, `hf_ptq.py` shards an MoE model's routed experts across the ranks (expert parallelism) instead of loading the whole model on every GPU. Each rank holds `num_experts / ep_degree` experts and calibrates on its own shard of the calibration data. An all-to-all sends every token to the rank that owns its expert, so each expert's quantizers see the tokens routed to it from all ranks. The dense (non-expert) part is replicated on every rank, or FSDP2-sharded when it does not fit one GPU. Each rank then writes its share of the unified HF checkpoint, without gathering the model onto one GPU.
+
+A model can be expert-parallel when its config declares a `base_model_ep_plan`. transformers ships one for several MoEs, and ModelOpt supplies the plan or an adapter for others (see the `ep_spec` entries in `modelopt/torch/models/*/specs.py`). Expert parallelism needs transformers 5.7 or newer, and some models a newer release; a run that cannot be expert-parallel is refused before loading.
+
+```bash
+torchrun --nproc_per_node=<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 54f805faabd..e67466fd166 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -23,8 +23,6 @@
 import warnings
 from collections.abc import Iterator
 from contextlib import contextmanager, nullcontext
-from dataclasses import dataclass
-from datetime import timedelta
 from pathlib import Path
 from typing import Any
 
@@ -65,39 +63,6 @@
 SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"]
 
 
-@dataclass
-class DistributedState:
-    """Example-local distributed state for model loading, dataloader sharding, and rank-0 output."""
-
-    rank: int
-    world_size: int
-    device: torch.device | str
-    is_main: bool
-
-
-def setup_distributed_args(args):
-    """Initialize and attach ``args.dist_state`` (single-process if FSDP2 off)."""
-    if getattr(args, "use_fsdp2", False):
-        # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can
-        # exceed it, and barrier() has no per-call timeout (it must be set at PG creation).
-        dist_utils.setup(timeout=timedelta(hours=2))
-        rank = dist_utils.rank()
-        args.dist_state = DistributedState(
-            rank=rank,
-            world_size=dist_utils.size(),
-            device=torch.device(f"cuda:{dist_utils.local_rank()}"),
-            is_main=rank == 0,
-        )
-    else:
-        args.dist_state = DistributedState(rank=0, world_size=1, device=args.device, is_main=True)
-
-
-def cleanup_distributed(args):
-    """Destroy the process group if ``--use_fsdp2`` set it up."""
-    if getattr(args, "use_fsdp2", False):
-        dist_utils.cleanup()
-
-
 def _awq_calibrations(args) -> set[str]:
     """The AWQ calibrations (``awq_lite``, ``awq_full``) the run's --recipe or --qformat uses."""
     if getattr(args, "recipe", None) is not None:
@@ -289,18 +254,18 @@ def is_nemotron_vl(model_or_config):
     return any("nemotron" in arch.lower() for arch in architectures)
 
 
-def create_vlm_calibration_loop(full_model, calib_dataloader):
-    """Create a calibration loop for VLM models that handles multimodal inputs.
+def create_vlm_calibration_forward(full_model):
+    """Return ``forward(batch)``, which runs one multimodal calibration batch through the VLM.
 
-    This function inspects the model's forward signature and filters batch kwargs
-    to only include supported parameters, then calls the appropriate forward method.
+    It filters the batch kwargs to the parameters the model's forward accepts, then calls the
+    appropriate forward method. :func:`create_vlm_calibration_loop` runs it over a dataloader; an
+    expert-parallel run hands it to ``distributed_forward_loop`` instead.
 
     Args:
         full_model: The full VLM model
-        calib_dataloader: DataLoader yielding multimodal batches
 
     Returns:
-        A calibration function that can be passed to mtq.quantize()
+        A function that runs one batch through ``full_model``
     """
     # Import here to avoid circular dependency
     from nemotron_vl_calib import safe_nemotron_vl_forward
@@ -324,37 +289,55 @@ def create_vlm_calibration_loop(full_model, calib_dataloader):
         or (is_nemotron_vl(full_model) and not hasattr(full_model, "extract_feature"))
     )
 
-    def calibrate_loop(_model):
-        # Inspect model's forward signature to determine what parameters it accepts
-        forward_params = inspect.signature(full_model.forward).parameters
-        accepts_kwargs = any(
-            p.kind == inspect.Parameter.VAR_KEYWORD for p in forward_params.values()
-        )
-        allowed_keys = set(forward_params.keys())
+    # Inspect model's forward signature to determine what parameters it accepts
+    forward_params = inspect.signature(full_model.forward).parameters
+    accepts_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in forward_params.values())
+    allowed_keys = set(forward_params.keys())
+
+    def forward(batch):
+        # For encoder-decoder models, rename input_ids → decoder_input_ids
+        # and disable KV caching to avoid tuple index errors in decoder layers
+        if is_enc_dec and "input_ids" in batch and "pixel_values" in batch:
+            batch["decoder_input_ids"] = batch.pop("input_ids")
+            if "attention_mask" in batch:
+                batch["decoder_attention_mask"] = batch.pop("attention_mask")
+            batch["use_cache"] = False
+
+        # Filter batch to only include parameters the model accepts
+        if accepts_kwargs:
+            call_kwargs = batch
+        else:
+            call_kwargs = {k: v for k, v in batch.items() if k in allowed_keys}
+        # Remove None values
+        call_kwargs = {k: v for k, v in call_kwargs.items() if v is not None}
+
+        if use_safe_forward:
+            safe_nemotron_vl_forward(full_model, call_kwargs)
+        else:
+            full_model(**call_kwargs)
+
+    return forward
+
 
+def create_vlm_calibration_loop(full_model, calib_dataloader):
+    """Create a calibration loop for VLM models that handles multimodal inputs.
+
+    It runs every batch through :func:`create_vlm_calibration_forward`.
+
+    Args:
+        full_model: The full VLM model
+        calib_dataloader: DataLoader yielding multimodal batches
+
+    Returns:
+        A calibration function that can be passed to mtq.quantize()
+    """
+
+    def calibrate_loop(_model):
+        forward = create_vlm_calibration_forward(full_model)
         full_model.eval()
         with torch.no_grad():
             for batch in calib_dataloader:
-                # For encoder-decoder models, rename input_ids → decoder_input_ids
-                # and disable KV caching to avoid tuple index errors in decoder layers
-                if is_enc_dec and "input_ids" in batch and "pixel_values" in batch:
-                    batch["decoder_input_ids"] = batch.pop("input_ids")
-                    if "attention_mask" in batch:
-                        batch["decoder_attention_mask"] = batch.pop("attention_mask")
-                    batch["use_cache"] = False
-
-                # Filter batch to only include parameters the model accepts
-                if accepts_kwargs:
-                    call_kwargs = batch
-                else:
-                    call_kwargs = {k: v for k, v in batch.items() if k in allowed_keys}
-                # Remove None values
-                call_kwargs = {k: v for k, v in call_kwargs.items() if v is not None}
-
-                if use_safe_forward:
-                    safe_nemotron_vl_forward(full_model, call_kwargs)
-                else:
-                    full_model(**call_kwargs)
+                forward(batch)
 
     return calibrate_loop
 
@@ -1125,6 +1108,13 @@ def recipe_layerwise_blocks(recipe) -> list[dict]:
     return blocks
 
 
+def layerwise_export_requested(blocks) -> bool:
+    """Whether a recipe's ``layerwise`` blocks (:func:`recipe_layerwise_blocks`) export layerwise."""
+    return any(
+        block.get("export_dir") is not None and block.get("enable", False) for block in blocks
+    )
+
+
 def _layerwise_checkpoint_dir(algorithm) -> str | None:
     """First ``layerwise.checkpoint_dir`` across the algorithm entries, or None."""
     return next(
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 753f4be5c5d..04a9d02cf4b 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -18,6 +18,7 @@
 import random
 import time
 import warnings
+from datetime import timedelta
 from pathlib import Path
 from typing import Any
 
@@ -35,7 +36,7 @@
     HF_PTQ,
     _prepare_quant_cfg,
     build_quant_cfg,
-    cleanup_distributed,
+    create_vlm_calibration_forward,
     create_vlm_calibration_loop,
     generate_excludes_prompt,
     get_model,
@@ -44,12 +45,12 @@
     is_nemotron_vl,
     is_trtllm_enc_dec_export,
     keep_the_experts_implementation,
+    layerwise_export_requested,
     mlflow_run,
     preview_generate_kwargs,
     recipe_layerwise_blocks,
     run_nemotron_vl_preview,
     save_source_config,
-    setup_distributed_args,
     validate_fsdp2_supported,
 )
 from torch.utils.data import DataLoader
@@ -101,17 +102,22 @@
     get_max_batch_size,
     get_supported_datasets,
 )
-from modelopt.torch.utils.distributed import DistributionPlan, abort, set_distribution_plan
+from modelopt.torch.utils.distributed import (
+    _reinstall_a2a_forwards,
+    distributed_forward_loop,
+    distributed_run,
+    distribution_plan,
+    local_rank,
+    shard_dataloader,
+)
 from modelopt.torch.utils.memory_monitor import launch_memory_monitor
 from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args
+from modelopt.torch.utils.moe import expert_parallel_mesh
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
     copy_non_model_files,
     ensure_local_checkpoint,
 )
-from modelopt.torch.utils.plugins.model_load_utils import (
-    load_pretrained,
-    parallel_load_and_prepare_fsdp2,
-)
+from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan
 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
@@ -194,6 +200,20 @@ def _to_device(value):
     return {k: _to_device(v) for k, v in batch.items()}
 
 
+def _shard_calib_dataloader(calib_dataloader, rank: int, world_size: int):
+    """This rank's share of the calibration data, in the loader type it came in."""
+    if isinstance(calib_dataloader, _DeviceDataLoader):
+        return _DeviceDataLoader(
+            shard_dataloader(calib_dataloader.dataloader, rank, world_size), calib_dataloader.device
+        )
+    if isinstance(calib_dataloader, DataLoader):
+        return shard_dataloader(calib_dataloader, rank, world_size)
+    raise NotImplementedError(
+        f"A {type(calib_dataloader).__name__} calibration loader cannot be split across this "
+        f"run's {world_size} ranks."
+    )
+
+
 def make_calib_dataloader(
     args: argparse.Namespace,
     language_model: torch.nn.Module,
@@ -275,12 +295,15 @@ def make_calib_dataloader(
             max_sample_length=args.calib_seq,
             device=device,
             include_labels=include_labels,
-            distributed=args.use_fsdp2,
-            sampler_kwargs=(
-                {"num_replicas": args.dist_state.world_size, "rank": args.dist_state.rank}
-                if args.use_fsdp2
-                else None
-            ),
+        )
+    world_size = distribution_plan(language_model).world_size
+    if world_size > 1 and calib_dataloader is not None:
+        # Every distributed calibration path is data-parallel over the world, so each rank calibrates
+        # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel
+        # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2
+        # shards the params but still feeds each rank its own data.
+        calib_dataloader = _shard_calib_dataloader(
+            calib_dataloader, args.dist_state.rank, world_size
         )
     return calib_dataloader, first_text_speech_dataset
 
@@ -323,37 +346,140 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None:
         )
 
 
+def _resolve_distribution(args: argparse.Namespace):
+    """Decide how this run distributes the model: ``(plan, hf_config)``.
+
+    ``resolve_distribution_plan`` decides, and --use_fsdp2 asks it for FSDP2 and no expert
+    parallelism. Without --use_fsdp2 a plan that FSDP2-shards the whole model is refused: the
+    FSDP2 run is the flag's, which checks what this script cannot do with it.
+    """
+    if args.use_fsdp2:
+        # Its checks read the config alone, and come first: deciding builds the model on the meta
+        # device, which some checkpoints they refuse (VILA, multimodal, pack-quantized) cannot do.
+        validate_fsdp2_supported(
+            args,
+            AutoConfig.from_pretrained(
+                args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
+            ),
+        )
+    plan, hf_config = resolve_distribution_plan(
+        args.local_checkpoint_path,
+        trust_remote_code=args.trust_remote_code,
+        ep_size=1 if args.use_fsdp2 else None,
+        dense="fsdp" if args.use_fsdp2 else None,
+    )
+    if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2:
+        raise RuntimeError(
+            f"{plan!r} FSDP2-shards the model, which this script does under --use_fsdp2. Pass it, "
+            "or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU; an MoE can instead shard its "
+            "experts (MODELOPT_EP_DEGREE above 1)."
+        )
+    if plan.world_size > 1 and _exports_tensorrt_llm_checkpoint(args, hf_model_type(hf_config)):
+        # That exporter reads the process count as Megatron's training tensor parallelism. Rank 0
+        # alone waits in its barrier for ranks that never arrive, and every rank would merge
+        # whole replicas as if they were shards. Refused before the load, like the layout above.
+        raise NotImplementedError(
+            f"{plan!r}: the TensorRT-LLM checkpoint export (encoder-decoder models, "
+            "--sparsity_fmt, int8_smoothquant) runs in a single process; launch hf_ptq without "
+            "torchrun for it."
+        )
+    if (plan.is_expert_parallel or plan.is_fsdp_dense) and args.vllm_fakequant_export:
+        # That exporter writes the whole model with save_pretrained from each process it runs in,
+        # and every rank of a sharded model exports: each would write its own part into the same
+        # directory. Rank 0 alone writes a replicated model; --use_fsdp2 refuses it while parsing.
+        raise NotImplementedError(
+            f"{plan!r}: --vllm_fakequant_export writes the whole model from one process, which "
+            "holds only its own part of a sharded model; launch hf_ptq without torchrun for it."
+        )
+    if plan.is_expert_parallel and args.experts_implementation is not None:
+        # The expert-parallel load runs the routed experts through torch_a2a, whatever backend the
+        # config names, so the flag would be dropped without a word.
+        raise NotImplementedError(
+            f"{plan!r}: --experts_implementation picks the experts backend a model is loaded with, "
+            "but the expert-parallel load runs the routed experts through torch_a2a whatever it "
+            "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a "
+            "single process."
+        )
+    if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
+        # Its search scores each rank's own calibration shard and reduces no score across ranks, so
+        # the ranks can choose different formats per layer: under expert parallelism, a checkpoint
+        # of mixed formats. --use_fsdp2 keeps main's warning for it.
+        raise NotImplementedError(
+            f"{plan!r}: AutoQuantize scores each rank's own calibration data and reduces no score "
+            "across ranks, so ranks could choose different formats per layer; launch hf_ptq "
+            "without torchrun for an AutoQuantize recipe."
+        )
+    if (
+        plan.world_size > 1
+        and args.recipe is not None
+        and layerwise_export_requested(recipe_layerwise_blocks(load_recipe(args.recipe)))
+    ):
+        # Every rank would write the same per-layer shard files. The exporter refuses that too, but
+        # only once calibration binds it, after the load.
+        raise NotImplementedError(
+            f"{plan!r}: the recipe's layerwise export writes its checkpoint shards from one "
+            "process; launch hf_ptq without torchrun for it."
+        )
+    if plan.world_size > 1 and not args.use_fsdp2:
+        # These configure loading one model in a single process: compressed as it loads, or
+        # spread over the node's GPUs, CPU and disk. An expert-parallel run never reaches those
+        # loaders, and a replicated one pins each rank's copy to the rank's GPU, which none of them
+        # works with. Refused before the load rather than ignored. --max_cpu_memory_gb and
+        # --max_gpu_memory_gb come only with --offload_folder.
+        for flag, given in (
+            ("--low_memory_mode", args.low_memory_mode),
+            ("--use_seq_device_map", args.use_seq_device_map),
+            ("--offload_folder", args.offload_folder is not None),
+        ):
+            if given:
+                raise NotImplementedError(
+                    f"{plan!r}: {flag} configures loading one model in a single process, but each "
+                    "rank of this run loads its own share or copy of it onto its own GPU. Drop it, "
+                    "or launch hf_ptq without torchrun to use it."
+                )
+    return plan, hf_config
+
+
+def _exports_tensorrt_llm_checkpoint(args: argparse.Namespace, model_type: str | None) -> bool:
+    """Whether the run exports through the deprecated TensorRT-LLM checkpoint exporter."""
+    return (
+        is_trtllm_enc_dec_export(model_type)
+        or args.sparsity_fmt != "dense"
+        or "int8_smoothquant" in args.qformat
+    )
+
+
 def load_model(args: argparse.Namespace):
     # If low memory mode is enabled, we compress the model while loading the HF checkpoint.
     calibration_only = False
-    if args.use_fsdp2:
-        hf_config = AutoConfig.from_pretrained(
-            args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
+    plan, hf_config = _resolve_distribution(args)
+    if plan.is_expert_parallel:
+        full_model = load_pretrained(
+            args.local_checkpoint_path,
+            distribution=plan,
+            hf_config=hf_config,
+            trust_remote_code=args.trust_remote_code,
+            attn_implementation=args.attn_implementation,
         )
+    elif plan.is_fsdp_dense:
         if args.experts_implementation is not None:
             # Set on the config, as build_meta_causal_lm sets a pre-fetched config's
             # attn_implementation: the meta model is built from it, with no kwarg of its own. The
             # setter carries it to the sub-configs, such as the text config a VLM's experts read.
             hf_config._experts_implementation = args.experts_implementation
-        validate_fsdp2_supported(args, hf_config)
-        full_model = parallel_load_and_prepare_fsdp2(
+        full_model = load_pretrained(
             args.local_checkpoint_path,
-            args.dist_state.device,
-            args.dist_state.rank,
-            args.dist_state.world_size,
+            distribution=plan,
+            hf_config=hf_config,
             trust_remote_code=args.trust_remote_code,
             cpu_offload=args.cpu_offload,
             attn_implementation=args.attn_implementation,
-            hf_config=hf_config,
-        )
-        # Calibration sizes its batch, and the export writes, by how the model is distributed.
-        set_distribution_plan(
-            full_model, DistributionPlan(args.dist_state.world_size, dense="fsdp")
         )
     elif args.specdec_offline_dataset is not None or not args.low_memory_mode:
         full_model = get_model(
             args.local_checkpoint_path,
             args.dist_state.device,
+            distribution=plan,
             gpu_mem_percentage=args.gpu_max_mem_percentage,
             trust_remote_code=args.trust_remote_code,
             use_seq_device_map=args.use_seq_device_map,
@@ -378,9 +504,6 @@ def load_model(args: argparse.Namespace):
         with init_quantized_weights(
             quant_cfg, gpu_mem_percentage=args.gpu_max_mem_percentage, quant_gemm=False
         ):
-            hf_config = AutoConfig.from_pretrained(
-                args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
-            )
             model_kwargs = {"trust_remote_code": args.trust_remote_code}
             if args.attn_implementation is not None:
                 model_kwargs["attn_implementation"] = args.attn_implementation
@@ -390,6 +513,7 @@ def load_model(args: argparse.Namespace):
                 args.local_checkpoint_path,
                 AutoModelForCausalLM,
                 model_type=hf_config.model_type,
+                distribution=plan,
                 **model_kwargs,
             )
         calibration_only = True
@@ -401,7 +525,7 @@ def load_model(args: argparse.Namespace):
     # model reports its own sub-config's type, which varies across transformers versions.
     model_type = hf_model_type(full_model)
 
-    if args.use_fsdp2:
+    if plan.is_fsdp_dense:
         device = args.dist_state.device
     else:
         device = full_model.device
@@ -568,7 +692,33 @@ def mono_quantize(
             # Image calibration batches contain multimodal kwargs (for example pixel_values).
             # They must be consumed by the complete VLM even when only a nested component is the
             # quantization target; the full forward still exercises that component's quantizers.
-            if args.calib_with_images:
+            if distribution_plan(full_model).is_expert_parallel:
+                # An expert-parallel run has to calibrate through distributed_forward_loop, not the
+                # plain one: it calls _reinstall_a2a_forwards first, which re-plants the
+                # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on
+                # transformers >= 5.16 (where the styles install a forward rather than a pre-hook).
+                # Without it the experts forward runs with process_group=None and raises
+                # "torch_a2a experts forward requires the EP `process_group`" at calibration step 0.
+                # On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op --
+                # which is why this is invisible on older transformers and fatal on newer. Other
+                # distributed runs (FSDP2, or a model replicated on every rank) keep the plain loop,
+                # which retries no batch that runs out of memory on a plan of several ranks: the
+                # error reaches distributed_run, which aborts the rank and with it the run. Image
+                # batches go through the VLM's own forward, as they do without expert parallelism:
+                # Nemotron-H Omni, for one, calibrates on images by default. Text batches go through
+                # the language model, but the loop re-plants over the model it is handed, so that is
+                # the full model: transformers keeps the expert-parallel plan on the model
+                # from_pretrained built, and a VLM's language model names no experts in its own.
+                calibrate_loop = distributed_forward_loop(
+                    full_model,
+                    calib_dataloader,
+                    device=args.dist_state.device,
+                    ep_mesh=expert_parallel_mesh(full_model),
+                    forward=create_vlm_calibration_forward(full_model)
+                    if args.calib_with_images
+                    else lambda batch: language_model(**batch),
+                )
+            elif args.calib_with_images:
                 calibrate_loop = create_vlm_calibration_loop(full_model, calib_dataloader)
             else:
                 calibrate_loop = create_forward_loop(
@@ -587,6 +737,12 @@ def mono_quantize(
             )
         else:
             language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop)
+        if distribution_plan(full_model).is_expert_parallel:
+            # The calibration loop re-plants the a2a forwards quantizer insertion drops (on
+            # transformers >= 5.16; a no-op before), but a config without static activation scales
+            # calibrates without it, and algorithm None (the MX formats) never calls it; the
+            # preview and the export forward through them.
+            _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model))
 
         # For VL models, update full_model to use the quantized language model
         if is_nemotron_vl_model and language_model is not full_model:
@@ -671,6 +827,11 @@ def export_quantized(
     language_model: torch.nn.Module,
     model_type: str | None,
 ):
+    plan = distribution_plan(full_model)
+    if not (args.dist_state.is_main or plan.is_fsdp_dense or plan.is_expert_parallel):
+        # Replicated over the ranks, the model is whole on each: rank 0 writes it, and the rest
+        # would race it into the same directory. A sharded model needs every rank for its part.
+        return
     # Not inference_mode: the FSDP2 path gathers full params in this context and
     # inference tensors break the subsequent state_dict() -> param.detach().
     with torch.no_grad():
@@ -679,6 +840,13 @@ def export_quantized(
         # Early exit for speculative decoding checkpoints
         # No tokenizer saving needed for spec ckpts
         if has_spec_opt(full_model):
+            if plan.is_fsdp_dense or plan.is_expert_parallel:
+                # export_speculative_decoding writes the whole checkpoint from one process and
+                # knows no plan: each rank of a sharded model would write it over the others.
+                raise NotImplementedError(
+                    f"Exporting a speculative-decoding model sharded as {plan!r} is not "
+                    "supported; quantize it in a single process."
+                )
             export_speculative_decoding(full_model, export_dir=export_path)
             args.checkpoint_exported = True
             print(f"Quantized speculative decoding checkpoint exported to: {export_path}")
@@ -686,17 +854,13 @@ def export_quantized(
 
         if is_multimodal_model(full_model):
             # Per-layer export writes its own config.json with quantization_config, which
-            # the source config would replace; it never writes a processor config.
-            if not args.layerwise_export:
+            # the source config would replace; it never writes a processor config. Rank 0's
+            # export writes config.json, so no other rank may write the source config after it.
+            if not args.layerwise_export and args.dist_state.is_main:
                 save_source_config(args, export_path)
 
         start_time = time.time()
-        is_tensorrt_llm_export = (
-            is_trtllm_enc_dec_export(model_type)
-            or args.sparsity_fmt != "dense"
-            or "int8_smoothquant" in args.qformat
-        )
-        if is_tensorrt_llm_export:
+        if _exports_tensorrt_llm_checkpoint(args, model_type):
             if (
                 args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1
             ) and args.qformat == "nvfp4_svdquant":
@@ -890,15 +1054,18 @@ def post_quantize(
                 raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e
             warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
     elif is_nemotron_vl_model and tokenizer is not None:
-        generated_ids_after_ptq = run_nemotron_vl_preview(
-            full_model,
-            tokenizer,
-            preview_input_ids,
-            args.local_checkpoint_path,
-            "after quantization",
-            allow_fallback=False,
-            trust_remote_code=args.trust_remote_code,
-        )
+        try:
+            generated_ids_after_ptq = run_nemotron_vl_preview(
+                full_model,
+                tokenizer,
+                preview_input_ids,
+                args.local_checkpoint_path,
+                "after quantization",
+                allow_fallback=False,
+                trust_remote_code=args.trust_remote_code,
+            )
+        except Exception as e:
+            raise RuntimeError(_PREVIEW_STOPS_THE_RUN) from e
     else:
         warnings.warn(
             "Llama4 Maverick generation after quantization has a bug. Skipping generation sample."
@@ -985,9 +1152,7 @@ def quantize_main(
     is_layerwise = any(cfg.get("enable", False) for cfg in layerwise_cfgs)
 
     # The value is a placeholder, replaced with --export_path below; presence is the switch.
-    args.layerwise_export = any(
-        cfg.get("export_dir") is not None and cfg.get("enable", False) for cfg in layerwise_cfgs
-    )
+    args.layerwise_export = layerwise_export_requested(layerwise_cfgs)
     if not args.layerwise_export and any(
         cfg.get("export_dir") is not None for cfg in layerwise_cfgs
     ):
@@ -1547,27 +1712,48 @@ def main(args: argparse.Namespace):
     random.seed(RAND_SEED)
     np.random.seed(RAND_SEED)
 
-    setup_distributed_args(args)
-
-    try:
-        # Entered inside the try: opening the run is fatal by design, and skipping
-        # cleanup_distributed would leave the other ranks blocked on the first collective
-        # until the NCCL timeout.
-        with mlflow_run(args):
-            # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole
-            # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local
-            # path.
-            args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(
-                args.pyt_ckpt_path
-            )
-
-            # launch a memory monitor to read the currently used GPU memory.
+    # A distributed launch owns its process group for the whole run, and the run takes this
+    # process's rank, size and device from it. Entered here rather than inside load_model so
+    # cleanup still runs if loading raises -- otherwise the surviving ranks block on the first
+    # collective until the NCCL timeout.
+    with distributed_run(timeout=timedelta(hours=3), device=args.device) as dist_state:
+        args.dist_state = dist_state
+        _main_inner(args)
+
+
+def _main_inner(args: argparse.Namespace):
+    with mlflow_run(args):
+        # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole
+        # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local
+        # path.
+        args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path)
+
+        # launch a memory monitor to read the currently used GPU memory. It reads every GPU on the
+        # node, so a distributed run starts one per node, on its local rank 0, rather than one per
+        # rank.
+        if args.dist_state.world_size == 1 or local_rank() == 0:
             launch_memory_monitor()
 
-            # Force eager execution for all model types.
-            torch.compiler.set_stance("force_eager")
+        # Force eager execution for all model types.
+        torch.compiler.set_stance("force_eager")
+
+        (
+            full_model,
+            language_model,
+            model_type,
+            calibration_only,
+            processor,
+            tokenizer,
+            device,
+        ) = load_model(args)
 
-            (
+        if args.sparsity_fmt != "dense":
+            # Sparse
+            sparsity_main(args, full_model, tokenizer, device)
+        else:
+            # Quantize
+            quantize_main(
+                args,
                 full_model,
                 language_model,
                 model_type,
@@ -1575,33 +1761,7 @@ def main(args: argparse.Namespace):
                 processor,
                 tokenizer,
                 device,
-            ) = load_model(args)
-
-            if args.sparsity_fmt != "dense":
-                # Sparse
-                sparsity_main(args, full_model, tokenizer, device)
-            else:
-                # Quantize
-                quantize_main(
-                    args,
-                    full_model,
-                    language_model,
-                    model_type,
-                    calibration_only,
-                    processor,
-                    tokenizer,
-                    device,
-                )
-    except BaseException:
-        # A rank that fails alone, such as in a sharded model's lockstep preview, leaves the others
-        # in a collective it never reaches, and cleanup's barrier and destroy_process_group would
-        # wait on them with the traceback unprinted. abort() prints it and exits, so the launcher
-        # stops the other ranks.
-        if args.dist_state.world_size > 1:
-            abort()
-        raise
-    finally:
-        cleanup_distributed(args)
+            )
 
 
 if __name__ == "__main__":
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index e4300892572..9181ea33ddf 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -26,7 +26,7 @@
 import yaml
 from _test_utils.mlflow import clean_env  # noqa: F401
 from _test_utils.torch.transformers_models import get_tiny_qwen3
-from transformers import Qwen3_5MoeConfig, Qwen3Config, Qwen3MoeConfig
+from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig
 
 from modelopt.recipe import load_recipe
 from modelopt.recipe.config import (
@@ -38,11 +38,7 @@
 from modelopt.torch.quantization import tensor_quant
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.utils import mlflow as mlflow_lib
-from modelopt.torch.utils.distributed import (
-    DistributionPlan,
-    distribution_plan,
-    set_distribution_plan,
-)
+from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
 from modelopt.torch.utils.mlflow import describe_run, run_tags
 
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
@@ -414,7 +410,8 @@ def distributed_load(checkpoint, *_args, **_kwargs):
         assert checkpoint == str(tmp_path)
         raise RuntimeError("distributed loader reached")
 
-    monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load)
+    # hf_ptq loads an FSDP2 plan through load_pretrained.
+    monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load)
 
     with pytest.raises(RuntimeError, match="distributed loader reached"):
         hf_ptq.load_model(args)
@@ -449,7 +446,7 @@ def distributed_load(*_args, hf_config, **_kwargs):
         assert hf_config.get_text_config()._experts_implementation == "eager"
         raise RuntimeError("distributed loader reached")
 
-    monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load)
+    monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load)
     with pytest.raises(RuntimeError, match="distributed loader reached"):
         hf_ptq.load_model(args)
 
@@ -475,7 +472,7 @@ def test_fsdp2_rejects_mtp_before_distributed_loading(monkeypatch, tmp_path, has
     def distributed_load(*args, **kwargs):
         raise RuntimeError("distributed loader reached")
 
-    monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load)
+    monkeypatch.setattr(hf_ptq, "load_pretrained", distributed_load)
     error, message = (
         (NotImplementedError, "Nemotron-H MTP")
         if has_mtp
@@ -485,27 +482,6 @@ def distributed_load(*args, **kwargs):
         hf_ptq.load_model(args)
 
 
-def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path):
-    """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an
-    FSDP2-sharded model without one is refused rather than taken for a single process's."""
-    Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path)
-    monkeypatch.setenv("RANK", "0")
-    hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2")
-    args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=2)
-    args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path)  # as main() sets them
-    model = torch.nn.Linear(1, 1)
-    monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *args, **kwargs: model)
-
-    def loaded(model):
-        raise RuntimeError("loaded")
-
-    monkeypatch.setattr(hf_ptq, "prepare_model_for_calibration", loaded)
-
-    with pytest.raises(RuntimeError, match="loaded"):
-        hf_ptq.load_model(args)
-    assert distribution_plan(model) == DistributionPlan(2, dense="fsdp")
-
-
 def test_mlflow_flag_defaults_the_experiment_name(monkeypatch):
     monkeypatch.setattr(getpass, "getuser", lambda: "tester")
     hf_ptq, args = _parse_hf_ptq_args(
@@ -1111,41 +1087,36 @@ def load_model(args):
     ]
 
 
-@pytest.mark.parametrize("world_size", [1, 2], ids=["single_process", "two_ranks"])
-def test_a_rank_that_fails_aborts_a_run_of_several_ranks(monkeypatch, world_size):
-    """The other ranks may wait in a collective the failed rank never reaches, where tearing the
-    process group down would wait on them with the error unprinted: a run of several ranks aborts
-    instead, which exits, so the launcher stops the others. A single process raises as before."""
-    hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "org/model")
-    events = []
-
-    class _AbortedError(BaseException):
-        """abort() exits the process; this stand-in raises instead, so the test sees it."""
-
-    def setup_distributed_args(args):
-        args.dist_state = SimpleNamespace(rank=0, world_size=world_size, is_main=True)
+@pytest.mark.parametrize(
+    ("rank", "node_rank", "launches"),
+    [(0, 0, True), (1, 1, False), (4, 0, True)],
+    ids=["rank_0", "other_rank", "first_rank_of_another_node"],
+)
+def test_one_rank_per_node_runs_the_gpu_memory_monitor(
+    monkeypatch, tmp_path, rank, node_rank, launches
+):
+    """The monitor reads every GPU on the node: one per rank would poll them all once per rank and
+    print the same report once per rank, and one on rank 0 alone would report node 0 alone."""
+    hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path))
+    args.dist_state = SimpleNamespace(is_main=rank == 0, rank=rank, world_size=8)
+    monkeypatch.setattr(hf_ptq, "local_rank", lambda: node_rank)
+    launched = []
 
-    def abort():
-        events.append("abort")
-        raise _AbortedError
+    class _LoadReachedError(Exception):
+        pass
 
     def load_model(args):
-        raise RuntimeError("this rank's load failed")
+        raise _LoadReachedError
 
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
     monkeypatch.setattr(torch.compiler, "set_stance", lambda *args, **kwargs: None)
-    monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: None)
-    monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", lambda path: (path, path))
-    monkeypatch.setattr(hf_ptq, "setup_distributed_args", setup_distributed_args)
-    monkeypatch.setattr(hf_ptq, "cleanup_distributed", lambda args: events.append("cleanup"))
-    monkeypatch.setattr(hf_ptq, "abort", abort)
+    monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: launched.append(True))
+    monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", lambda path: (None, path))
     monkeypatch.setattr(hf_ptq, "load_model", load_model)
 
-    with pytest.raises(_AbortedError if world_size > 1 else RuntimeError):
-        hf_ptq.main(args)
+    with pytest.raises(_LoadReachedError):
+        hf_ptq._main_inner(args)
 
-    # The real abort() exits before the teardown runs.
-    assert events == (["abort", "cleanup"] if world_size > 1 else ["cleanup"])
+    assert len(launched) == launches
 
 
 def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypatch, tmp_path):
diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
new file mode 100644
index 00000000000..62cc6ca6b2d
--- /dev/null
+++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
@@ -0,0 +1,963 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Expert-parallel PTQ wiring in ``examples/hf_ptq``.
+
+hf_ptq.py decides a distribution plan once, with ``resolve_distribution_plan``, and loads the model
+by it through ``load_pretrained``: expert-parallel, FSDP2-sharded under ``--use_fsdp2``, or whole
+through ``get_model``. Everything after reads the plan back from the model: every rank of a
+distributed plan calibrates on its own shard of the data, an expert-parallel one through the
+distributed forward loop, and a whole model is pinned to the rank's own GPU instead of ``"auto"``.
+
+Nothing here needs a checkpoint, and only the test of a VLM's all-to-all forwards needs a process
+group, of one gloo rank: elsewhere the process group's rank and size are answered the way a torchrun
+launch would, the configs are tiny real transformers configs served from memory, and the loaders are
+replaced by stand-ins that record how they were called.
+"""
+
+import contextlib
+import importlib
+import inspect
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+import transformers
+from _test_utils.torch.transformers_models import get_tiny_tokenizer
+from packaging.version import Version
+from torch.utils.data import DataLoader
+from transformers import AutoConfig, DeepseekV3Config, LlamaConfig
+
+import modelopt.torch.utils.distributed as dist
+from modelopt.torch.models import get_spec
+from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan
+from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS
+from modelopt.torch.utils.plugins import model_load_utils
+from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
+
+_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
+
+# Small enough that the dense part fits any GPU, so "auto" dense sharding always resolves to
+# replication whether or not the machine running the tests has one.
+_TINY = {
+    "hidden_size": 16,
+    "intermediate_size": 32,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 1,
+    "vocab_size": 32,
+}
+
+_CKPT = "/ckpt/model"
+
+
+class _LoaderReachedError(Exception):
+    """Raised by a loader stand-in once the call under test has been recorded."""
+
+
+class _RunFailedError(Exception):
+    """Raised by a stand-in for the body of a run that fails."""
+
+
+@pytest.fixture
+def hf_ptq(monkeypatch):
+    monkeypatch.syspath_prepend(str(_EXAMPLES_DIR))
+    return importlib.import_module("hf_ptq")
+
+
+@pytest.fixture(autouse=True)
+def _no_distribution_knobs(monkeypatch):
+    """A shell that exports the knobs would otherwise steer every decision below."""
+    monkeypatch.delenv("MODELOPT_EP_DEGREE", raising=False)
+    monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False)
+
+
+@pytest.fixture(autouse=True)
+def _checkpoint_is_local(monkeypatch):
+    """The checkpoint path names no files; resolving it must not reach the Hub."""
+    monkeypatch.setattr(
+        model_load_utils, "ensure_local_checkpoint", lambda path, group=None, **_: (None, path)
+    )
+
+
+def _gathered_alike(gathered, mine, group=None):
+    """Answer a gather as if every rank handed it what this rank did."""
+    gathered[:] = [mine] * len(gathered)
+
+
+@pytest.fixture
+def torchrun(monkeypatch):
+    """Run as rank ``rank`` of a ``world_size``-rank torchrun launch; returns its ``dist_state``.
+
+    Every rank acts as its group's first, deciding for itself, and reads the checkpoint as this one
+    does; how the ranks agree on what they decide is tested with the collectives.
+    """
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0)
+    monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None)
+    monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike)
+
+    def enter(world_size, rank=0):
+        monkeypatch.setattr(dist, "is_initialized", lambda: world_size > 1)
+        monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1)
+        monkeypatch.setattr(dist, "size", lambda group=None: world_size)
+        monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size)
+        monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0)
+        monkeypatch.setattr(dist, "local_rank", lambda: rank)
+        return _dist_state(world_size, rank)
+
+    return enter
+
+
+def _dist_state(world_size, rank=0):
+    """What distributed_run yields to ``rank`` of a ``world_size``-rank launch."""
+    return dist.DistributedState(
+        rank=rank, world_size=world_size, device=torch.device(f"cuda:{rank}")
+    )
+
+
+def _args(dist_state, **overrides):
+    """What load_model reads before it picks a loader, as main() leaves it."""
+    args = SimpleNamespace(
+        dist_state=dist_state,
+        use_fsdp2=False,
+        recipe=None,
+        local_checkpoint_path=_CKPT,
+        trust_remote_code=False,
+        attn_implementation=None,
+        experts_implementation=None,
+        sparsity_fmt="dense",
+        qformat="fp8",
+        low_memory_mode=False,
+        use_seq_device_map=False,
+        offload_folder=None,
+        vllm_fakequant_export=False,
+    )
+    vars(args).update(overrides)
+    return args
+
+
+def _serve(monkeypatch, hf_ptq, config):
+    """Answer every config read with ``config`` rather than a checkpoint on disk."""
+    monkeypatch.setattr(AutoConfig, "from_pretrained", lambda *args, **kwargs: config)
+    return config
+
+
+# DeepSeek-V3 is expert-parallel only from the transformers release the expert-parallel load needs;
+# below it, a multi-rank run of it is refused before the route the test exercises.
+needs_ep_transformers = pytest.mark.skipif(
+    Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS),
+    reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}",
+)
+
+
+def _deepseek_v3(n_routed_experts):
+    """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable.
+
+    Its one layer is an MoE layer: DeepSeek-V3's first layers are dense by default.
+    """
+    config = DeepseekV3Config(
+        architectures=["DeepseekV3ForCausalLM"],
+        n_routed_experts=n_routed_experts,
+        first_k_dense_replace=0,
+        moe_intermediate_size=8,
+        n_group=1,
+        topk_group=1,
+        num_experts_per_tok=1,
+        **_TINY,
+    )
+    config.base_model_ep_plan = None  # holds even if a later transformers ships a plan
+    return config
+
+
+# -------------------------------------------------------------------------------- the decision
+
+
+def test_use_fsdp2_plans_fsdp2_without_expert_parallelism(hf_ptq, monkeypatch, torchrun):
+    """--use_fsdp2 names its load, so even an explicit EP degree for an EP-capable MoE must not
+    divert the run to the expert-parallel one."""
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", "2")
+    _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8))
+    args = _args(torchrun(world_size=2), use_fsdp2=True)
+
+    assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(2, dense="fsdp")
+
+
+def test_use_fsdp2_refuses_a_checkpoint_before_building_it_to_decide(hf_ptq, monkeypatch, torchrun):
+    """--use_fsdp2's checks read the config alone. They come before the decision, which builds the
+    model on the meta device: some checkpoints they refuse cannot be built, and that build's error
+    would name neither --use_fsdp2 nor the way out."""
+    _serve(
+        monkeypatch,
+        hf_ptq,
+        LlamaConfig(
+            architectures=["LlamaForCausalLM"], quantization_config={"quant_method": "fp8"}, **_TINY
+        ),
+    )
+    monkeypatch.setattr(
+        model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built the model")
+    )
+
+    with pytest.raises(NotImplementedError, match="--use_fsdp2 does not support"):
+        hf_ptq._resolve_distribution(_args(torchrun(world_size=2), use_fsdp2=True))
+
+
+def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monkeypatch, torchrun):
+    """Without expert parallelism a plan that shards the dense part is an FSDP2 run, which this
+    script runs only under --use_fsdp2, the flag its FSDP2 checks hang off; so it must fail and point
+    there rather than load the model whole on every rank."""
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY))
+
+    with pytest.raises(RuntimeError, match="--use_fsdp2"):
+        hf_ptq._resolve_distribution(_args(torchrun(world_size=2)))
+
+
+@needs_ep_transformers
+@pytest.mark.parametrize("ep_degree", [2, 4])
+def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused(
+    hf_ptq, monkeypatch, torchrun, ep_degree
+):
+    """Below the world size each expert-parallel group is a replica, and FSDP2 shards the dense part
+    within each: a layout the export cannot write, so it must fail before the load and the
+    calibration rather than at the end. Expert parallelism over the whole world can still shard it."""
+    monkeypatch.setenv("MODELOPT_EP_DEGREE", str(ep_degree))
+    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+    _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8))
+    args = _args(torchrun(world_size=4))
+
+    if ep_degree == 4:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(4, ep_size=4, dense="fsdp")
+    else:
+        with pytest.raises(NotImplementedError, match=r"MODELOPT_EP_DEGREE.*MODELOPT_FSDP_SIZE=1"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    ("world_size", "options"),
+    [
+        (2, {"qformat": "int8_smoothquant"}),
+        (2, {"sparsity_fmt": "sparsegpt"}),
+        (1, {"sparsity_fmt": "sparsegpt"}),
+    ],
+    ids=["int8_smoothquant", "sparsity", "single_process"],
+)
+def test_a_tensorrt_llm_checkpoint_export_needs_a_single_process(
+    hf_ptq, monkeypatch, torchrun, world_size, options
+):
+    """That exporter reads the process count as Megatron's training tensor parallelism: launched
+    with torchrun, it hangs in a barrier on rank 0 alone, and on every rank it merges replicas as
+    shards. So it fails before the load and the calibration; a single process exports as before."""
+    _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY))
+    args = _args(torchrun(world_size=world_size), **options)
+
+    if world_size == 1:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+    else:
+        with pytest.raises(NotImplementedError, match="single process"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    ("world_size", "env", "use_fsdp2"),
+    [
+        pytest.param(2, {"MODELOPT_EP_DEGREE": "2"}, False, marks=needs_ep_transformers),
+        pytest.param(2, {}, False, marks=needs_ep_transformers),
+        (2, {}, True),
+        (1, {}, False),
+    ],
+    ids=["expert_parallel", "replicated", "use_fsdp2", "single_process"],
+)
+def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds(
+    hf_ptq, monkeypatch, torchrun, world_size, env, use_fsdp2
+):
+    """AutoQuantize scores each rank's own calibration shard and reduces no score across ranks, so
+    the ranks could choose different formats per layer. An expert-parallel or replicated run refuses
+    the recipe before the load; --use_fsdp2 keeps main's warning, and a single process runs it."""
+    for name, value in env.items():
+        monkeypatch.setenv(name, value)
+    _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8))
+    args = _args(
+        torchrun(world_size=world_size),
+        use_fsdp2=use_fsdp2,
+        recipe="general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits",
+    )
+
+    if world_size == 1 or use_fsdp2:
+        hf_ptq._resolve_distribution(args)
+    else:
+        with pytest.raises(NotImplementedError, match="AutoQuantize"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize("world_size", [2, 1], ids=["torchrun", "single_process"])
+def test_a_layerwise_export_recipe_is_refused_before_a_distributed_load(
+    hf_ptq, monkeypatch, torchrun, world_size
+):
+    """Every rank would write the same per-layer shard files, which the exporter refuses only once
+    calibration binds it, after the load; a launch of several ranks refuses the recipe before it.
+    A single process exports as before."""
+    _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY))
+    args = _args(
+        torchrun(world_size=world_size),
+        recipe="general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export",
+    )
+
+    if world_size == 1:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+    else:
+        with pytest.raises(NotImplementedError, match="layerwise export"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    "option",
+    [{"low_memory_mode": True}, {"use_seq_device_map": True}, {"offload_folder": "/offload"}],
+    ids=["low_memory_mode", "use_seq_device_map", "offload_folder"],
+)
+@pytest.mark.parametrize(
+    ("make_config", "world_size"),
+    [
+        pytest.param(lambda: _deepseek_v3(n_routed_experts=8), 2, marks=needs_ep_transformers),
+        (lambda: LlamaConfig(**_TINY), 2),
+        (lambda: LlamaConfig(**_TINY), 1),
+    ],
+    ids=["expert_parallel", "replicated", "single_process"],
+)
+def test_the_single_process_load_options_are_refused_on_the_routes_this_script_adds(
+    hf_ptq, monkeypatch, torchrun, make_config, world_size, option
+):
+    """Each rank of an expert-parallel or replicated run loads its own share or copy of the model
+    onto its own GPU, where these options of a single process's load would be ignored or would break
+    the load. Both routes refuse them before the load; a single process loads by them as before."""
+    _serve(monkeypatch, hf_ptq, make_config())
+    args = _args(torchrun(world_size=world_size), **option)
+
+    if world_size == 1:
+        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+    else:
+        (name,) = option
+        with pytest.raises(NotImplementedError, match=f"--{name}"):
+            hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    ("make_config", "world_size", "refused"),
+    [
+        pytest.param(
+            lambda: _deepseek_v3(n_routed_experts=8), 2, True, marks=needs_ep_transformers
+        ),
+        (lambda: LlamaConfig(**_TINY), 2, False),
+        (lambda: LlamaConfig(**_TINY), 1, False),
+    ],
+    ids=["expert_parallel", "replicated", "single_process"],
+)
+def test_the_vllm_fakequant_export_is_refused_for_a_sharded_model(
+    hf_ptq, monkeypatch, torchrun, make_config, world_size, refused
+):
+    """That exporter writes the whole model with save_pretrained from each process it runs in, so
+    every rank of an expert-parallel run would write its own experts into the same directory. Such a
+    run refuses it before the load; rank 0 alone writes a replicated model, and a single process
+    exports as before."""
+    _serve(monkeypatch, hf_ptq, make_config())
+    args = _args(torchrun(world_size=world_size), vllm_fakequant_export=True)
+
+    if refused:
+        with pytest.raises(NotImplementedError, match="--vllm_fakequant_export"):
+            hf_ptq._resolve_distribution(args)
+    else:
+        hf_ptq._resolve_distribution(args)
+
+
+@pytest.mark.parametrize(
+    ("make_config", "world_size", "use_fsdp2", "refused"),
+    [
+        pytest.param(
+            lambda: _deepseek_v3(n_routed_experts=8), 2, False, True, marks=needs_ep_transformers
+        ),
+        (lambda: _deepseek_v3(n_routed_experts=8), 2, True, False),
+        (lambda: LlamaConfig(**_TINY), 2, False, False),
+        (lambda: LlamaConfig(**_TINY), 1, False, False),
+    ],
+    ids=["expert_parallel", "use_fsdp2", "replicated", "single_process"],
+)
+def test_an_experts_implementation_is_refused_where_the_load_picks_the_backend(
+    hf_ptq, monkeypatch, torchrun, make_config, world_size, use_fsdp2, refused
+):
+    """The expert-parallel load runs the routed experts through torch_a2a whatever backend is asked
+    for, so an expert-parallel run refuses --experts_implementation before the load rather than drop
+    it. --use_fsdp2, a replicated run and a single process load by it."""
+    _serve(monkeypatch, hf_ptq, make_config())
+    args = _args(
+        torchrun(world_size=world_size), use_fsdp2=use_fsdp2, experts_implementation="eager"
+    )
+
+    if refused:
+        with pytest.raises(NotImplementedError, match="--experts_implementation"):
+            hf_ptq._resolve_distribution(args)
+    else:
+        hf_ptq._resolve_distribution(args)
+
+
+# ---------------------------------------------------------------------------------- load_model
+
+
+def _record_load_pretrained(hf_ptq, monkeypatch):
+    """Replace hf_ptq's load_pretrained and every other loader; return the recorded call."""
+    calls = []
+
+    def load_pretrained_stub(*args, **kwargs):
+        calls.append(inspect.signature(load_pretrained).bind(*args, **kwargs).arguments)
+        raise _LoaderReachedError
+
+    monkeypatch.setattr(hf_ptq, "load_pretrained", load_pretrained_stub)
+    monkeypatch.setattr(hf_ptq, "get_model", lambda *a, **k: pytest.fail("whole-model loader ran"))
+    return calls
+
+
+@pytest.mark.parametrize(
+    ("make_config", "knobs", "world_size", "options", "plan", "passed"),
+    [
+        pytest.param(
+            lambda: _deepseek_v3(n_routed_experts=8),
+            {"MODELOPT_FSDP_SIZE": "world"},
+            4,
+            {"trust_remote_code": True, "attn_implementation": "eager"},
+            DistributionPlan(4, ep_size=4, dense="fsdp"),
+            {"trust_remote_code": True, "attn_implementation": "eager"},
+            id="expert_parallel",
+            marks=needs_ep_transformers,
+        ),
+        pytest.param(
+            lambda: LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY),
+            {},
+            2,
+            {
+                "use_fsdp2": True,
+                "cpu_offload": True,
+                "specdec_offline_dataset": None,
+                "low_memory_mode": False,
+            },
+            DistributionPlan(2, dense="fsdp"),
+            {"trust_remote_code": False, "cpu_offload": True, "attn_implementation": None},
+            id="use_fsdp2",
+        ),
+    ],
+)
+def test_load_model_hands_load_pretrained_the_plan_and_the_config_it_was_decided_on(
+    hf_ptq, monkeypatch, torchrun, make_config, knobs, world_size, options, plan, passed
+):
+    """With an EP plan, or under --use_fsdp2, load_model hands load_pretrained the plan -- the
+    decided degree and dense sharding -- the config it was decided on, spec-prepared for EP, and
+    the caller's options, such as --use_fsdp2's CPU-offload choice."""
+    for name, value in knobs.items():
+        monkeypatch.setenv(name, value)
+    served = _serve(monkeypatch, hf_ptq, make_config())
+    calls = _record_load_pretrained(hf_ptq, monkeypatch)
+
+    with pytest.raises(_LoaderReachedError):
+        hf_ptq.load_model(_args(torchrun(world_size=world_size), **options))
+
+    (call,) = calls
+    assert call["ckpt_path"] == _CKPT
+    assert call["distribution"] == plan
+    hf_config = call["kwargs"].pop("hf_config")
+    assert call["kwargs"] == passed
+    if plan.is_expert_parallel:
+        # The spec's plan goes on the decision's own copy of the config.
+        assert hf_config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan
+        assert served.base_model_ep_plan is None
+    else:
+        assert hf_config is served
+
+
+def test_a_whole_model_loads_by_the_plan_through_get_model(hf_ptq, monkeypatch, torchrun):
+    """Data-parallel, every rank loads the whole model, by the plan, onto its own GPU."""
+    _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY))
+    loads = []
+
+    def get_model(ckpt_path, device, **kwargs):
+        loads.append((device, kwargs["distribution"]))
+        raise _LoaderReachedError
+
+    monkeypatch.setattr(hf_ptq, "get_model", get_model)
+    args = _args(
+        torchrun(world_size=4, rank=1),
+        specdec_offline_dataset=None,
+        low_memory_mode=False,
+        gpu_max_mem_percentage=0.8,
+        use_seq_device_map=False,
+        offload_folder=None,
+        max_cpu_memory_gb=None,
+        max_gpu_memory_gb=None,
+    )
+
+    with pytest.raises(_LoaderReachedError):
+        hf_ptq.load_model(args)
+
+    assert loads == [(torch.device("cuda:1"), DistributionPlan(4))]
+
+
+# ------------------------------------------------------------------ process group and dist_state
+
+
+def test_main_owns_the_process_group_around_the_whole_run(hf_ptq, monkeypatch):
+    """main() runs inside distributed_run, with a collective timeout above NCCL's 30-minute
+    default, and leaves it even when the run raises -- or the other ranks would wait in their next
+    collective until that timeout. The run reads its rank, size and device from the state
+    distributed_run yields, which puts a single process on --device."""
+    events, calls, seen = [], [], []
+    state = dist.DistributedState(rank=0, world_size=1, device=torch.device("cpu"))
+
+    @contextlib.contextmanager
+    def distributed_run(timeout=None, device="cuda"):
+        calls.append((timeout, device))
+        events.append("enter")
+        try:
+            yield state
+        finally:
+            events.append("exit")
+
+    def main_inner(args):
+        events.append("run")
+        seen.append(args.dist_state)
+        raise _RunFailedError
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(hf_ptq, "distributed_run", distributed_run)
+    monkeypatch.setattr(hf_ptq, "_main_inner", main_inner)
+
+    with pytest.raises(_RunFailedError):
+        hf_ptq.main(SimpleNamespace(use_fsdp2=False, device="cpu"))
+
+    assert events == ["enter", "run", "exit"]
+    timeout, device = calls[0]
+    assert timeout > timedelta(minutes=30)
+    assert device == "cpu"
+    assert len(seen) == 1
+    assert seen[0] is state
+
+
+# --------------------------------------------------------------------------------- calibration
+
+
+def test_distributed_calibration_gives_each_rank_its_own_shard(hf_ptq, monkeypatch):
+    """Every distributed path calibrates data-parallel: the ranks get equal, disjoint shards that
+    together cover the calibration set. A single process keeps the loader as built."""
+    loader = DataLoader(list(range(8)), batch_size=1)
+    monkeypatch.setattr(hf_ptq, "get_dataset_dataloader", lambda **kwargs: loader)
+    args = SimpleNamespace(
+        specdec_offline_dataset=None,
+        calib_with_images=False,
+        dataset=["cnn_dailymail"],
+        batch_size=1,
+        calib_size=[8],
+        calib_seq=16,
+    )
+    tokenizer = get_tiny_tokenizer()
+
+    def calib_loader(world_size, rank=0):
+        args.dist_state = _dist_state(world_size, rank)
+        model = torch.nn.Linear(2, 2)
+        set_distribution_plan(model, DistributionPlan(world_size))
+        built, _ = hf_ptq.make_calib_dataloader(
+            args, model, None, tokenizer, torch.device("cpu"), "llama"
+        )
+        return built
+
+    shards = [[int(b) for b in calib_loader(world_size=2, rank=r)] for r in range(2)]
+
+    assert len(shards[0]) == len(shards[1]) == 4
+    assert sorted(shards[0] + shards[1]) == list(range(8))
+    assert calib_loader(world_size=1) is loader
+
+
+@pytest.mark.parametrize("images", [False, True], ids=["text", "images"])
+@pytest.mark.parametrize(
+    ("plan", "loop"),
+    [(DistributionPlan(2, ep_size=2), "distributed"), (DistributionPlan(2, dense="fsdp"), "plain")],
+    ids=["expert_parallel", "fsdp2"],
+)
+def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_loop(
+    hf_ptq, monkeypatch, plan, loop, images
+):
+    """Only distributed_forward_loop re-plants the all-to-all forwards that quantizer insertion
+    clobbers, so an expert-parallel run must calibrate through it, over the mesh load_and_prepare_ep
+    loaded the experts over, image batches included: they run through the whole VLM's own forward
+    there (Nemotron-H Omni calibrates on images by default). The loop re-plants over the model it is
+    handed, the whole model, and text batches run through the language model. Without expert
+    parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry,
+    as on main."""
+    ep_mesh, vlm_forward = object(), object()
+    full_model, language_model = torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)
+    set_distribution_plan(full_model, plan)
+    chosen, quantized = [], []
+    monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False)
+    monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: True)
+    monkeypatch.setattr(hf_ptq, "expert_parallel_mesh", lambda full_model: ep_mesh)
+    monkeypatch.setattr(
+        hf_ptq,
+        "distributed_forward_loop",
+        lambda model, loader, **k: chosen.append((model, k["ep_mesh"], k["forward"])) or object(),
+    )
+    monkeypatch.setattr(hf_ptq, "create_vlm_calibration_forward", lambda model: vlm_forward)
+    monkeypatch.setattr(
+        hf_ptq,
+        "create_vlm_calibration_loop",
+        lambda model, loader: chosen.append("vlm") or object(),
+    )
+    monkeypatch.setattr(
+        hf_ptq, "create_forward_loop", lambda **k: chosen.append("plain") or object()
+    )
+    monkeypatch.setattr(
+        hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: quantized.append(forward_loop) or m
+    )
+    args = SimpleNamespace(
+        qformat="fp8",
+        calib_with_images=images,
+        dist_state=_dist_state(world_size=2),
+        layerwise_export=False,
+        specdec_offline_dataset=None,
+    )
+
+    hf_ptq.mono_quantize(
+        args,
+        {"algorithm": "max"},
+        full_model,
+        language_model,
+        model_type="qwen3_moe",
+        calibration_only=False,
+        calib_dataloader=object(),
+        is_nemotron_vl_model=False,
+    )
+
+    assert len(quantized) == 1
+    if loop == "plain":
+        assert chosen == ["vlm" if images else "plain"]
+    elif images:
+        assert chosen == [(full_model, ep_mesh, vlm_forward)]
+    else:
+        ((model, mesh, forward),) = chosen
+        assert (model, mesh) == (full_model, ep_mesh)
+        x = torch.ones(1, 2)
+        assert torch.equal(forward({"input": x}), language_model(x))
+
+
+@pytest.mark.parametrize(
+    "loop_handed", [False, True], ids=["calibration_free", "loop_never_called"]
+)
+@pytest.mark.parametrize(
+    ("plan", "replanted"),
+    [(DistributionPlan(2, ep_size=2), True), (DistributionPlan(2, dense="fsdp"), False)],
+    ids=["expert_parallel", "fsdp2"],
+)
+def test_an_expert_parallel_run_replants_the_all_to_all_forwards_after_quantizing(
+    hf_ptq, monkeypatch, plan, replanted, loop_handed
+):
+    """The calibration loop re-plants the all-to-all forwards that quantizer insertion drops only
+    when the calibration runs it: a config without static activation scales calibrates without the
+    loop, and algorithm None (the MX formats) never calls the loop it is handed. The preview and the
+    export forward through the experts either way, so an expert-parallel run re-plants them once
+    quantized; without expert parallelism there is nothing to re-plant."""
+    ep_mesh, full_model = object(), torch.nn.Linear(2, 2)
+    set_distribution_plan(full_model, plan)
+    events = []
+    monkeypatch.setattr(hf_ptq, "is_quantized", lambda m: False)
+    monkeypatch.setattr(hf_ptq, "need_calibration", lambda cfg: loop_handed)
+    monkeypatch.setattr(hf_ptq, "expert_parallel_mesh", lambda model: ep_mesh)
+    monkeypatch.setattr(hf_ptq, "distributed_forward_loop", lambda *args, **kwargs: object())
+    monkeypatch.setattr(hf_ptq, "create_forward_loop", lambda **kwargs: object())
+    monkeypatch.setattr(
+        hf_ptq.mtq, "quantize", lambda m, cfg, forward_loop: events.append("quantize") or m
+    )
+    monkeypatch.setattr(
+        hf_ptq,
+        "_reinstall_a2a_forwards",
+        lambda model, mesh: events.append(("replant", model, mesh)),
+    )
+    args = SimpleNamespace(
+        qformat="mxfp8",
+        calib_with_images=False,
+        dist_state=_dist_state(world_size=2),
+        layerwise_export=False,
+        specdec_offline_dataset=None,
+    )
+
+    hf_ptq.mono_quantize(
+        args,
+        {"algorithm": None},
+        full_model,
+        full_model,
+        model_type="qwen3_moe",
+        calibration_only=False,
+        calib_dataloader=object(),
+        is_nemotron_vl_model=False,
+    )
+
+    assert events == ["quantize"] + ([("replant", full_model, ep_mesh)] if replanted else [])
+
+
+@pytest.fixture
+def gloo_world_of_one(tmp_path):
+    """A one-rank gloo process group, for the all-to-all forwards to run on."""
+    torch.distributed.init_process_group(
+        "gloo", init_method=(tmp_path / "store").as_uri(), rank=0, world_size=1
+    )
+    yield torch.distributed.group.WORLD
+    torch.distributed.destroy_process_group()
+
+
+class _EpMesh:
+    """The expert-parallel submesh of a two-rank plan, over this process's one-rank group."""
+
+    def size(self):
+        return 2
+
+    def get_group(self):
+        return torch.distributed.group.WORLD
+
+
+def test_a_vlm_calibrated_on_text_gets_its_experts_all_to_all_forwards_back(
+    hf_ptq, gloo_world_of_one, monkeypatch
+):
+    """Text batches run through a VLM's language model, but transformers keeps the expert-parallel
+    plan on the model from_pretrained built, and the language model's own plan names no experts.
+    Re-planted over the language model, the experts' all-to-all forward that quantizer insertion
+    drops (transformers 5.16 and later) stayed dropped, and the first batch raised for the missing
+    EP process group."""
+    pytest.importorskip("transformers", minversion="5.16")
+    from transformers import Qwen3VLMoeConfig, Qwen3VLMoeForConditionalGeneration
+    from transformers.distributed.configuration_utils import DistributedConfig
+    from transformers.distributed.tensor_parallel import ALL_PARALLEL_STYLES, _get_parameter_tp_plan
+
+    from modelopt.torch.utils.moe import register_torch_a2a
+
+    register_torch_a2a()
+    config = Qwen3VLMoeConfig(
+        text_config={
+            **_TINY,
+            "head_dim": 8,
+            "num_experts": 4,
+            "num_experts_per_tok": 2,
+            "moe_intermediate_size": 8,
+            "decoder_sparse_step": 1,
+        },
+        vision_config={
+            "depth": 1,
+            "hidden_size": 16,
+            "intermediate_size": 16,
+            "num_heads": 2,
+            "out_hidden_size": 16,
+            "patch_size": 2,
+            "spatial_merge_size": 1,
+            "temporal_patch_size": 1,
+            "deepstack_visual_indexes": [],
+        },
+    )
+    # The model as load_and_prepare_ep leaves it: its plan's all-to-all styles on the language
+    # model's config and expert parallelism on the top one, the forwards from_pretrained installs
+    # for them, the experts dispatching through the all-to-all, and their seeded parallel state.
+    config.text_config.base_model_ep_plan = {
+        "layers.*.mlp.gate": "torch_a2a_router",
+        "layers.*.mlp.experts": "torch_a2a_experts",
+    }
+    config.distributed_config = DistributedConfig(enable_expert_parallel=True)
+    model = Qwen3VLMoeForConditionalGeneration(config).eval()
+    mesh = _EpMesh()
+    for name, module in model.named_modules():
+        style = _get_parameter_tp_plan(parameter_name=name, tp_plan=model.tp_plan, is_weight=False)
+        if style in ("torch_a2a_router", "torch_a2a_experts"):
+            ALL_PARALLEL_STYLES[style].install_forward(module, mesh)
+    model.config.text_config._experts_implementation = "torch_a2a"
+    model._device_mesh = mesh
+    set_distribution_plan(model, DistributionPlan(2, ep_size=2))
+    # One process cannot build the plan's two-rank expert-parallel row, which calibration checks the
+    # experts' group against, so the one-rank group answers as that row.
+    group_ranks = torch.distributed.get_process_group_ranks
+    monkeypatch.setattr(
+        torch.distributed,
+        "get_process_group_ranks",
+        lambda group: [0, 1] if group is gloo_world_of_one else group_ranks(group),
+    )
+    model_load_utils._seed_parallel_states(
+        model,
+        [name for name, _ in model.named_modules() if name.endswith("mlp.experts")],
+        dense_state=ParallelState(data_parallel_group=-1),
+        expert_state=ParallelState(
+            data_parallel_group=-1, expert_model_parallel_group=gloo_world_of_one
+        ),
+    )
+    language_model = model.model.language_model
+    args = SimpleNamespace(
+        qformat="int8",
+        calib_with_images=False,
+        dist_state=dist.DistributedState(rank=0, world_size=2, device=torch.device("cpu")),
+        layerwise_export=False,
+        specdec_offline_dataset=None,
+    )
+    batch = {"input_ids": torch.tensor([[1, 2, 3, 4]]), "attention_mask": torch.ones(1, 4)}
+
+    hf_ptq.mono_quantize(
+        args,
+        hf_ptq.mtq.INT8_DEFAULT_CFG,
+        model,
+        language_model,
+        model_type="qwen3_vl_moe",
+        calibration_only=False,
+        calib_dataloader=[batch],
+        is_nemotron_vl_model=False,
+    )
+
+    assert language_model.layers[0].mlp.experts.down_proj_input_quantizer.amax is not None
+
+
+def test_the_offline_specdec_loader_is_split_across_ranks_too(hf_ptq):
+    """It wraps its DataLoader to move each batch onto the device; each rank still calibrates on
+    its own share, as it would from the bare loader."""
+
+    def loader():
+        return DataLoader(list(range(8)), batch_size=1, collate_fn=lambda b: {"x": torch.tensor(b)})
+
+    wrapped = hf_ptq._DeviceDataLoader(loader(), torch.device("cpu"))
+
+    shard = hf_ptq._shard_calib_dataloader(wrapped, 1, 2)
+
+    assert isinstance(shard, hf_ptq._DeviceDataLoader)
+    assert [b["x"].tolist() for b in shard] == [
+        b["x"].tolist() for b in dist.shard_dataloader(loader(), 1, 2)
+    ]
+    with pytest.raises(NotImplementedError, match="cannot be split"):
+        hf_ptq._shard_calib_dataloader(iter(loader()), 1, 2)
+
+
+@pytest.mark.parametrize(
+    ("plan", "nemotron_vl", "raises"),
+    [
+        (DistributionPlan(4), False, False),
+        (DistributionPlan(4, dense="fsdp"), False, False),
+        (DistributionPlan(4, ep_size=4), False, True),
+        (DistributionPlan(4, ep_size=4), True, True),
+    ],
+    ids=["replicated", "fsdp2", "expert_parallel", "expert_parallel_nemotron_vl"],
+)
+def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run(
+    hf_ptq, monkeypatch, plan, nemotron_vl, raises
+):
+    """A rank that skipped an expert-parallel model's preview would leave the others decoding in
+    lockstep, blocked in the next step's collective, so the failure is raised, which aborts the run,
+    with the preview's error as its cause. A model every rank holds whole keeps the best-effort
+    skip and exports, and so does --use_fsdp2, as on main."""
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("the preview failed on this rank")
+
+    class _FailingModel(torch.nn.Module):
+        generate = fail
+
+    model = _FailingModel()
+    set_distribution_plan(model, plan)
+    exported = []
+    monkeypatch.setattr(hf_ptq, "export_quantized", lambda *args, **kwargs: exported.append(True))
+    monkeypatch.setattr(hf_ptq, "run_nemotron_vl_preview", fail)  # as vlm_utils raises in lockstep
+    ids = torch.zeros(1, 4, dtype=torch.long)
+    args = SimpleNamespace(
+        specdec_offline_dataset=None,
+        verbose=False,
+        local_checkpoint_path=_CKPT,
+        trust_remote_code=False,
+    )
+    tokenizer = SimpleNamespace()
+    run = lambda: hf_ptq.post_quantize(  # noqa: E731
+        args, model, model, "llama", tokenizer, None, ids, None, ids, nemotron_vl, None, None
+    )
+
+    if raises:
+        with pytest.raises(RuntimeError) as stopped:
+            run()
+        assert "failed on this rank" in str(stopped.value.__cause__)
+    else:
+        with pytest.warns(UserWarning, match="sanity check failed"):
+            run()
+    assert exported == ([] if raises else [True])
+
+
+# ------------------------------------------------------------------------------------- export
+
+
+def _planned(plan):
+    """A module carrying ``plan``, as the load records it."""
+    model = torch.nn.Linear(1, 1)
+    set_distribution_plan(model, plan)
+    return model
+
+
+@pytest.mark.parametrize(
+    ("plan", "rank", "exports"),
+    [
+        (DistributionPlan(4), 0, True),
+        (DistributionPlan(4), 1, False),
+        (DistributionPlan(4, dense="fsdp"), 1, True),
+        (DistributionPlan(4, ep_size=4), 1, True),
+    ],
+    ids=["replicated_rank0", "replicated_rank1", "fsdp2_rank1", "expert_parallel_rank1"],
+)
+def test_only_rank_0_writes_a_replicated_model(hf_ptq, monkeypatch, plan, rank, exports):
+    """A replicated model is whole on every rank, whose exports would write one checkpoint into the
+    same directory over each other; a sharded one needs every rank for its part. The source config
+    is rank 0's either way: another rank's, written after rank 0's export, would drop the
+    quantization config."""
+    written = []
+    monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: False)
+    monkeypatch.setattr(hf_ptq, "is_multimodal_model", lambda model: True)
+    monkeypatch.setattr(hf_ptq, "save_source_config", lambda args, path: written.append("config"))
+    monkeypatch.setattr(
+        hf_ptq, "export_hf_checkpoint", lambda model, export_dir: written.append("checkpoint")
+    )
+    args = SimpleNamespace(
+        dist_state=_dist_state(plan.world_size, rank),
+        export_path="/export",
+        layerwise_export=False,
+        sparsity_fmt="dense",
+        qformat="fp8",
+        inference_tensor_parallel=1,
+        inference_pipeline_parallel=1,
+        vllm_fakequant_export=False,
+    )
+
+    model = _planned(plan)
+    hf_ptq.export_quantized(args, model, model, "qwen3_moe")
+
+    assert written == ["config"] * (rank == 0) + ["checkpoint"] * exports
+
+
+@pytest.mark.parametrize(
+    "plan", [DistributionPlan(4, dense="fsdp"), DistributionPlan(4, ep_size=4)], ids=["fsdp2", "ep"]
+)
+def test_a_sharded_speculative_decoding_model_is_refused_at_export(hf_ptq, monkeypatch, plan):
+    """The speculative-decoding exporter writes the whole checkpoint from one process and has no
+    rank gate, so every rank of a sharded model would write it over the others: it is refused on
+    every rank instead."""
+    written = []
+    monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: True)
+    monkeypatch.setattr(
+        hf_ptq, "export_speculative_decoding", lambda model, export_dir: written.append(export_dir)
+    )
+    args = SimpleNamespace(dist_state=_dist_state(plan.world_size, 1), export_path="/export")
+
+    model = _planned(plan)
+    with pytest.raises(NotImplementedError, match="speculative-decoding model sharded"):
+        hf_ptq.export_quantized(args, model, model, "qwen3_moe")
+    assert written == []
diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
index 48165090687..20c86f2f1a8 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
@@ -20,6 +20,8 @@
 import pytest
 import torch
 
+from modelopt.torch.utils.distributed import DistributionPlan
+
 _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq"
 
 
@@ -65,6 +67,10 @@ def test_image_calibration_model_target_follows_recipe(
     )
 
     monkeypatch.setattr(hf_ptq, "get_model", lambda *args, **kwargs: full_model)
+    # A single process's plan, decided without reading the "dummy" checkpoint.
+    monkeypatch.setattr(
+        hf_ptq, "resolve_distribution_plan", lambda *args, **kwargs: (DistributionPlan(), None)
+    )
     # A VLM's extracted language model reports its own sub-config type; hf_ptq must keep the root's.
     monkeypatch.setattr(
         hf_ptq,
diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py
index a820158fa9a..72c27541b98 100644
--- a/tests/examples/hf_ptq/test_llm_ptq.py
+++ b/tests/examples/hf_ptq/test_llm_ptq.py
@@ -13,10 +13,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import dataclasses
+import json
+import os
 
 import pytest
+import torch
+import transformers
 from _test_utils.examples.hf_ptq_utils import PTQCommand
 from _test_utils.examples.models import MIXTRAL_PATH, T5_PATH, TINY_LLAMA_PATH
+from _test_utils.examples.run_command import run_example_command
+from _test_utils.torch.distributed.utils import get_free_port
+from _test_utils.torch.misc import minimum_gpu
+from _test_utils.torch.transformers_models import create_tiny_qwen3_moe_dir
+from packaging.version import Version
+from safetensors.torch import load_file
 
 
 def _run_sharing_the_gpu(command: PTQCommand, model_path: str):
@@ -85,3 +95,92 @@ def test_ptq_mixtral(command):
 @pytest.mark.usefixtures("shared_gpu")
 def test_ptq_llama(command):
     _run_sharing_the_gpu(command, TINY_LLAMA_PATH)
+
+
+@minimum_gpu(2)
+@pytest.mark.usefixtures("exclusive_gpu")
+@pytest.mark.parametrize(
+    ("launch_args", "launch_env", "min_transformers"),
+    [
+        # qwen3_moe's base_model_ep_plan ships in transformers 5.7.
+        ([], {"MODELOPT_EP_DEGREE": "2"}, "5.7"),
+        (["--use_fsdp2"], {}, None),
+    ],
+    ids=["expert_parallel", "fsdp2"],
+)
+def test_ptq_qwen3_moe_sharded(tmp_path, launch_args, launch_env, min_transformers):
+    """A torchrun launch shards a tiny Qwen3-MoE over two ranks, previews, calibrates and exports.
+
+    The model's experts are sharded over the ranks, or the whole model is with FSDP2. Each rank holds
+    and writes only its own share, so every expert of every layer must come out FP8-quantized: a rank
+    whose share is lost, or written under the wrong expert index, fails it. The previews generate
+    before and after calibration: a sharded preview broke in three ways (ranks finishing at different
+    steps hung in a collective, Transformers refused an FSDP2-wrapped class's experts switch, and the
+    switch decoded unquantized experts), and only a real launch reaches it.
+    """
+    if torch.cuda.get_device_capability() < (8, 9):
+        pytest.skip("Requires sm89 or higher")
+    if min_transformers and Version(transformers.__version__) < Version(min_transformers):
+        pytest.skip(f"Requires transformers {min_transformers}")
+    num_layers, num_experts = 2, 4
+    model_dir = create_tiny_qwen3_moe_dir(
+        tmp_path,
+        with_tokenizer=True,
+        hidden_size=128,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        intermediate_size=128,
+        moe_intermediate_size=128,
+        num_experts=num_experts,
+        num_experts_per_tok=2,
+        num_hidden_layers=num_layers,
+        max_position_embeddings=128,
+    )
+    # A local calibration set: no dataset download.
+    calib = tmp_path / "calib.jsonl"
+    calib.write_text(
+        "".join(
+            json.dumps({"text": f"Sample {i}: the quick brown fox jumps over the lazy dog."}) + "\n"
+            for i in range(8)
+        )
+    )
+    export_dir = tmp_path / "export"
+    run_example_command(
+        [
+            "torchrun",
+            "--nproc_per_node=2",
+            f"--master_port={get_free_port()}",
+            "hf_ptq.py",
+            "--pyt_ckpt_path",
+            str(model_dir),
+            "--recipe",
+            "general/ptq/fp8_default-kv_fp8_cast",
+            "--dataset",
+            str(calib),
+            "--calib_size",
+            "8",
+            "--batch_size",
+            "2",
+            "--calib_seq",
+            "64",
+            "--no-verbose",
+            "--export_path",
+            str(export_dir),
+            *launch_args,
+        ],
+        "hf_ptq",
+        env={**os.environ, **launch_env},
+    )
+
+    config = json.loads((export_dir / "config.json").read_text())
+    assert config["quantization_config"]["quant_algo"] == "FP8"
+    state_dict = {}
+    for shard in export_dir.glob("*.safetensors"):
+        state_dict.update(load_file(shard))
+    for layer in range(num_layers):
+        for expert in range(num_experts):
+            for proj in ("gate_proj", "up_proj", "down_proj"):
+                name = f"model.layers.{layer}.mlp.experts.{expert}.{proj}"
+                assert state_dict[f"{name}.weight"].dtype == torch.float8_e4m3fn, name
+                assert f"{name}.weight_scale" in state_dict, name
diff --git a/tests/gpu/torch/quantization/test_fsdp2_autoquant.py b/tests/gpu/torch/quantization/test_fsdp2_autoquant.py
index e2f45ee9507..3a259465a57 100644
--- a/tests/gpu/torch/quantization/test_fsdp2_autoquant.py
+++ b/tests/gpu/torch/quantization/test_fsdp2_autoquant.py
@@ -113,6 +113,9 @@ def _run_kv_autoquant_example(rank, size, checkpoint_dir, search_dir, composed):
         calib_with_images=False,
         specdec_offline_dataset=None,
         low_memory_mode=False,
+        sparsity_fmt="dense",
+        qformat="fp8",
+        vllm_fakequant_export=False,
         dataset=["local-test-data"],
         calib_size=[2 * size],
         batch_size=1,

From 0c1245306738ea6e4eff9028bce5eed4593d5f8c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:08:40 +0000
Subject: [PATCH 14/24] Say in the changelog what a multi-rank torchrun launch
 of hf_ptq now does

The Backward Breaking entry named only the --device refusal. A torchrun launch
of several ranks without --use_fsdp2 used to run one independent copy of
hf_ptq per process; it now runs one distributed job, expert-parallel for an
MoE that can shard its experts and replicated otherwise, and refuses the
options its routes cannot honor. The entry now says so.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index e9b5aa0d28d..1c068e7e34b 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -52,7 +52,7 @@ Changelog
 **Backward Breaking Changes**
 
 - ``QDQAutotunerBase.export_onnx(output_path=...)`` now returns ``None`` instead of serialized bytes. Omit ``output_path`` to keep the bytes-returning behavior, which remains subject to protobuf size limits.
-- ``examples/hf_ptq/hf_ptq.py`` launched by ``torchrun`` on more than one rank 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`` launched by ``torchrun`` on more than one rank without ``--use_fsdp2`` now runs one distributed job instead of one independent run per process: an MoE that can shard its experts runs expert-parallel, any other model is replicated with each rank calibrating its own share of the data, and only rank 0 exports a replicated model. Such a launch refuses ``--low_memory_mode``, ``--use_seq_device_map``, ``--offload_folder``, AutoQuantize recipes and, under expert parallelism, ``--vllm_fakequant_export`` and ``--experts_implementation``, and any multi-rank launch now refuses layerwise-export recipes, the TensorRT-LLM checkpoint export and a ``--device`` other than ``cuda``: launch without ``torchrun`` for them.
 - ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). It also raises when a calibration batch runs out of GPU memory, instead of retrying it in smaller pieces: set ``--batch_size`` to one that fits.
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later.

From cb704c21eb708202eba3cd463512a0e16b2a451d Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 06:00:37 +0000
Subject: [PATCH 15/24] Drop the calibration test's claim that FSDP2 keeps the
 out-of-memory retry

#2752 dropped the plain loop's split-on-OOM retry for a plan of several ranks,
so the docstring now says only that FSDP2 keeps the plain loop, as there is
nothing to re-plant without expert parallelism.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
index 62cc6ca6b2d..f5ade5dde92 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
@@ -599,8 +599,7 @@ def test_only_an_expert_parallel_run_calibrates_through_the_distributed_forward_
     loaded the experts over, image batches included: they run through the whole VLM's own forward
     there (Nemotron-H Omni calibrates on images by default). The loop re-plants over the model it is
     handed, the whole model, and text batches run through the language model. Without expert
-    parallelism there is nothing to re-plant: FSDP2 keeps the plain loop and its split-on-OOM retry,
-    as on main."""
+    parallelism there is nothing to re-plant, so FSDP2 keeps the plain loop."""
     ep_mesh, vlm_forward = object(), object()
     full_model, language_model = torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)
     set_distribution_plan(full_model, plan)

From 580028dbb37a7bd658452f61fcddc903235228bb Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 19:12:00 +0000
Subject: [PATCH 16/24] Drop the failed-preview test's --use_fsdp2 row, which
 [2/3] now tests

[2/3] stops the run when a sharded model's preview fails, --use_fsdp2
included, so this PR's expert-parallel-only raise and its message went down
into it. The row that expected --use_fsdp2 to skip the failure and export no
longer holds, and [2/3]'s test covers the case.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
index f5ade5dde92..729cfc37849 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
@@ -846,11 +846,10 @@ def loader():
     ("plan", "nemotron_vl", "raises"),
     [
         (DistributionPlan(4), False, False),
-        (DistributionPlan(4, dense="fsdp"), False, False),
         (DistributionPlan(4, ep_size=4), False, True),
         (DistributionPlan(4, ep_size=4), True, True),
     ],
-    ids=["replicated", "fsdp2", "expert_parallel", "expert_parallel_nemotron_vl"],
+    ids=["replicated", "expert_parallel", "expert_parallel_nemotron_vl"],
 )
 def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run(
     hf_ptq, monkeypatch, plan, nemotron_vl, raises
@@ -858,7 +857,7 @@ def test_a_failed_preview_of_an_expert_parallel_model_stops_the_run(
     """A rank that skipped an expert-parallel model's preview would leave the others decoding in
     lockstep, blocked in the next step's collective, so the failure is raised, which aborts the run,
     with the preview's error as its cause. A model every rank holds whole keeps the best-effort
-    skip and exports, and so does --use_fsdp2, as on main."""
+    skip and exports."""
 
     def fail(*args, **kwargs):
         raise RuntimeError("the preview failed on this rank")

From 4c306cb00466a55592b1aa1cef8be8795b1f5900 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:56:33 +0000
Subject: [PATCH 17/24] Move hf_ptq's distribution decision into example_utils

resolve_distribution, next to validate_fsdp2_supported, holds the
example-specific checks of a distributed hf_ptq run around the library's
resolve_distribution_plan, with shorter comments; hf_ptq.py keeps the PTQ
workflow. exports_tensorrt_llm_checkpoint moves with it. The tests call and
patch the names where they now live.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/example_utils.py              |  99 ++++++++++++++-
 examples/hf_ptq/hf_ptq.py                     | 114 +-----------------
 tests/examples/hf_ptq/test_hf_ptq_args.py     |   1 -
 .../hf_ptq/test_hf_ptq_expert_parallel.py     |  34 +++---
 .../examples/hf_ptq/test_hf_ptq_model_type.py |   3 +-
 .../hf_ptq/test_hf_ptq_vision_quantization.py |   4 +-
 6 files changed, 123 insertions(+), 132 deletions(-)

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index e67466fd166..a76611e8735 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -30,6 +30,7 @@
 import transformers
 from accelerate import infer_auto_device_map, init_empty_weights
 from accelerate.utils import get_max_memory
+from autoquant_utils import _recipe_is_auto_quantize
 from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static
 from safetensors import safe_open
 from transformers import (
@@ -56,7 +57,11 @@
 from modelopt.torch.utils import distributed as dist_utils
 from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
 from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run
-from modelopt.torch.utils.plugins.model_load_utils import is_fsdp2_text_only_vlm, load_pretrained
+from modelopt.torch.utils.plugins.model_load_utils import (
+    is_fsdp2_text_only_vlm,
+    load_pretrained,
+    resolve_distribution_plan,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -120,6 +125,98 @@ def validate_fsdp2_supported(args, config):
         )
 
 
+def exports_tensorrt_llm_checkpoint(args, model_type: str | None) -> bool:
+    """Whether the run exports through the deprecated TensorRT-LLM checkpoint exporter."""
+    return (
+        is_trtllm_enc_dec_export(model_type)
+        or args.sparsity_fmt != "dense"
+        or "int8_smoothquant" in args.qformat
+    )
+
+
+def resolve_distribution(args) -> tuple[DistributionPlan, Any]:
+    """Decide how hf_ptq distributes the model, refusing what that run cannot do: ``(plan, config)``.
+
+    :func:`resolve_distribution_plan` decides; ``--use_fsdp2`` asks it for FSDP2 without expert
+    parallelism.
+    """
+    if args.use_fsdp2:
+        # Before the decision, which builds the model on the meta device: some checkpoints these
+        # checks refuse (VILA, multimodal, pack-quantized) cannot be built there.
+        validate_fsdp2_supported(
+            args,
+            AutoConfig.from_pretrained(
+                args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
+            ),
+        )
+    plan, hf_config = resolve_distribution_plan(
+        args.local_checkpoint_path,
+        trust_remote_code=args.trust_remote_code,
+        ep_size=1 if args.use_fsdp2 else None,
+        dense="fsdp" if args.use_fsdp2 else None,
+    )
+    if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2:
+        raise RuntimeError(
+            f"{plan!r} FSDP2-shards the model, which this script does under --use_fsdp2. Pass it, "
+            "or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU; an MoE can instead shard its "
+            "experts (MODELOPT_EP_DEGREE above 1)."
+        )
+    if plan.world_size > 1 and exports_tensorrt_llm_checkpoint(args, hf_model_type(hf_config)):
+        # That exporter takes the process count for tensor parallelism.
+        raise NotImplementedError(
+            f"{plan!r}: the TensorRT-LLM checkpoint export (encoder-decoder models, "
+            "--sparsity_fmt, int8_smoothquant) runs in a single process; launch hf_ptq without "
+            "torchrun for it."
+        )
+    if (plan.is_expert_parallel or plan.is_fsdp_dense) and args.vllm_fakequant_export:
+        # It writes the whole model from every rank that exports, and every rank of a sharded
+        # model exports.
+        raise NotImplementedError(
+            f"{plan!r}: --vllm_fakequant_export writes the whole model from one process, which "
+            "holds only its own part of a sharded model; launch hf_ptq without torchrun for it."
+        )
+    if plan.is_expert_parallel and args.experts_implementation is not None:
+        # The expert-parallel load runs the experts through torch_a2a whatever the flag names.
+        raise NotImplementedError(
+            f"{plan!r}: --experts_implementation picks the experts backend a model is loaded with, "
+            "but the expert-parallel load runs the routed experts through torch_a2a whatever it "
+            "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a "
+            "single process."
+        )
+    if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
+        # Each rank would score its own data, and could pick other formats per layer.
+        raise NotImplementedError(
+            f"{plan!r}: AutoQuantize scores each rank's own calibration data and reduces no score "
+            "across ranks, so ranks could choose different formats per layer; launch hf_ptq "
+            "without torchrun for an AutoQuantize recipe."
+        )
+    if (
+        plan.world_size > 1
+        and args.recipe is not None
+        and layerwise_export_requested(recipe_layerwise_blocks(load_recipe(args.recipe)))
+    ):
+        # Every rank would write the same shard files.
+        raise NotImplementedError(
+            f"{plan!r}: the recipe's layerwise export writes its checkpoint shards from one "
+            "process; launch hf_ptq without torchrun for it."
+        )
+    if plan.world_size > 1 and not args.use_fsdp2:
+        # Options for loading one model in a single process, refused rather than ignored.
+        # --max_cpu_memory_gb and --max_gpu_memory_gb come only with --offload_folder.
+        for flag, given in (
+            ("--low_memory_mode", args.low_memory_mode),
+            ("--use_seq_device_map", args.use_seq_device_map),
+            ("--offload_folder", args.offload_folder is not None),
+        ):
+            if given:
+                raise NotImplementedError(
+                    f"{plan!r}: {flag} configures loading one model in a single process, but each "
+                    "rank of this run loads its own share or copy of it onto its own GPU. Drop it, "
+                    "or launch hf_ptq without torchrun to use it."
+                )
+    return plan, hf_config
+
+
 def preview_generate_kwargs(model) -> dict[str, Any]:
     """Extra ``generate`` kwargs that keep every rank of a sharded model decoding a preview.
 
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 04a9d02cf4b..11bf68d6bef 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -38,24 +38,23 @@
     build_quant_cfg,
     create_vlm_calibration_forward,
     create_vlm_calibration_loop,
+    exports_tensorrt_llm_checkpoint,
     generate_excludes_prompt,
     get_model,
     get_processor,
     get_tokenizer,
     is_nemotron_vl,
-    is_trtllm_enc_dec_export,
     keep_the_experts_implementation,
     layerwise_export_requested,
     mlflow_run,
     preview_generate_kwargs,
     recipe_layerwise_blocks,
+    resolve_distribution,
     run_nemotron_vl_preview,
     save_source_config,
-    validate_fsdp2_supported,
 )
 from torch.utils.data import DataLoader
 from transformers import (
-    AutoConfig,
     AutoModelForCausalLM,
     AutoProcessor,
     PreTrainedTokenizer,
@@ -117,7 +116,7 @@
     copy_non_model_files,
     ensure_local_checkpoint,
 )
-from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan
+from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
@@ -346,113 +345,10 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None:
         )
 
 
-def _resolve_distribution(args: argparse.Namespace):
-    """Decide how this run distributes the model: ``(plan, hf_config)``.
-
-    ``resolve_distribution_plan`` decides, and --use_fsdp2 asks it for FSDP2 and no expert
-    parallelism. Without --use_fsdp2 a plan that FSDP2-shards the whole model is refused: the
-    FSDP2 run is the flag's, which checks what this script cannot do with it.
-    """
-    if args.use_fsdp2:
-        # Its checks read the config alone, and come first: deciding builds the model on the meta
-        # device, which some checkpoints they refuse (VILA, multimodal, pack-quantized) cannot do.
-        validate_fsdp2_supported(
-            args,
-            AutoConfig.from_pretrained(
-                args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
-            ),
-        )
-    plan, hf_config = resolve_distribution_plan(
-        args.local_checkpoint_path,
-        trust_remote_code=args.trust_remote_code,
-        ep_size=1 if args.use_fsdp2 else None,
-        dense="fsdp" if args.use_fsdp2 else None,
-    )
-    if plan.is_fsdp_dense and not plan.is_expert_parallel and not args.use_fsdp2:
-        raise RuntimeError(
-            f"{plan!r} FSDP2-shards the model, which this script does under --use_fsdp2. Pass it, "
-            "or set MODELOPT_FSDP_SIZE=1 if the model does fit one GPU; an MoE can instead shard its "
-            "experts (MODELOPT_EP_DEGREE above 1)."
-        )
-    if plan.world_size > 1 and _exports_tensorrt_llm_checkpoint(args, hf_model_type(hf_config)):
-        # That exporter reads the process count as Megatron's training tensor parallelism. Rank 0
-        # alone waits in its barrier for ranks that never arrive, and every rank would merge
-        # whole replicas as if they were shards. Refused before the load, like the layout above.
-        raise NotImplementedError(
-            f"{plan!r}: the TensorRT-LLM checkpoint export (encoder-decoder models, "
-            "--sparsity_fmt, int8_smoothquant) runs in a single process; launch hf_ptq without "
-            "torchrun for it."
-        )
-    if (plan.is_expert_parallel or plan.is_fsdp_dense) and args.vllm_fakequant_export:
-        # That exporter writes the whole model with save_pretrained from each process it runs in,
-        # and every rank of a sharded model exports: each would write its own part into the same
-        # directory. Rank 0 alone writes a replicated model; --use_fsdp2 refuses it while parsing.
-        raise NotImplementedError(
-            f"{plan!r}: --vllm_fakequant_export writes the whole model from one process, which "
-            "holds only its own part of a sharded model; launch hf_ptq without torchrun for it."
-        )
-    if plan.is_expert_parallel and args.experts_implementation is not None:
-        # The expert-parallel load runs the routed experts through torch_a2a, whatever backend the
-        # config names, so the flag would be dropped without a word.
-        raise NotImplementedError(
-            f"{plan!r}: --experts_implementation picks the experts backend a model is loaded with, "
-            "but the expert-parallel load runs the routed experts through torch_a2a whatever it "
-            "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a "
-            "single process."
-        )
-    if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
-        # Its search scores each rank's own calibration shard and reduces no score across ranks, so
-        # the ranks can choose different formats per layer: under expert parallelism, a checkpoint
-        # of mixed formats. --use_fsdp2 keeps main's warning for it.
-        raise NotImplementedError(
-            f"{plan!r}: AutoQuantize scores each rank's own calibration data and reduces no score "
-            "across ranks, so ranks could choose different formats per layer; launch hf_ptq "
-            "without torchrun for an AutoQuantize recipe."
-        )
-    if (
-        plan.world_size > 1
-        and args.recipe is not None
-        and layerwise_export_requested(recipe_layerwise_blocks(load_recipe(args.recipe)))
-    ):
-        # Every rank would write the same per-layer shard files. The exporter refuses that too, but
-        # only once calibration binds it, after the load.
-        raise NotImplementedError(
-            f"{plan!r}: the recipe's layerwise export writes its checkpoint shards from one "
-            "process; launch hf_ptq without torchrun for it."
-        )
-    if plan.world_size > 1 and not args.use_fsdp2:
-        # These configure loading one model in a single process: compressed as it loads, or
-        # spread over the node's GPUs, CPU and disk. An expert-parallel run never reaches those
-        # loaders, and a replicated one pins each rank's copy to the rank's GPU, which none of them
-        # works with. Refused before the load rather than ignored. --max_cpu_memory_gb and
-        # --max_gpu_memory_gb come only with --offload_folder.
-        for flag, given in (
-            ("--low_memory_mode", args.low_memory_mode),
-            ("--use_seq_device_map", args.use_seq_device_map),
-            ("--offload_folder", args.offload_folder is not None),
-        ):
-            if given:
-                raise NotImplementedError(
-                    f"{plan!r}: {flag} configures loading one model in a single process, but each "
-                    "rank of this run loads its own share or copy of it onto its own GPU. Drop it, "
-                    "or launch hf_ptq without torchrun to use it."
-                )
-    return plan, hf_config
-
-
-def _exports_tensorrt_llm_checkpoint(args: argparse.Namespace, model_type: str | None) -> bool:
-    """Whether the run exports through the deprecated TensorRT-LLM checkpoint exporter."""
-    return (
-        is_trtllm_enc_dec_export(model_type)
-        or args.sparsity_fmt != "dense"
-        or "int8_smoothquant" in args.qformat
-    )
-
-
 def load_model(args: argparse.Namespace):
     # If low memory mode is enabled, we compress the model while loading the HF checkpoint.
     calibration_only = False
-    plan, hf_config = _resolve_distribution(args)
+    plan, hf_config = resolve_distribution(args)
     if plan.is_expert_parallel:
         full_model = load_pretrained(
             args.local_checkpoint_path,
@@ -860,7 +756,7 @@ def export_quantized(
                 save_source_config(args, export_path)
 
         start_time = time.time()
-        if _exports_tensorrt_llm_checkpoint(args, model_type):
+        if exports_tensorrt_llm_checkpoint(args, model_type):
             if (
                 args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1
             ) and args.qformat == "nvfp4_svdquant":
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 9181ea33ddf..0ea0dab0264 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -1135,7 +1135,6 @@ def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypa
     copies = []
     monkeypatch.setattr(hf_ptq, "has_spec_opt", lambda model: False)
     monkeypatch.setattr(hf_ptq, "is_multimodal_model", lambda model: False)
-    monkeypatch.setattr(hf_ptq, "is_trtllm_enc_dec_export", lambda model_type: False)
     monkeypatch.setattr(hf_ptq, "remove_hook_from_module", lambda *args, **kwargs: None)
     monkeypatch.setattr(hf_ptq, "export_tensorrt_llm_checkpoint", lambda *args, **kwargs: None)
     monkeypatch.setattr(hf_ptq, "copy_non_model_files", lambda *args: copies.append(args))
diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
index 729cfc37849..c668b63d240 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
@@ -193,7 +193,7 @@ def test_use_fsdp2_plans_fsdp2_without_expert_parallelism(hf_ptq, monkeypatch, t
     _serve(monkeypatch, hf_ptq, _deepseek_v3(n_routed_experts=8))
     args = _args(torchrun(world_size=2), use_fsdp2=True)
 
-    assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(2, dense="fsdp")
+    assert hf_ptq.resolve_distribution(args)[0] == DistributionPlan(2, dense="fsdp")
 
 
 def test_use_fsdp2_refuses_a_checkpoint_before_building_it_to_decide(hf_ptq, monkeypatch, torchrun):
@@ -212,7 +212,7 @@ def test_use_fsdp2_refuses_a_checkpoint_before_building_it_to_decide(hf_ptq, mon
     )
 
     with pytest.raises(NotImplementedError, match="--use_fsdp2 does not support"):
-        hf_ptq._resolve_distribution(_args(torchrun(world_size=2), use_fsdp2=True))
+        hf_ptq.resolve_distribution(_args(torchrun(world_size=2), use_fsdp2=True))
 
 
 def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monkeypatch, torchrun):
@@ -223,7 +223,7 @@ def test_dense_sharding_without_expert_parallelism_needs_use_fsdp2(hf_ptq, monke
     _serve(monkeypatch, hf_ptq, LlamaConfig(**_TINY))
 
     with pytest.raises(RuntimeError, match="--use_fsdp2"):
-        hf_ptq._resolve_distribution(_args(torchrun(world_size=2)))
+        hf_ptq.resolve_distribution(_args(torchrun(world_size=2)))
 
 
 @needs_ep_transformers
@@ -240,10 +240,10 @@ def test_dense_fsdp2_within_expert_data_parallel_replicas_is_refused(
     args = _args(torchrun(world_size=4))
 
     if ep_degree == 4:
-        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(4, ep_size=4, dense="fsdp")
+        assert hf_ptq.resolve_distribution(args)[0] == DistributionPlan(4, ep_size=4, dense="fsdp")
     else:
         with pytest.raises(NotImplementedError, match=r"MODELOPT_EP_DEGREE.*MODELOPT_FSDP_SIZE=1"):
-            hf_ptq._resolve_distribution(args)
+            hf_ptq.resolve_distribution(args)
 
 
 @pytest.mark.parametrize(
@@ -265,10 +265,10 @@ def test_a_tensorrt_llm_checkpoint_export_needs_a_single_process(
     args = _args(torchrun(world_size=world_size), **options)
 
     if world_size == 1:
-        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+        assert hf_ptq.resolve_distribution(args)[0] == DistributionPlan(1)
     else:
         with pytest.raises(NotImplementedError, match="single process"):
-            hf_ptq._resolve_distribution(args)
+            hf_ptq.resolve_distribution(args)
 
 
 @pytest.mark.parametrize(
@@ -297,10 +297,10 @@ def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds(
     )
 
     if world_size == 1 or use_fsdp2:
-        hf_ptq._resolve_distribution(args)
+        hf_ptq.resolve_distribution(args)
     else:
         with pytest.raises(NotImplementedError, match="AutoQuantize"):
-            hf_ptq._resolve_distribution(args)
+            hf_ptq.resolve_distribution(args)
 
 
 @pytest.mark.parametrize("world_size", [2, 1], ids=["torchrun", "single_process"])
@@ -317,10 +317,10 @@ def test_a_layerwise_export_recipe_is_refused_before_a_distributed_load(
     )
 
     if world_size == 1:
-        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+        assert hf_ptq.resolve_distribution(args)[0] == DistributionPlan(1)
     else:
         with pytest.raises(NotImplementedError, match="layerwise export"):
-            hf_ptq._resolve_distribution(args)
+            hf_ptq.resolve_distribution(args)
 
 
 @pytest.mark.parametrize(
@@ -347,11 +347,11 @@ def test_the_single_process_load_options_are_refused_on_the_routes_this_script_a
     args = _args(torchrun(world_size=world_size), **option)
 
     if world_size == 1:
-        assert hf_ptq._resolve_distribution(args)[0] == DistributionPlan(1)
+        assert hf_ptq.resolve_distribution(args)[0] == DistributionPlan(1)
     else:
         (name,) = option
         with pytest.raises(NotImplementedError, match=f"--{name}"):
-            hf_ptq._resolve_distribution(args)
+            hf_ptq.resolve_distribution(args)
 
 
 @pytest.mark.parametrize(
@@ -377,9 +377,9 @@ def test_the_vllm_fakequant_export_is_refused_for_a_sharded_model(
 
     if refused:
         with pytest.raises(NotImplementedError, match="--vllm_fakequant_export"):
-            hf_ptq._resolve_distribution(args)
+            hf_ptq.resolve_distribution(args)
     else:
-        hf_ptq._resolve_distribution(args)
+        hf_ptq.resolve_distribution(args)
 
 
 @pytest.mark.parametrize(
@@ -407,9 +407,9 @@ def test_an_experts_implementation_is_refused_where_the_load_picks_the_backend(
 
     if refused:
         with pytest.raises(NotImplementedError, match="--experts_implementation"):
-            hf_ptq._resolve_distribution(args)
+            hf_ptq.resolve_distribution(args)
     else:
-        hf_ptq._resolve_distribution(args)
+        hf_ptq.resolve_distribution(args)
 
 
 # ---------------------------------------------------------------------------------- load_model
diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py
index 34b5194e1c2..2539311313e 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py
@@ -93,7 +93,8 @@ def test_pre_quantize_preview_removes_padding(hf_ptq, ids, mask):
     ],
 )
 def test_is_trtllm_enc_dec_export(hf_ptq, model_type: str, expected: bool):
-    assert hf_ptq.is_trtllm_enc_dec_export(model_type) is expected
+    example_utils = importlib.import_module("example_utils")  # imported with hf_ptq
+    assert example_utils.is_trtllm_enc_dec_export(model_type) is expected
 
 
 def _config(model_type: str, is_encoder_decoder: bool) -> SimpleNamespace:
diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
index 20c86f2f1a8..dbf194c2b3f 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
@@ -68,9 +68,7 @@ def test_image_calibration_model_target_follows_recipe(
 
     monkeypatch.setattr(hf_ptq, "get_model", lambda *args, **kwargs: full_model)
     # A single process's plan, decided without reading the "dummy" checkpoint.
-    monkeypatch.setattr(
-        hf_ptq, "resolve_distribution_plan", lambda *args, **kwargs: (DistributionPlan(), None)
-    )
+    monkeypatch.setattr(hf_ptq, "resolve_distribution", lambda args: (DistributionPlan(), None))
     # A VLM's extracted language model reports its own sub-config type; hf_ptq must keep the root's.
     monkeypatch.setattr(
         hf_ptq,

From d12305a0df55d57b42db142f51a28363658fd792 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:56:42 +0000
Subject: [PATCH 18/24] Trim the comments hf_ptq's expert-parallel run adds

Each now says why in a line or two. The calibration loop's comment no longer
gives a second, contradicting version boundary for the a2a forwards'
re-plant: transformers 5.16, as everywhere else.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/hf_ptq.py | 42 ++++++++++-----------------------------
 1 file changed, 11 insertions(+), 31 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 11bf68d6bef..1b2c07a8fc5 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -297,10 +297,8 @@ def make_calib_dataloader(
         )
     world_size = distribution_plan(language_model).world_size
     if world_size > 1 and calib_dataloader is not None:
-        # Every distributed calibration path is data-parallel over the world, so each rank calibrates
-        # on its own per-rank shard (and the quantizer amax is synced across ranks): MoE expert-parallel
-        # a2a routes each token to its expert-owning rank; dense DP replicates the model per rank; FSDP2
-        # shards the params but still feeds each rank its own data.
+        # Every distributed plan is data-parallel over the world: each rank calibrates on its own
+        # shard, and the quantizers' amax is synced across the ranks.
         calib_dataloader = _shard_calib_dataloader(
             calib_dataloader, args.dist_state.rank, world_size
         )
@@ -589,22 +587,10 @@ def mono_quantize(
             # They must be consumed by the complete VLM even when only a nested component is the
             # quantization target; the full forward still exercises that component's quantizers.
             if distribution_plan(full_model).is_expert_parallel:
-                # An expert-parallel run has to calibrate through distributed_forward_loop, not the
-                # plain one: it calls _reinstall_a2a_forwards first, which re-plants the
-                # torch_a2a expert/router forward wrappers that quantizer insertion clobbers on
-                # transformers >= 5.16 (where the styles install a forward rather than a pre-hook).
-                # Without it the experts forward runs with process_group=None and raises
-                # "torch_a2a experts forward requires the EP `process_group`" at calibration step 0.
-                # On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op --
-                # which is why this is invisible on older transformers and fatal on newer. Other
-                # distributed runs (FSDP2, or a model replicated on every rank) keep the plain loop,
-                # which retries no batch that runs out of memory on a plan of several ranks: the
-                # error reaches distributed_run, which aborts the rank and with it the run. Image
-                # batches go through the VLM's own forward, as they do without expert parallelism:
-                # Nemotron-H Omni, for one, calibrates on images by default. Text batches go through
-                # the language model, but the loop re-plants over the model it is handed, so that is
-                # the full model: transformers keeps the expert-parallel plan on the model
-                # from_pretrained built, and a VLM's language model names no experts in its own.
+                # The distributed loop first re-plants the torch_a2a forwards that quantizer
+                # insertion clobbers on transformers >= 5.16. It is handed the full model, whose
+                # expert-parallel plan names the experts; text batches still run through the
+                # language model.
                 calibrate_loop = distributed_forward_loop(
                     full_model,
                     calib_dataloader,
@@ -634,10 +620,8 @@ def mono_quantize(
         else:
             language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop)
         if distribution_plan(full_model).is_expert_parallel:
-            # The calibration loop re-plants the a2a forwards quantizer insertion drops (on
-            # transformers >= 5.16; a no-op before), but a config without static activation scales
-            # calibrates without it, and algorithm None (the MX formats) never calls it; the
-            # preview and the export forward through them.
+            # A config that runs no calibration loop never re-planted the a2a forwards, which the
+            # preview and the export need.
             _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model))
 
         # For VL models, update full_model to use the quantized language model
@@ -1608,10 +1592,7 @@ def main(args: argparse.Namespace):
     random.seed(RAND_SEED)
     np.random.seed(RAND_SEED)
 
-    # A distributed launch owns its process group for the whole run, and the run takes this
-    # process's rank, size and device from it. Entered here rather than inside load_model so
-    # cleanup still runs if loading raises -- otherwise the surviving ranks block on the first
-    # collective until the NCCL timeout.
+    # The run owns a distributed launch's process group, and ends it even if loading raises.
     with distributed_run(timeout=timedelta(hours=3), device=args.device) as dist_state:
         args.dist_state = dist_state
         _main_inner(args)
@@ -1624,9 +1605,8 @@ def _main_inner(args: argparse.Namespace):
         # path.
         args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path)
 
-        # launch a memory monitor to read the currently used GPU memory. It reads every GPU on the
-        # node, so a distributed run starts one per node, on its local rank 0, rather than one per
-        # rank.
+        # launch a memory monitor to read the currently used GPU memory, one per node: it reads all
+        # of the node's GPUs.
         if args.dist_state.world_size == 1 or local_rank() == 0:
             launch_memory_monitor()
 

From 37a75639a73a81e87b6dfea8a60d5ce4f6f7ba73 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:57:43 +0000
Subject: [PATCH 19/24] Add a Slurm example of expert-parallel PTQ across nodes

multinode_ep_ptq.slurm follows multinode_fsdp2_ptq.slurm: one torchrun per
node, launched without --use_fsdp2, so hf_ptq.py shards the MoE's experts
over every rank. It quantizes Qwen3-235B-A22B to NVFP4 by default, and names
the two knobs that override the plan. The README's expert-parallel section
points to it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/README.md                    |  6 ++
 examples/hf_ptq/slurm/multinode_ep_ptq.slurm | 93 ++++++++++++++++++++
 2 files changed, 99 insertions(+)
 create mode 100644 examples/hf_ptq/slurm/multinode_ep_ptq.slurm

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 7db3517ae49..4815a5273fb 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -677,6 +677,12 @@ torchrun --nproc_per_node=<num_gpus_per_node> \
 
 To run across nodes, add the `--nnodes`, `--node_rank`, `--master_addr` and `--master_port` flags shown in [Manual (run on each node)](#manual-run-on-each-node). `--calib_size` is the total over all ranks.
 
+Under Slurm, [`slurm/multinode_ep_ptq.slurm`](./slurm/multinode_ep_ptq.slurm) is a ready-to-run example that quantizes Qwen3-235B-A22B to NVFP4 across nodes. Edit its `CONFIG` block and submit:
+
+```bash
+sbatch --nodes=2 slurm/multinode_ep_ptq.slurm
+```
+
 Two environment variables override the automatic distribution plan:
 
 | Variable | Default | Effect |
diff --git a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
new file mode 100644
index 00000000000..758ee429843
--- /dev/null
+++ b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
@@ -0,0 +1,93 @@
+#!/bin/bash
+
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Multi-node expert-parallel post-training quantization of an MoE model, ready to run under Slurm.
+#
+# Slurm allocates the nodes and launches one task per node; each task starts a `torchrun` that
+# spawns one process per GPU. Without `--use_fsdp2`, `hf_ptq.py` shards the MoE's routed experts
+# across all ranks (world_size = num_nodes * gpus_per_node), calibrates each rank on its own share
+# of the data, and has every rank write its share of the checkpoint. Expert parallelism needs
+# transformers 5.7 or newer; see the README's expert-parallel section for the supported models.
+#
+# The defaults below quantize Qwen/Qwen3-235B-A22B to NVFP4 with an FP8 KV cache. Edit the CONFIG
+# block for your cluster/model, then submit with e.g.:
+#
+#   sbatch --nodes=2 multinode_ep_ptq.slurm
+#
+# For a single-node run, `--nodes=1` works unchanged.
+
+#SBATCH --job-name=ep-ptq
+#SBATCH --account={account}
+#SBATCH --partition={partition}
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=1          # one torchrun launcher per node; it fans out to the GPUs
+#SBATCH --gpus-per-node=8
+#SBATCH --exclusive
+#SBATCH --time=04:00:00
+#SBATCH --output=%x_%j.log
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# CONFIG — edit these for your cluster and model. They are exported so `srun`
+# propagates them into the container task below.
+# ---------------------------------------------------------------------------
+export CONTAINER_IMAGE={container_image}                # e.g. nvcr.io#nvidia/pytorch:25.10-py3
+export MODELOPT_PATH={path_to_modelopt_repo}            # host clone of TensorRT-Model-Optimizer
+export MODEL_PATH=Qwen/Qwen3-235B-A22B                  # HF repo id (auto-downloaded) or a local dir
+export EXPORT_PATH={path_to_export_dir}                 # where the quantized checkpoint is written
+export HF_HOME={path_to_hf_cache}                       # HF cache; persist so a repo id isn't re-downloaded each run
+# export HF_TOKEN={hf_token}                            # required for gated repos (or `huggingface-cli login` on host)
+export RECIPE=general/ptq/nvfp4_default-kv_fp8          # built-in recipe name or /path/to/recipe.yaml
+export CALIB_SIZE=512                                   # total over all ranks
+export BATCH_SIZE=4
+# export MODELOPT_EP_DEGREE=16                          # expert-parallel degree (default: auto)
+# export MODELOPT_FSDP_SIZE=1                           # 1 replicates the dense part, world shards it (default: auto)
+
+# Rendezvous: node 0 is the master; all ranks meet at MASTER_ADDR:MASTER_PORT.
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1)
+export MASTER_PORT=29531
+
+# Mount the repo, export dir, and HF cache; add the model dir only when MODEL_PATH is a local path
+# (a repo id is downloaded into HF_HOME instead). Run from the hf_ptq example dir.
+CONTAINER_MOUNTS="${MODELOPT_PATH}:/modelopt,${EXPORT_PATH}:${EXPORT_PATH},${HF_HOME}:${HF_HOME}"
+if [ -d "${MODEL_PATH}" ]; then
+    CONTAINER_MOUNTS="${CONTAINER_MOUNTS},${MODEL_PATH}:${MODEL_PATH}"
+fi
+srun --container-image="${CONTAINER_IMAGE}" \
+     --container-mounts="${CONTAINER_MOUNTS}" \
+     --container-workdir=/modelopt/examples/hf_ptq \
+  bash -c '
+    set -euo pipefail
+    export PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+    # ModelOpt + its HF deps (transformers/accelerate/datasets/...) from the mounted source, then example extras.
+    pip install -q -e "/modelopt[hf]" --no-build-isolation
+    pip install -q -r requirements.txt
+
+    torchrun \
+      --nnodes="${SLURM_NNODES}" \
+      --node_rank="${SLURM_NODEID}" \
+      --nproc_per_node="$(nvidia-smi -L | wc -l)" \
+      --rdzv_backend=c10d \
+      --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \
+      hf_ptq.py \
+        --pyt_ckpt_path "${MODEL_PATH}" \
+        --recipe "${RECIPE}" \
+        --calib_size "${CALIB_SIZE}" \
+        --batch_size "${BATCH_SIZE}" \
+        --export_path "${EXPORT_PATH}"
+  '

From 00b2951c189f65b7b90433f1c7580a40a41e7767 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 21:14:22 +0000
Subject: [PATCH 20/24] Install the transformers expert parallelism needs in
 the EP Slurm example

The header says expert parallelism needs transformers 5.7 or newer, but the
install kept whatever the image ships when it meets ModelOpt's own floor,
such as 5.5, where the run would not shard the experts. Ask pip for 5.7 or
newer together with ModelOpt, and print the version into the job log.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/slurm/multinode_ep_ptq.slurm | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
index 758ee429843..1d0516bf35a 100644
--- a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
+++ b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
@@ -75,8 +75,10 @@ srun --container-image="${CONTAINER_IMAGE}" \
     set -euo pipefail
     export PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
     # ModelOpt + its HF deps (transformers/accelerate/datasets/...) from the mounted source, then example extras.
-    pip install -q -e "/modelopt[hf]" --no-build-isolation
+    # Expert parallelism needs transformers 5.7 or newer, which the image may not ship.
+    pip install -q -e "/modelopt[hf]" "transformers>=5.7" --no-build-isolation
     pip install -q -r requirements.txt
+    python -c "import transformers as t; print(t.__name__, t.__version__)"
 
     torchrun \
       --nnodes="${SLURM_NNODES}" \

From 2a9e94dc234374bfa135a69094f410a1f21b763f Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 23:39:22 +0000
Subject: [PATCH 21/24] Keep hf_ptq's expert-parallel run off private helpers

hf_ptq.py re-planted the experts' all-to-all forwards through the private
_reinstall_a2a_forwards. modelopt.torch.utils.distributed now exports
reinstall_a2a_forwards, a thin public wrapper for a caller that quantizes
without distributed_forward_loop, and hf_ptq calls it. example_utils.py
checks for an AutoQuantize recipe with the public recipe API instead of
importing autoquant_utils' private helper.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/example_utils.py              | 10 +++++++---
 examples/hf_ptq/hf_ptq.py                     |  4 ++--
 modelopt/torch/utils/distributed.py           | 19 ++++++++++++++++---
 .../hf_ptq/test_hf_ptq_expert_parallel.py     |  2 +-
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index a76611e8735..b9cc0158d36 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -30,7 +30,6 @@
 import transformers
 from accelerate import infer_auto_device_map, init_empty_weights
 from accelerate.utils import get_max_memory
-from autoquant_utils import _recipe_is_auto_quantize
 from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static
 from safetensors import safe_open
 from transformers import (
@@ -43,7 +42,7 @@
     ProcessorMixin,
 )
 
-from modelopt.recipe import load_recipe
+from modelopt.recipe import ModelOptAutoQuantizeRecipe, load_recipe
 from modelopt.recipe.presets import QUANT_CFG_CHOICES
 from modelopt.torch.export import has_spec_opt
 from modelopt.torch.export.model_utils import is_multimodal_model
@@ -183,7 +182,12 @@ def resolve_distribution(args) -> tuple[DistributionPlan, Any]:
             "names. Drop it, or quantize without expert parallelism: with --use_fsdp2, or in a "
             "single process."
         )
-    if plan.world_size > 1 and not args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe):
+    if (
+        plan.world_size > 1
+        and not args.use_fsdp2
+        and args.recipe is not None
+        and isinstance(load_recipe(args.recipe), ModelOptAutoQuantizeRecipe)
+    ):
         # Each rank would score its own data, and could pick other formats per layer.
         raise NotImplementedError(
             f"{plan!r}: AutoQuantize scores each rank's own calibration data and reduces no score "
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 1b2c07a8fc5..18ff07274d1 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -102,11 +102,11 @@
     get_supported_datasets,
 )
 from modelopt.torch.utils.distributed import (
-    _reinstall_a2a_forwards,
     distributed_forward_loop,
     distributed_run,
     distribution_plan,
     local_rank,
+    reinstall_a2a_forwards,
     shard_dataloader,
 )
 from modelopt.torch.utils.memory_monitor import launch_memory_monitor
@@ -622,7 +622,7 @@ def mono_quantize(
         if distribution_plan(full_model).is_expert_parallel:
             # A config that runs no calibration loop never re-planted the a2a forwards, which the
             # preview and the export need.
-            _reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model))
+            reinstall_a2a_forwards(full_model, expert_parallel_mesh(full_model))
 
         # For VL models, update full_model to use the quantized language model
         if is_nemotron_vl_model and language_model is not full_model:
diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 8a41448af03..5abceddf3f5 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -69,6 +69,7 @@
     "is_master",
     "rank",
     "recorded_distribution_plan",
+    "reinstall_a2a_forwards",
     "set_distribution_plan",
     "shard_dataloader",
     "size",
@@ -849,9 +850,11 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int:
     transformers installs its TP transforms by replacing ``module.forward`` on the instance
     (``install_forward``); for torch_a2a_experts that wrapper passes the EP ``process_group`` to the
     experts forward. Quantizer insertion rebinds ``module.forward``, dropping it, and the experts
-    forward then raises. Called after conversion, before the first forward. Only a forward quantizer
-    insertion rebound is planted again: a wrapper still in place is kept, since wrapping it again
-    would run its transforms twice. Without an ``ep_mesh`` there is nothing to re-plant.
+    forward then raises. Called after conversion, before the first forward: by the calibration
+    loop, and through :func:`reinstall_a2a_forwards` by a caller that quantizes without one. Only a
+    forward quantizer insertion rebound is planted again: a wrapper still in place is kept, since
+    wrapping it again would run its transforms twice. Without an ``ep_mesh`` there is nothing to
+    re-plant.
     """
     if ep_mesh is None:
         return 0
@@ -892,6 +895,16 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int:
     return n
 
 
+def reinstall_a2a_forwards(model, ep_mesh) -> int:
+    """Re-plant the routed experts' all-to-all forwards that quantizer insertion dropped.
+
+    For a caller that quantizes an expert-parallel model without :func:`distributed_forward_loop`,
+    which does it itself: a recipe that runs no calibration runs no loop. ``ep_mesh`` is the mesh the
+    experts are sharded over. Returns how many forwards it planted again.
+    """
+    return _reinstall_a2a_forwards(model, ep_mesh)
+
+
 def _release_cached_cuda_memory():
     """Hand the caching allocator's reserved-but-unused pool back to the driver.
 
diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
index c668b63d240..b382e0fe054 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
@@ -684,7 +684,7 @@ def test_an_expert_parallel_run_replants_the_all_to_all_forwards_after_quantizin
     )
     monkeypatch.setattr(
         hf_ptq,
-        "_reinstall_a2a_forwards",
+        "reinstall_a2a_forwards",
         lambda model, mesh: events.append(("replant", model, mesh)),
     )
     args = SimpleNamespace(

From 45304fb204e2d26b7e287d698f648e3cd7d36ebd Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 23:41:40 +0000
Subject: [PATCH 22/24] Warn about LSQ on a multi-rank hf_ptq run, and say
 which calibrations were tested

LSQ's first scale calibration is synced or refused as the algorithm it runs
is, but nothing has checked its learnable scales across ranks, so an
expert-parallel or replicated run now warns before the load. The calibration
methods a run uses come from one helper, which the AWQ check reads too. The
README's expert-parallel section says which calibrations have run on GPUs on
more than one rank, which are supported but not yet run there, and which are
refused.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/README.md                     |  2 ++
 examples/hf_ptq/example_utils.py              | 20 ++++++++++++---
 .../hf_ptq/test_hf_ptq_expert_parallel.py     | 25 +++++++++++++++++++
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 4815a5273fb..198fbb6df02 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -690,6 +690,8 @@ Two environment variables override the automatic distribution plan:
 | `MODELOPT_EP_DEGREE` | `auto` | Expert-parallel degree. `auto` picks the largest degree that divides both the routed-expert count and the world size, or 1 for a model whose spec refuses expert parallelism or needs a newer transformers. Below the world size, each group of `MODELOPT_EP_DEGREE` ranks holds a full set of experts, and the groups calibrate on different data. |
 | `MODELOPT_FSDP_SIZE` | `auto` | Whether to FSDP2-shard the dense part. `auto` shards it when it would take more than 60% of one GPU's memory; `world` always shards it; `1` replicates it. With `MODELOPT_EP_DEGREE` below the world size the export cannot write a sharded dense part yet, so such a run fails before loading. |
 
+On more than one rank, `max` and `mse` calibration have run on GPUs; `gptq`, `smoothquant`, `awq_clip` and `nvfp4_act_headroom` are supported but not yet run on GPUs; `awq_lite`, `awq_full`, `svdquant`, `local_hessian` and `imatrix`, and `max` or `mse` with `distributed_sync` off, are refused before quantizing; and an LSQ recipe runs with a warning, since nothing has checked its learnable scales across ranks yet.
+
 A run that cannot shard the model the way the plan asks fails before loading, rather than loading the whole model on every rank. `--use_fsdp2` skips expert parallelism and uses the [FSDP2 path](#multi-node-post-training-quantization-with-fsdp2) instead. Expert parallelism refuses GPT-OSS and Llama-4: their experts are stored transposed, which the expert-parallel export cannot write yet. `--use_fsdp2` refuses their released checkpoints too (GPT-OSS's carry an MXFP4 `quantization_config`, and Llama-4's are multimodal), so quantize those in a single process (without `torchrun`), which spreads the model over the node's GPUs.
 
 ## Evaluate Accuracy
diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index b9cc0158d36..f23d173d44a 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -55,6 +55,7 @@
 from modelopt.torch.quantization.mode import get_modelike_from_algo_cfg
 from modelopt.torch.utils import distributed as dist_utils
 from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan
+from modelopt.torch.utils.logging import warn_rank_0
 from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run
 from modelopt.torch.utils.plugins.model_load_utils import (
     is_fsdp2_text_only_vlm,
@@ -67,8 +68,8 @@
 SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"]
 
 
-def _awq_calibrations(args) -> set[str]:
-    """The AWQ calibrations (``awq_lite``, ``awq_full``) the run's --recipe or --qformat uses."""
+def _calibration_methods(args) -> set[str]:
+    """The calibration methods (``max``, ``awq_lite``, ...) the run's --recipe or --qformat uses."""
     if getattr(args, "recipe", None) is not None:
         recipe = load_recipe(args.recipe)
         configs = [getattr(recipe, "quantize", None)]
@@ -80,12 +81,16 @@ def _awq_calibrations(args) -> set[str]:
     else:
         preset = QUANT_CFG_CHOICES.get(getattr(args, "qformat", None) or "")
         algorithms = [preset.get("algorithm")] if preset is not None else []
-    methods = {
+    return {
         mode.removesuffix("_calibrate")
         for algorithm in algorithms
         for mode, _ in get_modelike_from_algo_cfg(algorithm)
     }
-    return methods & {"awq_lite", "awq_full"}
+
+
+def _awq_calibrations(args) -> set[str]:
+    """The AWQ calibrations (``awq_lite``, ``awq_full``) the run's --recipe or --qformat uses."""
+    return _calibration_methods(args) & {"awq_lite", "awq_full"}
 
 
 def validate_fsdp2_supported(args, config):
@@ -194,6 +199,13 @@ def resolve_distribution(args) -> tuple[DistributionPlan, Any]:
             "across ranks, so ranks could choose different formats per layer; launch hf_ptq "
             "without torchrun for an AutoQuantize recipe."
         )
+    if plan.world_size > 1 and not args.use_fsdp2 and "lsq" in _calibration_methods(args):
+        # Its scale calibration is the algorithm it runs first, synced or refused as that one is,
+        # but nothing has checked its learnable scales across ranks yet.
+        warn_rank_0(
+            f"{plan!r}: LSQ has not been validated on a run of several ranks: nothing checks that "
+            "the ranks' learnable scales match. Launch hf_ptq without torchrun for the tested path."
+        )
     if (
         plan.world_size > 1
         and args.recipe is not None
diff --git a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
index b382e0fe054..176d416f034 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
@@ -30,6 +30,7 @@
 import contextlib
 import importlib
 import inspect
+import warnings
 from datetime import timedelta
 from pathlib import Path
 from types import SimpleNamespace
@@ -303,6 +304,30 @@ def test_an_auto_quantize_recipe_is_refused_on_the_routes_this_script_adds(
             hf_ptq.resolve_distribution(args)
 
 
+@pytest.mark.parametrize(
+    ("world_size", "use_fsdp2", "warns"),
+    [(2, False, True), (2, True, False), (1, False, False)],
+    ids=["replicated", "use_fsdp2", "single_process"],
+)
+def test_an_lsq_recipe_warns_on_the_routes_this_script_adds(
+    hf_ptq, monkeypatch, torchrun, world_size, use_fsdp2, warns
+):
+    """Nothing has checked LSQ's learnable scales across ranks: a replicated or expert-parallel run
+    warns before the load; --use_fsdp2 and a single process run it as before."""
+    _serve(monkeypatch, hf_ptq, LlamaConfig(architectures=["LlamaForCausalLM"], **_TINY))
+    args = _args(
+        torchrun(world_size=world_size),
+        use_fsdp2=use_fsdp2,
+        recipe="general/qad/nvfp4_lsq-mse_init-fp8_kv",
+    )
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        hf_ptq.resolve_distribution(args)
+
+    assert any("LSQ" in str(warning.message) for warning in caught) is warns
+
+
 @pytest.mark.parametrize("world_size", [2, 1], ids=["torchrun", "single_process"])
 def test_a_layerwise_export_recipe_is_refused_before_a_distributed_load(
     hf_ptq, monkeypatch, torchrun, world_size

From 1a35fe13f5e993136c12051400ded31ed0f84d12 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 3 Nov 2026 01:16:52 +0000
Subject: [PATCH 23/24] Say the EP Slurm example needs shared storage, take
 turns on its install, and drop --node_rank

Every rank writes its share of the checkpoint to EXPORT_PATH, and rank 0
downloads a repo id into HF_HOME for the others to read, so both must be
on storage every node can see; the script and the README now say so.
Every node runs the editable install against the one mounted clone, which
the build writes into, so the nodes now take turns under flock where the
filesystem offers it. torchrun ignores --node_rank under the c10d
rendezvous, which assigns the node ranks itself.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 examples/hf_ptq/README.md                    |  2 +-
 examples/hf_ptq/slurm/multinode_ep_ptq.slurm | 13 ++++++++++---
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 198fbb6df02..42307b8126a 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -675,7 +675,7 @@ torchrun --nproc_per_node=<num_gpus_per_node> \
     --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.
+To run across nodes, add the `--nnodes`, `--node_rank`, `--master_addr` and `--master_port` flags shown in [Manual (run on each node)](#manual-run-on-each-node). `--calib_size` is the total over all ranks. The export directory and the Hugging Face cache (`HF_HOME`) must be on storage every node can see: every rank writes its share of the checkpoint, and rank 0 downloads a Hub model for the others to read.
 
 Under Slurm, [`slurm/multinode_ep_ptq.slurm`](./slurm/multinode_ep_ptq.slurm) is a ready-to-run example that quantizes Qwen3-235B-A22B to NVFP4 across nodes. Edit its `CONFIG` block and submit:
 
diff --git a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
index 1d0516bf35a..04fb0b61a9d 100644
--- a/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
+++ b/examples/hf_ptq/slurm/multinode_ep_ptq.slurm
@@ -45,6 +45,8 @@ set -euo pipefail
 # ---------------------------------------------------------------------------
 # CONFIG — edit these for your cluster and model. They are exported so `srun`
 # propagates them into the container task below.
+# EXPORT_PATH and HF_HOME must be on storage every node can see: every rank writes its share of
+# the checkpoint to EXPORT_PATH, and rank 0 downloads a repo id into HF_HOME for the others to read.
 # ---------------------------------------------------------------------------
 export CONTAINER_IMAGE={container_image}                # e.g. nvcr.io#nvidia/pytorch:25.10-py3
 export MODELOPT_PATH={path_to_modelopt_repo}            # host clone of TensorRT-Model-Optimizer
@@ -75,14 +77,19 @@ srun --container-image="${CONTAINER_IMAGE}" \
     set -euo pipefail
     export PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
     # ModelOpt + its HF deps (transformers/accelerate/datasets/...) from the mounted source, then example extras.
-    # Expert parallelism needs transformers 5.7 or newer, which the image may not ship.
-    pip install -q -e "/modelopt[hf]" "transformers>=5.7" --no-build-isolation
+    # Expert parallelism needs transformers 5.7 or newer, which the image may not ship. The editable
+    # install writes into the one clone every node mounts, so the nodes take turns where the
+    # filesystem offers flock.
+    mkdir -p /modelopt/build
+    (
+      flock 9 || true
+      pip install -q -e "/modelopt[hf]" "transformers>=5.7" --no-build-isolation
+    ) 9>/modelopt/build/.pip-install.lock
     pip install -q -r requirements.txt
     python -c "import transformers as t; print(t.__name__, t.__version__)"
 
     torchrun \
       --nnodes="${SLURM_NNODES}" \
-      --node_rank="${SLURM_NODEID}" \
       --nproc_per_node="$(nvidia-smi -L | wc -l)" \
       --rdzv_backend=c10d \
       --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \

From eed54f9114b268b9b3057f8e45f8dce4225b58df Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 3 Nov 2026 01:17:14 +0000
Subject: [PATCH 24/24] Make reinstall_a2a_forwards the one re-plant function,
 not an alias

The public name was a bare wrapper of _reinstall_a2a_forwards. Rename the
function itself, merge the two docstrings into one, and have
distributed_forward_loop and the tests call the public name.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py           | 27 +++++++------------
 .../utils/test_distributed_calibration.py     | 12 ++++-----
 2 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 5abceddf3f5..8f2f99158b2 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -844,17 +844,18 @@ def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoad
     )
 
 
-def _reinstall_a2a_forwards(model, ep_mesh) -> int:
-    """Re-plant the torch_a2a TP forward wrappers that quantizer insertion clobbered.
+def reinstall_a2a_forwards(model, ep_mesh) -> int:
+    """Re-plant the routed experts' all-to-all forwards that quantizer insertion dropped.
 
     transformers installs its TP transforms by replacing ``module.forward`` on the instance
     (``install_forward``); for torch_a2a_experts that wrapper passes the EP ``process_group`` to the
     experts forward. Quantizer insertion rebinds ``module.forward``, dropping it, and the experts
-    forward then raises. Called after conversion, before the first forward: by the calibration
-    loop, and through :func:`reinstall_a2a_forwards` by a caller that quantizes without one. Only a
-    forward quantizer insertion rebound is planted again: a wrapper still in place is kept, since
-    wrapping it again would run its transforms twice. Without an ``ep_mesh`` there is nothing to
-    re-plant.
+    forward then raises. Call this after conversion, before the first forward:
+    :func:`distributed_forward_loop` does, and so must a caller that quantizes an expert-parallel
+    model without that loop (a recipe that runs no calibration runs no loop). Only a forward
+    quantizer insertion rebound is planted again: a wrapper still in place is kept, since wrapping
+    it again would run its transforms twice. ``ep_mesh`` is the mesh the experts are sharded over;
+    without one there is nothing to re-plant. Returns how many forwards it planted again.
     """
     if ep_mesh is None:
         return 0
@@ -895,16 +896,6 @@ def _reinstall_a2a_forwards(model, ep_mesh) -> int:
     return n
 
 
-def reinstall_a2a_forwards(model, ep_mesh) -> int:
-    """Re-plant the routed experts' all-to-all forwards that quantizer insertion dropped.
-
-    For a caller that quantizes an expert-parallel model without :func:`distributed_forward_loop`,
-    which does it itself: a recipe that runs no calibration runs no loop. ``ep_mesh`` is the mesh the
-    experts are sharded over. Returns how many forwards it planted again.
-    """
-    return _reinstall_a2a_forwards(model, ep_mesh)
-
-
 def _release_cached_cuda_memory():
     """Hand the caching allocator's reserved-but-unused pool back to the driver.
 
@@ -954,7 +945,7 @@ def distributed_forward_loop(
         raise ValueError(f"{plan!r}: the ep_mesh spans {ep_mesh.size()} ranks, not {plan.ep_size}.")
 
     def calibrate(_inner_model):
-        _reinstall_a2a_forwards(model, ep_mesh)
+        reinstall_a2a_forwards(model, ep_mesh)
         with _disable_use_cache(model), torch.no_grad():
             # Hand the caching allocator's unused pool back: NCCL allocates its buffers outside it.
             _release_cached_cuda_memory()
diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py
index 95a954b9272..a0f750c7312 100644
--- a/tests/unit/torch/utils/test_distributed_calibration.py
+++ b/tests/unit/torch/utils/test_distributed_calibration.py
@@ -51,8 +51,8 @@
 from modelopt.torch.utils.distributed import (
     DistributionPlan,
     ParallelState,
-    _reinstall_a2a_forwards,
     distributed_forward_loop,
+    reinstall_a2a_forwards,
     set_distribution_plan,
     shard_dataloader,
 )
@@ -262,7 +262,7 @@ def test_sharded_loader_keeps_the_original_loader_settings(dataset):
         assert getattr(shard, setting) == getattr(loader, setting), setting
 
 
-# -------------------------------------------------------------------------- _reinstall_a2a_forwards
+# -------------------------------------------------------------------------- reinstall_a2a_forwards
 
 _TP_MODULE = "transformers.distributed.tensor_parallel"
 _LEGACY_TP_MODULE = "transformers.integrations.tensor_parallel"
@@ -398,7 +398,7 @@ def test_reinstall_is_a_noop_without_an_ep_mesh(monkeypatch):
     """A single-process or non-EP model is left untouched, even with the torch_a2a styles around."""
     styles = _install_fake_tp(monkeypatch)
 
-    assert _reinstall_a2a_forwards(_MoEModel(), None) == 0
+    assert reinstall_a2a_forwards(_MoEModel(), None) == 0
     assert not any(style.installed for style in styles.values())
 
 
@@ -423,7 +423,7 @@ def test_reinstall_skips_styles_whose_transforms_are_pre_hooks(monkeypatch, expe
     for layer in model.layers:  # quantizer insertion rebinds the planned modules' forwards
         _drop_tp_wrapper(layer.gate)
         _drop_tp_wrapper(layer.experts)
-    assert _reinstall_a2a_forwards(model, _FakeMesh()) == planted
+    assert reinstall_a2a_forwards(model, _FakeMesh()) == planted
 
 
 def test_reinstall_leaves_a_module_the_load_did_not_wrap(monkeypatch):
@@ -431,7 +431,7 @@ def test_reinstall_leaves_a_module_the_load_did_not_wrap(monkeypatch):
     the load did not wrap keeps running as it did, without the style's transforms."""
     styles = _install_fake_tp(monkeypatch)
 
-    assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0
+    assert reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0
     assert not any(style.installed for style in styles.values())
 
 
@@ -519,7 +519,7 @@ def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch,
     """
     events = []
     monkeypatch.setattr(
-        mto_dist, "_reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args))
+        mto_dist, "reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args))
     )
     monkeypatch.setattr(torch.cuda, "is_initialized", lambda: cuda_started)
     monkeypatch.setattr(mto_dist.gc, "collect", lambda: events.append("collect"))