[shengliangxu]

What does this PR do?

Type of change: New feature

First of 21 PRs that add expert-parallel (EP) PTQ for Hugging Face MoE models to examples/hf_ptq under torchrun. They merge in order:

• [1/21]-[3/21], the building blocks: this all-to-all expert transport, amax synchronization that holds under data-parallel calibration ([2/21]), and fused-experts quantization under grouped_mm ([3/21]).
• [4/21]-[12/21], the distribution plan and the loads it routes to: a DistributionPlan decided once per run, and load_pretrained loading by it, replicated, FSDP2-sharded, or with an MoE's experts sharded. After [12/21], an MoE whose expert-parallel plan ships with transformers loads and calibrates expert-parallel.
• [13/21]-[17/21], model specs and export: the specs ([13/21]) make an MoE expert-parallel when transformers ships no usable plan for it, such as DeepSeek-V3, and refuse the ones the path cannot handle yet; [14/21]-[17/21] write a sharded model's Hugging Face checkpoint without gathering it, expert-parallel included.
• [18/21]-[20/21], models that plug in through their specs: compressed-tensors experts decompressed as they load, and the Kimi models.
• [21/21] runs it all from hf_ptq.py.

This PR adds the transport.

Calibrating a large MoE has no good option today. Replicating it needs the whole model on every rank. Sharding it with FSDP2 alone gathers every expert on every rank for each forward, so the expert weights dominate both memory and time, while each rank still only sees its own slice of tokens.

DP=EP instead: dense layers and the router run data-parallel, the routed experts are sharded across ranks, and tokens go to the rank that owns their expert by an all-to-all dispatch and come back by the combine. Each expert sees every token routed to it from all ranks, so its amax is collected over the global token pool rather than one rank's share. That is the accuracy point, not just a memory one. The expert matmul stays on transformers' own _grouped_linear, where [3/21] attaches ModelOpt's per-expert quantizers as in a single process. register_torch_a2a() installs the transport as a transformers experts implementation, lazily rather than at import.

The transport is plain all_to_all_single over NCCL, so it works across nodes at any rank count. Three things it has to get right:

• The routing weights stay on the token's rank. The experts' results come back unweighted, and the rank that holds each token weights them in the routing weights' own dtype, as transformers' grouped_mm does: a router that returns fp32 weights (Mixtral's, DeepSeek-V3's) gets an fp32 product. The weights never travel, so no hop carries them, and their gradient comes from the local product.
• The combine sums each token's top-k in a fixed order. A sum whose order depends on how tokens were distributed, or on the atomics index_add_ uses on CUDA, rounds differently from one EP degree or run to the next, so activation scales would shift and a checkpoint would not reproduce. The results come back in (token, slot) order, and each token's top-k rows are weighted and summed with view(n, top_k, H).sum(1), as transformers' grouped_mm does, a slice of tokens at a time. Each transport buffer is freed as soon as its stage is done: a whole fp32 copy of the results, beside buffers kept to the end, reached GiBs at calibration batch sizes.
• Large transfers. all_to_all_single addresses the flattened buffer with element offsets, which overflow int32 once rows × hidden exceeds 2**31. The transport chunks along the hidden dimension, into slices as wide as the most rows any rank moves allows. Ranks agree on that row count with one all-reduce per layer; deciding it from local sizes desyncs the hop count and deadlocks.

A rank that receives more than 4x the most tokens any rank sent (a collapsed router piling the cluster's tokens onto a few experts) warns once, before the expert stage allocates for that count. A rank can receive at most world size times that, so the warning fires only above 4 ranks.

The incidents behind the transport's less obvious code:

Code	Why	Seen on
Hops over 2**31 elements go as hidden-column slices	Keeps every hop within all_to_all_single's int32 element offsets, for an extra hop or two past 2**31. Added after a hang in EP calibration that has not been reproduced since: with the slicing disabled, such hops complete exactly on NCCL 2.29, 2.30 and 2.31, including at 32 ranks over the network on 2.29 and 2.30	Qwen3.5-MoE EP calibration at batch 8 x seq 4096: rows x hidden = 2.68e9
The host waits for the count exchange before reading the counts	A read racing the collective's stream sized the dispatch from garbage, and the next all-to-all crashed	Qwen3.5-MoE EP calibration
Each token's top-k is summed in a fixed order	A bf16, order-dependent combine made exported activation scales EP-dependent	Qwen3-30B-A3B, top-8: 27% of input scales differed from single-GPU, by up to 28%

Usage

Internal to the expert-parallel loader ([9/21]); nothing to call directly.

Testing

• tests/unit/torch/utils/test_moe_a2a.py, under gloo on CPU:

  • transformers' fused Qwen3-MoE experts split over two ranks: through the registered torch_a2a forward they give the whole module's output and gradients, for the hidden states, the routing weights and each rank's own experts (mutation-checked);
  • hops split along the hidden dim, each within the per-hop limit, once that limit is shrunk below the transfer size, with a bit-identical result;
  • the collapse warning, its ratio lowered for two ranks: collapsed routing warns on the receiving rank only, and unequal batches with even routing warn nowhere (fails on the previous revision, which compared with the rank's own count);
  • with fp32 routing weights on bf16 experts, at EP=1 the torch_a2a forward gives grouped_mm's result bit for bit (fails on the previous revision, which rounded the weights to bf16);
  • the combine's temporaries stay within one slice of tokens, with the same result as unsliced (fails on the previous revision, whose fp32 [tokens, hidden] accumulator exceeds the slice);
  • a rank that sends nothing, its routing slots all -1 or no tokens at all, but serves its peer's tokens: its backward still runs the reverse all-to-all, so both ranks get their input gradients (the test's group times out in 10 s, so a regression fails instead of hanging);
  • experts whose biases the plan left whole (GPT-OSS's, with only the weights sharded) are refused before any collective.

• tests/gpu/torch/utils/test_moe_a2a.py: with fp32 routing weights, at EP=1 the torch_a2a forward gives grouped_mm's result bit for bit on CUDA. Run on GB200, against the previous revision's index_add_ combine and this one (elements that differ from grouped_mm, of 4,194,304):

  Combine	1 GPU	4 GPUs, each rank its own one-rank group
  Previous (index_add_)	223	217 to 232 per rank
  This revision	0	0 on every rank

• On a real model, run to run: two identical calibrations of Qwen3-30B-A3B in FP8 (fp8_default-kv_fp8), EP4 on 4 GB200 with the dense layers replicated, batch size 2, compared scale by scale (the previous combine's pair ran at two revisions that both had it, on transformers 5.13; this revision's on 5.9):

  Combine	Input scales that differ	KV-cache scales that differ	Weight scales that differ
  Previous (index_add_)	262 of 18,624	2 of 96	0 of 18,624
  This revision	0 of 18,624	0 of 96	0 of 18,624

• Peak memory of one layer's transport, on CPU, at 4096 tokens, hidden 4096, top-k 8, fp32 routing weights:

  Revision	Peak RSS growth
  Before the bounded combine and the early frees	2122 MiB
  This revision	706 MiB

• The two-rank tests run on cpu_dist_workers, the two gloo ranks main's unit tests keep up for the session (tests/unit/conftest.py), since spawning ranks cost each test about 14 s of the linux unit job. This PR makes those workers reseed after every job, so each job starts from the seed a freshly spawned rank gets; a job that draws a different amount of randomness on each rank otherwise leaves the next job's ranks out of step. Since both ranks then draw the same data, main's test_data_parallel now seeds its calibration data by rank: its group-MAX check would otherwise hold with or without the data-parallel amax sync. It fails when that sync is removed. test_autoquant's data-parallel test seeds its batches by rank for the same reason, and fails when the score sync is removed.

• The result keeps the dtype it is given: fp32 and fp16 experts through the transport, in a one-rank group, match a dense reference in their own dtype. Tokens travel in the model's dtype.

• The test of a hand-written expert function against a dense reference is dropped: the Qwen3-MoE experts test checks the same forward and gradients, and the dtype test runs a one-rank group. The drop is checked by mutation.

• Locally on this revision, the series on main 6cd9ced: pre-commit (every hook, mypy included) over each PR of the stack against main; the changed tests in this PR and in [2/21], [3/21] and [8/21], and at the top of the stack the unit test files that run this transport, on transformers 5.18 and 5.5. CI runs the full unit lanes and the docs build.

• The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.

• On GPUs: tests/gpu/torch/utils/test_expert_parallel_load.py ([9/21]) runs a tiny Qwen3-MoE through this transport on every GPU and matches the unsharded model's logits; the series' model runs exercise it at 4 to 64 ranks (see [21/21]).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive; registered only when the expert-parallel loader asks for it.
• 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?: N/A. The feature's entry lands with [21/21].
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 1 of 21 of expert-parallel PTQ for Hugging Face MoE models. #2733 has merged, so this PR is based on main and merges first; the series then merges in order, [1/21] through [21/21].

Summary by CodeRabbit

• New Features

  • Added support for running mixture-of-experts models across multiple devices, routing tokens to experts and combining their outputs.
  • Added support for per-expert biases and split handling of large transfers, with results accumulated in FP32 before conversion.
  • Added a warning when a device receives substantially more tokens than it sends.

• Tests

  • Added coverage for distributed outputs, gradients, transfer splitting, imbalance warnings, and quantization synchronization.
  • Improved consistency by reseeding distributed workers between jobs.

[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: e256dcc
• Latest commit, not reviewed: 858dd1b

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

Adds a lazily registered Torch all-to-all MoE expert-parallel path with Transformers integration. Adds distributed tests for dispatch, output and gradient correctness, chunking, and imbalance warnings. Updates worker reseeding and distributed quantization synchronization checks.

Changes: Torch A2A MoE expert parallelism

Layer / File(s)	Summary
Token dispatch and output combination modelopt/torch/utils/moe.py	Routes token–expert pairs to owning ranks, exchanges token counts, runs local experts, and combines outputs. Large transfers are chunked across the hidden dimension, and output accumulation uses FP32.
Transformers expert integration modelopt/torch/utils/moe.py	Adds the Transformers experts forward and lazy registration of router and expert parallel styles. The expert style injects the process group and bypasses output reduction.
Distributed validation and worker setup tests/_test_utils/torch/distributed/utils.py, tests/unit/conftest.py, tests/unit/torch/utils/test_moe_a2a.py	Reseeds distributed workers after teardown. Tests cover outputs and gradients, dtype preservation, chunking, imbalance warnings, and Qwen3-MoE integration.

Distributed quantization test adjustments

Layer / File(s)	Summary
Rank-specific calibration and synchronization checks tests/unit/torch/quantization/test_dist.py	Seeds each rank independently and completes quantizer reductions before reporting all local amax mismatches.

Priority: ➖ Normal

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant SourceRank
  participant AllToAll
  participant DestinationRank
  participant LocalExpertFunction
  SourceRank->>AllToAll: Exchange counts and routed token data
  AllToAll->>DestinationRank: Deliver token groups
  DestinationRank->>LocalExpertFunction: Run local expert computation
  LocalExpertFunction->>AllToAll: Return weighted expert outputs
  AllToAll->>SourceRank: Deliver expert outputs

Loading

Suggested reviewers: kevalmorabia97, danielkorzekwa

Merge Risk: 🟡 Moderate · up to e256d

The new expert transport cannot be registered with a supported Transformers version. Add a compatible router-class import before merging.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 44.12% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 34 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 names the DP=EP all-to-all expert transport, which is the main change: the new modelopt/torch/utils/moe.py module adds a2a_moe, torch_a2a_experts_forward, and register_torch_a2a. The title a…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		No security anti-pattern from the custom check was introduced. The pull request changes one modelopt Python file and test files only; it adds no examples Python files or dependency files. The added co…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Title check

Explanation

The title names the DP=EP all-to-all expert transport, which is the main change: the new modelopt/torch/utils/moe.py module adds a2a_moe, torch_a2a_experts_forward, and register_torch_a2a. The title also carries the ticket and stack position tag [OMNIML-6059] [1/21], which is acceptable. The test and seeding changes are secondary and need not appear in the title.

Full details: Security Anti-Patterns

Explanation

No security anti-pattern from the custom check was introduced. The pull request changes one modelopt Python file and test files only; it adds no examples Python files or dependency files. The added code contains no unsafe torch.load, numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, externally derived eval()/exec(), or # nosec comments.

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• 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.}

[codecov]

Codecov Report

❌ Patch coverage is 88.23529% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.11%. Comparing base (6cd9ced) to head (858dd1b).

Files with missing lines	Patch %	Lines
modelopt/torch/utils/moe.py	88.23%	16 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2738      +/-   ##
==========================================
- Coverage   79.80%   79.11%   -0.69%     
==========================================
  Files         660      661       +1     
  Lines       74792    74928     +136     
==========================================
- Hits        59688    59282     -406     
- Misses      15104    15646     +542     

Flag	Coverage Δ
examples-diffusers	19.90% <0.00%> (-0.04%)	⬇️
examples-gpt-oss	13.44% <0.00%> (-0.03%)	⬇️
examples-hf_ptq	23.34% <0.00%> (-0.08%)	⬇️
examples-llm_distill	13.52% <0.00%> (-0.03%)	⬇️
examples-llm_eval	17.17% <0.00%> (-0.04%)	⬇️
examples-llm_qat	17.34% <0.00%> (-0.07%)	⬇️
examples-llm_sparsity	16.76% <0.00%> (-0.04%)	⬇️
examples-megatron_bridge	26.39% <0.00%> (-0.17%)	⬇️
examples-specdec_bench	13.23% <0.00%> (-0.03%)	⬇️
examples-speculative_decoding	17.42% <0.00%> (-0.10%)	⬇️
examples-torch_onnx	21.22% <0.00%> (-0.06%)	⬇️
examples-torch_trt	15.15% <0.00%> (-0.03%)	⬇️
examples-vllm_serve	17.38% <0.00%> (-0.04%)	⬇️
gpu	58.91% <77.94%> (-1.17%)	⬇️
regression	15.04% <0.00%> (-0.04%)	⬇️
unit	61.11% <88.23%> (+0.04%)	⬆️

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.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-24 01:19 UTC

[shengliangxu]

/ok to test cd7e9a6

[shengliangxu]

/ok to test 69f3a98

[shengliangxu]

/ok to test e301775

[shengliangxu]

/ok to test 7267471

[meenchen]

One minor comment, approve to unblock

[shengliangxu]

/ok to test 6eca1d7

[sychen52]

A few comments on the transport at 6eca1d7. They come from reading the code; nothing was run.

• [P3] tests/unit/torch/quantization/test_autoquant.py:956 (outside the diff). _test_data_parallel_auto_quantize builds its data with model.get_input(), which draws from the global RNG, so both ranks get the same batches. Its "same costs and scores on all ranks" check then passes even without the data-parallel score sync. That was often true before too, because the workers start from seed 1234, but the reseed after every job now makes it certain. It's the same gap this PR closes in test_dist.py, and a torch.manual_seed(rank) right after SimpleLinear() would make the check meaningful there too. The other pooled data-parallel tests are fine as they are. test_data_parallel_moe_score_module and test_data_parallel_aumann_shapley use the same data on every rank on purpose, and they check the reduced values against a single-rank reference, so a dropped reduction still fails them.
• Comment style. Several new comments read like incident notes: CAPS for emphasis, "observed: … all_to_all ran the full 3h NCCL timeout", "deterministic on the qwen3_5_moe_text EP path", "27% of input_scale values differed … by up to 28%", "the validated production path byte-for-byte unchanged". They do explain why, which helps. But CONTRIBUTING.md asks for one or two lines per comment, and puts benchmark numbers and root-cause analysis in the PR description. The longest are the combine comment (lines 208-219) and the count-exchange comment (148-153).

[shengliangxu]

/ok to test 858dd1b

[shengliangxu]

@sychen52, on the two points in your review body:

• test_autoquant.py: fixed in 858dd1b. The batches are seeded by rank after the model is built, and with the data-parallel score sync removed the test now fails.
• Comment style: fcf41ae trims the transport's comments to the reason for each piece of code. The incidents behind them, with their numbers, are now in the PR description.