Repository navigation
[OMNIML-6461] Run AutoQuantize on expert-parallel and replicated models - #3250
Draft
shengliangxu wants to merge 2 commits into
Draft
shengliangxu wants to merge 2 commits into
shengliangxu wants to merge 2 commits into
Conversation
|
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. |
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: true
Tip 🎯 Customize your risk assessments
✅ Automatic approvals, with your rules
|
Collaborator
Author
|
/ok to test c9a3396 |
Contributor
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3250 +/- ##
==========================================
+ Coverage 72.89% 80.92% +8.03%
==========================================
Files 688 688
Lines 79957 79960 +3
==========================================
+ Hits 58284 64707 +6423
+ Misses 21673 15253 -6420
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
mtq.auto_quantize refused a model whose routed experts are sharded expert-parallel, and hf_ptq refused every AutoQuantize recipe under torchrun without --use_fsdp2. The searcher already sums the scores and costs over each module's TP/DP/EP groups and broadcasts the chosen formats. What was missing: - auto_quantize re-plants the routed experts' all-to-all forwards that quantizer insertion drops on transformers >= 5.16, before the search calibrates and scores. - get_cost no longer sums a DTensor weight's size over the TP/EP groups. From transformers 5.16 the expert-parallel load leaves the experts' weights as DTensors, whose size is already the whole tensor's, so the experts counted ep_size times their size in the effective-bits budget. - aumann_shapley reduces its loss and token count over a dense module's groups. A search over the experts alone reduced them over one expert's ranks, which hold only part of the data. - hf_ptq runs AutoQuantize recipes expert-parallel and replicated, its score_size spread over all ranks. A KV-cache search, whose checks need every rank to hold the same quantizers, stays refused under expert parallelism, now before the load. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Below transformers 5.16 the expert-parallel load keeps each rank's plain share of the experts behind forward pre-hooks, the other half of the cost sum and of the forwards AutoQuantize runs; the test skipped it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/dp-ep-autoquant
branch
from
November 7, 2026 01:27
c9a3396 to
e7dd863
Compare
Collaborator
Author
|
/ok to test e7dd863 |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: New feature
mtq.auto_quantizerefused a model whose routed experts are sharded expert-parallel, and atorchrunlaunch ofexamples/hf_ptq/hf_ptq.pyrefused every AutoQuantize recipe unless--use_fsdp2was passed. The searcher already sums each module's scores and costs over its TP/DP/EP groups and broadcasts the chosen formats, so ranks that each score their own data can agree. What kept it from running, and from picking the right formats:forwardand drops transformers' wrapper that hands the torch_a2a experts their process group (transformers >= 5.16).distributed_forward_loopre-plants it formtq.quantize;auto_quantizenever did, so its first forward failed. It now takes the mesh from the model's recorded plan before converting it, and re-plants the forwards right after.get_costsummed that size over the EP group, so the experts countedep_sizetimes in the effective-bits budget and the search picked other formats. A DTensor weight is no longer summed over TP/EP. A plain tensor is still this rank's shard, as with Megatron or transformers < 5.16.aumann_shapleyover the experts alone. Its loss and token count were reduced over the first module carrying a parallel state. A search whose only configurable groups are routed experts took an expert's groups, which span only the ranks holding the same experts, while its expert-parallel ranks score data of their own too. It now reduces over a dense module's groups.hf_ptq.pyruns AutoQuantize recipes on its expert-parallel and replicated routes. Each rank scores its own shard of the data, so a step scoresbatch_sizesamples per rank: the recipe'sscore_sizeis spread over every rank of any distributed plan, not only under--use_fsdp2.Still refused: a KV-cache search under expert parallelism. Its checks require every rank to hold the same quantizers, and each expert-parallel rank holds its own experts'. The refusal moves from
auto_quantize's entry into the KV-cache search, still before anything converts the model.hf_ptq.pynow refuses such a recipe before the load, rather than after the weight search it runs first. Rank-local calibrations (AWQ and the others) stay refused per candidate, as before.Usage
Testing
New
tests/unit/torch/quantization/test_autoquant.py::test_expert_parallel_auto_quantize_picks_one_processes_formats[gradient|aumann_shapley], on two CPU ranks (gloo). A tiny Qwen3-MoE is sharded as transformers' expert-parallel load leaves it (apply_tensor_parallelismwith the torch_a2a styles: DTensor experts behind forward wrappers), with each module's groups seeded as the load seeds them. Each rank searches on its own batch; one process searches on every rank's batch, one per step. Formats and costs must be equal, scores within 1e-2 relative (gradient: ~1e-7). Theaumann_shapleycase searches the experts alone. Each fix fails it when reverted:torch_a2a experts forward requires the EP process_groupaumann_shapleyid: scores 3-4x offAbout 2 s per id once the shared worker pool is up.
Updated: the library's EP refusal test now checks the KV-cache search's refusal.
test_hf_ptq_expert_parallel.pychecks that a weight recipe runs on every route and a KV-cache recipe is refused before the load under expert parallelism.test_autoquant_utils.pychecks the score steps on a replicated plan. Each fails when its change is reverted.Pre-commit, mypy included, passes against main.
On CPU with transformers 5.18 and torch 2.14, before rebasing onto current main,
tests/unit/torch/quantization,tests/unit/torch/utilsandtests/examples/hf_ptqfailed exactly the same ids at main and on this branch, all of them failures of that host's environment, and the transformers-free lane passed the changed unit files. Main now carries [OMNIML-6370] [1/5] Isolate weight-rewriting AutoQuantize candidates from the source weights #3154, which wraps the search in a block that hands back the unconverted model on failure; the re-plant now sits inside it. After the rebase, the changed test files plustest_autoquant_shapley.py,test_ggml_gptq.pyandplugins/test_fused_experts.pypass (548 passed).New GPU test
tests/gpu/torch/utils/test_moe_a2a.py::test_expert_parallel_auto_quantize_and_export_match_the_whole_model, three layouts: experts over every rank, over half of them (experts replicated across two groups) and with the dense part FSDP2-sharded. It usesload_pretrained's expert-parallel load in float32, compares the search with the whole model's, then compares the no-gather writer's mixed-format checkpoint, tensors andhf_quant_config.json, with the whole model's single-process export:The 5.13 run is the pre-5.16 path: the experts load as each rank's plain share behind forward pre-hooks. The rest of
test_moe_a2a.pypasses and skips the same tests on main and on this branch at 2 and 4 GPUs. The 2- and 4-GPU rows with transformers 5.18 were re-run at the rebased head with the same results; the 1-GPU and 5.13 rows are from this change before its rebase onto current main.Regression on GPUs, main vs this branch:
test_fsdp2_autoquant.pypasses the same 6 tests on both at 2 GPUs, and skips them at 1.Megatron, at the rebased head on 2 GPUs:
test_megatron.py'stest_auto_quantize,test_auto_quantize_moe_epandtest_auto_quantize_mamba_hybrid_ep_costpass. The last checks that Megatron's AutoQuantize costs agree with the experts on 1 rank and on 2, a comparison CI's 1-GPU runner never reaches. Its 2-GPU run takes longer than the suite's default per-test timeout, so it ran with that timeout off.Real model: Qwen3-30B-A3B through
hf_ptq.pywithgeneral/auto_quantize/nvfp4_fp8_at_5p4bitson 4 GPUs (transformers 5.18, torch 2.11), in three layouts. Each run finished the search and exported its checkpoint, and the layouts chose the same format for all but two of the quantized layers:The two that differ are the attention output projections of layers 35 and 36, which have the same shape and so the same cost. The run with the experts over all 4 ranks gave NVFP4 to layer 35 and FP8 to layer 36, and the other two runs did the reverse, so every run ends with the same number of layers in each format. Spreading the experts differently changes the order of the floating-point arithmetic in them, which can tip a near-tie between two equal-cost layers either way. The runs did not keep the scores that would show that margin. These runs used this change before its rebase onto current main (42570f21ace).
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).hf_ptq.pyloads the model. Single-process, FSDP2 and Megatron searches are unchanged: their weights are plain tensors or FSDP2 DTensors whose TP/EP groups are unset.CONTRIBUTING.md: N/Atorchrunlaunch refuses. Both sit in 0.49.0 for now, the only unreleased section on main. Since this targets 0.50.0, the entry moves to a 0.50.0 section once main has one, and [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py, with lockstep previews #3139's 0.49.0 entry stays as it is.Additional Information
Tracked by OMNIML-6461, for ModelOpt 0.50.0. Lifts the refusals from #2740 (library) and #3139 (
hf_ptq.py), and builds on the autograd all-to-all from #2738. Overlaps #3236 in that one CHANGELOG line, which both edit; whichever merges second keeps both edits.