Repository navigation
[OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally - #2750
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. |
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. ℹ️ Recent review info
Walkthrough
Merge Risk:
 |
| 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 | |
No listed security anti-pattern was introduced. The reviewed Python additions contain no torch.load(..., weights_only=False), numpy.load/np.load with allow_pickle=True, hardcoded trust_remote_code=Tru… |
| Title check | |
The title clearly and concisely describes the main change: folding fused experts' quantizers into weight shards locally. |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The reviewed Python additions contain no torch.load(..., weights_only=False), numpy.load/np.load with allow_pickle=True, hardcoded trust_remote_code=True, eval/exec calls, or # nosec comments. The pull request does not change pyproject.toml or requirements.txt.
✨ 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.
65d7878 to
e832c67
Compare
660d339 to
784e815
Compare
e832c67 to
52b4592
Compare
|
|
/ok to test 52b4592 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2750 +/- ##
==========================================
+ Coverage 78.18% 79.91% +1.72%
==========================================
Files 664 664
Lines 76452 76659 +207
==========================================
+ Hits 59776 61262 +1486
+ Misses 16676 15397 -1279
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:
|
52b4592 to
d314e75
Compare
|
/ok to test d314e75 |
shengliangxu
changed the title
d314e75 to
636aea0
Compare
57494b6 to
25176cb
Compare
636aea0 to
dbb570a
Compare
|
/ok to test dbb570a |
shengliangxu
changed the title
dbb570a to
e85bc68
Compare
53b85aa to
ba6939d
Compare
shengliangxu
changed the title
ba6939d to
dec8f8c
Compare
e85bc68 to
be5d536
Compare
dec8f8c to
954d8bf
Compare
be5d536 to
4e1039c
Compare
|
Thanks. The optional item is done in 26e9863: the paragraph of |
`<<` is a silent no-op on a DTensor, and the INT4, NF4 and MXFP4 packers, and the export's INT4 packer, still shifted with it. They now use torch.bitwise_left_shift, which gives the same bytes on a plain tensor. None of them packs a shard either way: each combines its input with plain tensors, which DTensor refuses deep inside, or, on INT4's fallback, gathers it whole on every rank. So each now refuses a DTensor up front, and NVFP4 refuses one that comes with a plain per-block scale, such as a static quantizer's, which covers the whole weight. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The window keys a parameter its root owns by its bare name, and _param_key took any falsy module name for the root. A module the window's name map does not know (None) then took the root's key, and the window would have handed FSDP2 that module's weight as the root's own. Only "" names the root now; an unknown module keeps its old None. key, which no parameter has, so the window skips it as before. keep_fused_experts is documented in a docstring, like the fields after it, as the marker of the no-gather export. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…he window's root key, into the keep-fused fold keep_fused_experts' docstring, from [1/2], also says what the field means for fused experts here. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test fa69b30 |
Each rank folds the experts it holds on its GPU, at their global offsets, and the folds, gathered over NCCL, split into the per-expert export, an uneven shard included. Under expert parallelism the ranks decide on each projection together: their quantizer states go over Gloo in the cpu:gloo,cuda:nccl group the export runs on, and over NCCL in a group of NCCL alone. Every collective runs on a group that gives up after a minute. The test skips below two GPUs. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test ec2741c |
|
Following up on the second sign-off item, that the expert-parallel
GB200, torch 2.11.0a0 (NGC 26.02), CUDA 13.1, NCCL 2.29.2, transformers 5.9.0:
The mutant ignores the gathered states. It still refuses, but names the rank's local experts ("quantizers [1] of 2") instead of the global ones ("[3] of 4" on two GPUs, "[7] of 8" on four). So the test pins that the refusal comes from the states gathered across the ranks. |
…PU test _process_quantized_modules enters held_fsdp2_weight_windows around its dispatch loop, but its context leaves keep_fused_experts unset and its assert keeps sharded models out, so no window opens there yet. Its comment now says so instead of describing a live property. The GPU test's comment names what ends a rank left waiting in a window: pytest's per-test timeout, after which the worker pool kills its ranks. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…nto the keep-fused fold Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 6f1d1da |
[1/2] merged into main as one squashed commit, so this branch takes main in a merge, as it took [1/2]'s heads before, and keeps its own commits. registry.py and unified_export_hf.py conflicted where this PR extends [1/2]'s text, and take this PR's side. Main holds [1/2]'s files exactly as its last head left them, so the PR's diff against main is the diff it had against [1/2], byte for byte. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 830d704 |
| return sorted(prefixes) | ||
|
|
||
|
|
||
| def split_fused_experts_state_dict(state_dict: dict) -> dict: |
There was a problem hiding this comment.
Question (non-blocking): the per-expert keys are numbered from 0 over whatever tensors come in. The docstring says "one module's shard at a time", but the tests split full_tensor() results. Is the contract whole tensors, or a local shard with the caller adding the offset? A rank-1 shard would otherwise be numbered from 0. Worth one line in the docstring, since #2751/#2752 call it.
There was a problem hiding this comment.
Plain tensors, numbered by their row, with the caller owning the global numbers. The docstring now says so (fc6d998):
split_fused_experts_state_dicttakesn_experts = first_w.shape[0]and emits keys{prefix}.{expert_idx}.*forexpert_idx in range(n_experts). So a module's experts are numbered from 0 by their row in whatever it is given.- Whole tensors get their global numbers. The tests hand it
full_tensor()results, and the single-process layout is whole. - For a rank's local shard, the caller shifts the numbers. In [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752,
_split_local_fused_modulepasses theto_local()rows, then adds the shard's global offset (_dtensor_dim0_offset, plus the expert-parallel block base when the weight holds only its group's experts) to each index. - So a rank-1 shard is numbered from 0 only inside that call, and the writer never sees those numbers.
- "One module's shard at a time" meant one module's tensors, to bound memory. It read as a rank's shard, so it now says "tensors".
| return DTensor.from_local(local, mesh, placements, run_check=False, shape=shape, stride=stride) | ||
|
|
||
|
|
||
| class _LocalRows: |
There was a problem hiding this comment.
Suggestion (non-blocking): the proxy self works, but it depends on which attributes the accessor reads, and the docstring notes a bound method would read the whole DTensor. get_expert_weight is this repo's own method, so could it take an optional tensor instead? Also, native FP8 is now refused up front, so what does the proxy still cover that local_w[idx].to(dtype) doesn't?
There was a problem hiding this comment.
On what the proxy covers that local_w[idx].to(dtype) does not: the weight domain the per-expert export packs.
_export_fused_expertsreads each expert throughmodule.get_expert_weight(proj, idx, dtype). The base_QuantFusedExperts.get_expert_weightignoresdtypeand returnsgetattr(self, name)[idx]as stored.- The export derives the NVFP4 per-block scales from that weight, and leaves an unquantized projection in its stored dtype.
.to(dtype)would round each expert to the export dtype first. I swapped both of the fold's accessor reads for it, and 12 of the byte-for-byte tests fail: NVFP4, the uncalibrated fallback, the unquantized-projection cases, W4A8 AWQ's split, and the two-rank fold.
I'd keep the proxy rather than add a tensor argument to get_expert_weight, for now:
- Who defines it: [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 and [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 redefine both accessors,
_QuantFusedExperts.get_expert_weightand_QuantFP8Experts.get_expert_weight, to index through_local_experts. A new argument has to be added and honored in all of them, and whichever PR merges second has to carry it again. That couples this PR to theirs for a call only the fold makes. - What it reaches: for each of those accessors the rows stand in exactly where an argument would. Each reads the projection once as
getattr(self, name)and hands it to its helper as an argument (self._local_experts(getattr(self, name))). The docstring's caveat is for a helper that readsself.<proj>on its own, which none does. - Coverage: the keep-fused tests fold through the real accessors and compare byte for byte with the per-expert export. They use
_QuantFusedExpertson the gated layout and_QuantNonGatedFusedExpertson transformers'NemotronHExperts, and the two-rank job folds rows a rank holds at a non-zero offset.
Once #3074 and #3067 land, the argument would fit as one follow-up that changes every definition at once.
It numbers each module's experts from 0 by their row in the tensors it receives, so whole tensors get their global numbers, and a caller that splits a rank's local shard shifts each index by the shard's offset. "One module's shard at a time" meant one module's tensors, which read as a rank's shard. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test fc6d998 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| for proj_name, weight, ws, ins, ws2, row_slice in projections: | ||
| base = f"{prefix}.{expert_idx}.{proj_name}" | ||
| # .clone() every emitted tensor so no two keys share storage. gate/up share the | ||
| # per-tensor scale object (and slices alias the fused parent); without cloning, | ||
| # safetensors/save_pretrained shared-tensor dedup silently drops the duplicate | ||
| # (e.g. up_proj.weight_scale/input_scale would go missing). | ||
| out[f"{base}.weight"] = weight.detach().clone().contiguous() | ||
| sliced_ws = _slice_weight_scale(ws, expert_idx, row_slice, first_w.shape[1]) | ||
| if sliced_ws is not None: | ||
| out[f"{base}.weight_scale"] = sliced_ws.detach().clone().contiguous() | ||
| # weight_scale_2 is a per-expert scalar (shared by gate/up): _slice_weight_scale's | ||
| # dim<=1 branch returns ws2[expert_idx] regardless of row_slice. NVFP4 only; None for FP8. | ||
| sliced_ws2 = _slice_weight_scale(ws2, expert_idx, row_slice, first_w.shape[1]) |
There was a problem hiding this comment.
[SUGGESTION] _slice_weight_scale's rows denominator is first_w.shape[1] for every projection in the tuple list, including down_proj — whose own fused row count is down_w.shape[1].
It is correct today only by short-circuit: down_proj (and the ungated up_proj) always pass row_slice=None, so line 718 returns before rows is read. Nothing in the signature or the comment says the argument is only meaningful for the gate|up cut, so the first caller that ever passes a row_slice for down_proj gets proportional math against the wrong denominator — and for the common HIDDEN_DIM != 2 * INTERMEDIATE_DIM case that is a silently mis-sliced block scale rather than the NotImplementedError at line 725.
Cheapest fix is to carry the projection's own fused total in the tuple instead of re-deriving it at the call site:
projections = [
("gate_proj", first_w[expert_idx, :inter, :], first_ws, first_in, first_ws2, slice(0, inter), first_w.shape[1]),
("up_proj", first_w[expert_idx, inter:, :], first_ws, first_in, first_ws2, slice(inter, None), first_w.shape[1]),
("down_proj", down_w[expert_idx], down_ws, down_in, down_ws2, None, down_w.shape[1]),
]
...
for proj_name, weight, ws, ins, ws2, row_slice, fused_rows in projections:
...
sliced_ws = _slice_weight_scale(ws, expert_idx, row_slice, fused_rows)
sliced_ws2 = _slice_weight_scale(ws2, expert_idx, row_slice, fused_rows)There was a problem hiding this comment.
The comment now says it, in 945271d: rows is the fused gate|up's row count, and only the gate and up cuts read it. That is the case today too. down_proj and an ungated up_proj pass row_slice=None, which returns before rows is read, and the only row slices are the two halves of first_w, which the loop builds itself.
I didn't add a per-projection total to the tuples. It would be a field no current entry reads, and this PR sits at 499 of its 500 source lines.
There was a problem hiding this comment.
Claude review — keep-fused expert fold, [2/2] (round 4)
Scope: full review (the trigger comment carried no scoping instructions). 6 files changed (+1498/−9). I read the whole modelopt/ diff (4 files, 506 insertions) and both test files. The branch is a single squashed commit against main, so instead of diffing against the earlier rounds I re-traced the cross-file pieces: _fsdp2_weight_window / held_fsdp2_weight_windows / _export_weight (hf_export_handlers.py), _process_quantized_modules (unified_export_hf.py), _export_fused_experts and _delete_fused_moe_source_attrs (moe_utils.py), the two get_expert_weight implementations in quantization/plugins/huggingface.py, and the dist_workers / cpu_dist_workers / need_2_gpus fixtures.
Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2 (both posted inline, both non-blocking)
What I verified this round
Fold/split equivalence with the reference path. I traced _export_fused_experts_keep_fused + split_fused_experts_state_dict against _export_fused_experts stage by stage. The two agree where it matters:
- The uncalibrated fallback uses the whole fused
(2I, H)expert slice, so gate and up share one amax and therefore oneweight_scale_2— the same invariant that the shared-amax step atmoe_utils.py:473-491exists to preserve for the W1/W3 fusion in vLLM. The non-gated anddown_projcases fall through to the per-projection slice, matching the fallback at line 557. - Quantizing the fused slice once and slicing the resulting scale is equivalent to the reference order (slice the amax, then quantize) for per-row and per-block amax, since each row scale depends only on its own amax; and trivially so for per-tensor amax, which the
numel() > 1guard leaves unsliced. - The proportional cut in
_slice_weight_scaleis exact for the layouts claimed: NVFP4(2I, H/16), INT4 and W4A8(2I, H/group), FP8 2-D(ceil(2I/128), ceil(H/128)). The gpt-oss-shaped case where a 128-row block straddles the gate|up cut raises rather than rounding.
Rank agreement. The decision loop runs every check and both all_gather_object calls for both projections before the fold loop touches anything, so a refusal on the second projection cannot leave the first folded — the "refused with the module untouched" claim holds. _refuse_unfoldable_shards reaches the same verdict on every rank: the global shape and the mesh are replicated, and the only rank-dependent input, an uneven EP quantizer count, is refused on all ranks by the num_quantizers * shards check. shards can accumulate at most one mesh dim, because a multi-Shard placement is refused first. The zero-expert guard matches the chunking torch does (ceil(E/size)), so 3-over-2 passes and 7-over-8 is refused.
Previously raised, now closed.
- The round-1 hazard where
_fused_experts_prefixesalso matched the BMM experts:_export_bmm_expertsnow raises before itsQUANTIZATION_NONEearly return, and it is registered by class name, so it is dispatched for everyGptOssExperts/Llama4TextExpertswhether quantized or not. A model with transposed BMM experts can no longer reach the keep-fused export at all, so the split never meets one. - The round-2 held-window module-list dedup is in place at
hf_export_handlers.py:89-90(identity check). - The all-gather staging-device suggestion now has real coverage rather than an argument:
test_fused_experts_keep_fused_gpu.pyparametrizes the group overcpu:gloo,cuda:nccland over NCCL alone, so the object gather runs on both, withtorch.cuda.set_devicedone by the worker pool.
The new GPU test (tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py, 211 lines) is the main thing added since the last round, and it is well targeted: the sharded fold checked byte for byte against the plain fold at the rank offset, round-tripped through full_tensor() into the per-expert export, plus the uneven 2*W-1 shard, the no-expert-rank refusal, EP-local quantizer lists, and the joint decision of the ranks on a partly-disabled projection. One soft spot: the EP comparison at lines 180-183 iterates the folded state dict and indexes plain[key], so an extra key fails the test but a missing one does not — the assert set(fused) == set(plain) used at line 138 would close that.
Risk
Low to moderate. The feature is dead code until #2751/#2752 (nothing sets keep_fused_experts, nothing calls split_fused_experts_state_dict), every refusal is a loud NotImplementedError rather than a silent fallback, and the non-keep-fused path through _export_fused_experts_module and _export_bmm_experts is unchanged. The residual risk is the one the description states plainly: the EP object gather had not run over NCCL until the GPU test added today, which the linux CI job will not execute.
Both suggestions harden split_fused_experts_state_dict against inputs the fold itself cannot produce. Neither blocks.
🤖 Generated with Claude Code
…U test's keys The split numbers each module's experts by the first projection, so a down_proj holding more experts lost the extra ones unwritten, and one holding fewer failed deep in the loop. It is refused up front, naming the module. The scale slicer's comment says that only the gate|up cut reads its row count. The GPU test's expert-parallel comparison checks the folded keys too, as the two-rank unit test does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 945271d |
|
Thanks. All three points are addressed in 945271d:
The PR is at 499 source lines. Pre-commit passes against main. |
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 [14/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.
[2/2] of this PR's split (2026-10-28), made at its review's request because it was over the 500-line budget. [1/2], #3099, with the NVFP4 DTensor packing and the held FSDP2 weight-update windows, merged on 2026-10-29. This PR keeps the fused-experts fold.
Base:
main. The diff below is this PR's alone (499 source lines). Review fixes are listed at the end of Testing.Series layout
mainWhat does this PR do?
Type of change: New feature
The fused-experts export splits each 3-D expert weight into full per-expert tensors before it quantizes them. Under FSDP2 that first all-gathers the weight, so every rank ends up holding every expert: a 235B MoE OOMed at about 157 GB per rank, however many nodes it ran on.
keep_fused_expertsset on the export context, the handler folds each expert's quantizer into the fused weight in place:split_fused_experts_state_dictlater cuts the folded weight and its stacked scales into exactly the per-expert keys the single-process export writes. It takes one module's shard at a time, not a consolidated state dict: its input keeps the fused tensors while the per-expert copies are made._fsdp2_weight_window). The window hands FSDP2 the folded weights and reshards on the way out; folded outside it, the next reshard would put FSDP2's own unquantized shards back.The held FSDP2 windows and the NVFP4 packing of DTensor shards are #3099's, [1/2] of this split.
Nothing sets
keep_fused_expertsyet. The writer that consumes the fold arrives in #2751 and #2752.The fold drops the source checkpoint's expert scales (
*_scale_inv,*_activation_scale) as the per-expert export does, and refuses on every rank a layout that leaves a rank no experts, a fused weight sharded on a dim other than the expert axis or over more than one mesh dim, and an unsharded weight whosenum_experts, expert rows and weight quantizers disagree. The shard checks run for every sharded projection, quantized or not. It also refuses native FP8 expert weights in every projection, quantized or not, since it drops each projection'sscale_inv(dequantize them first,unpack_weight()): a native-FP8 projection whose experts a recipe leaves unquantized would otherwise be kept as raw FP8 bytes without it. The split cuts a packed or block scale in proportion to the weight rows it cuts, so INT4 blockwise, W4A8 AWQ and FP8 2-D blockwise gate and up get their own scales, and refuses a block scale a gate|up cut would split (128-row blocks over gpt-oss's 2880-row gate, say) instead of rounding.Usage
N/A: internal to the distributed export (#2752).
Testing
New
tests/unit/torch/export/test_fused_experts_keep_fused.py, on CPU under gloo, one and two ranks:_export_fused_experts, byte for byte, for gated and ungated experts in FP8, FP8 weight-only and NVFP4;num_expertsshort of the weight's experts refused with the module untouched, a native-FP8 projection whose experts are all unquantized included;Shard(0)experts, each with its own quantizer at its global offset (_dtensor_dim0_offset), against the unsharded fold, an uneven 3-over-2 shard included; gathered whole withfull_tensor(), the two ranks' folds split into exactly the per-expert export;(Shard(2),)) or over two mesh dims ((Shard(0), Shard(0))) refused, on one gloo rank;The NVFP4 DTensor packing and held-window tests are [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's (
test_sharded_weight_packing.py);_split_local_fused_module's tests went to [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 with it.Mutation-checked: each test fails under at least one deliberate break of the code, the handlers' two plan reads, the shard-local quantizer pairing, the quantizer-count, native-FP8 (an all-unquantized native-FP8 projection's too),
num_experts, placement and split refusals, the ranks' joint decision on a projection, and the window's unit included.The export comparisons are one table, the one-rank and faked uneven shards fold into the real two-rank job, and the NVFP4 zero-block tests move beside main's scale-clamping tests. Each drop is checked by mutation.
The per-expert comparison covers source scales, and INT4 blockwise, W4A8 AWQ and FP8 2-D blockwise experts at 128-wide dims; one expert over two ranks is refused on both.
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/uniton 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: the fold's gather has run on two gloo ranks only, not yet over NCCL in an expert-parallel export on GPUs. (It has since 2026-10-29, in the GPU test at the end of this list.)
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 fold reads each expert through the module's own accessor, in the weight domain the per-expert export packs; one FSDP2 weight window per unit rather than per linear; a fused-experts module not named
expertsis refused before anything is folded. 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: a unit's FSDP2 window is held only inside
held_fsdp2_weight_windows, which the quantized-module export enters, and each handler packs its module in one call (the BMM experts, which pack two weights, refuse the keep-fused export); the fold puts its stacked scales on the weight's device; the fold hands the expert accessor the module with only the projection swapped, which [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's accessor needs (with [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged in, 26 keep-fused and fused-experts unit tests fail without it and pass with it). 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 3: the fold writes each packed expert into one preallocated stack, so it holds the packed projection once (table below); a zero NVFP4 block scale is replaced with a scalar rather than a tensor of ones; a window on the model itself reuses the module names the export context holds; comments say what orders an expert-parallel rank's quantizers and which experts a rank folds inside an FSDP2 weight window. The byte-for-byte comparisons against
_export_fused_expertspass unchanged; pre-commit, mypy included, and this PR's own unit tests pass at this head. Peak memory the fold adds over its input, one gated module (hidden 1024, intermediate 1024), CPU, torch 2.14: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).
Review suggestions, 2026-10-27 evening: the keep-fused fold refuses a projection whose weight quantizers use a GGML format, before anything is folded.
Review suggestion, 2026-10-27 night:
_local_expert_weight's docstring says which reads of the projection the local rows stand in for. The keep-fused tests' ungated case runs transformers' realNemotronHExperts, so it passes whichever wrapper quantizes Nemotron-H's experts ([OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 gives them one of their own).Review fixes, 2026-10-27, last round: the BMM experts (
GptOssExperts,Llama4TextExperts) refuse the keep-fused export: a GPT-OSS checkpoint without aquantization_config, loaded with FSDP2 and exported without gathering, would have had its fusedgate_up_projsplit on the wrong axis. The new test fails without the change. Docstring and wording fixes beside it.Review suggestions, 2026-10-27, last round: comments name why the GGML refusal agrees across ranks and where the replicated
input_scaleis synchronized; two unreachableNoneguards are removed.Review fix, 2026-10-27, last round: the
gloo_world_of_onefixture destroys its process group in afinally, also when its mesh fails to build.Review suggestion, 2026-10-28:
test_meta_dcp_load.py's module-localgloo_world_of_oneis deleted in favor of the shared conftest fixture, the same deletion [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 makes.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 main1ca8bd1fc4d(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 the fused-experts and FSDP2 helper tests pass (266), and with [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 merged in, which also changes_QuantFusedExperts, the keep-fused, fused-experts and all-to-all tests pass (244). pre-commit passes against main.Split on 2026-10-28, late night, at the review's request (meenchen's review bot): the NVFP4 DTensor packing and the held FSDP2 weight-update windows moved to [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099 ([1/2]) with their tests, and this PR merges [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099 into its history, so its own commits stay as they were.
_split_local_fused_module, whose only caller is [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's writer, moves to [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 with its two tests (it lands there at [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's next rebase; until then [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's base still carries it). The module's local imports moved to the top, theDTensor = Nonefallback going with them (2504934ab24).keep_fused_expertskeeps its name: it switches the no-gather export, as its comment says, and the FSDP2 windows are part of that path. Rebased onto mainbdc45368d10after [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 merged. The fold's tests, the export tests andtest_fused_experts.pypass (730; 7test_export_diffusers.pycases fail the same way without this PR in that environment), the fold with [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 merged in passes its tests (238), and pre-commit passes against main.2026-10-29: merges [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's first review round (
ebfb6a7b050): a module joins its unit's held window once, however many weights it packs; comments cut to their constraint; module-scope imports. No change of its own.tests/unit/torch/export,test_nvfp4_tensor.py,test_meta_dcp_load.pyandquantization/test_utils.pypass (620; the 7test_export_diffusers.pycases fail as before in that environment), and pre-commit passes against main.2026-10-29, later:
_export_bmm_experts's docstring no longer says a rank packs its experts without a gather. Every path that reaches it now packs whole weights, and the keep-fused export refuses these experts (d34142120fd, at [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's review). Docstring only; pre-commit passes against main.Review (meenchen's review bot), 2026-10-29: the two-rank fold test is back (
3a86980a421), without_split_local_fused_module, which is [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's now. Each rank's fold is checked against the unsharded one at its offset, and the ranks' folds, gathered whole withfull_tensor(), against the per-expert export (gloo cannot gather FP8, so one-byte values move as bytes). Each of four deliberate breaks fails it: quantizers paired by position for every layout, the ranks' quantizer-state gather skipped, the rewrap inferring the global shape, and the no-expert refusal dropped. The native-FP8 test importsPretrainedConfigat module scope;FP8Expertsstays behindimportorskip, since transformers 5.5 imports triton in its module. The keep-fused tests pass (48) on transformers 5.18 and 5.5.4, and pre-commit passes against main. The description no longer claims [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's windows and NVFP4 packing.2026-10-29, later: merges [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's second review round (
54f8632c78b): a unit that owns a weight directly keeps its packed weight in the window, and the packing has a two-rank test. No change of its own (a322d417b4b). The keep-fused, sharded-packing, NVFP4 and param-index tests pass (79), and pre-commit passes against main.2026-10-29, later: merges [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's GPU test of the sharded packing (
51f187fb029). No change of its own (0759477e213); pre-commit passes against main.2026-10-29, evening (meenchen's review bot, optional item):
_export_fused_experts_module's docstring drops a paragraph from before the keep-fused fold, which said FSDP2 ranks pack their own experts kept fused (26e9863b87d). Docstring only; pre-commit passes against main.2026-10-29, evening: merges [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's third review round (
6033bb8e2a8): the INT4, NF4 and MXFP4 packers refuse a DTensor shard, NVFP4 refuses one with a static quantizer's scale, and a module the weight window cannot name keeps off the root's key (fa69b309c15). Its only change of its own iskeep_fused_experts' new docstring, which here also says the flag keeps fused experts fused and sharded. Pre-commit passes against main.2026-10-29, evening: a GPU test of the fold,
tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py(ec2741cdef3), for meenchen's review bot's question about NCCL. On every visible GPU, and skipped below two:cpu:gloo,cuda:ncclgroup the export runs on, and over NCCL in a group of NCCL alone.GPU run, 2026-10-29, at
ec2741cdef3:test_fused_experts_keep_fused_gpu.py::test_ranks_fold_their_own_experts_over_nccl, both cases (objects-over-gloo,objects-over-nccl). GB200, torch 2.11.0a0 (NGC 26.02), CUDA 13.1, NCCL 2.29.2, transformers 5.9.0. A mutant whose ranks each decide alone, ignoring the gathered quantizer states, fails both cases: it still refuses, but names the rank's local experts ("quantizers [1] of 2") instead of the global ones ("[3] of 4" on two GPUs, "[7] of 8" on four).ec2741cdef3)2026-10-29, evening: merges [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's fourth review round (
46742ffd24e), comments only (6f1d1da4b2b). Pre-commit passes against main.2026-10-29, night: [1/2] ([OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099) merged into main as
6c475d716cf, and this PR now merges main (830d7044438).registry.pyandunified_export_hf.pyconflicted where this PR extends [1/2]'s text, and take this PR's side. The PR's diff against main is byte for byte its diff against [1/2]'s last head. Pre-commit passes against main.Review (Fridah-nv), 2026-10-29, night:
split_fused_experts_state_dict's docstring says how it numbers the experts it is given, from 0 by their row in plain tensors, which a caller splitting a rank's local shard shifts by the shard's offset (fc6d99829d8, docstring only). The suggestion to handget_expert_weighta tensor instead of the fold's proxy is answered in its thread. Swapping the fold's accessor reads forlocal_w[idx].to(dtype)fails 12 of the byte-for-byte tests. Pre-commit passes against main; 497 source lines.Review (Claude's review, approved), 2026-10-29, night (
945271deb04):test_split_refuses_projections_holding_different_expert_countsfails without the check.Before your PR is "Ready for review"
CONTRIBUTING.md: N/Afc6d99829d8; its two suggestions are in945271deb04.Additional Information
Part of expert-parallel PTQ for Hugging Face MoE models.
The Codex re-review item "Validate quantization consistency across EP ranks before folding" is handled in this PR, by the fold's gather, as well as by #2753's model-wide pre-check.
Summary by CodeRabbit
(Edited after the split: the NVFP4 fixes this summary listed are #3099's.)