[shengliangxu]

Note

[3/3] of this PR's split (2026-10-30, on review), cut to the 500-line budget: [1/3] #3133 library pieces → [2/3] #3139 hf_ptq's previews and whole-model placement → [3/3] this PR, hf_ptq's expert-parallel run. They merge in that order. The series is a tree of PRs (re-cut 2026-10-27 and 2026-10-29), so PRs that do not depend on each other can be reviewed and merged in parallel; this PR was [21/21] of the original stack.

Base: shengliangx/hfep-14b-hf-ptq-previews, [2/3]'s branch, which sits on shengliangx/hfep-base-14-hf-ptq, a branch (not a PR) that merges #2746, #2742, #2753 and #3133. The diff below is this PR's alone.

Series layout

This PR's chain only; the other PRs of the series sit on the same helper branches.

• main ([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088, [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099, [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122, [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131, [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751, [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134 and [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 merged)
  • [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 → [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 load routing
  • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H experts
  • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
  • hfep-base-6a-ep-loader (merges [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
    • [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133 [1/3] library pieces for hf_ptq's distributed runs
    • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 expert-parallel loader, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding, [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 → [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 expert parallelism in the model specs (each split into slices too)
  • hfep-base-12c-nogather-export (merges [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
    • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export
      • [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export
  • hfep-base-14-hf-ptq (merges [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742, [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 and [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133)
    • [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 [2/3] previews in lockstep, whole-model loads pinned by plan
      • [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 [3/3] expert-parallel PTQ from hf_ptq.py (this PR)

What does this PR do?

Type of change: New feature

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.

• 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, which the load sets up ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779) and distributed_run destroys 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, with resolve_distribution_plan ([OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780, [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746) wrapped in example_utils.resolve_distribution, which adds this script's checks, 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. load_pretrained records the plan, so the script records none itself: [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's record after the --use_fsdp2 load is gone. Everything after reads the plan back from the model, not the flag or the process group: the calibration device, the data sharding, the calibration loop, and get_model's GPU placement.
  • MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE override the automatic choices.
  • --use_fsdp2 asks the decision for FSDP2 and no expert parallelism. Its --experts_implementation ([OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074) still goes on the config the FSDP2 load builds from.
  • --use_fsdp2 now loads through load_pretrained's FSDP2 route, which [1/3] [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133 opens to the text-only Qwen3.5 and Qwen3.5-MoE VLMs main's FSDP2 path calibrates since [OMNIML-5987] [3/3] Support Qwen text-only FSDP2 AutoQuant loading and export #2676.
• Refusals. A plan that shards the whole model with FSDP2 raises without --use_fsdp2, the flag this script's FSDP2 checks hang on, which now run before anything is set up. resolve_distribution_plan refuses a plan that would quietly replicate an expert-parallel-capable MoE, and a model whose spec refuses expert parallelism (GPT-OSS, Llama-4, [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746) raises unless --use_fsdp2 rules expert parallelism out. Left to decide, the plan replicates a dense part it cannot FSDP2-shard within several expert-parallel groups, which the export cannot write yet; asking for that layout raises before loading and names the settings that avoid it. An AutoQuantize recipe raises on the expert-parallel and replicated routes, since its search scores each rank's shard without reducing the scores; --use_fsdp2 keeps main's warning. --low_memory_mode, --use_seq_device_map and --offload_folder (with it --max_cpu_memory_gb / --max_gpu_memory_gb) configure a single process's load, so an expert-parallel or replicated run refuses them before loading: an expert-parallel run never reaches those loaders, and a replicated one pins each rank's copy to its GPU, which they cannot work with. A single process and --use_fsdp2 keep main's handling. An expert-parallel run also refuses --vllm_fakequant_export before loading, since 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. It refuses --experts_implementation ([OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074) before loading too, since 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, image batches through the VLM's own forward by the loop's forward= ([1/3] [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133). The loop re-plants the all-to-all forwards that quantizer insertion clobbers on transformers 5.16 and later. From transformers 5.16, quantizer insertion drops the routed experts' all-to-all forward wrapper, and a recipe whose calibration runs no loop (configs that need no calibration, such as int8_weight_only or w4a16_nvfp4, and the MX presets, whose algorithm is None) would leave it dropped, so the preview after PTQ would fail before the export. Once quantized, an expert-parallel run therefore re-plants it. 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, so 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, which lets a module with no routed experts (the vision tower 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 #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. The previews keep a sharded model's ranks decoding together, keep generate from switching a MoE's experts to batched_mm, and stop the run when a sharded model's preview fails, rather than leave the other ranks in the next step's collective ([2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139); the error says the export was not reached and names --skip_generate. Here that covers the expert-parallel runs, Nemotron-VL's previews included.
• GPU placement. A replicated run loads through get_model by its plan, which pins each rank's whole copy to its own GPU on every load path ([2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139). --use_seq_device_map and --offload_folder are refused before loading there, as above.
• Docs. README section "Expert-Parallel Post-Training Quantization for MoE Models": how to launch, what the knobs override, which models are refused, and that GPT-OSS and Llama-4 (refused by expert parallelism, and by --use_fsdp2 for their released MXFP4 / multimodal checkpoints) quantize in a single process.
• Memory monitor. Each node's first rank starts it, since it reads every GPU on its node. The NVML probe fix it needs is [1/3] [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133.

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

Testing

• New tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py (55 cases; the decision itself is tested in [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 and [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746; the placement and preview tests moved with their code to [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139):

  • 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 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;
  • the all-to-all forwards re-planted after a quantization that runs no calibration loop, for an expert-parallel plan only. A 2-GPU A/B on an earlier revision of this branch (tiny Qwen3-MoE, EP 2 with the dense part replicated, MXFP8 on weight and input quantizers with algorithm None, previews on): on transformers 5.18 the run without the re-plant failed in the post-quantization preview on both ranks ("torch_a2a experts forward requires the EP process_group") and wrote no checkpoint, and with it the run completed, restoring 8 forward wrappers, and exported; on transformers 5.9.0, where there is nothing to re-plant (up to 5.15 the TP transforms are hooks that quantizer insertion leaves alone), both runs exported byte-identical files;
  • a tiny Qwen3-VL-MoE, as the expert-parallel load leaves it, calibrated on text over a one-rank gloo group (fails without the fix on transformers 5.18; skipped below 5.16). On GPUs at this revision, a tiny Kimi-K3, a VLM MoE (bf16, 32 experts), on 2 GPUs at EP 2 with transformers 5.18, NVFP4 with an FP8 KV cache, 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;
  • 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): 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 is the tests' one real sharded run of the script; the file above covers the wiring with stand-ins. 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

    An earlier, expert-parallel-only version of the smoke test passed on 2 GB200s in 29.5 s at f4bb395488 (2026-10-16).

• tests/examples/hf_ptq/test_hf_ptq_args.py: the memory monitor starts on each node's first rank only.

• Plan tests the library's and get_model's tests cover are dropped, and the load-call, calibration-loop and memory-monitor pairs are one parametrized test each. Each drop is checked by mutation.

• test_hf_ptq_expert_parallel.py also covers: the offline speculative-decoding loader split across ranks; 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 --use_fsdp2 case is tested in [2/3]); and an AutoQuantize recipe refused on the routes this script adds.

• The DeepSeek-V3 multi-rank cases of test_hf_ptq_expert_parallel.py skip below transformers 5.7, where [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 refuses expert parallelism before the route they exercise; their single-process and --use_fsdp2 cases still run there.

• Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.

• The GPU results below were measured on earlier revisions of this branch, 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 #2780); this revision's GPU tests run in CI only. The series' GPU tests on this branch: 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;
  • the distributed export.

• Against the series' base, on one GPU node:

  • nox -s gpu (all of tests/gpu, 1 GPU, CI's PR-runner shape): every test failing on this branch also fails at the base, and the only new tests are the series' own.
  • The 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: 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, which 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 0.94 / 0.93 on gsm8k (5-shot, 100 samples, flexible / strict), against the FP8 source's 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.

• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.

• Review fixes, 2026-10-27: the Whisper calibration loader stays whole on every rank, so the preview matches its printed transcript. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.

• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). hf_ptq keeps main's removal of the FSDP2 KV-cache AutoQuantize refusal ([OMNIML-5987] [2/3] Support DP/FSDP2 KV AutoQuant and checkpoint resume #2664); main's new test that such a run reaches the distributed load now stubs load_pretrained, which this PR loads every distributed plan through, and an import that only the removed refusal test used is gone.

• Review fixes, 2026-10-27 (Claude's review after the rebase): --use_fsdp2's checks (validate_fsdp2_supported) run on the checkpoint's config before the plan is decided, as they ran before this PR, so a checkpoint they refuse (VILA, multimodal, pack-quantized) gets their message instead of the meta-device build's error. The new test fails without the change.

• Review suggestions, 2026-10-27 evening: hf_ptq refuses a layerwise-export recipe on a multi-rank launch, before the load; the memory monitor runs on each node's first rank; the unreachable Whisper loader exemption is gone; the off-GPU warning covers a replicated rank's device.

• Review suggestions, 2026-10-27 night: one helper, layerwise_export_requested, decides both the pre-load refusal and the layerwise export switch; a comment names the load paths without a device map (Nemotron VL), where each rank builds a whole host copy.

• Re-cut on 2026-10-29 onto the new base ([OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 2221490, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 aec25d4, [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 a5f06c4): the 57 commits rebased with three conflicts, all mechanical (load_model's replaced --use_fsdp2 block against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's --experts_implementation lines, a test import list, the changelog), and seven commits on top:

  • --experts_implementation ([OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074) goes on the config the FSDP2 load builds from again, where [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 set it in the block this PR replaces; [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's test now stubs load_pretrained;
  • an expert-parallel run refuses --experts_implementation before loading (new test);
  • get_model compares its device by type: the run state's device is always a torch.device now ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), which never equals a string such as "cpu" (new test case);
  • load_pretrained's FSDP2 route loads the text-only Qwen3.5 and Qwen3.5-MoE VLMs that main's --use_fsdp2 supports since [OMNIML-5987] [3/3] Support Qwen text-only FSDP2 AutoQuant loading and export #2676, instead of refusing them as non-causal-LMs (new test);
  • the VLM text-calibration test's one-rank group answers as the plan's two-rank expert-parallel row, which [OMNIML-6071] Calibrate a distributed model on per-rank data #2740's group check now reads;
  • the calibration comment no longer says FSDP2 keeps the split-on-OOM retry, which [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 drops for a plan of several ranks.
  • main's 2-GPU test_fsdp2_kv_autoquant_example ([OMNIML-5987] [2/3] Support DP/FSDP2 KV AutoQuant and checkpoint resume #2664) hands load_model the options it now reads (--sparsity_fmt, --qformat, --vllm_fakequant_export, and [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's --experts_implementation). No PR runner has two GPUs; on CPU, with the FSDP2 loader stubbed, its arguments raised AttributeError before the change and reach the loader after it.

  Each new test fails with its fix reverted. Pre-commit, mypy included, passes against main, and this PR's unit tests with the neighbouring tests/unit/torch/utils ones pass on CPU with transformers 5.18 and 5.5.4.

• Docs, 2026-10-29: the hf_ptq README's out-of-memory note covers every torchrun launch of several ranks, not only --use_fsdp2: the expert-parallel and replicated runs this PR adds neither retry a calibration batch that runs out of memory ([OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752) nor probe the batch size they estimate. Pre-commit passes against main.

• Split on 2026-10-30 into [1/3] [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133, [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139 and this PR (meenchen's bot review): the three diffs together are this PR's former diff; tests moved with their code; this PR's tests with their neighbours pass on CPU (transformers 5.18: 562 passed; 5.5.4: 520 passed, 42 skipped), and pre-commit, mypy included, passes against main. The CHANGELOG's Backward Breaking entry now lists every change to a torchrun launch of several ranks without --use_fsdp2, as the review asked.

• Rebased on 2026-10-30 onto [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139, itself rebased onto main 99c3834's rebuilt parents: no conflicts, own diff unchanged (the experts_implementation=None line of the FSDP2 KV-AutoQuantize GPU test now comes from [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074). Pre-commit passes against main; this PR's tests with their neighbours pass on CPU (transformers 5.18: 565 passed; 5.5.4: 523 passed, 42 skipped).

• Review fixes, 2026-10-30 (keval): the calibration test's docstring no longer says FSDP2 keeps a split-on-OOM retry ([OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 removed it), and the smoke-test bullet above names test_ptq_qwen3_moe_sharded with both its cases. Pre-commit passes against main.

• Restacked on 2026-10-30 onto [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139's review fix (get_model refuses a plan that shards the model; this PR loads such plans through load_pretrained): no conflicts, own diff unchanged. Pre-commit passes against main; the hf_ptq example tests that run without a GPU pass on CPU (transformers 5.18: 306 passed, 17 skipped; 5.5.4: 295 passed, 28 skipped).

• Restacked on 2026-10-30 onto [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139's lockstep fix: [2/3] now stops the run on a failed preview of any sharded model, so this PR's expert-parallel-only raise and its message moved down into it, and the failed-preview test drops its --use_fsdp2 row, which [2/3] tests. Pre-commit passes against main; the same hf_ptq example tests pass on CPU (transformers 5.18: 306 passed, 17 skipped; 5.5.4: 295 passed, 28 skipped).

• Review fixes, 2026-10-30 (meenchen, Claude): the distribution decision and its checks moved to example_utils.resolve_distribution, next to validate_fsdp2_supported; the comments this PR adds to hf_ptq.py are trimmed; new Slurm example slurm/multinode_ep_ptq.slurm, following multinode_fsdp2_ptq.slurm (not run on a cluster yet). Restacked onto [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139's abort fix: main here runs inside distributed_run, which aborts a failed rank itself, so [2/3]'s except and its test go. Pre-commit passes against main; the hf_ptq example tests that run without a GPU pass on CPU (transformers 5.18: 306 passed, 17 skipped; 5.5.4: 295 passed, 28 skipped).

• Review fix, 2026-10-30 (Claude): the expert-parallel Slurm example installs the transformers 5.7 or newer it says expert parallelism needs, and prints the version into the job log. Pre-commit passes against main.

• Rebased on 2026-10-30 onto main 059588b, through [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139: main's [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128 added an AWQ check to example_utils.py that reads the process group's size, so this PR keeps the dist_utils import, and main now brings the load_recipe import; the rest of the own diff is unchanged. Pre-commit passes against main; the hf_ptq example tests that run without a GPU pass on CPU (transformers 5.18: 310 passed, 17 skipped; 5.5.4: 299 passed, 28 skipped).

• Review suggestions, 2026-10-30 (Fridah-nv): hf_ptq.py re-plants the experts' all-to-all forwards through modelopt.torch.utils.distributed.reinstall_a2a_forwards, a new thin public wrapper, and example_utils.py checks for an AutoQuantize recipe through the public recipe API instead of a private helper of autoquant_utils; an expert-parallel or replicated run warns before the load when the recipe uses LSQ (new test fails without it); the README says which calibrations have run on GPUs on more than one rank. Rebased onto main d464d30 through [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139: no conflicts, own diff otherwise unchanged. Pre-commit passes against main; the hf_ptq example tests that run without a GPU pass on CPU (transformers 5.18: 313 passed, 17 skipped; 5.5.4: 302 passed, 28 skipped).

• Rebased on 2026-10-31 onto main e352cac through [2/3] [OMNIML-6085] [2/3] Keep hf_ptq's previews in lockstep on a sharded model, and pin whole-model loads by plan #3139: no conflicts, own diff unchanged. Pre-commit passes against main; the hf_ptq example tests that run without a GPU pass on CPU (transformers 5.18: 313 passed, 17 skipped; 5.5.4: 302 passed, 28 skipped), on a base that now refuses a sharded load inside init_quantized_weights ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779).

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, and refuses the single-process load options, AutoQuantize and layerwise-export recipes and the TensorRT-LLM export there; listed under Backward Breaking Changes in the CHANGELOG. Single-process and --use_fsdp2 runs behave as before, except that 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), where main's --use_fsdp2 ignored it and ran each rank on cuda:<local_rank>.
• 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?: ✅ Approved on 2026-10-27 and 2026-10-29; a re-review is requested after the split.

Additional Information

[3/3] of this PR's split, and the last PR of expert-parallel PTQ for Hugging Face MoE models. It merges after [2/3] #3139, which merges after [1/3] #3133, #2746, #2742 and #2753.

Summary by CodeRabbit

• New Features
  • Added distributed, expert-parallel post-training quantization for Hugging Face MoE models, with calibration data distributed across ranks and checkpoint shards written in parallel.
  • Added options to control expert-parallel degree and dense-layer FSDP2 sharding, with validation for unsupported configurations.
  • Added synchronized preview generation for distributed runs.
• Bug Fixes
  • Improved GPU memory monitoring when NVML initialization or memory reporting is unavailable.
• Documentation
  • Added setup guidance, launch examples, requirements, and distribution options for expert-parallel quantization.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 78f56fe
• Latest commit, not reviewed: 1a8e007

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: a3be3d46-b98f-4ccd-8e07-1df4498d6504

📥 Commits

Reviewing files that changed from the base of the PR and between 12e1719 and 78f56fe.

📒 Files selected for processing: 9

• examples/hf_ptq/example_utils.py
• examples/hf_ptq/hf_ptq.py
• examples/hf_ptq/vlm_utils.py
• modelopt/torch/utils/distributed.py
• tests/examples/hf_ptq/test_hf_ptq_args.py
• tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py
• tests/examples/hf_ptq/test_hf_ptq_model_type.py
• tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py
• tests/unit/torch/utils/test_distributed_calibration.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The HF PTQ example adds expert-parallel and dense-FSDP distribution support for model loading, calibration, previews, and checkpoint export. The changes also update NVML initialization around the memory probe and add tests and documentation for these behaviors.

Changes: Hugging Face PTQ distribution

Layer / File(s)	Summary
Distribution planning and model loading examples/hf_ptq/hf_ptq.py, examples/hf_ptq/example_utils.py, tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py, tests/examples/hf_ptq/test_hf_ptq_args.py, tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py, examples/hf_ptq/README.md, CHANGELOG.rst	The example resolves distribution plans, checks supported layouts, and passes plans to pretrained-loading paths. The guide describes expert-parallel configuration and launch options. Tests cover plan selection and loading behavior.
Distributed execution and calibration examples/hf_ptq/hf_ptq.py, examples/hf_ptq/example_utils.py, modelopt/torch/utils/distributed.py, tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py, tests/examples/hf_ptq/test_llm_ptq.py, tests/unit/torch/utils/test_distributed_calibration.py	The run uses a distributed context with a timeout above 30 minutes. Supported calibration data loaders are sharded across ranks, and expert-parallel plans use distributed_forward_loop. The integration test runs two-rank Qwen3-MoE PTQ and checks the exported FP8 expert weights and scales.
Synchronized previews and checkpoint export examples/hf_ptq/example_utils.py, examples/hf_ptq/hf_ptq.py, examples/hf_ptq/vlm_utils.py, tests/examples/hf_ptq/test_hf_ptq_expert_parallel.py, tests/examples/hf_ptq/test_hf_ptq_model_type.py	Generation kwargs enable synchronized previews for sharded models. Replicated models export on rank 0; dense-FSDP and expert-parallel models export on all ranks. Tests cover preview kwargs and export participation.

NVML memory probe

Layer / File(s)	Summary
NVML probe lifecycle modelopt/torch/utils/memory_monitor.py, tests/unit/torch/utils/test_memory_monitor.py	The memory monitor initializes NVML before probing GPU memory and shuts it down after the probe. Tests cover supported and unsupported memory reporting.

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Ranks as torchrun ranks
  participant PTQ as HF PTQ
  participant Loader as load_pretrained
  participant Data as Calibration loader
  participant Forward as distributed_forward_loop
  participant Export as Checkpoint export
  Ranks->>PTQ: Resolve distribution plan
  PTQ->>Loader: Load model with distribution plan
  PTQ->>Data: Shard calibration data across ranks
  Data->>Forward: Provide rank-local calibration batches
  Forward->>Export: Complete expert-parallel calibration
  Ranks->>Export: Export checkpoint shards

Loading

Merge Risk: Minimal · up to 78f56

No confirmed issue blocks merging. The distributed speculative-decoding calibration path remains unverified.

Pre-merge checks: 5

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		No CRITICAL security anti-pattern was introduced. The changed modelopt and examples Python code adds no unsafe torch.load, numpy.load(..., allow_pickle=True), eval/exec, or # nosec bypas…
Title check		The title clearly and concisely identifies the main change: adding expert-parallel PTQ execution to hf_ptq.py.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

No CRITICAL security anti-pattern was introduced. The changed modelopt and examples Python code adds no unsafe torch.load, numpy.load(..., allow_pickle=True), eval/exec, or # nosec bypass. Transformer loading receives caller-controlled trust_remote_code values, with defaults set to False; the literal True values occur only in newly added test fixtures. The pull request changes no pyproject.toml or requirements file.

Finishing Touches

Generate docstrings

• Commit to this branch
• Create a new PR

Fix failing CI checks

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

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

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2754/
Built to branch gh-pages at 2026-10-27 00:05 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 62.69%. Comparing base (ca98606) to head (481197b).

Additional details and impacted files

@@                       Coverage Diff                        @@
##           shengliangx/hfep-11c-kimi-k3    #2754      +/-   ##
================================================================
+ Coverage                         62.63%   62.69%   +0.06%     
================================================================
  Files                               668      668              
  Lines                             74485    74490       +5     
================================================================
+ Hits                              46654    46703      +49     
+ Misses                            27831    27787      -44     

Flag	Coverage Δ
unit	62.69% <100.00%> (+0.06%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test 9d9b154

[shengliangxu]

/ok to test 487a8e8

[shengliangxu]

/ok to test 22ac8fc

[shengliangxu]

/ok to test 5b2a422

[shengliangxu]

/claude review

[claude]

[SUGGESTION] The version boundary here contradicts the two other places that state it. Line 699 (three lines above) and line 742 both say the clobber exists on transformers >= 5.16; _reinstall_a2a_forwards' own docstring and body say the same ("install_forward() -- and the clobber it causes -- only exist on transformers >=5.16. Older releases install their TP input/output transforms as forward PRE-HOOKS"), and the PR description says "On 5.15 and earlier the TP transforms are forward hooks". So 5.14 and 5.15 are also pre-hook releases, and <= 5.13 reads as if 5.14/5.15 were a third, unexplained case.

Impact is only on the next reader — the code keys on hasattr(style, "install_forward") per style, not on a version — but a comment that disagrees with the comment three lines up is the kind of thing that sends someone version-hunting.

Suggested change

	# On <= 5.13 the wrappers are pre-hooks that survive, so the reinstall is a no-op --
	# On <= 5.15 the wrappers are pre-hooks that survive, so the reinstall is a no-op --

[shengliangxu]

Fixed in c2dbf28: the trimmed comment gives only the 5.16 boundary, as the other places do; the <= 5.13 sentence is gone.

[claude]

Claude review — [3/3] Run expert-parallel PTQ from hf_ptq.py

Findings: CRITICAL 0 · IMPORTANT 0 · SUGGESTION 1 — approving.

Scope

Full review — bare /claude review, no scoping instructions. 10 changed files (+1453/-190), two-dot diff against origin/shengliangx/hfep-14b-hf-ptq-previews. I read every source hunk with surrounding context (examples/hf_ptq/hf_ptq.py, example_utils.py, README.md, modelopt/torch/utils/distributed.py) plus CHANGELOG.rst, and traced the base-branch callees this slice leans on: resolve_distribution_plan (including its config read and the spec/plan broadcasts), load_pretrained, distributed_run / DistributedState / local_rank, distribution_plan / set_distribution_plan, shard_dataloader, distributed_forward_loop / _reinstall_a2a_forwards, expert_parallel_mesh, preview_generate_kwargs / keep_the_experts_implementation, validate_fsdp2_supported, get_model, layerwise_export_requested. Of tests/ I read the four small diffs in full plus the fixtures, helper namespaces, calibration/preview/export cases and test-name index of the new 961-line test_hf_ptq_expert_parallel.py. I could not execute tests in this environment, so pass/fail evidence is the description's.

Delta since the 2026-10-30 approval

One commit, 5b2a4221 ("Drop the failed-preview test's --use_fsdp2 row"), plus the restack onto [2/3]'s get_model review fix. Both check out:

• The dropped row is right. The base tip 7772c0df moved the raise down into [2/3] and widened it to every sharded plan, so preview_generate_kwargs returns synced_gpus — and the raise fires — for dense="fsdp" too; a row asserting --use_fsdp2 warns and still exports would now fail. The three surviving rows still discriminate what this slice owns: DistributionPlan(4) yields lockstep == {}, so the replicated row genuinely exercises warn-and-export rather than passing tautologically.
• The restack strands no load path. get_model now raises for is_expert_parallel or is_fsdp_dense, and load_model's branch order means only a replicated plan reaches it: EP first, then is_fsdp_dense, with is_fsdp_dense and not is_expert_parallel and not --use_fsdp2 already refused in _resolve_distribution. The hybrid EP+FSDP-dense plan takes the EP branch.

The one finding (non-blocking)

hf_ptq.py:702 — the a2a comment says the wrappers are pre-hooks on transformers <= 5.13, while line 699 three lines above, line 742, and _reinstall_a2a_forwards' own docstring all put the boundary at >= 5.16. No behavioural effect (the code keys on hasattr(style, "install_forward") per style, not on a version), but two comments in the same function disagree.

What I checked this round and found sound

• Dropping the low-memory branch's own config read is safe. resolve_distribution_plan reads the config unconditionally inside all_reduce("checkpoint read", ...) and returns a PretrainedConfig, never None, so hf_config.model_type there cannot hit an attribute error. That branch is single-process-only anyway (--low_memory_mode is refused above world size 1 without --use_fsdp2, and by validate_fsdp2_supported with it).
• The seven pre-load refusals are rank-identical and fire before anything is allocated. validate_fsdp2_supported reads only args and a config, ahead of the first collective; everything after derives from the broadcast plan plus rank-identical args. The short-circuit order is right: is_expert_parallel and args.experts_implementation is not None never reads the flag under --use_fsdp2 (which forces ep_size=1), while (is_expert_parallel or is_fsdp_dense) and args.vllm_fakequant_export does read it there — which is why test_fsdp2_autoquant.py had to grow vllm_fakequant_export. That namespace now carries every attribute _resolve_distribution reads on the --use_fsdp2 path.
• Calibration shards cannot desync step counts. make_calib_dataloader reads distribution_plan(language_model).world_size — resolvable on a VLM sub-model since the plan is stamped on every module — and shards with args.dist_state.rank, the same size()/rank() the plan was built from. Map-style loaders pad via DistributedSampler(drop_last=False), the VLM stream drops its trailing partial group, every loader is num_workers=0, and no dp_size is passed, so there is exactly one shard. _shard_calib_dataloader covers both shapes the builder can return, keeps _DeviceDataLoader's device move outermost, tests the non-DataLoader subclass first, and refuses anything else loudly. The Whisper branch is unreachable above one rank.
• The a2a lifetime is covered end to end. The loop re-plants first and asserts is_expert_parallel == (ep_mesh is not None) and ep_mesh.size() == plan.ep_size, so a mismatched mesh fails loudly instead of running with process_group=None; the post-quantize re-plant sits outside the use_calibration guard, so a recipe that runs no loop keeps its wrappers for the preview and export; and the re-plant skips a wrapper already in place, so the two calls cannot stack transforms. The EP text lambda closes over the language_model variable, which mtq.quantize rebinds only after returning, so during calibration it is the in-place-converted module.
• Export gating. The early return filters exactly the non-main ranks of a replicated plan — for which _hf_shard_flags reports no sharding, so rank 0 issues no collective they skipped — save_source_config is is_main-gated and ordered ahead of rank 0's own export, and the sharded spec-dec refusal raises on every rank with nothing written, leaving checkpoint_exported at False.
• main / _main_inner. distributed_run(timeout=3h, device=args.device) owns the group for the whole run and is exited even when the body raises; args.dist_state is assigned as its first statement and read only within it. The monitor gate world_size == 1 or local_rank() == 0 short-circuits before local_rank() can warn in a single process, and local_rank() falls back to the global rank rather than raising when LOCAL_RANK is absent.
• Tests assert behaviour, not mocks. No checkpoint needed, one gloo rank only for the test that must run a real all-to-all forward, autouse fixtures clear MODELOPT_EP_DEGREE/MODELOPT_FSDP_SIZE so a developer's shell cannot steer the decisions, and the 2-GPU smoke test requires every expert of every layer to come out float8_e4m3fn with its weight_scale, so a lost or misindexed rank share fails it.
• No prompt-injection or review-steering text anywhere in the diff.

Still open from earlier rounds (non-blocking, not re-filed)

The config read twice under --use_fsdp2 and the recipe parsed twice; the private _reinstall_a2a_forwards import and the doubled expert_parallel_mesh call; the post-quantize re-plant keyed on the quantize block rather than the plan; the EP image route building its VLM forward pre-conversion and skipping full_model.eval(); the plan-ungated Nemotron-VL re-raise (not a regression — that path raised before the series, this only replaces the message); and the README knob table's auto row for a spec that refuses EP.

Wording point for the PR body, not a defect: a multi-rank --use_fsdp2 run of a spec-dec checkpoint without --specdec_offline_dataset is now refused at export, where before every rank raced export_speculative_decoding into the same directory. A fix, but a third item (with --sparsity_fmt <non-dense> and --device cuda:N) the "--use_fsdp2 runs behave as before" answer does not cover.

Risk

Low. The delta is one test-only commit plus a restack, and neither disturbs the invariants this series rests on: one plan decided once, broadcast, and read back off the model everywhere after; refusals before the load from rank-identical inputs, so a rejection leaves nothing allocated; rank 0 writes a replicated model and every rank writes a sharded one; equal per-rank step counts. Residual exposure is where the code says it is — multi-node EP paths CI cannot run, backed by the described end-to-end runs and the 2-GPU smoke test — plus a preview failure on one rank alone parking its peers in generate's collectives until the timeout.

Claude review passed — no blocking issues found. LGTM

🤖 Generated with Claude Code

[meenchen]

Could we move _resolve_distribution into example_utils.py, alongside validate_fsdp2_supported? It primarily contains example-specific CLI validation, and extracting it would keep hf_ptq.py focused on the PTQ workflow. The generic planning logic should remain in the existing library resolve_distribution_plan.

[shengliangxu]

Moved in 50a28e9: resolve_distribution now sits in example_utils.py after validate_fsdp2_supported, with exports_tensorrt_llm_checkpoint, which it needs, and load_model calls it. The generic planning stays in the library's resolve_distribution_plan. Its comments are shorter, and the tests call and patch the names where they now live.

[meenchen]

Could you also add examples of EP for slurm? similar to https://github.com/NVIDIA/Model-Optimizer/blob/5b2a4221ba65e85eb84589e393adb61cb04ac690/examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm

[meenchen]

Could you trim the comments?

[shengliangxu]

Trimmed in c2dbf28: this 16-line comment is 4 lines now, and the other comments this PR adds to hf_ptq.py say why in a line or two. The moved checks' comments were cut down in 50a28e9.

[shengliangxu]

/ok to test be8d1df

[shengliangxu]

/claude review

[shengliangxu]

Added in be8d1df: examples/hf_ptq/slurm/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 with an FP8 KV cache by default, and names MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE for overriding the plan. The README's expert-parallel section links it, with the sbatch line. I have not run the script on a cluster yet; the expert-parallel runs in the description were launched with torchrun directly.

[claude]

[SUGGESTION] The script's header (line 24) states expert parallelism needs transformers 5.7 or newer, but nothing in the install path enforces it: pyproject.toml pins transformers>=5.5,<5.19 and examples/hf_ptq/requirements.txt pins no version, so pip install -e "/modelopt[hf]" leaves a transformers 5.5 that the container already ships (the pin's own comment notes TensorRT-LLM 1.3.0 ships 5.5) in place.

Why it matters for this example specifically: on 5.5 _ep_degree("auto", "qwen3_moe") returns 1, so the spec is never applied, the 235B MoE is read as dense, dense="auto" picks fsdp because it does not fit one GPU, and resolve_distribution raises "FSDP2-shards the model, which this script does under --use_fsdp2. Pass it, or set MODELOPT_FSDP_SIZE=1…". It fails closed, which is right, but the message points at --use_fsdp2/MODELOPT_FSDP_SIZE and never mentions the transformers version that actually diverted the run off the expert-parallel path — a confusing first experience for a ready-to-run example. (On a model small enough to fit, the same mismatch silently gives a replicated run instead of an expert-parallel one, with no error at all.)

Cheapest fix is to make the install satisfy the script's own stated requirement:

    pip install -q -e "/modelopt[hf]" --no-build-isolation
    pip install -q -r requirements.txt
    pip install -q -U "transformers>=5.7"   # expert parallelism needs 5.7+; images may ship 5.5

Printing python -c "import transformers; print(transformers.__version__)" before torchrun would at least put the version in the job log.

[shengliangxu]

Fixed in fc2c6f6: the install asks pip for transformers>=5.7 together with ModelOpt, so an image that ships 5.5 gets upgraded within ModelOpt's own pin rather than kept, and the job prints the transformers version before torchrun.

[claude]

[SUGGESTION] This file is a near-verbatim copy of the existing slurm/multinode_fsdp2_ptq.slurm: of 93 lines, only the job name, the two prose paragraphs, MODEL_PATH, RECIPE, the two commented MODELOPT_* knobs, the # total over all ranks comment and the dropped --use_fsdp2 differ. Everything load-bearing — the #SBATCH block, the CONFIG names, the MASTER_ADDR/MASTER_PORT rendezvous, the mount construction, the srun/pip/torchrun body — is duplicated verbatim.

CONTRIBUTING.md's simplicity guidance (prefer extending an existing mechanism over adding a parallel one) applies here, and the cost is already visible in this very PR: the README's out-of-memory note had to be reworded because it now covers both launches, and a future fix to the mount logic or the rendezvous has to be applied twice or it silently drifts.

Consider one script driven by a config switch, e.g. in the CONFIG block:

export USE_FSDP2=0   # 1 shards the whole model with FSDP2 instead of the MoE's experts

and in the torchrun body:

        --export_path "${EXPORT_PATH}" \
        ${USE_FSDP2:+$([ "${USE_FSDP2}" = 1 ] && echo --use_fsdp2)}

(or simply $EXTRA_ARGS), with the README pointing at the one script for both sections. Non-blocking — a self-contained copy-pasteable script per workflow is a defensible UX choice; flagging it so the duplication is a deliberate call rather than an accident.

[shengliangxu]

Kept as a separate script on purpose: meenchen asked for an EP example similar to multinode_fsdp2_ptq.slurm, so it mirrors that script, and each workflow stays one self-contained script to copy and edit, with its own model, recipe and knobs. A switch in one script would also change the FSDP2 script that users and the README already point at.

[claude]

Claude review — [3/3] Run expert-parallel PTQ from hf_ptq.py

Findings: CRITICAL 0 · IMPORTANT 0 · SUGGESTION 2 — approving.

Scope

Full review — bare /claude review, no scoping instructions. 12 changed files (+1519/-230), two-dot diff against origin/shengliangx/hfep-14b-hf-ptq-previews. I read every source hunk in full (examples/hf_ptq/hf_ptq.py, example_utils.py, README.md, the new slurm/multinode_ep_ptq.slurm, modelopt/torch/utils/distributed.py) plus CHANGELOG.rst, re-read the surrounding hf_ptq.py regions the diff only partially shows (make_calib_dataloader, export_quantized, the post_quantize preview block), and the five smaller test diffs. For the delta since the previous approval I additionally read the sibling slurm/multinode_fsdp2_ptq.slurm on the base, modelopt_recipes/general/ptq/nvfp4_default-kv_fp8.yaml, examples/hf_ptq/requirements.txt, the transformers pin in pyproject.toml, and resolve_distribution_plan's contract. I did not re-read the 961-line test_hf_ptq_expert_parallel.py in full — it is unchanged since the round that reviewed it. I cannot execute tests here, so pass/fail evidence is the description's.

Delta since the 2026-10-30 approval

The new commit be8d1df2 adds examples/hf_ptq/slurm/multinode_ep_ptq.slurm (+93) and the three README paragraphs that point at it; the rest of the diff shift is the restack onto [2/3]'s lockstep fix, already reviewed. Both suggestions below are on the new file. What checks out:

• The Slurm script is launch-correct. MASTER_ADDR from scontrol show hostnames | head -1, --rdzv_backend=c10d with --nnodes="${SLURM_NNODES}" and --nproc_per_node="$(nvidia-smi -L | wc -l)", one task per node fanning out to the GPUs — identical to the FSDP2 sibling that already runs, minus --use_fsdp2, which is exactly the switch that selects the expert-parallel plan in resolve_distribution. set -euo pipefail in both shells, and the conditional MODEL_PATH mount correctly distinguishes a local dir from a Hub ID downloaded into the mounted HF_HOME.
• The recipe it names is a legal choice on the expert-parallel route. general/ptq/nvfp4_default-kv_fp8 resolves (modelopt_recipes/general/ptq/nvfp4_default-kv_fp8.yaml) and its algorithm: max is in the supported row of the description's table — not one of the per-rank-statistics algorithms _validate_recipe_calibration refuses, not AutoQuantize, and recipe_layerwise_blocks finds no layerwise export, so none of the seven pre-load refusals fire. It is also the non-_cast variant, right for a BF16 source where the sibling's _cast is for an MXFP4 one. The commented MODELOPT_EP_DEGREE=16 is consistent with the --nodes=2 x --gpus-per-node=8 header and divides Qwen3-235B-A22B's routed-expert count.
• The README additions match the code. --calib_size really is the total over all ranks — make_calib_dataloader builds the full loader from args.calib_size and only then shards it with _shard_calib_dataloader(..., args.dist_state.rank, world_size) — the #manual-run-on-each-node anchor resolves to the existing #### Manual (run on each node) heading, the two knob rows match resolve_distribution_plan's documented ep_size/dense semantics including the 60%-of-one-GPU threshold, and widening the out-of-memory note from --use_fsdp2 to any multi-rank torchrun is right now that [2/3] drops the split-on-OOM retry for every plan of several ranks.
• Re-verified independently on the source side (not taking the earlier rounds on trust): the export_quantized early return filters exactly the non-main ranks of a replicated plan and leaves args.checkpoint_exported at False there, which only feeds tracked_run's exported= on the is_main rank, so no rank reports a write it did not do; save_source_config is is_main-gated and ordered ahead of rank 0's own config.json write; the sharded spec-dec refusal raises on every rank with nothing written; _shard_calib_dataloader keeps _DeviceDataLoader's device move outermost and refuses an unsplittable loader loudly, with the Whisper branch unreachable above one rank (whisper is in is_trtllm_enc_dec_export's list, so exports_tensorrt_llm_checkpoint refuses it first); the post-quantize _reinstall_a2a_forwards sits outside the use_calibration guard so a no-calibration recipe keeps its wrappers for the preview and export; and test_fsdp2_autoquant.py's namespace now carries every attribute resolve_distribution reads on the --use_fsdp2 path — recipe, experts_implementation, sparsity_fmt, qformat, vllm_fakequant_export — with the unread ones correctly short-circuited.
• No prompt-injection or review-steering text anywhere in the diff, the new Slurm script included.

The two findings (both non-blocking, both on the new Slurm script)

1. multinode_ep_ptq.slurm:78 — the header says expert parallelism needs transformers 5.7+, but nothing installs it: the [hf] extra allows >=5.5 and requirements.txt pins nothing, so an image already shipping 5.5 keeps it. The run then fails closed, but with a message pointing at --use_fsdp2 / MODELOPT_FSDP_SIZE rather than the version that actually diverted it off the expert-parallel path.
2. multinode_ep_ptq.slurm:33 — the file is a near-verbatim copy of multinode_fsdp2_ptq.slurm; every load-bearing line is duplicated, so a future fix to the mounts or rendezvous has to land twice.

Still open from earlier rounds (non-blocking, not re-filed)

The config read twice under --use_fsdp2 and the recipe parsed twice; the _reinstall_a2a_forwards comment at hf_ptq.py:702 putting the boundary at <= 5.13 where its neighbours say >= 5.16; the private _reinstall_a2a_forwards import and the doubled expert_parallel_mesh call; the EP image route building its VLM forward pre-conversion and skipping full_model.eval(); the plan-ungated Nemotron-VL re-raise; and the README knob table's auto row for a spec that refuses EP.

Risk

Low. The delta is one new documentation artifact — a Slurm submit script and the README text pointing at it — with no behavioural code change, and it is a minimal edit of a script that already runs for the FSDP2 path. The invariants the series rests on are undisturbed: one plan decided once, broadcast, and read back off the model everywhere after; refusals before the load from rank-identical inputs, so a rejection leaves nothing allocated; rank 0 writes a replicated model and every rank writes a sharded one; equal per-rank step counts. Residual exposure stays where the description says it is — multi-node expert-parallel paths CI cannot run, backed by the described end-to-end runs and the 2-GPU smoke test.

Claude review passed — no blocking issues found. LGTM

🤖 Generated with Claude Code

[shengliangxu]

/ok to test fc2c6f6

[shengliangxu]

/ok to test a909d4e

[Fridah-nv]

Suggestion (non-blocking): lsq is "not assessed" and isn't refused, so under an EP plan it would run, possibly with unsynced scales. Could it be refused like the others, or at least warn? Same for the "supported, not yet run on GPUs" row (gptq, smoothquant, nvfp4_act_headroom): worth a line in the README's EP section so users know what was tested.

[shengliangxu]

Done in d8bbfee. An expert-parallel or replicated run, the routes this PR adds, now warns before the load when the recipe uses LSQ. Its first scale calibration is synced or refused as the algorithm it runs is, but nothing has checked its learnable scales across ranks yet. --use_fsdp2 and a single process run it as before. New test test_an_lsq_recipe_warns_on_the_routes_this_script_adds; it fails without the warning.

The README's expert-parallel section now says which calibrations have run on GPUs on more than one rank (max, mse), which are supported but not yet run there (gptq, smoothquant, awq_clip, nvfp4_act_headroom), and which are refused before quantizing (awq_lite, awq_full, svdquant, local_hessian, imatrix, and max or mse with distributed_sync off).

[Fridah-nv]

Suggestion (non-blocking): hf_ptq.py imports and calls the underscore-prefixed _reinstall_a2a_forwards, and example_utils.py imports _recipe_is_auto_quantize from another example module. This PR's own docstring edit now documents that hf_ptq calls the first one. Could they be made public (or wrapped), so an example doesn't depend on internals?

[shengliangxu]

Done in 45e9c81. 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.py calls it. example_utils.py checks for an AutoQuantize recipe with the public recipe API, isinstance(load_recipe(...), ModelOptAutoQuantizeRecipe), instead of importing autoquant_utils' private helper. The private names hf_ptq.py already imported from autoquant_utils on main are unchanged.

[Fridah-nv]

Approve to unblock

[shengliangxu]

/ok to test d8bbfee

[shengliangxu]

/claude review

[shengliangxu]

/ok to test 1a8e007