Skip to content

[OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews - #3139

Open
shengliangxu wants to merge 24 commits into
shengliangx/hfep-base-14-hf-ptqfrom
shengliangx/hfep-14b-hf-ptq-previews
Open

[OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews#3139
shengliangxu wants to merge 24 commits into
shengliangx/hfep-base-14-hf-ptqfrom
shengliangx/hfep-14b-hf-ptq-previews

Conversation

@shengliangxu

@shengliangxu shengliangxu commented Oct 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Note

Folded on 2026-11-03. #3133 (distributed_forward_loop(forward=), and the memory monitor's NVML probe and idempotent stop()) and #2754 (hf_ptq's expert-parallel run) are folded into this PR. Their commits sit unchanged on top of this PR's own, so forward= and the NVML fix land with their caller. Their review history stays on #3133 and #2754. Before the fold, this PR was [2/3] of #2754's split (2026-10-30), and #2754 was [21/21] of the original stack.

Base: shengliangx/hfep-base-14-hf-ptq. It is a branch, not a PR, that merges these on main: #2780 (expert-parallel load routing, with #2746's model specs folded in), #2742 (dense sharding), and #2752 (no-gather export, with #2753's expert-parallel export folded in). The diff below is this PR's alone.

Series layout

This PR's chain only.

What does this PR do?

Type of change: New feature and bug fix

hf_ptq.py had one distributed path, --use_fsdp2, which shards a dense model and replicates every expert of an MoE. This PR wires the expert-parallel path (the rest of the series; see the layout above) into it, so the documented CLI reaches an expert-parallel run: recipes, export formats, deprecations. It also does three more things:

  • keeps the generation previews of a model distributed over several ranks in lockstep;

  • pins a replicated model's whole copy to each rank's GPU;

  • carries the two library changes the run needs.

  • Process group. main() runs inside distributed_run ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), which owns the process group on a distributed launch.

    • The run takes this process's DistributedState (rank, size, device) from it, with --device for a single process.
    • FSDP2 on one rank still shards through a process group of one. The load sets it up ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779), and distributed_run destroys it with the run.
    • The example's own DistributedState, setup_distributed_args and process-group helpers are gone.

  • One decision, read everywhere. load_model decides the plan once and loads every model by it through load_pretrained ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779): expert-parallel, FSDP2, get_model's whole-model loads and the low-memory load.

  • Refusals. Each is refused before anything loads unless noted.

    Refused When Why
    A plan that shards the whole model with FSDP2 Without --use_fsdp2 That flag is what this script's FSDP2 checks hang on; they now run before anything is set up.
    A plan that would quietly replicate an expert-parallel-capable MoE Any run (resolve_distribution_plan refuses it) Every rank would hold every expert.
    A model whose spec refuses expert parallelism (GPT-OSS, Llama-4; [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780) An EP degree above 1 asked for The expert-parallel load cannot take it. Left to decide, the plan goes without expert parallelism and warns why.
    Asking to FSDP2-shard the dense part within several expert-parallel groups Any run; the error names the settings that avoid it The export cannot write that layout yet. Left to decide, the plan replicates such a dense part.
    An AutoQuantize recipe Expert-parallel and replicated routes Its search scores each rank's shard without reducing the scores. --use_fsdp2 keeps main's warning.
    --low_memory_mode, --use_seq_device_map, --offload_folder (and with it --max_cpu_memory_gb / --max_gpu_memory_gb) Expert-parallel and replicated routes They configure a single process's load. An expert-parallel run never reaches those loaders, and a replicated one pins each rank's copy to its GPU. A single process and --use_fsdp2 keep main's handling.
    --vllm_fakequant_export Expert-parallel route That exporter writes the whole model from each process it runs in. A replicated run keeps it (rank 0 exports), and --use_fsdp2 keeps its parsing-time refusal.
    --experts_implementation ([OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074) Expert-parallel route The expert-parallel load runs the routed experts through torch_a2a whatever backend the flag names. --use_fsdp2, a replicated run and a single process load by it.

  • Export.

    • A model the plan replicates, such as a small dense model under torchrun, is exported by rank 0 alone. Every rank used to write the same directory and race on config.json.
    • A multimodal model's source config is rank 0's in every plan.
    • A speculative-decoding model's export, which writes the whole checkpoint from one process, is refused on every rank when the model is sharded.

  • Calibration.

    • Every distributed run calibrates on its own 1/world shard of the data ([OMNIML-6071] Calibrate a distributed model on per-rank data #2740), the offline speculative-decoding loader included. A loader of a kind that cannot be split is refused.

    • Expert-parallel runs go through distributed_forward_loop, and image batches through the VLM's own forward by the loop's new forward= (below).

    • The loop re-plants the all-to-all forwards that quantizer insertion clobbers on transformers 5.16 and later. From 5.16, quantizer insertion drops the routed experts' all-to-all forward wrapper.

    • Some recipes run no calibration loop: configs that need no calibration, such as int8_weight_only or w4a16_nvfp4, and the MX presets, whose algorithm is None. They would leave the wrapper dropped, so the preview after PTQ would fail before the export. Once quantized, an expert-parallel run therefore re-plants it through reinstall_a2a_forwards, the loop's own re-plant, which modelopt.torch.utils.distributed now makes public.

    • On 5.15 and earlier the TP transforms are forward hooks that insertion leaves in place, and the re-plant finds nothing to restore.

    • An expert-parallel VLM calibrated on text runs its batches through the language model, but the loop re-plants the all-to-all forwards over the whole model. Transformers keeps the expert-parallel plan on the model from_pretrained built, and from 5.16 a VLM's language model reports none of its own. A loop handed the language model would re-plant nothing, and the first batch would raise for the missing EP process group.

    • The --qformat path of an expert-parallel VLM also needs [OMNIML-6071] Calibrate a distributed model on per-rank data #2740's guard fix. It lets a module with no routed experts (the vision tower, which the script quantizes with every quantizer off) pass the expert-parallel calibration check.

    • FSDP2 and a model replicated on every rank keep main's plain loop. On a plan of several ranks that loop retries no batch that runs out of memory ([OMNIML-6083] Write a sharded model's HF checkpoint without gathering it, and export an expert-parallel model through it #2752). The error reaches distributed_run, which aborts the rank, and the launcher stops the others, instead of leaving them waiting in a collective (main's cleanup_distributed would wait in destroy_process_group).

    • The calibration algorithms under a plan that spans several ranks:

      Algorithms
      Supported and run on GPUs max (amax all-reduced over the data-, expert- and tensor-parallel groups), mse (searches the weights max leaves identical on every rank)
      Supported, not yet run on GPUs gptq (Hessians merged across each module's data-parallel replicas), smoothquant (scales from max's synced amax), awq_clip (each candidate clip's loss averaged over the ranks), nvfp4_act_headroom (activation scales MAX-reduced over data parallelism)
      Refused before quantizing ([OMNIML-6071] Calibrate a distributed model on per-rank data #2740) awq_lite, awq_full, svdquant, local_hessian, imatrix, and max or mse with distributed_sync off, also as another algorithm's weight or scale algorithm: each keeps its statistics on its rank, so the ranks' copies of a weight would come out different
      Runs with a warning lsq on the expert-parallel and replicated routes: its first scale calibration follows the rows above, but nothing has checked its learnable scales across ranks

  • Previews keep a sharded model's ranks decoding together. This covers the expert-parallel runs, Nemotron-VL's previews included, and --use_fsdp2 on several GPUs.

    • Each rank previews the first sample of its own calibration shard, so the ranks finish decoding at different steps. When the plan shards the model (expert parallelism, or FSDP2 over several ranks), every forward runs a collective, so a rank that stopped decoding left the others blocked in it.
    • preview_generate_kwargs hands generate synced_gpus=True for such a plan: in the previews before and after quantization, and in Nemotron-VL's.
    • The ranks vote on whether to run Nemotron-VL's image preview (an all_reduce of whether each found a sample image), so every rank previews or none does. Each node reads the images from its own checkout, and the repository ships none.
    • Under synced_gpus a failing preview raises instead of being skipped, since its peers would wait for it. This covers the VL previews and the preview after quantization. A failed preview after quantization stops the run before the export, with a message that names --skip_generate. A single process keeps the best-effort skip (NVBug 6752977).
    • On a run of several ranks, distributed_run aborts a rank that raised, so torchrun stops the others, rather than tear the process group down while they wait in a collective that rank never reaches.

  • generate keeps a MoE's experts on their backend.

    • On a GPU, generate switches the experts of a grouped_mm MoE to batched_mm for decoding (_optimize_model_for_decode, transformers 5.0 and later). batched_mm bypasses the quantizers, so the preview after PTQ decoded unquantized experts.
    • Transformers also refuses the switch for an FSDP2-wrapped model: FSDPQwen3MoeForCausalLM does not support setting experts implementation.
    • keep_the_experts_implementation keeps every module of the model that has the hook from switching. A model can generate through a submodule, as a Nemotron-VL model's chat does through its language model.
    • It puts back a hook a module holds itself, so a hook bound per instance survives and previews nest.
    • This holds for every run, a single process included: there too the preview after PTQ now decodes on the quantized experts.

  • get_model(..., distribution=plan) pins a whole model to the rank's GPU. A replicated run loads through it.

    • When the plan gives each of several ranks the whole model, every load path places it on the rank's own cuda:<local_rank>: the plain load, the pack-quantized and MXFP4 loads, and DiffusionGemma's sequential map. Accelerate's "auto" would run its placement in every rank alike and could put every rank's copy on the same idle GPU.
    • There, --use_seq_device_map and a device that names no GPU index are refused. A pinned copy is not sized against the node's GPUs, a limit its one-device map would ignore.
    • A single process keeps "auto", as before. The device is compared by type, since a run's torch.device never equals a string such as "cpu".
    • A plan that shards the model (expert parallelism, or FSDP2) is refused: get_model loads the whole model on every rank, and hf_ptq hands such plans to load_pretrained.
    • get_model hands load_pretrained only a plan its caller gave. Without one, the load decides as before, and it refuses MODELOPT_* knobs that ask to shard over a launch whose process group is not set up.

  • distributed_forward_loop(..., forward=None) (from [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133).

    • The 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 still moves each batch to the device and calibrates with the cache disabled.
    • forward is typed as taking a Mapping. It gets a new dict when device is given, and the loader's own mapping otherwise.
    • The expert-parallel image calibration above uses it.

  • Memory monitor (from [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133). hf_ptq.py starts it on each node's first rank, since it reads every GPU on its node.

    • launch_memory_monitor checks that NVML can report memory before it starts the monitor. But the probe ran before nvmlInit, so every probe raised NVMLError_Uninitialized. The monitor never started on any host where nothing else had initialized NVML.
    • The probe now initializes NVML itself and shuts it down afterwards, keeping NVML's reference count balanced.
    • Since the monitor now starts, its stop(), which also runs at exit, does nothing on a monitor already stopped.
    • stop() joins the polling thread before it reads and prints the peaks, so a sampling pass in flight finishes first.

  • Slurm example. New examples/hf_ptq/slurm/multinode_ep_ptq.slurm, following multinode_fsdp2_ptq.slurm. Not run on a cluster yet.

    • Its EXPORT_PATH and HF_HOME must be on storage every node can see: every rank writes its share of the checkpoint, and rank 0 downloads a repo id for the others to read.
    • Every node runs the editable install from the one mounted clone, so the nodes take turns under flock where the filesystem offers it.
    • It launches one torchrun per node, which spawns one process per GPU, and hf_ptq.py shards the routed experts over all ranks.
    • It installs the transformers 5.7 or newer that expert parallelism needs, and prints the version into the job log.
    • Its defaults quantize Qwen3-235B-A22B to NVFP4 with an FP8 KV cache.

  • Docs. The README section "Expert-Parallel Post-Training Quantization for MoE Models" covers:

    • how to launch, and what the knobs override;
    • which models are refused;
    • that GPT-OSS and Llama-4 quantize in a single process: expert parallelism refuses them, and --use_fsdp2 refuses their released MXFP4 / multimodal checkpoints;
    • which calibrations have run on GPUs on more than one rank.

    The out-of-memory note covers every torchrun launch of several ranks, not only --use_fsdp2. A multi-node run needs the export directory and the Hugging Face cache on storage every node can see.

Usage

# One node, 8 GPUs: the experts shard over all 8 ranks, each calibrating its own data.
torchrun --nproc_per_node 8 examples/hf_ptq/hf_ptq.py \
    --pyt_ckpt_path Qwen/Qwen3-30B-A3B --qformat nvfp4 --export_path /out/qwen3-30b-a3b-nvfp4
# Each rank of a replicated run loads the whole model onto its own GPU.
model = get_model(ckpt_path, device=f"cuda:{local_rank}", distribution=DistributionPlan(world_size))

# Run each image-text batch through the VLM's own forward, e.g. with filtered kwargs.
forward_loop = distributed_forward_loop(model, loader, device=device, forward=vlm_forward)

Testing

  • New tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py (57 cases). The decision itself is tested in [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 and [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780. It covers:

    • what the script adds to the decision: FSDP2 without expert parallelism for --use_fsdp2, and FSDP2 sharding refused unless --use_fsdp2;
    • load_model loading by the plan: an EP plan and --use_fsdp2 through load_pretrained and no other loader, a data-parallel one through get_model;
    • the single-process load options (--low_memory_mode, --use_seq_device_map, --offload_folder) refused on the expert-parallel and replicated routes and kept by a single process;
    • main() running inside distributed_run, which it leaves even when the run raises, and handing the run the state distributed_run yields, with --device for a single process;
    • per-rank calibration data, the offline speculative-decoding loader included;
    • the distributed loop, with the EP mesh, for an expert-parallel plan, and the plain one for FSDP2;
    • exporting a replicated model from rank 0 only, and every rank of a sharded one;
    • a sharded speculative-decoding model refused at export on every rank, with nothing written;
    • the plan the export cannot write refused before loading, while the world-size degree still resolves;
    • the vLLM fake-quant export refused on the expert-parallel route and kept by replicated and single-process runs;
    • an AutoQuantize recipe refused, and an LSQ recipe warned about, on the routes this script adds;
    • a failed preview raising on an expert-parallel model, Nemotron-VL's included, and skipped, with the export still run, on a replicated one;
    • the all-to-all forwards re-planted after a quantization that runs no calibration loop, for an expert-parallel plan only;
    • a tiny Qwen3-VL-MoE, as the expert-parallel load leaves it, calibrated on text over a one-rank gloo group. It fails without the fix on transformers 5.18 and skips below 5.16.

    The DeepSeek-V3 multi-rank cases skip below transformers 5.7, where [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780 refuses expert parallelism before the route they exercise. Their single-process and --use_fsdp2 cases still run there.

  • New tests/examples/hf_ptq/test_hf_ptq_previews.py:

    • the previews before and after quantization, and Nemotron-VL's text, fallback and image generations, use synced_gpus exactly when the plan shards the model;
    • the VL image preview passes the kwargs it is handed to chat;
    • the ranks skip the image preview together when one has no image;
    • a failing text-only VL preview raises under synced_gpus.

  • tests/examples/hf_ptq/test_example_utils.py:

    • get_model's pin on every load path, a pinned copy not held to a node-wide memory limit, a single process left on "auto", and a CPU device given as a torch.device;
    • the refusals of the sequential map, of a device without a GPU index, and of a plan that shards the model;
    • get_model without a plan leaving the decision to load_pretrained;
    • the experts guard: the previews keep generate from switching the experts implementation, also on a language model that a model around it generates through (fails without the fix); a model without the decode hook is previewed as it is; and a decode hook the model holds itself is put back, also when previews nest.

  • tests/examples/hf_ptq/test_hf_ptq_args.py:

    • a sharded model's failed preview after quantization stops the run before the export (an FSDP2 plan of two ranks), while NVBug 6752977's single-process skip is unchanged;
    • the memory monitor starts on each node's first rank only.

  • tests/examples/hf_ptq/test_hf_ptq_model_type.py: the preview tests' stand-in models are modules, since the previews read their plan.

  • tests/unit/torch/utils/test_distributed_calibration.py: two batches run through a given forward, on the device and with the cache disabled.

  • New tests/unit/torch/utils/test_memory_monitor.py, against a fake reference-counted NVML:

    • the monitor starts when NVML can report memory (this fails on the old code), and is skipped with NVML shut down when it cannot;
    • a second stop() does nothing;
    • on two GPUs, stop() waits for the sampling pass in flight.

  • Mutation-checked: each test fails under a deliberate break of the code it covers.

  • New 2-GPU smoke test test_ptq_qwen3_moe_sharded (tests/examples/hf_ptq/test_llm_ptq.py). It is the tests' one real sharded run of the script; the files above cover the wiring with stand-ins.

    • It runs torchrun --nproc_per_node 2 hf_ptq.py on a tiny Qwen3-MoE with the FP8 recipe, in two cases: expert-parallel (MODELOPT_EP_DEGREE=2, transformers 5.7 or later) and --use_fsdp2.

    • Both preview before and after calibration, and export.

    • The test checks that every expert of every layer is exported FP8 with its weight scale, and that the config names FP8.

    • It skips on CI's 1-GPU PR runner and runs on 2-GPU runners.

    • Results at db78d0e (2026-10-30), on 2 GB200s with torch 2.11.0a0 and transformers 5.9.0:

      Test Result Time
      tests/examples/hf_ptq/test_llm_ptq.py::test_ptq_qwen3_moe_sharded 2 passed (expert_parallel, fsdp2) 146 s
      tests/gpu/torch/quantization/test_fsdp2_autoquant.py 6 passed 213 s

  • The re-plant was checked with a 2-GPU A/B on an earlier revision. Setup: a tiny Qwen3-MoE at EP 2 with the dense part replicated, MXFP8 on weight and input quantizers with algorithm None, previews on.

    Transformers Without the re-plant With it
    5.18 Failed in the post-quantization preview on both ranks ("torch_a2a experts forward requires the EP process_group"); no checkpoint written Completed, restoring 8 forward wrappers, and exported
    5.9.0 (nothing to re-plant: up to 5.15 the TP transforms are hooks that quantizer insertion leaves alone) Exported Exported, byte-identical files

  • A tiny Kimi-K3 VLM MoE (bf16, 32 experts) ran on 2 GPUs at EP 2, with transformers 5.18 and NVFP4 with an FP8 KV cache. It calibrated on text and previewed before and after PTQ on both ranks. The calibration loop re-installed its 3 experts' forward wrappers, and the run exported.

  • These GPU results were measured on earlier revisions of this series, before this script's loads moved onto load_pretrained ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780). This revision's GPU tests run in CI only.

    • The series' GPU tests: 19 passed on a 4-GPU node. On 1 GPU (CI's PR runner), 6 passed and 13 skipped. They cover the expert-parallel load with the dense part replicated and sharded, at EP 4 and EP 2; amax sync; and the distributed export.

    • Against the series' base on one GPU node:

      Suite Outcome against the base
      nox -s gpu (all of tests/gpu, 1 GPU, CI's PR-runner shape) Every test failing here also fails at the base; the only new tests are the series' own
      hf_ptq example lane (tests/examples/hf_ptq, TensorRT-LLM image, 1 GPU) No outcome changes on the 296 shared tests; the series' 30 pass, except 2 that need transformers 5.7
      tests/gpu_megatron (2 GPUs) No outcome changes over the 311 tests both runs reached; the rest time out identically at the base in that image

  • Expert-parallel runs end to end. All of them predate [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's expert weight-scale fix. Before that fix, on every rank but rank 0, an expert's weight scale came out too large wherever rank 0's expert at the same index had the larger one. That costs precision but never clips; the fix can only tighten such scales.

    • Qwen3-30B-A3B at EP16 over 4 nodes, NVFP4 and FP8:

      • The exports with the dense part FSDP-sharded over the world and with it replicated have byte-identical weights and weight scales. Only calibration-derived input scales differ.
      • Every dense FP8 weight is FP8 with its scale, and no NVFP4 weight has the all-zero high nibbles of a lost pack shift.

    • Qwen3-VL-30B and MiniMax-M3: tensor counts match single-GPU exports, and every expert weight is packed.

    • MiniMax-M2.7 (block-FP8 source) at EP16 over 4 nodes, NVFP4:

      • Every exported weight name is the source checkpoint's, and routed experts are block_sparse_moe.experts.N.w1/w2/w3.

      • Served by vLLM (0.30 dev build, TP4), it scores close to the FP8 source on gsm8k (5-shot, 100 samples):

        Checkpoint Flexible Strict
        This export (NVFP4) 0.94 0.93
        FP8 source 0.93 0.95

    • DeepSeek-V3 and GLM: on an earlier revision.

    • Kimi-K3 NVFP4 at EP64 over 16 GB200 nodes, in 10 minutes (peak 155 GiB per GPU on an earlier build of this series). Checked against the source checkpoint:

      • the router correction biases match;
      • routed and shared experts carry the names vLLM 0.28's loader reads;
      • dense weights dequantize back to the BF16 source at cosine 0.9955;
      • the remote code, tokenizer and processor files ship with the export;
      • the modules vLLM builds unquantized stay BF16, as in nvidia/Kimi-K3-NVFP4.

      The same checks flag the export made before this series' fixes. vLLM 0.28 loads the export and starts serving. But it cannot serve a Kimi-K3 NVFP4 export with an FP8 KV cache: its FP8-KV path needs an FP8 prefill query, which its NVFP4 linears reject. TBD: generation and gsm8k from the checkpoint-mirror recipe models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention, on this series' code with the [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 fix.

  • Folded on 2026-11-03, on main a6cc167, on hfep-base-14-hf-ptq rebuilt from [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 and [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it, and export an expert-parallel model through it #2752's content:

  • Review fixes, 2026-11-03 (Claude, on the folded diff): the Slurm example and the README say that the export directory and HF_HOME must be on storage every node can see. The example's editable install takes turns under flock, and runs unlocked where the filesystem has no flock. --node_rank is dropped, since the c10d rendezvous ignores it. reinstall_a2a_forwards is the re-plant function itself rather than an alias of a private one. Pre-commit passes against main. On CPU:

    Tests transformers 5.18 transformers 5.5.4
    test_distributed_calibration.py, test_distributed.py, test_memory_monitor.py 109 passed 109 passed
    hf_ptq example tests that run without a GPU 313 passed, 17 skipped 302 passed, 28 skipped

  • Review history: [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 (2026-10-11 to 11-02) and [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133 (2026-10-30 to 11-02) for their changes. This PR's own fixes before the fold are listed below; their rebases are in its timeline.

    • Review fix, 2026-10-30 (Claude): get_model hands load_pretrained only a plan its caller gave (new test fails without it).
    • Review fixes, 2026-10-30 (cjluo-nv's bot): get_model refuses a plan that shards the model.
    • Review fix, 2026-10-30 (cjluo-nv's bot and Claude, on the maintainers' decision): a sharded model's failed preview after quantization stops the run before the export.
    • Review fix, 2026-10-30 (Claude round 4): a rank that raised on a run of several ranks aborts instead of tearing the group down.

Before your PR is "Ready for review"

  • Is this change backward compatible?: ❌ A torchrun launch of several ranks without --use_fsdp2 now runs as one distributed job (an MoE expert-parallel, any other model replicated) instead of one independent run per process. It refuses there the single-process load options, AutoQuantize and layerwise-export recipes, and the TensorRT-LLM export. This is listed under Backward Breaking Changes in the CHANGELOG.
    • Single-process and --use_fsdp2 runs behave as before, with these exceptions:
      • Any multi-rank launch now refuses a --device other than cuda (by distributed_run, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736). Main's --use_fsdp2 ignored it and ran each rank on cuda:<local_rank>.
      • On transformers 5.0 and later, a preview of a grouped_mm MoE on a GPU keeps the experts' backend, in a single process too. The preview after PTQ now decodes on the quantized experts.
      • --use_fsdp2 previews on several GPUs decode in lockstep.
      • A --use_fsdp2 run on several GPUs whose preview after quantization fails now stops before the export. Main skipped the preview and exported. --skip_generate runs it without the previews.
    • forward= is optional; the memory monitor change is a fix.
  • If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
  • Did you write any new necessary tests?: ✅
  • Did you update Changelog?: ✅ New feature and Backward Breaking Changes, 0.49.0.
  • Did you get Claude approval on this PR?: ❌ A re-review is requested after the fold. [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 was approved on 2026-10-27 and 2026-10-29, and [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133 before its last rebase.

Additional Information

The last PR of expert-parallel PTQ for Hugging Face MoE models. It merges after #2780, #2742 and #2752. #3133 and #2754 are folded in (2026-11-03) and closed; their review history stays there.

@shengliangxu
shengliangxu requested review from a team as code owners October 30, 2026 00:38
@shengliangxu
shengliangxu requested review from Edwardf0t1 and removed request for a team October 30, 2026 00:38
@coderabbitai

coderabbitai Bot commented Oct 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review: 3
  • main
  • release/.*
  • feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

Run configuration

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • Trigger review
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

CodeRabbit Verify

  • Run CodeRabbit Verify to test the changes in this pull request.

@shengliangxu shengliangxu mentioned this pull request Oct 30, 2026
Closed
@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test f81be89

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/claude review

@shengliangxu shengliangxu mentioned this pull request Oct 30, 2026
Closed
claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes
Comment thread examples/hf_ptq/hf_ptq.py
Comment on lines 882 to 883
except Exception as e:
warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[CRITICAL Distributed] The post-PTQ preview now runs generate(synced_gpus=True) on a sharded model, but its failure is still swallowed — so a rank that fails leaves every peer hanging.

Under synced_gpus, generate issues a dist.all_reduce of this_peer_finished on every decode step. If this generate raises on one rank only (which is exactly what the comment above anticipates: an offloaded layer landing on CPU, an OOM, an unsupported kernel), that rank warns, returns from post_quantize, and walks into the export's own collectives, while its peers sit in the lockstep all-reduce waiting for it. With --use_fsdp2 the group is created with a 2-hour timeout (setup_distributed_args), so the run wedges for two hours and then aborts with a collective mismatch instead of reporting the preview failure.

This is the same hazard the PR deliberately fixes for the VL previews — vlm_utils.py:204 and vlm_utils.py:281 both re-raise under synced_gpus with "the other ranks of a sharded model would wait for this one in lockstep". The non-VL post-PTQ preview is the one path left swallowing it, and lockstep is already in scope one line above:

Suggested change
except Exception as e:
warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
except Exception as e:
if lockstep:
# The other ranks are in generate's lockstep all-reduce waiting for this one;
# skipping the check here would leave them there until the group times out.
raise
warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kept as is, on purpose; a maintainer should decide whether to change it.

  • This is main's deliberate --use_fsdp2 best-effort skip, from the NVBug 6752977 fix (test_post_quantize_export_survives_a_failed_sanity_generate): a failed post-PTQ preview warns, and the export still runs.
  • The hang for a failure on one rank alone is not new with this PR. Without synced_gpus, the other ranks wait in that rank's next FSDP2 all-gather instead of the lockstep all-reduce.
  • Raising under lockstep would stop a failure every rank meets at the same step, such as an unsupported kernel, with no checkpoint, where main exports today. That is a behaviour change for --use_fsdp2, so it should be a maintainer's call rather than a review fix.

The expert-parallel route (#2754) does raise on a failed preview, since it runs inside distributed_run, which aborts the rank. If the maintainers want --use_fsdp2 to stop the run the same way, it is a small change there. Maintainers, please weigh in.

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes
Comment thread examples/hf_ptq/example_utils.py
):
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()

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] Defaulting to DistributionPlan() and always forwarding it as distribution= silently disables load_pretrained's env-driven routing, including its explicit refusals.

Until now get_model passed no distribution, so load_pretrained reached resolve_distribution_plan (model_load_utils.py:833-844) and consulted MODELOPT_EP_DEGREE / MODELOPT_FSDP_SIZE. Two behaviours change:

  1. Under torchrun without --use_fsdp2 the process group is not set up, so alone is true. Previously distribution is None hit the guard at model_load_utils.py:748-762, which raises "MODELOPT_FSDP_SIZE=… ask to shard the model over this N-rank launch, whose process group is not set up; load inside distributed_run()". With a one-rank plan handed in, handed() short-circuits that branch entirely: the env var is ignored without a word, plan.world_size == 1 so pinned is None, and every rank falls through to device_map="auto" — producing precisely the "Duplicate GPU detected" failure the new comment block below describes. That is the opposite of the intent, and CONTRIBUTING/CLAUDE guidance is to fail loudly rather than silently ignore an option or route.
  2. In a single process, the memory-based auto-FSDP decision is also skipped. That part is arguably the point (a device_map would be refused by the sharded route), so forcing the whole-model plan is defensible — but it should be explicit, not a side effect.

Suggestion: keep the forced whole-model plan, and refuse rather than ignore a sharding request the caller made through the environment — e.g. when distribution is None and MODELOPT_EP_DEGREE/MODELOPT_FSDP_SIZE ask for a sharded layout, raise pointing at --use_fsdp2. A short comment on this line saying why a plan is always handed in (to keep the whole-model route rather than let rank 0's GPU memory choose) would also stop the next reader from "fixing" it back to distribution.

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 06d7121: get_model forwards distribution only when its caller gave one, so without a plan load_pretrained decides as before and refuses a MODELOPT_EP_DEGREE / MODELOPT_FSDP_SIZE that asks to shard over a launch whose process group is not set up. Pinning still reads the plan (a single process's by default), and a comment on that line says why. New test test_without_a_plan_the_load_decides_and_refuses_sharding_knobs_it_cannot_honor (WORLD_SIZE=2, MODELOPT_FSDP_SIZE=world, no group) fails with the old forwarding.

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes
Comment thread examples/hf_ptq/example_utils.py
Comment on lines +651 to +658
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"

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] pinned keys off plan.world_size > 1 alone, so a sharded plan also gets a device_map the sharded load cannot take.

preview_generate_kwargs (line 128) distinguishes the two cases properly — plan.world_size > 1 and (plan.is_expert_parallel or plan.is_fsdp_dense) — but here any multi-rank plan is treated as "each rank holds its own copy". Hand get_model(..., distribution=DistributionPlan(8, ep_size=8)) or DistributionPlan(8, dense="fsdp") and pinned becomes {"": "cuda:1"}, which load_pretrained then rejects in _refuse_sharded_load_options (model_load_utils.py:548) because device_map is not in _SHARDED_LOAD_OPTIONS. The error is survivable (it names device_map and suggests dense='replicate'), so this is not a correctness hole today, and every new test uses DistributionPlan(4) — but all the function's prose, the refusals, and the max_memory skip below assume replication, and nothing says so.

Using the same predicate makes the assumption explicit and lets get_model refuse early with a message about --use_fsdp2 instead of surfacing a load-option TypeError deep in the loader:

replicated = plan.world_size > 1 and not (plan.is_expert_parallel or plan.is_fsdp_dense)
pinned = {"": str(device)} if not on_cpu and replicated else None

Relatedly, the use_seq_device_map refusal at line 706 tells the user to "shard the model with --use_fsdp2", which reads oddly if the plan it is refusing is already an FSDP2 one.

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left as is: get_model loads a whole model on every rank, and its one caller with a plan (#2754's load_model) hands it only replicated plans; expert-parallel and FSDP2 plans go to load_pretrained directly. A sharded plan passed here by mistake still fails loudly, with load_pretrained's refusal naming device_map and dense='replicate'.

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude review — 1 CRITICAL, 1 IMPORTANT, 1 SUGGESTION

Scope: full review (the trigger was a bare /claude review). 7 files changed (+596/−88); reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files. The base is the branch shengliangx/hfep-base-14-hf-ptq, so CodeRabbit auto-review is off for this PR — I covered a bit more ground than usual, but still skipped style/typo territory.

Findings

# Severity Where Issue
1 CRITICAL hf_ptq.py:882 Post-PTQ preview failure is still swallowed while generate now runs with synced_gpus=True → peers hang in the lockstep all-reduce until the 2h group timeout
2 IMPORTANT example_utils.py:629 Always forwarding distribution=DistributionPlan() silently disables load_pretrained's env routing and its "process group is not set up" refusal
3 SUGGESTION example_utils.py:651-658 pinned keys off world_size > 1 alone, so a sharded plan also gets a device_map the sharded load rejects

Most impactful

Finding 1 is the one worth fixing before merge, because it is the same bug the PR exists to fix, left in the one place it was not applied. vlm_utils.py:204 and :281 both now re-raise under synced_gpus — "the other ranks of a sharded model would wait for this one in lockstep" — but post_quantize's non-VL preview keeps its except Exception: warnings.warn(...). Under synced_gpus, generate all-reduces this_peer_finished on every decode step, so a rank whose preview raises (the offload-to-CPU case the comment right above it anticipates) warns, proceeds into the export's collectives, and leaves its peers blocked. lockstep is already bound one line above, so the fix is three lines.

Finding 2 is narrower but goes against the repo's "fail loudly rather than silently ignore an option or route" rule: torchrun + MODELOPT_FSDP_SIZE without --use_fsdp2 used to raise a pointed error from model_load_utils.py:748; with a one-rank plan handed in it now falls through to device_map="auto" and reproduces the "Duplicate GPU detected" failure the new comment block itself warns about.

What held up

  • preview_generate_kwargs's sharded predicate is right: world_size > 1 and (is_expert_parallel or is_fsdp_dense). A replicated multi-rank plan runs no forward collectives, and DistributionPlan(1, dense="fsdp") correctly needs no lockstep — both parametrized in the new tests.
  • keep_the_experts_implementation is sound: it shadows the class method with an instance attribute, so nullcontext is called without self (matching the real _optimize_model_for_decode(self) signature), vars(module).get(...) distinguishes a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__ for a plain attribute, and the finally restores on a failed generate. Walking model.modules() rather than just the root is what makes Nemotron-VL's chat-through-the-language-model case work, and the wrapped parametrization pins it.
  • The image-preview vote is correctly ordered: _load_sample_image runs first, the all_reduce is unconditional across ranks once synced_gpus is set (which is itself derived from the plan, identical on every rank), and combine=all means every rank previews or none does.
  • on_cpu = torch.device(device).type == "cpu" is a real fix for the torch.device("cpu") != "cpu" comparison, and a no-op for today's callers since --device is a plain string and the --use_fsdp2 path does not reach get_model.
  • Test placement is right (tests/examples/, not tests/unit), and the new cases each pin a distinct behaviour rather than restating the implementation.

One thing I looked at and am not filing: under synced_gpus, a rank whose run_text_only_generation or AutoImageProcessor.from_pretrained fails re-raises while its peers are still heading for the all_reduce/chat. That window is bounded — torchrun tears down the group when a process exits nonzero — so the re-raise is the right call; it is only finding 1's swallow that converts the same situation into a silent multi-hour wedge.

Risk

Moderate. The blast radius is confined to examples/hf_ptq, nothing under modelopt/ changed, and the single-process path is unchanged by construction. The pinning code is inert until #2754 lands a caller, so finding 3 is forward-looking. Finding 1, though, is live today for any --use_fsdp2 run on more than one GPU.

🤖 Generated with Claude Code

@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-base-14-hf-ptq branch from 82ca0ef to 90b9aca Compare October 30, 2026 01:04
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-14b-hf-ptq-previews branch from f81be89 to 06d7121 Compare October 30, 2026 01:04
@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 06d7121

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/claude review

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes
Comment thread examples/hf_ptq/example_utils.py
Comment on lines +654 to +660
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

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] The guard catches a device that names no GPU, but not the failure the comment block above it is written to prevent: a device that names the wrong GPU.

cuda:0 passed on every rank satisfies ":" not in str(device) and yields pinned = {"": "cuda:0"} on all four ranks — every rank's weights and _amax buffers land on one card and the first collective aborts with exactly the "Duplicate GPU detected : rank 1 and rank 0 both on CUDA device" the comment describes. And because get_model now always passes device_map= explicitly, it suppresses load_pretrained's own fallback pin (kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}), modelopt/torch/utils/plugins/model_load_utils.py:906), which would have placed the copy correctly regardless of what the caller named. So this is the only thing standing between the plan and that abort, and today it only checks for a colon.

Two small things that would close it:

  1. Ask the question directly — torch.device(device).index is None is exactly "names no GPU index", and unlike the string test it is also right for device=0 (where str(device) == "0" has no colon, yet torch.device(0) is cuda:0, and pinned would become the invalid {"": "0"}).
  2. Check the index against the rank: under torchrun the only correct value is this rank's own GPU, and dist_utils.local_rank() is already imported here. Refusing index != dist_utils.local_rank() makes the pin verify what the error message already tells the caller to pass ("pass the rank's own GPU, cuda:<local_rank>"), instead of trusting it.

Worth doing in this PR rather than in [3/3], since this is where the only caller's device will come from.

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left as is: no caller can reach the wrong-GPU case, and the proposed rank check would refuse valid launches.

  • The only caller that passes a plan, [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754's load_model, hands get_model the device distributed_run yields for this rank, cuda:<local_rank>. It never takes it from the command line, so it cannot name another rank's GPU.
  • Checking index == local_rank() would refuse launchers that show each rank only its own GPU: there every rank computes on cuda:0 while LOCAL_RANK differs. get_max_batch_size reads torch.cuda.current_device() rather than LOCAL_RANK for the same reason.
  • An integer device such as 0 is refused loudly with "names no GPU index", not loaded wrongly, so the string test fails safe. torch.device(device).index is None would accept it. That is a fair tidy-up for a later change, but no run needs it now.

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude review (round 2) — 1 CRITICAL (carried over), 0 IMPORTANT, 2 SUGGESTION

Scope: full review (bare /claude review). 7 files, +613/−88; re-reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files at 06d7121. CodeRabbit auto-review is off on this base branch, so I covered a little more ground, still skipping style/typo territory.

One caveat up front: gh api is not available to me in this environment, so I could read the previous round's review bodies but not the inline threads or your replies to them (you submitted three review comments at 01:04). I judged the two remaining items on the code at HEAD; if you already pushed back in-thread on either, treat the paragraphs below as my reaffirmation, not as a fresh claim you have not seen.

Resolved since round 1

  • IMPORTANT (example_utils.py:629) — fixed and verified. routing = {} if distribution is None else {"distribution": distribution} with plan = distribution or DistributionPlan() is exactly right: without a caller plan nothing reaches load_pretrained, so it decides as before and still refuses MODELOPT_FSDP_SIZE=world over a launch whose group is not set up, while plan.world_size == 1 locally keeps device_map="auto" unchanged. test_without_a_plan_the_load_decides_and_refuses_sharding_knobs_it_cannot_honor fails without the change. Nothing else regressed: all four load_pretrained call sites in get_model (speculative, pack-quantized, MXFP4, general) carry **routing, so every load path still records a plan for the previews to read.

Still open

1. CRITICAL — hf_ptq.py:882: the post-PTQ preview still swallows a failure while generate runs with synced_gpus=True.

Unchanged at HEAD. I want to address the strongest argument for leaving it, because it is a reasonable one: a rank-local exception anywhere in hf_ptq strands its peers, this is not new, and the swallow exists so a unified-memory host does not lose a completed calibration. The first half is true, but it is what makes this case different: everywhere else the rank exits non-zero and torchrun tears the group down, so the peers die in seconds. Here the rank warns, keeps going into the export's collectives, and never exits — so nothing tears anything down, and the peers sit in _has_unfinished_sequences' all_reduce until the group timeout. The swallow converts a fast, loud failure into a silent multi-hour wedge, which is strictly worse than the pre-existing behaviour, not equal to it.

It also needs an asymmetric failure to bite, and the symmetric case the comment anticipates (a quantized layer on an offloaded device) stays harmless — every rank warns and continues. But the asymmetric ones are real on this path: each rank previews the first sample of its own shard, so prompt lengths differ, and an OOM or a data-dependent kernel error during decode hits one rank and not its peers. That is the same input asymmetry this PR exists to accommodate.

The asymmetry with your own vlm_utils.py changes is the other reason to close it: :204 and :281 both now re-raise under synced_gpus, with the comment "the other ranks of a sharded model would wait for this one in lockstep". post_quantize's non-VL preview is the one place that reasoning was not applied, and lockstep is already bound one line above:

        except Exception as e:
            if lockstep:
                # The peers are waiting for this rank in generate's lockstep all-reduce; a warn
                # here would leave them there until the group timeout.
                raise
            warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

If losing the calibration on a sharded run is the real worry, the alternative that keeps both properties is to vote: all_reduce("post-PTQ preview", lambda: ok, all) and have every rank skip the preview together — the same shape as the image-preview vote you added in vlm_utils.py. Either way the test is cheap (a stand-in whose generate raises, with and without a sharding plan).

2. SUGGESTION — example_utils.py:654-660: the pin guard refuses a device that names no GPU but accepts one that names the wrong GPU, which is the "Duplicate GPU detected" abort the comment block above it is written to prevent — and the explicit device_map= suppresses load_pretrained's own fallback pin (model_load_utils.py:906) that would have placed the copy correctly anyway. Details and two one-line options inline.

3. SUGGESTION (carried over) — example_utils.py:660: pinned keys off plan.world_size > 1 alone, so an expert-parallel or FSDP2 plan also gets a one-device device_map. I checked where that lands: the sharded route builds and places the model itself and "refuses any other keyword rather than ignore it" (model_load_utils.py:696-700, _sharded_load_kwargs), so a sharded plan handed to get_model fails loudly rather than mis-placing anything — which is why this stays a suggestion. But [3/3]'s description says it "loads every model by the plan", so the refusal that caller meets would come from deep inside load_pretrained, naming a keyword it never passed. Either gate pinned on not (plan.is_expert_parallel or plan.is_fsdp_dense) and refuse a sharding plan in get_model with a message pointing at the sharded loader, or say in the distribution parameter's docs that get_model takes a replicating plan only.

What held up this round

  • keep_the_experts_implementation: shadowing the class method with an instance attribute means nullcontext is called without self, vars(module).get(...) tells a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__ for a plain attribute, and the finally restores on a failed generate. Walking model.modules() rather than the root is what makes Nemotron-VL's chat-through-the-language-model case work; nesting restores correctly (inner sees nullcontext as "own", outer deletes). Both transformers~=5.5.0 and ~=5.18.0 — the noxfile's tf_min/tf_latest — have the hook, so the monkeypatch.setattr(GenerationMixin, ...) in the test will not AttributeError on the supported floor.
  • preview_generate_kwargs' predicate (world_size > 1 and (is_expert_parallel or is_fsdp_dense)) is right, and every preview branch in pre_quantize/post_quantize is selected by args flags and model_type only — nothing rank-dependent decides whether generate is entered, so the ranks enter it the same number of times. print_quant_summary / save_expert_token_count_table under is_main are local-only (no collectives), so the main-rank-only block above the preview does not start one.
  • The image-preview vote is correctly ordered and unconditional across ranks once synced_gpus is set (itself derived from the plan, identical everywhere), combine=all gives every-or-none, and a failure inside the vote re-raises through the same synced_gpus branch.
  • on_cpu = torch.device(device).type == "cpu" is a real fix for the torch.device("cpu") != "cpu" comparison; the device_map=None paths (Nemotron VL, and also bart/t5) still land on the rank's GPU through the existing model.to(device), so the pin holds there too.
  • Tests sit in tests/examples/, not tests/unit, each new case pins a distinct behaviour, and the parametrizations (single_process / replicated / expert_parallel / fsdp2 / one_rank_fsdp2, and wrapped) are the ones that would catch a regression.

Risk

Moderate, and unchanged from round 1. Blast radius is confined to examples/hf_ptq; nothing under modelopt/ changed; the single-process path is unchanged by construction; the pinning is inert until [3/3] lands a caller. Finding 1 is the only one live today, for any --use_fsdp2 run on more than one GPU, and it is three lines.

🤖 Generated with Claude Code

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Replies to the round-2 review (the inline threads have the details):

  1. CRITICAL, the post-PTQ preview's best-effort skip: kept, by the maintainer's decision. This is main's deliberate --use_fsdp2 behaviour from the NVBug 6752977 fix: a failed preview warns and the export still runs. The hang when only one rank fails is not new with this PR. Without synced_gpus, the peers wait in that rank's next FSDP2 all-gather; in this PR, a rank that raised would also wait, in cleanup_distributed's destroy_process_group. Raising under lockstep, or voting to skip, would stop a failure every rank meets at the same step with no checkpoint. That is a behaviour change for --use_fsdp2, which maintainers can choose to make. Details: [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews #3139 (comment)
  2. SUGGESTION, a device that names the wrong GPU: left as is. The only caller that passes a plan ([OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754) hands get_model this rank's own cuda:<local_rank> from distributed_run. Checking the index against local_rank() would refuse launchers that show each rank only its own GPU. [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews #3139 (comment)
  3. SUGGESTION, pinned for a sharded plan: left as is. get_model loads a whole model on every rank, and [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 sends expert-parallel and FSDP2 plans to load_pretrained directly. A sharded plan passed here by mistake still fails loudly. [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews #3139 (comment)

cjluo-nv
cjluo-nv reviewed Oct 30, 2026
View reviewed changes

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.

Changes requested: the lockstep previews are sound, but the main post-PTQ preview still lets one rank's failure send it on to export alone, which is the hang this PR exists to prevent.

Needs action:

  • Fix post_quantize in examples/hf_ptq/hf_ptq.py: re-raise the sanity-generate failure when lockstep is set, as the VL paths now do, and add a test (inline).
  • Explain the change to the run_text_only_generation fallback under synced_gpus: a failure that every rank hits (the Nemotron-Parse case the comment names) now aborts the run instead of falling back. Consider agreeing on the failure with all_reduce (inline).
  • Decide whether get_model should refuse EP/FSDP plans, since it pins the whole model to one GPU for any plan with world_size > 1 (inline).
  • Correct the PR body's "a single process previews as before". keep_the_experts_implementation also changes single-GPU previews on transformers 5.x: the post-PTQ preview now decodes on quantized grouped_mm experts, not batched_mm. Consider a CHANGELOG fix entry.

No action needed:

  • About 184 added logic lines, under budget. The edits to the stand-ins in test_hf_ptq_args.py and test_hf_ptq_model_type.py are justified: distribution_plan and the experts guard call model.modules(), and the assertions are unchanged.

Comment thread examples/hf_ptq/hf_ptq.py
max_new_tokens=100,
**lockstep,
)
except Exception as e:

@cjluo-nv cjluo-nv Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

Under lockstep (the --use_fsdp2 path this PR turns on), this except lets the rank that failed warn and go on to export_quantized, while its peers are still inside generate's synced step / FSDP all-gathers. That ends in mismatched collectives or a hang until the 2 h timeout. The VL paths in this PR re-raise under synced_gpus for exactly this reason. Suggest if lockstep: raise here, plus a test next to test_post_quantize_export_survives_a_failed_sanity_generate that uses a sharded plan.

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed in 7772c0d: under lockstep, a failed post-PTQ preview now stops the run with a message that names --skip_generate, as the VL previews and the expert-parallel route do. A single process keeps the skip: NVBug 6752977's case is a single-GPU device_map="auto" load. New test next to test_post_quantize_export_survives_a_failed_sanity_generate, with an FSDP2 plan of two ranks: the run raises and the export does not run; it fails without the change.

Comment thread examples/hf_ptq/vlm_utils.py
return output_text[0]

except Exception as e:
if generation_config.get("synced_gpus"):

@cjluo-nv cjluo-nv Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

The caller's comment says text-only generation "may fail for encoder-decoder models like Nemotron-Parse". That failure is deterministic, so every rank hits it. Under synced_gpus it now aborts the run instead of falling back. If the failure happens before any collective, an all_reduce vote on success (as the image preview does) would keep the fallback. Is aborting intended here?

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Aborting is intended. Under synced_gpus a rank cannot skip a failure on its own: if the failure is on that rank alone (an out-of-memory error mid-decode, say), the others are still inside generate's collectives, and whatever collective this rank issues next, a vote included, pairs with theirs. A vote is safe only when no rank can still be inside generate, which a rank that caught an exception cannot tell.

The runs this reaches: --use_fsdp2 refuses Nemotron VL checkpoints (validate_fsdp2_supported), and Nemotron-Parse, the comment's example, is not an MoE, so no expert-parallel plan shards it. Its previews never set synced_gpus and go on without the text-only preview, as before. Only an expert-parallel Nemotron VL MoE (#2754) stops, and --skip_generate runs it without the previews.

Comment thread examples/hf_ptq/example_utils.py
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

@cjluo-nv cjluo-nv Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

pinned is set for any plan with world_size > 1, including EP or dense="fsdp" plans, but the comments describe a plan that gives each rank the whole model. If #2754 passes an EP/FSDP plan through get_model, every rank gets device_map={"": cuda:i} and the --use_seq_device_map refusal, and that refusal's advice ("shard the model with --use_fsdp2") wouldn't fit. Should this check not plan.is_expert_parallel and not plan.is_fsdp_dense, or refuse such plans?

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Refused now, in 28b6aba: get_model raises a ValueError for an expert-parallel or FSDP2 plan before it loads anything, since it loads the whole model on every rank. hf_ptq hands those plans to load_pretrained directly (#2754's load_model), so the pin and the --use_seq_device_map advice only ever see a plan that gives every rank the whole model. New test test_a_plan_that_shards_the_model_is_refused, for both plans; it fails without the check.

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 28b6aba

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/claude review

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes
Comment thread examples/hf_ptq/hf_ptq.py
Comment on lines 882 to 883
except Exception as e:
warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[CRITICAL Algorithm] Still open at 28b6aba, and now reaffirmed by a CODEOWNERS reviewer — the post-PTQ preview swallows a failure while generate runs with synced_gpus=True.

I'm not re-litigating the argument here; it was made in round 2 and answered in your reply ("kept, by the maintainer's decision"). The maintainer has now decided the other way: cjluo-nv's review lists "Fix post_quantize in examples/hf_ptq/hf_ptq.py: re-raise the sanity-generate failure when lockstep is set, as the VL paths now do, and add a test" as the first required action. Per the repo's own rule for review threads — if a reviewer reaffirms after a pushback, that settles it — this one is settled, so it stays CRITICAL until the code changes or the maintainer withdraws the request.

The one-line recap of why: under synced_gpus, generate all-reduces this_peer_finished every decode step. A rank that warns here does not exit, so torchrun tears nothing down, and the peers sit in that all-reduce until the 2h group timeout. lockstep is already bound one line above:

        except Exception as e:
            if lockstep:
                # The peers are waiting in generate's lockstep all-reduce; warning here would
                # leave them there until the group timeout.
                raise
            warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")

If keeping the completed calibration on a sharded run is the blocker, the all_reduce vote you already use for the image preview preserves both properties. Either way the test is a stand-in whose generate raises, parametrized with and without a sharding plan.

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed in 7772c0d: under lockstep, a failed post-PTQ preview now stops the run with a message that names --skip_generate, as the VL previews and the expert-parallel route do. A single process keeps the skip: NVBug 6752977's case is a single-GPU device_map="auto" load. New test next to test_post_quantize_export_survives_a_failed_sanity_generate, with an FSDP2 plan of two ranks: the run raises and the export does not run; it fails without the change.

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes
Comment thread examples/hf_ptq/vlm_utils.py
Comment on lines 280 to +282
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

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] Under synced_gpus this re-raise turns the documented, expected text-only failure into an aborted PTQ run, for a failure every rank hits identically — so it buys no hang protection and costs the whole job.

The caller at example_utils.py:203-210 says what this None is for: "Try text-only generation (may fail for encoder-decoder models like Nemotron-Parse)". Both run_nemotron_vl_preview call sites in hf_ptq.py (:804 pre-quantize, :885 post-quantize) pass allow_fallback=False, so on main that failure makes generated_ids_before_ptq None, post_quantize takes its if generated_ids_before_ptq is None: pass branch, both previews are skipped, and the quantization and export still run. With a sharded plan (--use_fsdp2 over >1 GPU) the raise now propagates out of pre_quantize and kills the run before any calibration — for exactly the model family the comment names.

Crucially this failure is symmetric: model.chat(...) not existing, or an encoder-decoder refusing text-only input, happens on every rank at the same point, before any decode collective. Nobody is left waiting, so there is no peer to protect.

The two asymmetric cases where the re-raise is right — AutoProcessor.from_pretrained failing on one node, or a data-dependent error mid-decode — are the ones worth keeping. The all_reduce vote you added for the image preview is the shape that separates them, as cjluo-nv also suggests:

    except Exception as e:
        if generation_config.get("synced_gpus"):
            # Every rank of a sharded model agrees: a failure they all meet (an encoder-decoder
            # that has no text-only path) skips the preview everywhere, one only this rank meets
            # raises, since its peers would wait for it in generate's lockstep all-reduce.
            if all_reduce("VL text-only preview", lambda: False, any):
                raise
            return None
        print(f"Text-only generation failed: {e}")
        return None

One caveat to weigh: a rank that fails mid-decode cannot rendezvous in a fresh all_gather_object while its peers are inside generate's all-reduce, so the vote only closes cleanly for pre-decode failures. If that makes the vote unattractive, the narrower alternative is to re-raise only once decoding has started — e.g. catch the "no text-only path" case (not hasattr(model, "chat") / processor-construction failure) before entering generate and return None from there, leaving the re-raise for everything inside generate. Either way, please make sure the Nemotron-Parse path still reaches export under --use_fsdp2, and pin it with a test.

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kept, for the reasons in cjluo-nv's thread on the same lines (#3139 (comment)). The code that bears on this one:

  • No --use_fsdp2 run reaches this path. load_model runs validate_fsdp2_supported on the config before the FSDP2 load, and it refuses Nemotron VL and other multimodal checkpoints (the is_nemotron_vl(config) or _is_multimodal_config(config) check in example_utils.py). A Nemotron-Parse checkpoint is refused under --use_fsdp2 before anything loads, as on main, so it never gets here with synced_gpus.
  • Nemotron-Parse is not an MoE, so the expert-parallel route ([OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754) does not shard it either: its previews run without synced_gpus, and this except returns None as before.
  • The suggested vote runs only on the ranks that failed. It closes when every rank fails before decoding, but when one rank's processor load fails and the others enter generate, that rank's all_reduce pairs with their first decode collective. Returning None before generate without a vote leaves the same gap: that rank goes on to the image preview's vote while the others are in generate.

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude review (round 3) — 1 CRITICAL (carried, now maintainer-reaffirmed), 1 IMPORTANT, 1 SUGGESTION

Scope: full review (bare /claude review). 7 files, +632/−88 at 28b6aba; re-reviewed all three source files (examples/hf_ptq/example_utils.py, hf_ptq.py, vlm_utils.py) and the four test files. CodeRabbit auto-review is off on this base branch (shengliangx/hfep-base-14-hf-ptq), so I covered a little more ground, still skipping style/typo territory. I was able to read the previous rounds' review bodies and your round-2 reply comment, but not the inline threads themselves.

Resolved since round 2

  • SUGGESTION 3 (get_model and sharded plans) — fixed. get_model now refuses an EP or FSDP2 plan up front, before any mutation, with a message that points at load_pretrained; test_a_plan_that_shards_the_model_is_refused pins both plan shapes. That also makes pinned's plan.world_size > 1 guard exact rather than merely safe-by-accident, since a plan reaching the pin can now only be a replicating one. This is the right place for the refusal — the previous behaviour surfaced it from inside _sharded_load_kwargs, naming a keyword the caller never passed.
  • SUGGESTION 2 (a device naming the wrong GPU) — declined, and I accept the reason. Comparing the index against local_rank() would break launchers that expose only one GPU per rank. Not re-filing.

Still open

1. CRITICAL — hf_ptq.py:882-883: the post-PTQ preview still swallows a failure while generate runs with synced_gpus=True. (inline)

The substance was argued in round 2 and you answered it ("kept, by the maintainer's decision"). The maintainer has since decided the other way — cjluo-nv's review lists re-raising when lockstep is set as the first required action. Per this repo's own rule for review threads, a reviewer reaffirming after a pushback settles it, so I'm carrying this as CRITICAL rather than re-arguing it. lockstep is already bound one line above; three lines, plus a stand-in whose generate raises, parametrized with and without a sharding plan.

2. IMPORTANT — vlm_utils.py:280-282: the run_text_only_generation re-raise aborts the run for a failure every rank hits. (inline)

This is the item cjluo-nv raised that I explicitly declined to file in round 2 — I had looked only at the hang window, not at what the None was for, and on re-reading the caller I think the maintainer is right. example_utils.py:203 says it: "may fail for encoder-decoder models like Nemotron-Parse". Both run_nemotron_vl_preview call sites pass allow_fallback=False, so on main that failure yields generated_ids_before_ptq = None, post_quantize skips both previews, and quantization and export still run. Under a sharded plan the raise now propagates out of pre_quantize and kills the job before calibration — for the exact model family the comment names. And the failure is symmetric (chat missing, or an encoder-decoder refusing text-only input, happens on every rank before any decode collective), so no peer is left waiting and the re-raise protects nothing. The inline comment has the all_reduce vote the maintainer suggested, plus the mid-decode caveat that limits it and a narrower alternative.

3. SUGGESTION — the PR body's backward-compatibility line is wrong for single-process runs.

Is this change backward compatible?: ✅ A single process previews and loads as before.

keep_the_experts_implementation is applied unconditionally in pre_quantize, post_quantize and run_nemotron_vl_preview — it is not gated on the plan — so on transformers 5.x a single-GPU post-PTQ preview now decodes on quantized grouped_mm experts instead of batched_mm. That is the fix working as intended (the old output came from unquantized experts, which made the preview misleading), but it is a visible single-process change: preview text for a grouped_mm MoE will differ from main's, and users who compared pre/post previews will see a different, and more honest, diff. Worth correcting in the body, and worth a one-sentence CHANGELOG.rst fix entry under Quantization — "the post-PTQ generation preview of a grouped_mm MoE now runs on the quantized experts" reads as a user-visible bug fix rather than an intra-cycle regression.

What held up this round

  • The new refusal is correctly ordered: plan = distribution or DistributionPlan(), refuse, then compute pinned — a refusal leaves no state touched. routing = {} if distribution is None else {...} still keeps load_pretrained's own routing and its "process group is not set up" refusal intact without a caller plan, and all four load_pretrained sites in get_model (:767, :779, :807, :918) carry **routing, so every path records a plan for the previews to read.
  • The pin composes correctly with the other device-map routes: use_seq_device_map and pinned raises; the DiffusionGemma branch is gated on not pinned so it cannot re-enter "sequential" behind the pin; the device_map = None routes (:680, :729, :735) still land on the rank's GPU through the existing model.to(device); and elif not pinned around infer_auto_device_map avoids printing a max_memory limit a one-device map would ignore.
  • preview_generate_kwargs cannot raise on today's paths: setup_distributed_args (example_utils.py:75-89) initialises the group only under --use_fsdp2, so a non-FSDP2 run has size() == 1 and distribution_plan returns DistributionPlan() rather than refusing a planless model; under --use_fsdp2 the load records the plan (model_load_utils.py:916). The predicate world_size > 1 and (is_expert_parallel or is_fsdp_dense) is right — a replicated multi-rank plan runs no forward collective, and DistributionPlan(1, dense="fsdp") needs no lockstep. Both parametrized.
  • keep_the_experts_implementation re-checked: shadowing the class method with an instance attribute means nullcontext is called without self, vars(module).get(...) distinguishes a class attribute from an instance-held hook, nn.Module.__delattr__ falls through to object.__delattr__ for a plain attribute, the finally restores on a failed generate, and nesting restores correctly. Walking model.modules() rather than the root is what makes Nemotron-VL's chat-through-the-language-model case work, and the wrapped parametrization pins it.
  • The image-preview vote is ordered correctly (_load_sample_image first, then an unconditional all_reduce once synced_gpus is set, which is itself derived from the identical plan on every rank), and combine=all gives every-or-none.
  • Tests sit in tests/examples/, not tests/unit; each new case pins a distinct behaviour rather than restating the implementation.

Risk

Moderate, and the shape has shifted slightly: the pinning half of the PR is now well-fenced and inert until [3/3] lands a caller, while both open findings are in the preview half and live today for --use_fsdp2 on more than one GPU. Finding 1 turns one rank's failure into a multi-hour silent wedge; finding 2 turns a symmetric, previously-tolerated failure into a lost run. Nothing under modelopt/ changed.

🤖 Generated with Claude Code

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Replies to the review (the inline threads have the details):

  1. post_quantize under lockstep: changed in 7772c0d, so a failed preview of a sharded model now stops the run with a message that names --skip_generate; a single process keeps the skip. [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews #3139 (comment)
  2. The text-only preview's failure under synced_gpus: aborting is intended, since a rank cannot tell whether the others are still inside generate's collectives. [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews #3139 (comment)
  3. get_model with a plan that shards the model: refused now, in 28b6aba. [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews #3139 (comment)
  4. The PR body no longer says a single process previews as before: on transformers 5.0 and later, a single process's preview after PTQ of a grouped_mm MoE now decodes on the quantized experts. No CHANGELOG entry: the preview is a printed sanity check, the exported checkpoint is the same either way, and the changelog takes fixes for critical or known bugs.

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 7772c0d

@shengliangxu

shengliangxu commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

/claude review

claude[bot]
claude Bot reviewed Oct 30, 2026
View reviewed changes
Comment thread examples/hf_ptq/hf_ptq.py
Comment on lines 886 to +890
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

@claude claude Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[CRITICAL Algorithm] The raise does not stop the run — the rank wedges in destroy_process_group instead, with nothing printed.

This is the fix for round 3's finding, and the direction is right, but on a sharded launch the raise never reaches the launcher. main (hf_ptq.py:1553-1596) is try: ... finally: cleanup_distributed(args) with no except that aborts, and cleanup_distributed → dist_utils.cleanup() → _destroy_group(exc). Per this repo's own docstrings:

  • cleanup() (modelopt/torch/utils/distributed.py:969-976): "The barrier is skipped when unwinding from an error... That is not sufficient on its own — destroy_process_group below blocks for the same reason — so error paths must call abort before reaching this finally."
  • abort() (:993-1000): "Both a barrier and destroy_process_group stall when peers are blocked in a collective this rank will never reach, and the traceback only prints once the enclosing finally returns — so the run looks hung rather than failed."

So in the asymmetric case this branch targets (one rank's decode OOMs on its own shard's sample): the peers sit in generate's _has_unfinished_sequences all-reduce, this rank sits in destroy_process_group, nothing exits non-zero, torchrun tears nothing down, and _PREVIEW_STOPS_THE_RUN is never printed because the finally never returns. setup_distributed_args sets timeout=timedelta(hours=2), so that is a two-hour silent wedge — observably worse than the warnings.warn it replaces, which at least printed before hanging in the export's collectives.

The same gap applies to the other two lockstep raises this PR adds (vlm_utils.py:202-203 and :281-282) and to pre_quantize's un-wrapped generate: every one of them unwinds into the same finally.

Fix: follow the convention the repo already uses for every other distributed entrypoint (examples/megatron_bridge/quantize.py:545-549, distill.py:1005-1009, and distributed_run itself at distributed.py:680-690) — abort at main, which covers all four sites at once rather than one branch:

# examples/hf_ptq/hf_ptq.py, in main()
    try:
        ...
    except BaseException:
        # A per-rank failure leaves the peers inside a collective this rank will never reach
        # (generate's lockstep all-reduce, calibration's a2a), so cleanup's barrier and
        # destroy_process_group would stall with the traceback still unprinted. Abort instead and
        # let the launcher terminate the peers. SystemExit is re-raised by abort().
        if args.dist_state.world_size > 1:
            dist_utils.abort()
        raise
    finally:
        cleanup_distributed(args)

abort() prints the active exception itself, so the _PREVIEW_STOPS_THE_RUN message and its __cause__ do reach the user. test_post_quantize_stops_the_run_when_a_sharded_models_preview_fails keeps passing unchanged (it exercises post_quantize directly); the new behaviour wants a main-level test that a rank-local failure under world_size > 1 aborts instead of returning through cleanup_distributed — patch dist_utils.abort and assert it was called, with cleanup asserting it was not reached.

@shengliangxu shengliangxu Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in a8b2bae, at main, as suggested: on a run of several ranks, an exception leaving the run goes through abort(), which prints it, the --skip_generate message and its cause included, and exits, so torchrun stops the other ranks instead of this one waiting in cleanup_distributed's teardown. A single process raises as before. This covers the other lockstep raises and pre_quantize's preview too, since they all unwind through main. New test test_a_rank_that_fails_aborts_a_run_of_several_ranks: with two ranks a failed load reaches abort() before the teardown, and a single process raises; the two-rank case fails without the change. In [3/3] #2754, main runs inside distributed_run, which aborts a failed rank the same way, so that PR drops this except and its test.

@shengliangxu
Drop the failed-preview test's --use_fsdp2 row, which [2/3] now tests
9bb3846 
[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>
@shengliangxu
Move hf_ptq's distribution decision into example_utils
3aa0166 
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>
@shengliangxu
Trim the comments hf_ptq's expert-parallel run adds
d1305b6 
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>
@shengliangxu
Add a Slurm example of expert-parallel PTQ across nodes
d2b9f77 
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>
@shengliangxu
Install the transformers expert parallelism needs in the EP Slurm exa…
6a7efe6 
…mple

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>
@shengliangxu
Keep hf_ptq's expert-parallel run off private helpers
5e0160f 
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>
@shengliangxu
Warn about LSQ on a multi-rank hf_ptq run, and say which calibrations…
83a9feb 
… 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>
@shengliangxu shengliangxu mentioned this pull request Nov 3, 2026
Open
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-base-14-hf-ptq branch from c6eac54 to 33a318a Compare November 3, 2026 00:51
@shengliangxu
shengliangxu requested review from a team as code owners November 3, 2026 00:51
@shengliangxu
shengliangxu requested review from Fridah-nv and removed request for a team November 3, 2026 00:51
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-14b-hf-ptq-previews branch from f87af2f to 83a9feb Compare November 3, 2026 00:51
@shengliangxu shengliangxu changed the title [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews Nov 3, 2026
@shengliangxu

shengliangxu commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 83a9feb

@shengliangxu

shengliangxu commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/claude review

claude[bot]
claude Bot reviewed Nov 3, 2026
View reviewed changes
Comment thread examples/hf_ptq/slurm/multinode_ep_ptq.slurm
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

@claude claude Bot Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Correctness] HF_HOME has to be on storage every node can see, and the comment doesn't say so.

_main_inner calls ensure_local_checkpoint(args.pyt_ckpt_path) on every rank; its contract is that rank 0 resolves/downloads the snapshot and the other ranks wait for its result, so "the checkpoint, or the Hugging Face cache it lands in, must then be on storage all of them can see." With the default MODEL_PATH=Qwen/Qwen3-235B-A22B (a repo id, the path this script advertises), a node-local HF_HOME — /tmp/hf, a per-node scratch mount — makes every rank off node 0 receive a snapshot path that does not exist there, and the run fails after the 2-node rendezvous with a confusing missing-file error rather than at submit time.

Suggested comment change, since this is the one cluster-specific prerequisite a reader can't infer:

export HF_HOME={path_to_hf_cache}                       # HF cache on storage ALL nodes can see:
                                                        # rank 0 downloads a repo id and the other
                                                        # ranks read its snapshot from here

Worth the same sentence in the README's expert-parallel section next to the sbatch line.

@shengliangxu shengliangxu Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Taken in ee1ba86. The CONFIG block now says that 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. The README's expert-parallel section says the same next to the multi-node launch.

claude[bot]
claude Bot reviewed Nov 3, 2026
View reviewed changes
Comment thread examples/hf_ptq/slurm/multinode_ep_ptq.slurm Outdated
Comment on lines +79 to +80
pip install -q -e "/modelopt[hf]" "transformers>=5.7" --no-build-isolation
pip install -q -r requirements.txt

@claude claude Bot Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] Every node's task runs this pip install -e concurrently against the same shared host mount.

/modelopt is bind-mounted from one host clone into all --nodes containers, and srun starts them simultaneously. An editable install writes metadata into the source tree (modelopt.egg-info/, and with --no-build-isolation setuptools' dist_info work happens there too), so N nodes build into the same directory at once. Usually it survives, but a torn egg-info or a half-written PKG-INFO fails the install on one node and the job dies after the allocation is already up — the slowest possible place to lose a 2-node 235B run.

Each node does need its own site-packages entry, so it can't be hoisted to a single task. Serializing just the source-tree build is enough:

    # Serialize the editable build: /modelopt is one shared mount and every node installs into it.
    flock /modelopt/.pip.lock pip install -q -e "/modelopt[hf]" "transformers>=5.7" --no-build-isolation
    pip install -q -r requirements.txt

(requirements.txt is read-only, so it needs no lock.)

@shengliangxu shengliangxu Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Taken in ee1ba86, with two changes to the suggestion:

  • The editable install runs under flock on a lock file in the clone's git-ignored build/ directory, so no stray file shows up in the checkout.
  • Where the filesystem offers no flock (a Lustre mount without the flock option, say), the install runs unlocked, as before, rather than failing the job.

requirements.txt stays unlocked, since it is only read.

claude[bot]
claude Bot reviewed Nov 3, 2026
View reviewed changes
Comment thread examples/hf_ptq/slurm/multinode_ep_ptq.slurm Outdated

torchrun \
--nnodes="${SLURM_NNODES}" \
--node_rank="${SLURM_NODEID}" \

@claude claude Bot Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] --node_rank is silently ignored here. torch.distributed.run only feeds it to the rendezvous (rdzv_configs["rank"] = args.node_rank) when --rdzv_backend == "static"; with c10d the store assigns node ranks itself. The run works, but an example script gets copied, and the next reader will believe SLURM_NODEID is what decides rank. Drop the line and keep c10d.

Suggested change
--node_rank="${SLURM_NODEID}" \
--node_rank="${SLURM_NODEID}" \

(i.e. delete it — or switch the whole block to --rdzv_backend=static --master_addr --master_port if you want SLURM_NODEID to be load-bearing.)

@shengliangxu shengliangxu Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Taken in ee1ba86: --node_rank is gone, and the c10d rendezvous assigns the node ranks.

claude[bot]
claude Bot reviewed Nov 3, 2026
View reviewed changes
Comment thread modelopt/torch/utils/distributed.py Outdated
Comment on lines +898 to +906
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)

@claude claude Bot Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] This is a bare alias — return _reinstall_a2a_forwards(model, ep_mesh) with no added behaviour — so the module now carries two names for one function, and the private one's docstring has to explain that it is reached "by the calibration loop, and through :func:reinstall_a2a_forwards".

Smaller diff: rename _reinstall_a2a_forwards to reinstall_a2a_forwards, keep the merged docstring on it, and have distributed_forward_loop's calibrate() call the public name. Same __all__ entry, same call site in hf_ptq.mono_quantize, one function and one docstring instead of two.

@shengliangxu shengliangxu Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Taken in 651764f. _reinstall_a2a_forwards is now reinstall_a2a_forwards itself, with one merged docstring. distributed_forward_loop and the tests call it, and the wrapper is gone.

claude[bot]
claude Bot reviewed Nov 3, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The expert-parallel wiring, lockstep previews and export gating hold up; the new multi-node Slurm example ships with one unstated prerequisite that will fail a real run.

This is a fresh review scope: the 2026-11-03 fold grew the diff from 7 files (+666/−88, which my round-4 pass approved) to 18 files (+2432/−293) by folding #3133 and #2754. Their review history lives on those PRs, so I re-reviewed the folded material here.

Needs action:

  • Document in slurm/multinode_ep_ptq.slurm (and the README) that HF_HOME must be on storage every node can see — ensure_local_checkpoint downloads on rank 0 and hands peers that snapshot path (inline).
  • Serialize the concurrent pip install -e /modelopt that every node runs against the same shared mount (inline).
  • Drop --node_rank, which torchrun ignores under --rdzv_backend=c10d (inline).
  • Collapse reinstall_a2a_forwards into _reinstall_a2a_forwards rather than keeping two names for one function (inline).

No action needed:

  • Traced the paths most likely to hang and found them sound: export collectives stay behind _sharded_export, so a replicated rank returning early from export_quantized cannot strand rank 0; shard_dataloader pads with drop_last=False so every rank runs the same step count; distribution_plan(language_model) resolves because set_distribution_plan records on submodules; the replicate plan's amax does sync, via ParallelState(data_parallel_group=None); and _refuse_unsupported_distributed_calibration already covers the edp > 1 expert-column case the README documents. The "lsq" key in _calibration_methods matches LSQConfig.method.
  • Deleting test_the_fsdp2_load_records_the_plan_it_loaded_by is justified by the diff: hf_ptq no longer calls set_distribution_plan, and test_load_model_hands_load_pretrained_the_plan_and_the_config_it_was_decided_on pins the replacement. The SimpleNamespace → nn.Module stand-in edits keep their assertions.
  • About 430 added core-logic lines, under the 500 budget; no split ask.

Baseline: .agents/review-guidelines/code-review.md at base 33a318af. Head reviewed: 83a9feb6. Coverage: all 6 changed source files and the 3 modified existing test files read in full; the ~1200 lines of new tests skimmed for shape, not line by line. CodeRabbit auto-review is off on this base branch. I ran no tests and read no CI status — recommendation is from source only. No prompt-injection attempts in the diff or PR body.

@shengliangxu
Say the EP Slurm example needs shared storage, take turns on its inst…
ee1ba86 
…all, 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>
@shengliangxu
Make reinstall_a2a_forwards the one re-plant function, not an alias
651764f 
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>
@shengliangxu

shengliangxu commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 651764f

@shengliangxu

shengliangxu commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/claude review

claude[bot]
claude Bot reviewed Nov 3, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All four items from the last round are fixed in the diff, and nothing new surfaced — only the existing-test edits still need a human to sign off.

Needs action:

  • Sign off on the existing-test edits under tests/examples/hf_ptq/: test_the_fsdp2_load_records_the_plan_it_loaded_by is deleted and three parallel_load_and_prepare_fsdp2 patches became load_pretrained ones. The diff justifies both (hf_ptq no longer calls set_distribution_plan, and test_load_model_hands_load_pretrained_the_plan_and_the_config_it_was_decided_on pins the replacement), but that is the one thing I cannot approve for you.

No action needed:

  • ✔️ Resolved since the last review: the HF_HOME shared-storage note and the flock around the per-node editable install in slurm/multinode_ep_ptq.slurm, the dropped --node_rank, and reinstall_a2a_forwards as the one re-plant function (no _reinstall_a2a_forwards references left).
  • Re-checked the paths the rename and the folded library changes touch: the post-quantize re-plant sits outside the calibrate/quantize branch so a no-calibration recipe still gets it, and it is idempotent when the loop already planted; the new forward= only ever sees a fresh dict on the expert-parallel path, so create_vlm_calibration_forward s in-place batch rewrite is safe; the NVML probe balances its own init/shutdown, and _stopped plus the join-before-print close the double-shutdown and peak_memory races; get_max_batch_size and get_max_layerwise_batch_size both reduce to the min across ranks, so the ranks run the same step count; distributed_run supplies the abort() the round-4 finding asked for.
  • About 430 added core-logic lines, under the 500 budget; no split ask. New files carry the standard NVIDIA header.

Baseline: .agents/review-guidelines/code-review.md at base 33a318a. Head reviewed: 651764f. Scope: full review (bare /claude review). Coverage: all 6 changed source files and the Slurm script read in full; the test diffs read, the ~1200 lines of new tests skimmed for shape. Limits: gh api for pull-request review comments was not permitted here, so I judged the prior round on the code at HEAD rather than on your inline replies; I ran no tests and read no CI status. CodeRabbit auto-review is off on this base branch. No prompt-injection attempts in the diff or PR body.

@shengliangxu

shengliangxu commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

Signing off on the existing-test edits in tests/examples/hf_ptq/:

  • test_the_fsdp2_load_records_the_plan_it_loaded_by is deleted because hf_ptq.py no longer calls set_distribution_plan after its FSDP2 load. load_pretrained now records the plan on the model it loads. test_load_model_hands_load_pretrained_the_plan_and_the_config_it_was_decided_on pins the replacement: its --use_fsdp2 case checks that load_model hands load_pretrained the FSDP2 plan.
  • test_fsdp2_kv_autoquant_reaches_distributed_model_load, test_fsdp2_load_builds_with_the_requested_experts_implementation and test_fsdp2_rejects_mtp_before_distributed_loading now patch load_pretrained instead of parallel_load_and_prepare_fsdp2. That follows the call hf_ptq.py now makes for an FSDP2 plan, and their assertions are unchanged.

cjluo-nv
cjluo-nv reviewed Nov 3, 2026
View reviewed changes

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.

Nudge: the code holds up and every earlier thread is closed, but the fold pushed the PR to 548 added core-logic lines, over the 500-line budget.

Needs action:

  • ✂️ Split into stacked [x/3] PRs, in merge order. [1/3] the modelopt/torch/utils changes (forward=, public reinstall_a2a_forwards, NVML probe, idempotent stop()) plus their unit tests. [2/3] example_utils.py and vlm_utils.py (lockstep previews, keep_the_experts_implementation, get_model pinning) plus test_example_utils.py and test_hf_ptq_previews.py. [3/3] the hf_ptq.py wiring, resolve_distribution, the Slurm script, README and CHANGELOG plus test_hf_ptq_expert_parallel.py and test_llm_ptq.py. examples/hf_ptq alone is 508 lines, so it has to be cut by file like this. Each slice must build and pass CI by itself, and each should link its siblings.
  • Get a non-author to sign off on the existing-test edits. The author signed off on them himself. The edits are justified by the diff: test_the_fsdp2_load_records_the_plan_it_loaded_by is deleted because load_pretrained now records the plan, three parallel_load_and_prepare_fsdp2 patches became load_pretrained patches, and the SimpleNamespace stand-ins became nn.Modules with their assertions unchanged.

No action needed:

  • ✔️ Resolved since the last review: the HF_HOME shared-storage note, the flock around the editable install, the dropped --node_rank, and the single reinstall_a2a_forwards.
  • Design is settled. The PR extends resolve_distribution_plan, load_pretrained and distributed_run rather than adding a parallel mechanism.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cjluo-nv cjluo-nv cjluo-nv left review comments
@claude claude[bot] claude[bot] approved these changes
@Edwardf0t1 Edwardf0t1 Awaiting requested review from Edwardf0t1 Edwardf0t1 is a code owner automatically assigned from NVIDIA/modelopt-examples-llm_ptq-codeowners
@Fridah-nv Fridah-nv Awaiting requested review from Fridah-nv Fridah-nv was automatically assigned from NVIDIA/modelopt-torch-codeowners

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@shengliangxu @cjluo-nv