From 20d7eee85e213948b321cd265f4f318c5214bab2 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:08:02 +0000
Subject: [PATCH 1/6] 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 3c5ffb9b63979b3abd9670c5c40da747d2ab58c2 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:08:08 +0000
Subject: [PATCH 2/6] 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 69047781d957a687e5a8b984fb8b6cdfe947f78c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:57:00 +0000
Subject: [PATCH 3/6] 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 b9bd7c342b2ab39c64d6854dd5bc1b9ba8776dab Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 18:59:39 +0000
Subject: [PATCH 4/6] 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 2d54de5bebcfe67e0a0418a083df37882203b621 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 19:10:08 +0000
Subject: [PATCH 5/6] 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 739b1857d75ca9329af94edf8e8f7e86691c73cc Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 20:50:51 +0000
Subject: [PATCH 6/6] 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(