[shengliangxu]

Note

The expert-parallel PTQ series is now a tree instead of one stack (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [8/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.

Base: #2736. The diff below is this PR's alone, and it merges after #2736.

Re-cut: this PR's diff is unchanged; only its base moved. Review fixes since are listed at the end of Testing.

Series layout

• main
  • [OMNIML-6073] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 fused experts under grouped_mm
    • [OMNIML-6073] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 calibration on grouped_mm
      • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
  • [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 meta-DCP checkpoint load (merged)
  • [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 distribution plan, run group, collectives
    • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data (this PR)
    • [OMNIML-6081] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold
    • [OMNIML-6082] Add the pieces of a distributed HF checkpoint writer #2751 distributed-writer pieces
      • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] Add the pieces of a distributed HF checkpoint writer #2751, [OMNIML-6081] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6054] Route load_pretrained by the distribution plan #2779 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
        • [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export
    • [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 plan decision
      • [OMNIML-6054] Route load_pretrained by the distribution plan #2779 load routing
        • [OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741 expert-parallel loader (on [OMNIML-6054] 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-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
            • [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
              • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
                • [OMNIML-6076] Describe expert parallelism in the model specs #2746 expert parallelism in the model specs
                  • [OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747 on-the-fly decompression
                  • [OMNIML-6078] Add the Kimi-K2.5 and Kimi-VL expert-parallel adapters #2748 Kimi-K2.5 and Kimi-VL adapters
                  • [OMNIML-6079] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter
                  • [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] Describe expert parallelism in the model specs #2746 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)

What does this PR do?

Type of change: New feature

The pieces a calibration loop needs once the model is spread across ranks.

• shard_dataloader gives each rank its own shard of the data.
  • Map-style datasets use a DistributedSampler with drop_last off, so every rank runs the same number of steps; a rank that exits early hangs the others on the per-step collectives.
  • An IterableDataset gets a strided copy of the stream, which must yield the same samples in the same order on every rank (a stream that shuffles needs a fixed seed of its own). A streaming loader with worker processes is refused: a dataset that splits itself across workers, as a Hugging Face streaming dataset does, would lose samples silently. It is not main's vlm_dataset_utils._ShardedIterable, which pads each rank to a count its loader knows: this wraps any loader's stream, which may have no length.
• distributed_forward_loop builds the forward_loop that mtq.quantize calls. It moves each batch to the device it is given and runs the full distributed forward, so the expert-parallel collectives fire inside it, with KV caching off as on the single-process path. Each batch is a mapping of the forward's keyword arguments; anything else is refused with a TypeError at the first batch.
  • A model whose distribution plan ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736) is expert-parallel passes the ep_mesh its experts are sharded over, and the loop re-plants the torch_a2a forward wrappers that quantizer insertion clobbers on transformers >= 5.16. A module whose forward is still the wrapper planted there is skipped, so passes do not stack wrappers; one whose forward was rebound since (another conversion, an export) is planted again. Any other model passes none. A mesh missing for an expert-parallel model, or given for another, raises.
  • It returns the allocator's reserved pool before the first forward and after the last. NCCL's buffers live outside the torch allocator, so a full reserved pool starves them. A process that never started CUDA skips this, and the garbage collection it costs.
• Every rank runs the same batch size. Equal step counts need equal batch sizes, and hf_ptq's default --batch_size 0 sizes calibration with get_max_batch_size on each rank: it measured every visible GPU, and on an out-of-memory error halved the batch and probed again, a different number of collective-bearing forwards per rank. For a model whose recorded plan spreads it over several ranks it is now collective: each rank measures the GPU it computes on (torch.cuda.current_device()) from one forward of the single-sample input, without probing the size it arrives at (a rank running out of memory partway through a probe's forward would leave its peers waiting in that forward's collectives), and every rank takes the smallest size any rank allows, reduced with all_reduce(..., min) ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736). A single process, and a model with no plan recorded (a Megatron model, say), are unchanged.
• Calibration that keeps its statistics on the rank is refused. Every rank calibrates on its own data, so an algorithm must share what it gathers across the ranks or the ranks calibrate their copies, or their shards, of a weight differently. max all-reduces its amax, GPTQ merges its Hessians across data-parallel replicas, and MSE and SmoothQuant work from max's synced amax. AWQ does not: it averages its activation scales over data parallelism but picks each layer's scale from its own rank's search, and a 2-rank probe showed the replicas' weights and pre_quant_scale diverging. local_hessian and imatrix share nothing either, and neither do max and MSE with distributed_sync off. mtq.quantize, mtq.calibrate and mtq.auto_quantize refuse these (awq_lite, awq_clip, awq_full, svdquant, which runs awq_lite, local_hessian, imatrix, and any config that turns distributed_sync off, also as another algorithm's weight_scale_algorithm or scale_algorithm) on a model whose recorded plan spans several ranks, before converting anything (auto_quantize in its format check). A model with no plan recorded, such as a Megatron model, is unchanged. On an expert-parallel plan, any calibration is refused unless every module holding routed experts carries the experts' expert-parallel group and, as its data-parallel group, just the plan's edp ranks, those holding the same experts: each rank's per-expert quantizers hold its own experts, and max's and GPTQ's data-parallel sync would otherwise merge different experts' scales. The modules holding routed experts are those carrying an expert-parallel group and, once quantization has converted them, the quantized modules under a container named experts, as transformers' expert-parallel plans name it; unseeded experts beside a seeded router are refused once converted, before they calibrate. [OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741's expert-parallel loader gives the routed experts those groups. A module with no routed experts and no expert-parallel group, such as a VLM's vision tower that hf_ptq quantizes on its own with every quantizer disabled, passes.
• hf_ptq's --use_fsdp2 load records its plan. The batch sizing above agrees across ranks only for a model that records its plan. hf_ptq loads FSDP2 through parallel_load_and_prepare_fsdp2, which records none, until [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754 moves it onto load_pretrained; it now records the FSDP2 plan right after that load. On several GPUs, --use_fsdp2 therefore now refuses the calibrations above (AWQ, SVDQuant, local Hessian, imatrix, also as an AutoQuantize candidate), which calibrated the ranks' shards of a weight differently; the README's FSDP2 section and the changelog's Backward Breaking Changes say so.

shard_dataloader refuses a loader that picks its own samples (a custom batch_sampler, or a sampler other than sequential, random or distributed, also inside a plain BatchSampler or with auto-batching off) instead of rebuilding it, and keeps a plain BatchSampler's batch size and drop_last in the rank's shard; a loader with auto-batching off (batch_size=None) shards one item per step, as it loads. With one rank, shard_dataloader returns the loader itself, its sampler and shuffling included. distributed_forward_loop runs its forwards without gradients. get_max_batch_size reads the recorded plan, as the calibration guard does. distributed_forward_loop refuses an ep_mesh whose size is not the plan's ep_size.

Usage

from modelopt.torch.utils.distributed import distributed_forward_loop, rank, shard_dataloader, size

loader = shard_dataloader(calib_loader, rank(), size())
# ep_mesh: the mesh an expert-parallel model's experts are sharded over; None for any other model
forward_loop = distributed_forward_loop(model, loader, device=device, ep_mesh=ep_mesh)
model = mtq.quantize(model, quant_cfg, forward_loop=forward_loop)

Testing

• New tests/unit/torch/utils/test_distributed_calibration.py (42 cases), on CPU, in one process except where noted. It covers:
  • strided shards that cover the data exactly once, with equal batch counts on uneven splits, for map-style and streaming datasets; a streaming loader with workers refused; map-style and streaming loaders with auto-batching off sharded one item per step;
  • every loader setting carried over, a plain BatchSampler's batch size and drop_last included;
  • the wrapper re-planting, on the planned modules only, without stacking, and again after another conversion drops it;
  • the loop running this rank's shard through the top-level model, with use_cache off and the allocator cache freed around it in a process that started CUDA, and not in one that did not; a batch that is not a mapping refused;
  • the ep_mesh refused when missing for an expert-parallel plan, given for another, or over the whole world for an ep_size=2 plan;
  • rank-local calibration refused through quantize, calibrate and auto_quantize (before conversion), also when another algorithm runs it for its weight scales or max is told not to sync, with the model left as it was;
  • on two gloo ranks, max calibration on an expert-parallel plan run when the experts carry their groups, and refused with no expert-parallel group anywhere or with the experts' data-parallel group over every rank (before conversion), and with the group on the router only (once the experts are converted, before they calibrate);
  • a module without routed experts calibrating under an expert-parallel plan, with every quantizer disabled (hf_ptq's vision-tower call) and quantized.
• distributed_forward_loop runs on real GPU ranks in [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742's expert-parallel calibration test, which checks each rank's per-expert scales against the whole model; this PR's own tests run it in one process.
• tests/examples/hf_ptq/test_hf_ptq_args.py: hf_ptq's --use_fsdp2 load, stubbed, leaves the model recording its FSDP2 plan.
• tests/unit/torch/utils/test_dataset_utils.py, on two real gloo ranks (cpu_dist_workers) with only the CUDA memory queries faked per rank: each rank measures only its own GPU, probes once, and both take the smaller batch size. A rank keeping its own size, measuring every GPU, or skipping the reduction fails it, the last after the group's 10 s timeout rather than by hanging; main's single-process retry test still passes. The two-rank test also runs a model with no plan recorded, which each rank measures on every visible GPU, as a single process does, without a refusal.
• Mutation-checked: deciding the reduction or the mesh other than by the plan fails a test (5 mutants), and so does dropping either refusal call, the nested-algorithm check, the distributed_sync check, hf_ptq's plan recording, the expert-parallel group refusal, or the ep_mesh size check. Refusing every expert-parallel calibration, or reading the plan strictly in get_max_batch_size, fails a test too. Dropping the check on converted experts or the data-parallel size check, keeping a replant flag once set, or planting on every pass each fail a test. Rebuilding a lone rank's loader, refusing a module without routed experts, or passing any module with no expert-parallel group each fail a test too.
• The shard tests are one table over dataset kind and world size, the two replant tests are one, and cases a remaining test covers are dropped. Each drop is checked by mutation.
• A subset sampler is refused, also under a plain BatchSampler and with auto-batching off, and so is a custom batch sampler; a lone rank keeps each such loader as it is. The sampler, batching, non-mapping-batch and auto_quantize cases each fail without their fix (mutation-checked).
• 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.
• GPU tests were not re-run on this revision.
• 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 a2a forward wrappers are re-planted only where quantizer insertion rebound the forward, so a wrapper still in place is never wrapped twice (found by review: the double wrap ran a module's input transforms twice, which no real run has shown to fail; on 4 GPUs at transformers 5.16.1 the expert-parallel load tests pass with either behaviour, and the CPU calibration-loop test is what tells them apart); the install probe checks each style it calls. 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.
• Review fixes, 2026-10-27, round 2: the batch size is measured on the GPU the rank computes on (torch.cuda.current_device()); a streaming loader with workers is refused; the docstring and the hf_ptq README say why the measurement does not probe. 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). The hf_ptq README's FSDP2 AutoQuantize note follows main's KV-cache AutoQuantize support ([OMNIML-5987] [2/3] Support DP/FSDP2 KV AutoQuant and checkpoint resume #2664): weight and KV-cache AutoQuantize with --use_fsdp2 are experimental, and refused when a candidate format is one of those this PR refuses.
• Review fix, 2026-10-27 night: shard_dataloader refuses a RandomSampler that does not draw the whole dataset once (other num_samples, or with replacement), as it refuses other samplers.
• Review suggestion, 2026-10-27, last round: the comment over the sync guard names the algorithms' defaults it relies on.

Before your PR is "Ready for review"

• Is this change backward compatible?: ❌ for hf_ptq.py --use_fsdp2 on several GPUs with a calibration that keeps its statistics per rank (AWQ, SVDQuant, local Hessian, imatrix, also as an AutoQuantize candidate), now refused: these calibrated the ranks' shards of a weight differently (changelog, README). Otherwise additive; get_max_batch_size changes only for a model whose recorded plan spans several ranks.
• 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?: ✅ A Backward Breaking Changes entry for --use_fsdp2's refusal; the feature's entry lands with [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 8 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2739, which has merged.

Summary by CodeRabbit

• New Features
  • Distributed calibration now shards map-style and streaming datasets across ranks, keeping batch counts aligned and preserving data loader settings.
  • Calibration supports expert- and tensor-parallel models with improved device handling, cache management, and compatibility with parallel model wrappers.
  • Distributed batch-size detection probes each rank’s local GPU and selects the smallest supported batch size across ranks.
• Bug Fixes
  • Calibration and quantization now reject selected rank-local calibration methods on distributed models before modifying model state.

[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]

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

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: shengliangx/hfep-5-distributed-calibration
• File: .coderabbit.yaml
• Commit: b2acf58
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: dc76c383-9843-487e-a1b9-1b7e8104163b

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

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: 60b267d
• Latest commit, not reviewed: b2acf58

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

Walkthrough

The changes add per-rank batch-size probing, distributed calibration data-loader sharding, and a calibration forward loop. They also reject selected rank-local calibration methods for distributed models.

Changes: Distributed Calibration

Layer / File(s)	Summary
Per-rank batch-size probing modelopt/torch/utils/dataset_utils.py, tests/unit/torch/utils/test_dataset_utils.py	Each rank probes its local GPU and contributes its regulated batch size to a minimum reduction. Device-specific probes do not retry smaller batches.
Calibration data sharding modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_calibration.py	Adds map-style and iterable dataset sharding. The rebuilt loader preserves its settings and handles uneven or short inputs as tested.
Calibration forward loop and wrapper restoration modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_calibration.py	Adds expert-parallel mesh validation, A2A wrapper restoration, temporary cache disabling, CUDA allocator cleanup, optional device movement, and calibration forwards.
Rank-local calibration guard modelopt/torch/quantization/model_quant.py, modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_calibration.py	Rejects selected rank-local calibration methods when the model has a recorded multi-rank plan. The guard runs before calibration or quantization changes the model.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant distributed_forward_loop
  participant DataLoader
  participant Model
  Caller->>distributed_forward_loop: Create calibration callback
  distributed_forward_loop->>DataLoader: Read calibration batches
  distributed_forward_loop->>Model: Reinstate wrappers and run batch forward
  Model-->>distributed_forward_loop: Return from forward

Loading

Suggested reviewers: kevalmorabia97

Merge Risk: Moderate · previous review 60b26

Some distributed calibration configurations can produce inconsistent scales, and automatic batch sizing fails for an unrecorded FSDP2 model. Address these paths before merging.

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		PASS. The pull request adds no listed security anti-pattern. Changed Python code does not add torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded `trust_remote_code=…
Title check		The title clearly summarizes the primary change: distributed-model calibration using per-rank data. It is concise and directly related to the pull request changes.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The pull request adds no listed security anti-pattern. Changed Python code does not add torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval()/exec() on external input, or # nosec. No pyproject.toml or requirements.txt file changes add dependencies.

✨ 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-2740/
Built to branch gh-pages at 2026-10-26 23:21 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 99.23664% with 1 line in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (shengliangx/hfep-4-meta-dcp-load@0e0b9f3). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
modelopt/torch/utils/distributed.py	98.79%	1 Missing ⚠️

Additional details and impacted files

@@                         Coverage Diff                         @@
##             shengliangx/hfep-4-meta-dcp-load    #2740   +/-   ##
===================================================================
  Coverage                                    ?   61.69%           
===================================================================
  Files                                       ?      649           
  Lines                                       ?    72627           
  Branches                                    ?        0           
===================================================================
  Hits                                        ?    44808           
  Misses                                      ?    27819           
  Partials                                    ?        0           

Flag	Coverage Δ
unit	61.69% <99.23%> (?)

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 963e78a

[shengliangxu]

/ok to test 935bd7f

[shengliangxu]

/ok to test 96159d7

[shengliangxu]

/ok to test 540716e

[shengliangxu]

/ok to test b2acf58