[OMNIML-6071] Calibrate a distributed model on per-rank data by shengliangxu · Pull Request #2740 · NVIDIA/Model-Optimizer

shengliangxu · 2026-10-12T00:56:50Z

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.

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] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 fused experts under grouped_mm (merged)
  - [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 [1/2] expert-parallel fused experts on each rank's own experts (merged)
    - [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 [2/2] calibration on grouped_mm (merged)
      - [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-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data (this PR)
- [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 (merged)
  - [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold (merged)
  - [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 distributed-writer pieces
    - [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6054] [2/2] 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 (merged)
    - [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 load routing
      - [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 expert-parallel loader (on [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-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] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 expert parallelism in the model specs
        
        [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 on-the-fly decompression
        [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 Kimi-K2.5 and Kimi-VL adapters
        [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter
        
        [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #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. The planted wrappers stay on the model, so its later forwards (a preview's, an export's) work once the loop has run. mtq.auto_quantize calls its forward_step directly and would plant nothing, so it refuses a model whose recorded plan shards its experts, before converting anything. 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-Lite 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. awq_clip now averages each candidate clip's loss over data parallelism before the ranks pick one, in both branches: each rank accumulates its own losses, and after the forward loop every module's losses and token count are all-reduced once, in module order, so the collective count does not depend on how many forwards a rank ran a module (an MoE expert's is each rank's own). Its static-block branch averaged per forward before, and its per-tensor branch, which NVFP4 takes, not at all; this also makes Megatron data-parallel replicas agree on the per-tensor clip. 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_full, svdquant, which runs awq_lite, local_hessian, imatrix, and any config that turns distributed_sync off, with False or None, 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, as its expert-parallel group, the rank's row of the (edp, ep) mesh the expert-parallel load lays the ranks out in, ep innermost (in row order, which picks each rank's expert shard), and, as its data-parallel group, its column, the plan's edp ranks 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] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #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 path is not changed here. Its load (parallel_load_and_prepare_fsdp2) records no plan, so the batch sizing and the refusal above do not apply to it yet. Recording the plan there lands with [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752, the first PR whose export reads it on that path, together with stopping a planned run from retrying a batch after a rank-local out-of-memory error; [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 then moves the load onto load_pretrained, which records the plan itself. nvfp4_awq_clip's data-parallel clip (below) needs no plan, so it applies to --use_fsdp2 runs now.

shard_dataloader refuses a loader that picks its own samples (a custom batch_sampler, a sampler other than sequential, random or distributed, or one of those over another source than the loader's dataset, 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/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.
Rebased on 2026-10-28, late night, onto [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 at cf4597dcbeb, on main 1ca8bd1fc4d (after [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged), without conflicts; this PR's own diff is unchanged. Its tests with [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736's and the hf_ptq argument tests pass on transformers 5.18 (220), and pre-commit passes against main.
Rebased on 2026-10-28, late night, onto main bdc45368d10 after [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 merged; [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736's squash has the tree of its last head, so this PR's tree is unchanged.
Review fixes, 2026-10-29: the two-rank batch-size test makes its bounded group with bounded_collectives, whose barrier keeps a rank from destroying the group while its peer is still connecting (one CI run had failed so); the batch sizing's distributed helpers are imported at module top; the rank-local refusal reads each algorithm's options through its mode's schema, so a distributed_sync of "false" is refused; the experts' groups are checked by the ranks they span, an expert-parallel group of ep_size ranks and a data-parallel group of edp ranks meeting it only at the rank itself (a test answers four ranks' group queries, which the two-rank pool cannot build). Each behaviour fix fails its test without it (mutation-checked). Pre-commit passes against main; this PR's tests, test_distributed.py, test_quantize_cpu.py and the hf_ptq argument tests pass on CPU with transformers 5.18 and 5.5.4.
Review fix, 2026-10-29, second round: awq_clip's per-tensor branch averages each candidate clip's loss over data parallelism, through the helper its static-block branch now uses too (same group, same op, every rank in the same order), and awq_clip is allowed on a distributed model; this also makes Megatron data-parallel replicas agree on the per-tensor clip. A two-rank CPU test of the static-block branch fails without the averaging (mutation-checked); a two-GPU NVFP4 test of the per-tensor branch is in tests/gpu (NVFP4 fake quantization needs CUDA; not run here). The hf_ptq README and the changelog say nvfp4_awq_clip works with --use_fsdp2 on several GPUs. Pre-commit passes against main; test_calib.py and this PR's tests pass on CPU.

GPU check, 2026-10-29, of the NVFP4 per-tensor branch on GB200 (torch 2.11.0a0, transformers 5.9.0), with the second round's per-forward averaging; without the data-parallel mean the even and odd ranks pick different per-tensor clips:

Run	1 GPU	2 GPUs	4 GPUs
`test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank`	skipped (needs 2+)	passed, 23.2 s	passed, 24.1 s
The same test with the averaging reverted	–	failed: clips 95 vs 100	–

GPU check, 2026-10-29, at cdb4e952468, with the fourth round's once-per-module averaging after the forward loop, on GB200 (torch 2.11.0a0, transformers 5.9.0); with that averaging disabled, the even and odd ranks pick different per-tensor clips again:

Run at `cdb4e952468`	1 GPU	2 GPUs	4 GPUs
`test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank`	skipped (needs 2+)	passed, 20.7 s	passed, 24.5 s
The same test with `sync_across_dp` disabled (`return` as its first statement)	–	failed: clips 95 vs 100	–

GPU check, 2026-10-29, at 01a1df717e3, with the fifth round's empty-input guard, on GB200 (torch 2.11.0a0, transformers 5.9.0). The test now also runs a second linear on no tokens on the odd ranks; without the guard, that linear's clip stays 0 on every replica:

Run at `01a1df717e3`	1 GPU	2 GPUs	4 GPUs
`test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank`	skipped (needs 2+)	passed, 21.1 s	passed, 25.0 s
The same test without the empty-input guard in `_clip_search`	–	failed on `amax[1] > 0`: amax = [100, 0]	–

Review fixes, 2026-10-29, third round: a distributed_sync of None, which skips the sync as False does, is refused too, nested scale algorithms included; the experts' groups are checked against the (edp, ep) mesh, the expert-parallel group as the rank's row in row order and the data-parallel group as its column, so diagonal data-parallel groups and a reordered expert-parallel group are refused (the four-rank test adds both). Each fails its test without the change (mutation-checked). Pre-commit passes against main; this PR's tests pass on CPU.
Review fixes, 2026-10-29, fourth round: awq_clip averages its clip search over the data-parallel ranks once per module after the forward loop, with the token count, in both branches, instead of once per forward, which hung when ranks ran an MoE expert a different number of times (the two-rank test now does; it fails without the change, on its bounded group's timeout) (the same round accepted the experts' groups of a mesh with either axis innermost; the eighth round reverts that). Pre-commit passes against main; test_calib.py and this PR's tests pass on CPU, and the quantization unit tests pass with the change.
Review fix, 2026-10-29, fifth round: awq_clip skips its clip search for an input with no tokens, which an MoE expert a rank routes none to can still get. Before, the per-tensor branch added NaN losses, which the post-loop sum spread to the replicas that had tokens (a clip of 0 for them), and the static-block branch failed to reshape the empty input. The two distributed tests now run an expert on no tokens on one rank, and check that its clips are positive and those a single process picks from the other ranks' tokens; each fails without the change (mutation-checked; the GPU one in the table above). Pre-commit passes against main; test_calib.py and this PR's tests pass on CPU.
Review fix, 2026-10-29, sixth round: hf_ptq's --use_fsdp2 plan recording, its README notes, changelog entry and argument test leave this PR for [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 (see above); the new GPU test tests/gpu/torch/quantization/test_fsdp2.py::test_planned_batch_size_then_max_calibration (2+ GPUs) sizes an FSDP2 model's batch by its recorded plan with the real get_max_batch_size, rank 0's own estimate far above the others', checks every rank takes the smallest, then calibrates each rank's shard with max through distributed_forward_loop and checks the steps and scales agree. Its logic ran on two CPU ranks (FSDP2 on gloo, the CUDA memory queries faked): it passes, and without the batch size's reduction it fails at once on the sizes ([40, 4]); GPU results in the table below. Rebased onto main 2b900c8fa41; this PR's own diff is otherwise unchanged. Pre-commit passes against main.

GPU check, 2026-10-29, at 10b3c1f0e35, of the new FSDP2 test on GB200 (torch 2.11.0a0, CUDA 13.1, NCCL 2.29.2, transformers 5.9.0); without the cross-rank minimum, rank 0 keeps its own estimate while the others take theirs:

Run at `10b3c1f0e35`	1 GPU	2 GPUs	4 GPUs
`test_fsdp2.py::test_planned_batch_size_then_max_calibration`	skipped (needs 2+)	passed, 22.3 s	passed, 25.0 s
The same test without the minimum's `all_reduce` in the planned branch	–	failed: sizes [36, 4]	–

Review fix, 2026-10-29, seventh round: stacked on [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 (rebased from main 2b900c8fa41 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's ec9e09c3fc0; this PR's own diff is unchanged). New GPU test tests/gpu/torch/utils/test_moe_a2a.py::test_expert_parallel_calibration_matches_the_whole_model (2+ GPUs, transformers ≥ 5.16): a tiny Qwen3-MoE loaded by transformers' own expert-parallel load on the torch_a2a styles and experts forward, quantized with FP8 through distributed_forward_loop, each rank on its own batch; every expert's weight scales and every activation scale match the whole model calibrated on every rank's batch, with bounded collectives. Its logic ran on two CPU ranks with transformers 5.18: it passes, re-planting the forward wrappers quantizer insertion drops, and fails without the re-plant, or with the conversion forcing eager as before [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074. GPU results in the table below. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 and main.

GPU check, 2026-10-29, at d90a479b6b5, of the new expert-parallel test on GB200 (torch 2.11.0a0, CUDA 13.1, NCCL 2.29.2, transformers 5.18.0): the expert-parallel calibration matches the whole model's, and needs both the forward re-plant and [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074:

Run at `d90a479b6b5`	1 GPU	2 GPUs	4 GPUs
`test_moe_a2a.py::test_expert_parallel_calibration_matches_the_whole_model`	skipped	passed, 21.6 s	passed, 26.7 s
The same test with the re-plant disabled (`_reinstall_a2a_forwards` returns 0)	–	failed: "torch_a2a experts forward requires the EP `process_group`"	–
The same test without [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 (`normalize_experts_impl_on_the_fly` forcing eager, main's behaviour)	–	failed: "Qwen3MoeExperts.forward() takes 4 positional arguments but 5 were given"	–

Restacked on 2026-10-29 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 at f68837967d3, now [2/2] on [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 [1/2] after [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's split; no conflicts, and this PR's own diff is unchanged. The new expert-parallel test's logic still passes on two CPU ranks on this base. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's head.
GPU check, 2026-10-29, at 54f0c1fd8e9 (on [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 f68837967d3, on [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 33736db61fa), on GB200 (torch 2.11.0a0, transformers 5.18.0): the expert-parallel calibration still matches the whole model with [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122's per-rank expert calibration underneath:

Run at 54f0c1fd8e9 2 GPUs 4 GPUs

test_moe_a2a.py::test_expert_parallel_calibration_matches_the_whole_model passed, 23.5 s passed, 27.0 s
Review fixes, 2026-10-29, eighth round (Claude's review of 54f0c1fd8e9): the experts' groups are checked against the ep-innermost (edp, ep) mesh only, the layout the expert-parallel load builds: accepting either axis innermost let transposed groups pass when edp == ep_size (a new four-rank case, groups-transposed, fails with that check). awq_clip reduces each module's clip search in one all-reduce, the token count riding in the losses' buffer (the ninth round reduces the losses in place instead), and every module's reduction runs before any weight gather. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074; test_calib.py and this PR's tests pass on CPU, and dropping the token count's sync fails the two-rank test.
On GPUs, with the rest of the expert-parallel series on top (transformers 5.13), before this round's change to awq_clip:
- Qwen3-30B-A3B at EP 4 and EP 2 on 4 GPUs, deterministic, batch size 1: every exported scale is bit-identical between the two. In the EP 2 run (edp 2), the same row/column check, reading the groups from the field that stack's loader seeds, accepted the loader's real (edp, ep) groups.
  
  Exported scale Differing, EP 4 vs EP 2
  
  weight_scale 0 of 18,624
  
  input_scale 0 of 18,624
  
  k_scale 0 of 48
  
  v_scale 0 of 48
- Qwen3-0.6B with nvfp4_awq_clip and --use_fsdp2 on 2 GPUs: awq_clip ran and the sharded checkpoint was exported, end to end. That run had the per-forward averaging this round replaces.
Review fixes, 2026-10-29, ninth round (Claude's review of 2dfeb9c23f0): awq_clip's candidate losses are views of one buffer, which the sync all-reduces and averages in place, with the token count in its own int64 all-reduce: no copy of the losses, where the packed buffer held up to three, and every module's reductions still run before any weight gather. The refusal names distributed_sync only for max and MSE, the calibrations that have it. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074; test_calib.py and this PR's tests pass on CPU, and the two-rank test fails without either all-reduce, or with the candidates' losses copied out of the buffer.
Review follow-up, 2026-10-29, tenth round (Claude's review of b517c7581b6): a docstring fix (07affa0fe76). The new expert-parallel test's job, run on two CPU ranks with transformers 5.18, shows that its dense linears are QuantLinear with the seeded data-parallel group and no tensor-parallel group (Qwen3-MoE's expert plan shards no dense module). Seeding them without that group fails the test: o_proj's input amax is off by 0.293 (relative).
Review follow-up, 2026-10-29, eleventh round (Claude's review of 07affa0fe76): a docstring note (e546bcca0b1) that every rank still reads and decodes every sample of a sharded stream. The experts' check that quantize() runs before conversion reads the groups a loader seeds. [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741's load is the first to seed them; until it lands, the unit tests and the expert-parallel GPU test, which seed them by hand, exercise that check.
Rebased on 2026-10-30 onto main 99c3834f62c through [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 1e74c31b49a and [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 3f3544435d8 (from [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's f5d36d4f794), as fd036d44f28; no conflicts, and this PR's own diff is unchanged. Pre-commit passes against main.
GPU check, 2026-10-30, at fd036d44f28, on GB200 (torch 2.11.0a0, transformers 5.9.0): awq_clip's in-place data-parallel sync (the ninth round's change) still gives every rank the same clip:

Run at fd036d44f28 2 GPUs 4 GPUs

test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank passed, 18 s passed, 24 s
Restacked on 2026-10-30 after [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 merged: onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 39697ad985b, which is on main 55898ee4e3b (from 3f3544435d8), as 3f56acfbb76; no conflicts, and this PR's own diff is unchanged. Pre-commit passes against main. This PR's unit tests (test_distributed_calibration.py, test_calib.py, test_dataset_utils.py, test_distributed.py) pass on CPU with transformers 5.18 and 5.5.4, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.
Rebased on 2026-10-30 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 a6758d77eb3, on main 059588b89a5, which has [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128 (from 3f56acfbb76), as 186efae02e5. [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's AWQ guard (_refuse_awq_on_a_sharded_fsdp2_model) and this PR's (_refuse_unsupported_distributed_calibration) met at the same three call sites in model_quant.py: calibrate(), quantize() and auto_quantize's candidate check. Both stay, [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's first. The changelog and the README keep both sides, and 186efae02e5 makes this PR's nvfp4_awq_clip README note its own paragraph, saying why it is not refused. This PR's own diff is otherwise unchanged. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py (with [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's guard test) and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.
Rebased on 2026-10-30 onto main d464d30b120 after [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 merged ([OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's squash; from 186efae02e5), as 9a87d153c89: [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's own commits drop out; no conflicts, and this PR's own diff is unchanged, with [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's guard still first at the three call sites. This PR now targets main, so the full CI runs, GPU included. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.
Review fixes, 2026-10-30, twelfth round (cjluo-nv's review of 9a87d153c89): mtq.auto_quantize refuses a model whose recorded plan shards its experts, before converting it (fbbe2275cf8); a two-rank test with seeded experts and an INT8 candidate checks the module types and state are unchanged, and fails without the refusal. shard_dataloader compares a sampler's source with the loader's dataset (e1976572627, typed in 9ac1486cdab); a SequentialSampler, RandomSampler or DistributedSampler over a shorter source is refused, and the three new cases fail with the old check. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.
Review fix, 2026-10-31, thirteenth round (Claude's review of 9ac1486cdab): the forward re-plant now plants only over a bound method, the forward quantizer insertion leaves, as its docstring says (aba16882660); a planned module whose forward the load never wrapped is left alone, which a new test checks and the old guard fails. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py and this PR's unit tests pass, and the expert-parallel GPU test's job still re-plants the two experts' wrappers and passes on two CPU ranks with transformers 5.18.
Rebased on 2026-10-31 onto main e352caca3c3 (from aba16882660, adding Run fused-MoE expert GPTQ updates on all of this process's GPUs #3002, Collect imatrix statistics with one fused float32 reduction #3044 and [OMNIML-6273] Let the Claude review bot submit its review again #3156), as dea9b793553; no conflicts (Collect imatrix statistics with one fused float32 reduction #3044 and Run fused-MoE expert GPTQ updates on all of this process's GPUs #3002 touch model_calib.py, which merges cleanly), and this PR's own diff is unchanged. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py, the ggml GPTQ and imatrix unit tests and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.

Before your PR is "Ready for review"

Is this change backward compatible?: ✅ Additive: get_max_batch_size and the calibration refusals change only for a model whose recorded plan spans several ranks. awq_clip's per-tensor (NVFP4) clip changes under data parallelism, where its replicas now agree (changelog, Bug Fixes).
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 Bug Fixes entry for AWQ-Clip's data-parallel clip; the feature's entry lands with [OMNIML-6085] [3/3] 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 · 2026-10-12T00:56:55Z

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 · 2026-10-12T00:56:59Z

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

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.

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

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

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

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 · 2026-10-12T02:31:16Z

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-31 04:35 UTC. Preview will be ready when the GitHub Pages deployment is complete.

codecov · 2026-10-12T03:23:29Z

Codecov Report

❌ Patch coverage is 99.42197% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 80.78%. Comparing base (2bb126f) to head (dea9b79).
⚠️ Report is 1 commits behind head on main.

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2740      +/-   ##
==========================================
+ Coverage   80.29%   80.78%   +0.48%     
==========================================
  Files         675      663      -12     
  Lines       77278    74300    -2978     
==========================================
- Hits        62054    60025    -2029     
+ Misses      15224    14275     -949

Flag	Coverage Δ
examples	`?`
examples-diffusers	`19.84% <16.76%> (-0.02%)`	⬇️
examples-gpt-oss	`13.46% <13.29%> (+<0.01%)`	⬆️
examples-hf_ptq	`23.47% <28.32%> (-0.04%)`	⬇️
examples-llm_distill	`13.53% <13.29%> (-0.01%)`	⬇️
examples-llm_eval	`17.26% <26.58%> (+<0.01%)`	⬆️
examples-llm_qat	`17.38% <16.76%> (-0.04%)`	⬇️
examples-llm_sparsity	`16.72% <13.29%> (-0.01%)`	⬇️
examples-megatron_bridge	`26.35% <16.76%> (-0.28%)`	⬇️
examples-specdec_bench	`13.23% <13.29%> (+<0.01%)`	⬆️
examples-speculative_decoding	`17.49% <16.76%> (-0.07%)`	⬇️
examples-torch_onnx	`21.10% <18.49%> (-0.04%)`	⬇️
examples-torch_trt	`15.20% <16.76%> (+<0.01%)`	⬆️
examples-vllm_serve	`17.47% <16.76%> (-0.01%)`	⬇️
gpu	`54.84% <23.12%> (-4.69%)`	⬇️
regression	`14.98% <13.29%> (-0.01%)`	⬇️
unit	`62.69% <99.42%> (+0.08%)`	⬆️

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 · 2026-10-12T03:24:22Z

/ok to test 963e78a

shengliangxu · 2026-10-12T10:39:32Z

/ok to test 935bd7f

shengliangxu · 2026-10-13T01:55:58Z

/ok to test 96159d7

shengliangxu · 2026-10-13T19:58:36Z

/ok to test 540716e

A stream that splits itself across the loader's workers by get_worker_info(), as a Hugging Face streaming dataset does, was split again by rank and silently lost samples: the shards stayed disjoint and the step counts equal, so nothing failed. shard_dataloader now refuses a streaming loader with workers, and the rank's shard no longer splits its groups across workers it cannot have. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shard_dataloader refuses a streaming loader with workers, so asking the stream to be the same in every worker, and explaining that torch seeds each worker differently, described a case that cannot arise. _StridedIterableShard's docstring carried the same clause. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shard_dataloader let any RandomSampler through, but one drawing fewer samples than the dataset (num_samples) or with replacement picks which and how many samples the loader reads, as a subset does. The rank's DistributedSampler then replaced it over the whole dataset, so calibration silently ran on all of it instead of the few asked for. Such a sampler is now refused like the other pickers; a sampler that shuffles the whole dataset is still sharded as the rank's unshuffled share. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The refusal reads the algorithm's options as the caller spelled them, which holds while every registered default syncs over the ranks; the comment names those defaults, so one that stops syncing is known to need refusing here. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The test created its own gloo group and destroyed it at once when the job ran no collective on it, as the no-plan case does. gloo's new_group can return on one rank while a peer is still connecting, so the peer's new_group failed with "Connection closed by peer" (seen once in CI). bounded_collectives makes the same bounded group and passes a barrier on it first. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Nothing defers them: modelopt.torch.utils.logging, which dataset_utils imports, already imports modelopt.torch.utils.distributed at its top. The hf_ptq argument test's new imports move to the top too. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The rank-local refusal read the algorithm dict as the caller spelled it, so {"method": "max", "distributed_sync": "false"}, which MaxCalibConfig coerces to False, calibrated without the sync. Each mode's options are now validated by its config class first, nested scale algorithms included, so coerced values and defaults count as calibration reads them. The test spells the sync off as "false", and gives nvfp4_act_headroom its local_hessian weight-scale algorithm in the form its schema accepts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…arallel size A data-parallel group of the plan's edp ranks passed the expert-parallel check even when it was the experts' own expert-parallel group, whose ranks hold other experts: at EP 2 over four ranks, max's and GPTQ's syncs would then merge different experts' scales. A routed expert's ParallelState now passes only with an expert-parallel group of ep_size ranks and a data-parallel group of edp ranks that meets it at this rank alone. The two-rank pool cannot build such groups, so the new test answers four ranks' process-group queries. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…t distributed awq_clip's static-block branch already all-reduced each candidate clip's loss over the module's data-parallel group; its per-tensor branch, which NVFP4 takes, kept each rank's own, so the ranks' copies of a weight took different clips. Both branches now average through one helper, issued on every rank in the same order, and awq_clip is no longer refused on a model distributed over several ranks. This also makes Megatron data-parallel replicas agree on the per-tensor clip. Tests: a two-rank CPU test of the static-block branch, and a two-GPU NVFP4 test of the per-tensor branch, each with ranks whose own data would pick different clips. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The schema-validated refusal read distributed_sync with "is False", but max and MSE accept None too, which skips their sync as False does, so {"method": "max", "distributed_sync": None} calibrated the ranks apart. Any falsy value is now refused, nested scale algorithms included. The refusal test gains an MSE weight-scale algorithm with distributed_sync None. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

… columns The previous check passed any expert-parallel group of ep_size ranks and data-parallel group of edp ranks meeting it only at this rank, so diagonal data-parallel groups ([0, 3] and [1, 2] at EP 2 over four ranks, where ranks 0 and 2 hold the same shard) still merged different experts' scales. A routed expert's groups must now be this rank's row of the (edp, ep) mesh the expert-parallel load lays the ranks out in, as its expert-parallel group in row order, and its column as its data-parallel group. The four-rank test adds the diagonal data-parallel group and an expert-parallel group in another order. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

… Fixes It is not backward breaking and not specific to hf_ptq: it also changes the amax values that data-parallel runs (Megatron included) export. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…the forwards The data-parallel mean ran inside each forward, in both branches (the static-block one on main too), so a module issued its all-reduces once per forward it ran. An MoE expert's forward count is each rank's own, so ranks that ran it a different number of times paired mismatched collectives and hung. Each rank now accumulates its own losses, and after the forward loop every module's candidate losses and token count are reduced once over its data-parallel group, in module order on every rank: the same numbers, as the steps are equal, with a collective count the model fixes. That also replaces up to one all-reduce per candidate, step and output block with one per module. The two-rank test now runs an expert a different number of times on each rank; it fails without the change, on the bounded group's timeout. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…rmost The row and column check took the ep axis as the mesh's innermost, as the expert-parallel load builds it today, and nothing ties the two together. It now accepts this rank's row and column of the mesh with either axis innermost: in both, the data-parallel group holds this rank's shard. Any other grouping is still refused. The four-rank test adds a mesh with ep outermost. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

An MoE expert a rank routes no tokens to can still be called, on an empty input. The per-tensor search then added a NaN loss (a mean over no outputs) for every candidate, which the post-loop sum spread to the replicas that had tokens, leaving them a clip of 0, and the static-block search failed to reshape the empty input. The search now returns at once for an empty input, leaving the losses and token count as they were, and the post-loop reduction still runs on every rank. Both distributed tests now run an expert on no tokens on one rank and check its clips are positive and those a single process picks from the other ranks' tokens: the CPU one (static-block branch) fails without the change on the reshape, the GPU one (NVFP4, per-tensor) would on a clip of 0. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…eads it hf_ptq --use_fsdp2 recorded the FSDP2 plan after its load, so that calibration sized its batch collectively and refused rank-local calibrations there. That path still retries a batch after a rank-local out-of-memory error, and the fix for that belongs where hf_ptq moves onto the run's abort path. The recording now lands with the no-gather export, the first change whose export needs the plan on that path, together with that fix. The README notes and the changelog entry on the refusal go with it; the AWQ-Clip ones stay, as its sync reads no plan. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…t, on GPUs A two-or-more-GPU test of the planned path end to end: a model under fully_shard records its plan, get_max_batch_size measures each rank's GPU in one forward (rank 0 holding a fortieth of its free memory per sample, every other rank a sixth) and every rank takes the smallest size, not its own; then each rank calibrates its shard with max through distributed_forward_loop. Every rank runs the same three steps and ends with the same scales. The batch size's reduction and the test's gathers give up in seconds on a bounded group. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

A two-or-more-GPU test of the expert-parallel calibration lifecycle on stock transformers: a tiny Qwen3-MoE loaded with its experts sharded over every rank by transformers' own expert-parallel load on the torch_a2a styles, seeded with the experts' groups, quantized with FP8 through distributed_forward_loop, each rank on its own batch. Every expert's weight scales and every activation scale match the whole model calibrated on every rank's batch. On transformers >= 5.16 quantizer insertion drops the transport's forward wrappers, which the loop re-plants: the test fails without the re-plant, and without the conversion keeping torch_a2a (#3074). The scales' syncs and the test's gathers give up in seconds on bounded groups. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Accepting the (edp, ep) mesh with either axis innermost left the check open when edp equals ep_size: there the two layouts' (row, column) pairs are each other's transpose, so experts whose expert-parallel and data-parallel groups were swapped passed, and max's and GPTQ's syncs merged different experts' scales. The check is exact again for the layout the expert-parallel load builds, ep innermost, and refuses any other. The four-rank test's transposed groups, which the either-way check accepted, are now refused. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…y weight gather The token count now rides in the fp32 losses' buffer, exact for any calibration set, so each module issues one all-reduce instead of two. And every module's reduction runs in its own pass ahead of the postprocess pass, whose weight access an FSDP2 model gathers for, instead of alternating with those gathers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The data-parallel sync stacked every candidate clip's loss into a new tensor, packed the token count after it and divided out of place, so while it ran it held up to three copies of the losses beyond the module's own. Each candidate's loss is now a view of one buffer, which the sync reduces and averages in place, and the token count goes in its own int64 all-reduce, exact for any count. Every module's reductions still run ahead of the weight gathers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The refusal advised calibrating with max, MSE, GPTQ, SmoothQuant or awq_clip "with distributed_sync on", but only max and MSE have that option. And local_hessian, refused although it has the option, says why where it is listed: its distributed_sync syncs only the max amax it starts from. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…uantizing The calibration guard runs before quantization converts anything, so on an expert-parallel plan it reads the routed experts' ParallelState from the caller, as the expert-parallel load seeds it, and checks their quantized modules again before they calibrate. The docstring now says so instead of implying that only the converted modules are checked. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

_StridedIterableShard takes one sample in world_size from the stream each rank builds, so every rank still reads and decodes all of them. Its docstring now says so, and that a map-style dataset, or a source that shards itself, saves that work for a stream costly to decode. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Main now refuses AWQ-Lite and AWQ-Full on an FSDP2 model sharded over several ranks, and its README note on that sits right above this PR's note on nvfp4_awq_clip. The two notes are now separate paragraphs, and this PR's says why nvfp4_awq_clip is not refused: AWQ-Clip folds no smoothing scale into the weights, and its clip search is averaged over the ranks. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

AutoQuantize calibrates and scores through forward_step on the model it has converted, and conversion drops the expert-parallel forwards that distributed_forward_loop restores. On an expert-parallel model the first of those forwards would fail with the experts forward's missing process_group error, after the model was converted. auto_quantize now refuses a model whose recorded plan shards its experts, before it changes anything, as quantize and calibrate refuse what they cannot calibrate. A two-rank test with seeded experts and an INT8 candidate checks the module types and state are left as they were. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shard_dataloader replaces a loader's sampler with the rank's DistributedSampler over the loader's dataset. It checked a RandomSampler's sample count against the sampler's own source, so a SequentialSampler, RandomSampler or DistributedSampler over a shorter source than the dataset passed, and the shard then calibrated on samples the loader never read. The coverage is now compared with the loader's dataset; a sampler that orders or shuffles the whole dataset is still replaced as before. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Only the three known samplers reach len(source), and each has a data_source or a dataset; the empty default only keeps the expression typed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

_reinstall_a2a_forwards planted the torch_a2a wrapper on any planned module whose instance forward was missing or a bound method. The clobber it repairs always leaves a bound method: a module with no instance forward is one the load never wrapped, and planting there would add the style's transforms to a module that ran without them. It now plants only over a bound method, as its docstring says. The pre-hook test now starts from rebound forwards, and a new test checks that a module the load did not wrap is left alone. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu · 2026-10-31T04:28:57Z

/ok to test dea9b79

shengliangxu added this pull request to stack #2755 October 12, 2026 00:57

shengliangxu force-pushed the shengliangx/hfep-4-meta-dcp-load branch from 1b34e71 to c99d63a Compare October 12, 2026 02:24

shengliangxu force-pushed the shengliangx/hfep-5-distributed-calibration branch from 7f91f94 to 963e78a Compare October 12, 2026 02:24

shengliangxu force-pushed the shengliangx/hfep-4-meta-dcp-load branch from c99d63a to d0d9b5e Compare October 12, 2026 08:39

shengliangxu force-pushed the shengliangx/hfep-5-distributed-calibration branch from 963e78a to 935bd7f Compare October 12, 2026 08:39

shengliangxu changed the title ~~[5/19] Calibrate a distributed model on per-rank data~~ [OMNIML-6071] [5/19] Calibrate a distributed model on per-rank data Oct 12, 2026

shengliangxu force-pushed the shengliangx/hfep-5-distributed-calibration branch from 935bd7f to c3ccb86 Compare October 13, 2026 00:53

shengliangxu force-pushed the shengliangx/hfep-4-meta-dcp-load branch from d0d9b5e to 64f8588 Compare October 13, 2026 00:53

shengliangxu force-pushed the shengliangx/hfep-5-distributed-calibration branch from c3ccb86 to 96159d7 Compare October 13, 2026 01:32

shengliangxu force-pushed the shengliangx/hfep-4-meta-dcp-load branch from 64f8588 to 7da05a2 Compare October 13, 2026 01:32

shengliangxu removed this pull request from stack #2755 October 13, 2026 17:42

shengliangxu added this pull request to stack #2775 October 13, 2026 17:43

shengliangxu changed the title ~~[OMNIML-6071] [5/19] Calibrate a distributed model on per-rank data~~ [OMNIML-6071] [5/20] Calibrate a distributed model on per-rank data Oct 13, 2026

shengliangxu force-pushed the shengliangx/hfep-5-distributed-calibration branch from 96159d7 to 540716e Compare October 13, 2026 19:38

shengliangxu force-pushed the shengliangx/hfep-4-meta-dcp-load branch from 7da05a2 to 00e5f72 Compare October 13, 2026 19:38

shengliangxu removed this pull request from stack #2775 October 13, 2026 19:38

shengliangxu added this pull request to stack #2781 October 13, 2026 19:40

shengliangxu changed the title ~~[OMNIML-6071] [5/20] Calibrate a distributed model on per-rank data~~ [OMNIML-6071] [6/21] Calibrate a distributed model on per-rank data Oct 13, 2026

shengliangxu removed this pull request from stack #2781 October 14, 2026 00:13

shengliangxu force-pushed the shengliangx/hfep-4-meta-dcp-load branch from 00e5f72 to 42a2de6 Compare October 14, 2026 00:14

shengliangxu force-pushed the shengliangx/hfep-5-distributed-calibration branch from 540716e to 177c692 Compare October 14, 2026 00:14

shengliangxu added this pull request to stack #2786 October 14, 2026 00:14

shengliangxu added 29 commits October 31, 2026 04:20

Give AWQ-Clip's data-parallel clip its own changelog entry, under Bug…

b1766f1

… Fixes It is not backward breaking and not specific to hf_ptq: it also changes the amax values that data-parallel runs (Megatron included) export. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Give mypy a sized default for the sampler's source

ad6e140

Only the three known samplers reach len(source), and each has a data_source or a dataset; the empty default only keeps the expression typed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Navigation Menu

Repository navigation

[OMNIML-6071] Calibrate a distributed model on per-rank data - #2740

[OMNIML-6071] Calibrate a distributed model on per-rank data#2740
shengliangxu wants to merge 60 commits into
mainfrom
shengliangx/hfep-5-distributed-calibration

shengliangxu commented Oct 12, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Oct 12, 2026

Uh oh!

coderabbitai Bot commented Oct 12, 2026 •

edited

Loading

Reviews paused

Sequence Diagram(s)

Uh oh!

github-actions Bot commented Oct 12, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-10-31 04:35 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

codecov Bot commented Oct 12, 2026 •

edited

Loading

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

shengliangxu commented Oct 13, 2026

Uh oh!

shengliangxu commented Oct 13, 2026

Uh oh!

shengliangxu commented Oct 31, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Exported scale	Differing, EP 4 vs EP 2
`weight_scale`	0 of 18,624
`input_scale`	0 of 18,624
`k_scale`	0 of 48
`v_scale`	0 of 48

Conversation

shengliangxu commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Summary by CodeRabbit

Uh oh!

copy-pr-bot Bot commented Oct 12, 2026

Uh oh!

coderabbitai Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Sequence Diagram(s)

Uh oh!

github-actions Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Built to branch gh-pages at 2026-10-31 04:35 UTC. Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

codecov Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

shengliangxu commented Oct 13, 2026

Uh oh!

shengliangxu commented Oct 13, 2026

Uh oh!

shengliangxu commented Oct 31, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shengliangxu commented Oct 12, 2026 •

edited

Loading

coderabbitai Bot commented Oct 12, 2026 •

edited

Loading

github-actions Bot commented Oct 12, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-10-31 04:35 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

codecov Bot commented Oct 12, 2026 •

edited

Loading