[shengliangxu]

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 ([1/21]-[20/21]) into it, so the documented CLI reaches an expert-parallel run: recipes, export formats, deprecations.

• Process group. main() runs inside distributed_run ([4/21]), 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 ([6/21]) 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 ([5/21], [12/21], [13/21]), and loads every model by it through load_pretrained ([6/21]): expert-parallel, FSDP2, get_model's whole-model loads and the low-memory load. 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.
• 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, [13/21]) raises unless --use_fsdp2 rules expert parallelism out. A plan that FSDP2-shards the dense part within several expert-parallel groups (the automatic degree when the world size does not divide the expert count), which the export cannot write yet, raises before loading and names the settings that avoid 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.
• Calibration.

  • Every distributed run calibrates on its own 1/world shard of the data ([8/21]).

  • Expert-parallel runs go through distributed_forward_loop, which re-plants the all-to-all forwards that quantizer insertion clobbers on transformers 5.16 and later.

  • FSDP2 keeps main's plain loop, with its split-on-OOM retry.

  • 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), nvfp4_act_headroom (activation scales MAX-reduced over data parallelism)
    Refused before quantizing ([8/21])	awq_lite, awq_clip, 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
    Not assessed	lsq's learnable scales (its first scale calibration follows the rows above)

• Previews keep a sharded model's ranks together. Each rank previews the first sample of its own calibration shard, so the ranks finish decoding at different steps. Under expert parallelism or FSDP2 every forward runs a collective, and a rank that stopped decoding left the others blocked in it: Qwen3-30B-A3B at NVFP4 and EP4 hung in the post-PTQ preview. The previews, Nemotron-VL's included, now generate with synced_gpus when the model's plan shards it.
• GPU placement. When the plan gives every one of several ranks the whole model, each rank pins it to its own GPU. A single process keeps spreading a model too big for one GPU, as before.
• Docs. README section "Expert-Parallel Post-Training Quantization for MoE Models": how to launch, what the knobs override, which models are refused.
• Memory report fix. launch_memory_monitor probed NVML before initializing it. Every probe raised Uninitialized, so the per-GPU peak-memory report never started on any host where nothing else had initialized NVML. The probe now initializes NVML itself.

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 (31 cases; the decision itself is tested in [5/21], [12/21] and [13/21]):
  • 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;
  • per-rank GPU pinning by the plan, and single-process placement left alone;
  • 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.
  • the previews, before and after quantization and Nemotron-VL's, generating with synced_gpus exactly when the plan shards the model (EP or FSDP2 over several ranks);
  • exporting a replicated model from rank 0 only, and every rank of a sharded one;
  • the plan the export cannot write refused before loading, while the world-size degree still resolves;
  • Mutation-checked: each test fails under a deliberate break of the code it covers.
• New 2-GPU smoke test test_ptq_qwen3_moe_expert_parallel (tests/examples/hf_ptq/test_llm_ptq.py): torchrun --nproc_per_node 2 hf_ptq.py with MODELOPT_EP_DEGREE=2 on a tiny Qwen3-MoE and the FP8 recipe. It checks that every expert of both ranks is exported FP8 with its weight scale, and that the config names FP8. It is the tests' one real expert-parallel run of the script; the file above covers the wiring with stand-ins. It passed on 2 GB200s in 29.5 s at f4bb395488 (2026-10-16), which carries this test on a revision before the latest loader and calibration fixes ([7/21], [8/21]). It skips on CI's 1-GPU PR runner and runs on 2-GPU runners.
• tests/examples/hf_ptq/test_hf_ptq_model_type.py (main's): the preview tests' stand-in models are modules now, since the previews read their distribution plan.
• 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.
• 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.
• Locally on this revision, the reordered series rebased on main 200011e: the full unit suite on transformers 5.18, and on 5.5 apart from the tests whose ONNX, TensorRT or diffusers packages that environment lacks; CI's transformers-free unit lane; pre-commit over the whole PR against main; and the Sphinx docs build. The same at every PR of the stack.
• The GPU results below were measured on earlier revisions of this branch, before this script's loads moved onto load_pretrained ([6/21], [12/21]); 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 [3/21]'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 [3/21] fix.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Single-process and --use_fsdp2 runs behave as before. A torchrun launch without --use_fsdp2 now runs an MoE expert-parallel.
• 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, 0.49.0.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 21 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [20/21].

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: fb135b1

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.

Priority: ➖ Normal

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 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 56.63% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 83 functions across 12 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

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 '[OMNIML-6085] [21/21] Run expert-parallel PTQ from hf_ptq.py' names the main change. The PR adds expert-parallel PTQ to hf_ptq.py, and the changeset adds distribution planning, expert-paral…
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.

Full details: Title check

Explanation

The title '[OMNIML-6085] [21/21] Run expert-parallel PTQ from hf_ptq.py' names the main change. The PR adds expert-parallel PTQ to hf_ptq.py, and the changeset adds distribution planning, expert-parallel model loading, calibration, previews, and rank-aware export there. The ticket and series tags are noise, but the core change is clear and specific.

✨ Finishing Touches: 💡 2

📝 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

_{Comment @coderabbitai help to get the list of available commands.}

[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-24 01:30 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (shengliangx/hfep-11c-kimi-k3@15287ee). Learn more about missing BASE report.

Additional details and impacted files

@@                       Coverage Diff                       @@
##             shengliangx/hfep-11c-kimi-k3    #2754   +/-   ##
===============================================================
  Coverage                                ?   62.28%           
===============================================================
  Files                                   ?      667           
  Lines                                   ?    74107           
  Branches                                ?        0           
===============================================================
  Hits                                    ?    46159           
  Misses                                  ?    27948           
  Partials                                ?        0           

Flag	Coverage Δ
unit	62.28% <100.00%> (?)

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 fb135b1