[shengliangxu]

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 and worker (a stream that shuffles needs a fixed seed of its own). The copy is split across the loader's worker processes so that each sample is still seen once. 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, and splits it across that loader's workers.
• 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 ([4/21]) 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 its own GPU from one forward of the single-sample input, without probing the size it arrives at, and every rank takes the smallest size any rank allows, reduced with all_reduce(..., min) ([4/21]). 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. [9/21]'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 [21/21] 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, including the split across DataLoader workers, run in-process with the worker info faked rather than by forking workers; 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 [10/21]'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.
• The worker-split test fails on the code before that fix.
• 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 0638a27: pre-commit over the whole PR against main, at every PR of the stack; each fix's own tests at its PR; and at the stack's tip the unit tests of the changed areas (tests/unit/torch/utils, export, models, the fused-experts quantization tests and the hf_ptq example tests) on transformers 5.18 and 5.5 and on torch 2.8. They pass apart from the diffusers-export and GGML tests, which fail the same way on main in that environment. CI runs the full unit suite.
• GPU tests were not re-run on this revision.

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 [21/21].
• 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 [7/21].

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]

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: 29fcc46

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

Failed checks: 1 warning

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

Passed checks: 5 passed

Check name	Status	Explanation
Title check		The title "[OMNIML-6071] [8/21] Calibrate a distributed model on per-rank data" names the main change. The PR adds distributed calibration utilities: sharding of map-style and iterable datasets across…
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=…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Title check

Explanation

The title "[OMNIML-6071] [8/21] Calibrate a distributed model on per-rank data" names the main change. The PR adds distributed calibration utilities: sharding of map-style and iterable datasets across ranks, a distributed forward loop, collective batch-size selection, and a guard that refuses rank-local calibration algorithms on distributed models with a recorded plan. The title is specific to the primary change. The ticket prefix and series marker "[8/21]" add noise but do not make the title unrelated or vague.

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-25 06:17 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.
✅ Project coverage is 61.54%. Comparing base (f9d1502) to head (29fcc46).

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.47%   61.54%   +0.06%     
====================================================================
  Files                                   649      649              
  Lines                                 72347    72469     +122     
====================================================================
+ Hits                                  44477    44598     +121     
- Misses                                27870    27871       +1     

Flag	Coverage Δ
unit	61.54% <99.23%> (+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 963e78a

[shengliangxu]

/ok to test 935bd7f

[shengliangxu]

/ok to test 96159d7

[shengliangxu]

/ok to test 540716e

[shengliangxu]

/ok to test 29fcc46