Repository navigation
[OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder - #3067
Open
shengliangxu wants to merge 14 commits into
Open
[OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder#3067shengliangxu wants to merge 14 commits into
shengliangxu wants to merge 14 commits into
Conversation
shengliangxu
requested review from
chadvoegele
and removed request for
a team
October 27, 2026 20:37
Contributor
|
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.
 |
| Layer / File(s) | Summary |
|---|---|
Fused-expert wrapper dispatch modelopt/torch/quantization/plugins/huggingface.py, modelopt/torch/export/moe_utils.py, tests/unit/torch/quantization/plugins/test_fused_experts.py |
The generic wrapper no longer detects non-gated experts by matching an up_proj shape. Module detection now recognizes registered wrappers that subclass _QuantFusedExperts. Documentation describes non-gated subclasses. Tests remove the synthetic non-gated wrapper suite and fixtures, and run the eager native FP8 test for gated and non-gated configurations. |
Nemotron-H wrapper and validation modelopt/torch/models/nemotron_h/modeling_ptq.py, tests/unit/torch/models/test_nemotron_h_modeling_ptq.py, tests/gpu/torch/models/test_modeling_ptq_registration.py |
NemotronHExperts is registered with a wrapper that uses up_proj as its first projection and marks the layout non-gated. Tests check conversion, projection quantizers, eager and grouped_mm output equivalence, calibration statistics, GPTQ Hessians, export, implementation normalization, and conditional registry availability. |
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant NemotronHExperts
participant QuantModuleRegistry
participant QuantNemotronHExperts
participant Projections
NemotronHExperts->>QuantModuleRegistry: Resolve registered wrapper
QuantModuleRegistry->>QuantNemotronHExperts: Select Nemotron-H wrapper
QuantNemotronHExperts->>Projections: Configure up_proj and down_proj quantizers
Merge Risk: 
70994
-
No blocking issue was established for Nemotron-H expert quantization; the PR is mergeable after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: adding Nemotron-H expert quantization for grouped_mm with a model-specific wrapper. 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 PR adds no torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/…Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The PR adds no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# nosecline. The changed-file inventory includes nopyproject.tomlor requirements file, so no dependency was added. The existingtrust_remote_code=Truereference inhuggingface.pyis unchanged by this PR.
-
Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
shengliangxu
requested review from
kevalmorabia97 and
meenchen
and removed request for
a team
October 27, 2026 20:37
This was referenced Oct 27, 2026
Merged
Open
[OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank
#2751
Merged
Collaborator
Author
|
/ok to test 23ab7f4 |
shengliangxu
force-pushed
the
shengliangx/hfep-8a-grouped-mm-calibration
branch
from
October 30, 2026 17:06
27c5f77 to
39697ad
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-8b-nemotron-h-experts
branch
from
October 30, 2026 17:06
94a0e71 to
bda6d27
Compare
Collaborator
Author
|
/ok to test bda6d27 |
shengliangxu
added a commit
that referenced
this pull request
Oct 30, 2026
… its LM head (#3131) > [!NOTE] > **[1/3] of #2741, split on review:** #2741 grew past the 500-line budget, so it is now three PRs, merged in order. This one sits on `main` and merges on its own. #3136 [2/3] (the expert-parallel loader) builds on it, and #2741 itself becomes [3/3] (the loader's post-load refusals). <details><summary>Series layout</summary> - `main` - **#2741 [1/3] the loader's helpers (this PR)** - #2746 [1/3] model-spec EP sections - `shengliangx/hfep-base-6a-ep-loader`, a branch that merges #2779, #3067, #2740 and this PR - #3136 [2/3] expert-parallel loader - #2741 [3/3] the loader's post-load refusals - #2742 dense sharding → #2743 block-FP8 dequantization - #2780 expert-parallel load routing (merges after #2741 [3/3]) - #2747 on-the-fly decompression </details> ### What does this PR do? Type of change: New feature It adds two helpers the expert-parallel loader (#2741 [2/3]) and the model specs (#2746) use. Neither changes how anything loads today. - **`moe.require_transformers(min_version, feature)`** raises if the installed transformers is older than the release a feature needs. The error names the feature, the installed release and the one to upgrade to. Expert parallelism relies on APIs only newer releases have, and an older release often runs without an error while doing the wrong thing: transformers 5.5 shards no expert. `transformers` is imported where it is checked, so `moe` still imports without it. - **`_auto_class(hf_config, require_lm_head=True, trust_remote_code=False)`** can build a model without its LM head. - With `require_lm_head=False`, a model whose architecture reads as a VLM's, or an encoder-decoder model, builds through `AutoModel`, which leaves the head out. - A causal LM keeps its head either way. - `require_lm_head` is the new second parameter. `_decide_distribution`, the one caller, now passes `trust_remote_code` by keyword. ### Usage ```python from modelopt.torch.utils.moe import require_transformers require_transformers("5.7", "expert-parallel PTQ") # raises on transformers < 5.7 ``` ### Testing - `test_a_transformers_below_a_features_floor_is_refused` (`test_moe_a2a.py`): 5.6.2 is refused against a 5.7 floor, 5.7.0 is not. - `test_the_class_that_builds_the_whole_model` (`test_model_load_utils.py`): gains the `require_lm_head` column and two headless rows, `vlm-backbone` and `encoder-decoder-backbone`. - 2026-10-30: - Pre-commit, mypy included, passes against main. - On CPU (torch 2.14), `test_moe_a2a.py` and `test_model_load_utils.py` pass on transformers 5.18, and on 5.5.4, where the version-gated cases skip. - Review fixes, 2026-10-30 (Claude, CodeRabbit): with `require_lm_head=False`, `_auto_class` returns `AutoModel` only when it can build the config (a native config, or a trusted remote one registering `AutoModel`), else the class that carries the head; two new table rows pin both sides, each failing without its half of the fix. The docstrings say the head is not required rather than absent, and that `require_transformers` refuses a pre-release of its floor on purpose. Pre-commit, mypy included, passes against main; `test_moe_a2a.py` and `test_model_load_utils.py` pass on CPU on transformers 5.18 and 5.5.4. - Rebased 2026-10-30 onto main (99c3834) with no conflict; its own diff is unchanged. Pre-commit, mypy included, passes against main. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ `_auto_class` is private; its one caller passes `trust_remote_code` by keyword. - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A. The feature's entry lands with #2754. - Did you get Claude approval on this PR?: ❌ Requested. ### Additional Information Part of expert-parallel PTQ for Hugging Face MoE models. The code comes from #2741 unchanged, except that `packaging` moved to the module top. It was reviewed there; Claude approved that PR on 2026-10-28. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Model loading can select a base model without a task-specific output head for supported configurations, including when trusted remote model code registers a base-model class. * If a feature requires a newer Transformers version, the app reports the minimum required version and provides an upgrade command. * **Bug Fixes** * Improved model-class selection across supported vision-language, encoder-decoder, and remote-code configurations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
added a commit
that referenced
this pull request
Oct 30, 2026
… each tensor's owner rank (#2751) > [!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 [15/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night. > > **[3/3] of this PR's split** (2026-10-29), at its review's request: [1/3] #3100 (the reverse conversion's `unmatched_merges` hook), [2/3] #3101 (`_materialize_owned`, ownership and transport), and this PR (reverse conversion and shard writing). Merge in that order. [1/3] #3100 merged on 2026-10-29 (`8ab63386152`) and [2/3] #3101 the same day (`4bef85279e4`); this PR now targets `main`. > > **Base:** `main`, on #3101's squash `4bef85279e4`. The diff below is this PR's alone. > > **Re-cut:** the `gloo_world_of_one` fixture its tests use is in #3101 now, in the same text #3099 adds, so either order merges cleanly. Review fixes since are listed at the end of Testing. <details><summary>Series layout</summary> - `main` - #2745 fused experts under grouped_mm (merged) - #3074 calibration on grouped_mm - #3067 Nemotron-H's experts under grouped_mm - #2739 meta-DCP checkpoint load (merged) - #2736 distribution plan, run group, collectives (merged) - #2740 calibration on per-rank data - #2750 keep-fused fold - #3100 conversion hook, [1/3] (merged) - #3101 ownership and transport, [2/3] (merged) - **#2751 shard writing, [3/3] (this PR)** - #2752 no-gather export (on #2751, #2750, #2779 and #2740) - #2753 expert-parallel export - #2737 plan decision - #2779 load routing - #2741 expert-parallel loader (on #2779, #3067 and #2740) - #2742 dense sharding - #2743 block-FP8 dequantization - #2780 expert-parallel load routing - #2746 expert parallelism in the model specs - #2747 on-the-fly decompression - #2748 Kimi-K2.5 and Kimi-VL adapters - #2749 Kimi-K3 adapter - #2754 expert-parallel PTQ from hf_ptq.py (on #2746 and #2753) </details> ### What does this PR do? Type of change: New feature Main's FSDP2 export already writes each rank's share, but it first unshards each FSDP2 unit in turn, an all-gather of that unit to every rank. It also refuses conversion mappings that merge or split tensors, since one unit at a time cannot regroup them, and it does not handle expert-parallel shards. In this writer every tensor moves once, whole, to one owner rank and is written from there (`modelopt/torch/export/distributed_writer.py`). After the split this PR holds the last of its pieces: `_materialize_owned` is in #3101, and the conversion's `unmatched_merges` hook is in #3100. - **`_materialize_owned`** (now in #3101) gives each tensor exactly one owner rank and makes it whole there: - replicated and plain values need no transport; - dim-0 shards travel point-to-point, batched and bucketed per peer (one bucket per key deadlocked at 16 ranks); - each group of keys the reverse conversion merges into one tensor (`_merge_groups`) goes to one owner, with its weights' quantization state, so the merge, or its refusal of a quantized group, happens there. Placements it cannot assemble, and shards that do not tile the global shape, raise rather than risk a silently mis-assembled checkpoint; a zero-element tensor in such a layout is refused like a non-empty one. Experts arrive whole and plain, one disjoint set per rank: an expert that is still a DTensor is refused on every rank before anything moves. - **`_revert_whole_sd`** runs the reverse weight conversion on those whole tensors, so they go out under the checkpoint's names. transformers' conversion ops cannot express a placement. Ownership keeps each merge group whole on one rank, so no rank sees part of one. The reverse runs inside the failure-carrying `all_reduce`, so a refusal on any rank, such as a merge's refusal of a quantized or incomplete group, makes every rank keep its in-memory names together. - **`_write_even_shards`**: - bin-packs each rank's set into evenly sized files numbered into one global range; - refuses a name two ranks both hold, on every rank, in its first collective, before any file is written; - writes the files with safetensors, both names of a tensor two keys share included (one is copied, as `ShardedSafetensorsWriter` does), which `save_file` would otherwise refuse; - unions the index; - never overwrites a checkpoint: a directory already holding `model.safetensors`, its index, or `model-NNNNN-of-NNNNN.safetensors` files is refused on every rank in the first collective, before anything is written. Other files, auxiliary weights such as `mtp.safetensors` included, are left alone; - writes every file under a temporary name and renames them into place only once every rank has written all of its own, the index last, so no index names a shard that is not in place. - **Failing together.** The export directory's creation with the file counts and tensor names, the shard writes and rank 0's index write each run inside the failure-carrying `all_reduce` (#2736), so one rank's failure raises on every rank with that rank's error rather than leaving the others in an unmatched collective. A tensor over `max_shard_size` gets a file of its own, as in HF, and the rest still pack. The entry point that drives them arrives in #2752. **Why not a DCP writer, or main's streaming writer (`_StreamingShardWriter`, `name_shards_and_write_index`)?** - DCP writes its own format, not safetensors shards under the hub's names with an index. vLLM and TRT-LLM would need a consolidation pass over it, and that pass is the gather this writer avoids. - The streaming writer serves main's FSDP2 export, which unshards each unit, an all-gather of it to every rank, and refuses conversion mappings that merge or split tensors. This writer moves each tensor once, whole, to one owner, so a merge can run there; that part has nothing to extend. - The file writing itself could share code with `_StreamingShardWriter`, which already stages to temporary files, given three things this writer needs: - every step fails on every rank together, by running in the failure-carrying `all_reduce`, where the streaming path names its shards on rank 0 after a barrier and a `gather_object`; - a name two ranks both hold is refused before anything is written; - each rank's owned set is bin-packed into even files. That sharing is deferred to a follow-up. It would change main's streaming export, a path this series does not otherwise touch, and no checkpoint format, so it can land on its own once this writer is in, without growing a PR already split for size. The metadata collective carries each sharded tensor's dtype, every dense tensor's size and the merge units, which owners are assigned from; ranks that disagree on any of them are refused on every rank. The send buffers are freed once the transfers finish, before a rank assembles what it owns (a `P2POp` and its request each keep their tensor, so the transfer runs in a helper that drops both); on top of its own shards a rank holds at most about twice its owned set. On two gloo ranks with 128 MiB of shards each: | Send buffers | Peak RSS above the shards | |---|---| | kept while a rank assembles (previous revision) | 257 MiB | | freed before it assembles (this revision) | 193 MiB | ### Usage N/A: internal to the distributed export (#2752). ### Testing - New `tests/unit/torch/export/test_distributed_writer.py`, on CPU under gloo, drives each piece directly: - HF shard-size parsing, a unit in any case read as bytes, as the streaming export reads it (transformers reads a lowercase `b` as bits); - binning: balanced, bins added until each fits, an oversized tensor isolated, no empty bins; - DTensor sizes and the shard-0 placement check; - ownership making every tensor whole and plain, and rejecting placements it cannot assemble; - the reverse conversion renaming in place and keeping in-memory names when it fails; - shard files and an index covering every tensor once; - two ranks where one rank's directory or write fails, or two ranks hold one name: both raise, and the name collision leaves no shard file. Run on a 10 s gloo group, these fail in seconds on the code before the fix rather than waiting out the process group's timeout. - Mutation-checked: each test fails under at least one deliberate break of the code. - Each helper's tests are one table, and checks a remaining test covers are dropped; the two-rank jobs use the shared `bounded_collectives`. Each drop is checked by mutation. - On two ranks, `_materialize_owned` moves an uneven shard, an empty one and two dtypes to their owners, and a dtype disagreement, or an expert that only rank 1 still holds as a DTensor, is refused on every rank. An empty `Shard(1)` tensor is refused, and two names for one CPU tensor are both written. The expert, collision, empty-tensor and shared-tensor cases each fail without their fix (mutation-checked). - On two ranks, what a rank sent is freed before it stores an assembled tensor, a rank that received nothing included; a dense size that differs across ranks is refused on every rank; and `_write_even_shards` with uneven file counts numbers the files across ranks and indexes both ranks' tensors. Each fails without its 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. - 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: each merge group goes to one owner rank, so the reverse conversion merges it; the shard size is read with the streaming export's parser; the mesh memo's comment gives its real reason. 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 merge groups are sorted before the units are built, so every rank derives the same owners whatever order its keys come in; a new case shows a quantized merge group, refused by the rank holding its scales, falling back on every rank. 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: a merged weight's quantization state stays in its merge group, so the rank that merges a quantized group refuses it even when the merge pattern does not match the scales; groups that share a key form one unit whatever their order; the ranks check in the metadata collective that they formed the same units. 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 #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 refusal of ranks that group the keys differently lists each rank's first differing merge units. - Review fix, 2026-10-27 night: when the shard or index step fails, every rank removes the files it wrote, a partial one included, so a retry starts from the directory as it was. The one-rank-failure test checks the directory is left empty; its shard and index cases fail without the change. - Review suggestions, 2026-10-27, last round: a failed write's cleanup ignores an `unlink` error, so the error naming the failing rank is the one raised, and a comment says a retry rebuilds the state dict. The new test fails without the change. - Review suggestion, 2026-10-27, last round: files a failed write's cleanup cannot remove are named in one warning, and the write's own error is still the one raised. The extended test fails without the change. - The same `gloo_world_of_one` fix as #2750's, in this PR's copy of the fixture, so the two stay identical. - The same `test_meta_dcp_load.py` fixture deletion as #2750's, so the two stay identical. - Review suggestion, 2026-10-28, second round: each tensor is dropped right after its host copy, so a file's tensors are no longer held on the device while it is written; the new test fails without the change. - Review fix, 2026-10-28, seventh round: the comment by the drop now says it frees each tensor once its host copy exists. It does not lower the device peak, which is the allocation at entry with or without it, so a peak-memory measurement reads the same either way by construction and is no evidence against it; the memory still held when each file is written is what moves. Measured on 4 GPUs, in a distributed HF export of tiny FP8 and NVFP4 models that runs this writer on CUDA tensors (torch 2.15 pre-release, transformers 5.18; this PR's own writer test runs on CPU, where device memory does not move). The table gives the device bytes allocated (`torch.cuda.memory_allocated`) on each rank just before `save_file`. Each rank wrote one file, so a file was its whole owned set; a rank writing several files was not measured. | Model | Tensors per rank | Held before this change | Held with it | Released before the write | |---|---|---|---|---| | FP8, dense, tied embeddings | 5 to 20 | 167,424 to 167,936 | 109,056 | 58,368 to 58,880 (35%) | | FP8, MoE | 39 to 53 | 526,848 to 528,384 | 275,456 | 251,392 to 252,928 (48%) | | NVFP4, dense, tied embeddings | 6 to 25 | 131,072 to 131,584 | 97,280 | 33,792 to 34,304 (26%) | | NVFP4, MoE | 49 to 71 | 371,712 to 376,320 | 220,672 | 151,040 to 155,648 (41%) | - Rebased on 2026-10-28, late night, onto #2736 at `cf4597dcbeb`, on main `1ca8bd1fc4d` (after #3088 and #2745 merged), without conflicts; this PR's own diff is unchanged. Its tests pass (143), and pre-commit passes against main. - Rebased on 2026-10-28, late night, onto main `bdc45368d10` after #2736 merged; #2736's squash has the tree of its last head, so this PR's tree is unchanged. - Review (meenchen's review bot), 2026-10-29, all items: - **Split in three:** #3100, #3101 and this PR. This branch merges #3101, so its own commits stay; source lines are about 40, 347 and 254. - **Rollback:** a failed export no longer deletes a checkpoint already in the directory. Every file is written under a temporary name, and the ranks rename theirs into place only once every rank has written all of its own, the index last. A failure before then removes only the temporary files. The one-rank-failure test runs over an existing checkpoint too and checks it byte for byte; its shard and index cases fail without the change. A new test re-exports over a checkpoint. - **Imports:** the writer's, the fixture's and a test's imports moved to module scope. - **Docstrings:** the long ones keep their constraints; the narrative moved here. - **GPU tests:** `tests/gpu/torch/export/test_distributed_writer_gpu.py` runs from two GPUs. It moves BF16, FP8 and uint8 shards of uneven and empty extent over NCCL and writes several files per rank. The GB200 runs and measurements follow below. - **Results:** 148 tests pass, and pre-commit passes against main. - **Narrative moved from the docstrings:** bucketing the transfers by peer replaced one operation per (key, peer), which at 235B was about 12.7k operations and 660 s for a few GB of payload. - GPU runs on GB200 of `tests/gpu/torch/export/test_distributed_writer_gpu.py` (its two tests), ran at `f577fc05af0`. The code is unchanged at this head; only [2/3]'s docstrings and tests changed since. | GPUs | Result | Time | |---|---|---| | 1 | 2 skipped (need two GPUs) | 0.5 s | | 2 | 2 passed | 17.6 s | | 4 | 2 passed | 23.5 s | - Measured on GB200 at `f577fc05af0` with a synthetic sharded state dict: 600 keys of shape (8741, 4096), 360 BF16, 150 FP8 e4m3 and 90 uint8, 32.0 GiB in total. 8741 rows divide by neither 4 nor 8, so the dim-0 shards are uneven. Timed `_materialize_owned(sd, set(), rank, world)`, then `_write_even_shards(sd, out_dir, "5GB", rank, world)`. Rank 0 is shown; all ranks were near-identical. | GPUs | Input per rank | Owned per rank | Transport | Peak device memory in the transport, above its input | Write | Files per rank | Bytes written per rank | Peak host RSS (process lifetime) | |---|---|---|---|---|---|---|---|---| | 4 (one node) | 8.01 GiB | 8.00 GiB | 6.22 s | 12.01 GiB | 4.36 s | 2 | 8.00 GiB | 8.68 GiB | | 8 (two nodes × 4) | 4.00 GiB | 4.00 GiB | 14.15 s | 7.01 GiB | 2.57 s | 1 | 4.00 GiB | 8.49 GiB | - The device peak is `torch.cuda.max_memory_allocated()` after `reset_peak_memory_stats`, minus the allocation before the transport. During the transfer a rank holds one send buffer and one receive buffer per (peer, dtype). The send buffers copy its shards of the tensors other ranks own, (w-1)/w of its input; the receive buffers hold the other ranks' shards of the tensors it owns, (w-1)/w of its owned set. Their sum is the measured peak at both sizes. In this run the state dict held the only reference to each input shard, so a rank freed its inputs as it assembled. - Host RSS is `ru_maxrss`, the maximum over the process's life. It was 2.1 GiB before the write, so the write's own peak is at most the number shown. - The 8 GPUs span two nodes, so that transport time includes cross-node NCCL traffic and does not scale like-for-like against the one-node run. - Re-measured on GB200 at `239260c7d8d` with the input DTensors kept alive through the write, as the model keeps its parameters in a real export. The writer's code is the same as at `f577fc05af0` (docstrings differ) and is unchanged at `2470bc4a4f7`. Same 32.0 GiB state dict; rank 0 is shown, and all ranks were near-identical. | GPUs | Owned per rank | Transport | Peak device memory above the allocation before the transport (transport and write) | Write | Files per rank | Peak host RSS during the write | |---|---|---|---|---|---|---| | 4 (one node) | 8.00 GiB | 6.83 s | 14.01 GiB | 4.12 s | 2 | 8.65 GiB | | 8 (two nodes × 4) | 4.00 GiB | 14.76 s | 7.51 GiB | 2.52 s | 1 | 8.21 GiB | - With the inputs alive, the peak comes from the assembly: the receive buffers, (w-1)/w of the owned set, plus the assembled owned tensors. Their sum is the measured peak at both sizes, within `_materialize_owned`'s bound of about twice the owned set. - The write never raised the peak above the transport's: it copies each tensor to the host and drops it, so it adds time, not device memory. - Host RSS here is windowed to the write (the high-water mark is reset just before it); the first table's is the process-lifetime `ru_maxrss`. - 2026-10-29: merges [2/3]'s review round (`46e42e10a74`, rebased onto #3100's `b9fbb242300`). That brings [1/3]'s docstring and test, plus [2/3]'s merge-group and strided-shard tests and trimmed docstrings. No change of its own. Pre-commit passes against main; `test_distributed_writer.py`, `test_quant_aware_conversion.py` and `test_meta_dcp_load.py` pass (149). - 2026-10-29, later (`2470bc4a4f7`, test only): `test_the_index_goes_in_place_only_after_every_shard` pins the commit order. Rank 1 fails to rename its shard into place after every file is written, and the index must stay as it was: absent in an empty directory, byte for byte over an existing checkpoint. So an index in the directory never names missing shards. With the shard and index commits swapped, both cases fail, and every other writer test stays green. The writer's code is unchanged since `239260c7d8d`. `test_distributed_writer.py` passes (44), and pre-commit passes against main. - 2026-10-29 (meenchen's review bot): why the writer is neither a DCP writer nor main's streaming writer, and why sharing its file writing with the latter is deferred, is now in the description above, moved from the review reply. Description only. - Review (meenchen's review bot), 2026-10-29, second round (`285c30bf08e`): once every rank has committed, rank 0 removes an earlier checkpoint's `model-NNNNN-of-NNNNN.safetensors` files the new index does not name, and a `model.safetensors`, which transformers loads ahead of the index. The replacement test runs over an earlier one-shard, three-shard and single-file checkpoint, with a `config.json` and an off-index weight file that must stay; the three-shard and single-file cases fail without the removal. The failure tests' existing checkpoint gained a `model.safetensors`, so a removal on a failure path fails them (checked by running it in the failure handler). The docstring states the commit's limit: a rename failing partway through the final commit can leave a checkpoint already there mixed; it is documented rather than fixed, for the reasons in the review thread. `test_distributed_writer.py` passes (46), and pre-commit passes against main. Source lines: 272. - 2026-10-29, later: merges [2/3]'s second review round (`055b57bd7cf`): mixed-device shards and whole-tensor disagreements refused on every rank, each transfer round's buffers freed when it ends. No change of its own (`1feacd040c9`). The writer, conversion and meta-DCP tests pass (158), and pre-commit passes against main. - Review fix, 2026-10-29 (`03ae6f5d5ff`, meenchen's review bot): a reverse conversion that fails with an exception carrying no message now counts as a failure. Before, `RuntimeError()` read as success, so rank 0 committed hub names while rank 1 kept in-memory ones (reproduced on two gloo ranks with tiny Mixtral). The one-rank-failure test gains that case, which fails without the fix. `test_distributed_writer.py` passes (52), and pre-commit passes against main. - GPU run, 2026-10-29, ran at `1feacd040c9` (the same code as the current head). GB200, torch 2.11.0a0 (NGC 26.02), transformers 5.9.0: | `tests/gpu/torch/export/test_distributed_writer_gpu.py` (2 tests) | 1 GPU | 2 GPUs | 4 GPUs | |---|---|---|---| | Result | 2 skipped | 2 passed, 20.3 s | 2 passed, 22.5 s | - Review (meenchen's review bot), 2026-10-29, third round (`d9bf5aafbbb`, the user's call): the writer refuses, on every rank and before writing anything, a directory that already holds a checkpoint (`model.safetensors`, its index, or `model-NNNNN-of-NNNNN.safetensors`), so a commit failure can no longer mix one. The refusal is decided in the first collective, so one rank's view of the directory stops them all. This replaces the documented commit limit and the stale-file cleanup of `285c30bf08e`, which is removed with its tests; the staged `.tmp` writes stay. A two-rank test checks the refusal for a one-shard, three-shard and single-file earlier checkpoint, each left byte for byte, and fails without the refusal. Another exports beside a `config.json` and an auxiliary weights file, and fails if the pattern also matches the auxiliary file. `test_distributed_writer.py` passes (49), and pre-commit passes against main. Source lines: 270. - 2026-10-29, later: merges [2/3] rebased onto main `2b900c8fa41` with its third review round (`b73c052e45f`): every rank initializes NCCL before a transfer only some ranks take part in, with a fresh-group GPU test, and the transport's memory bound names what a rank sends. No change of its own (`8b6f3a275f7`). The writer and conversion tests pass (128), and pre-commit passes against main. - GPU run, 2026-10-29, at `8b6f3a275f7`. GB200, torch 2.11.0a0 (NGC 26.02), CUDA 13.1, NCCL 2.29.2, transformers 5.9.0: | `tests/gpu/torch/export/test_distributed_writer_gpu.py` | 2 GPUs | 4 GPUs | |---|---|---| | Result | 2 passed, 1 skipped (the fresh-group test needs 3 or more), 20.4 s | 3 passed, 46.1 s | - 2026-10-29, later: merges [2/3]'s docstring note that owners are weighed by bytes (`764036cd64e`). No change of its own (`2a65d9fa364`); pre-commit passes against main. - Rebased on 2026-10-29 onto #3101's squash `4bef85279e4` after it merged, with this PR's own diff unchanged. The history is now linear: one commit for the pieces as they stood when the split was made, then this PR's four commits since. #3100's and #3101's pre-squash commits are dropped. Pre-commit passes against main; the writer and conversion tests pass (128). - Claude review, 2026-10-29 (approved, two suggestions; `55b62f6fdf4`, docstring and message only): `_write_even_shards`' docstring says that a failure during the final renames leaves the shards already renamed, which a retry is refused over until they are removed, and the refusal's message names that cause. It also says the caller drops tied weights by name. Pre-commit passes against main; `test_distributed_writer.py` passes (49). - Rebased on 2026-10-30 onto main `99c3834f62c` (after #3099 and #2750 merged, and #2737), own diff unchanged and no conflicts: main's new commits touch none of this PR's files. Pre-commit passes against main. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Additive; nothing calls these yet. - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A. The no-gather export's entry lands with #2752. - Did you get Claude approval on this PR?: ❌ Not yet. ### Additional Information Part 15 of 21 of expert-parallel PTQ for Hugging Face MoE models; its [1/3] #3100 and [2/3] #3101 have merged. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added distributed checkpoint export in Hugging Face safetensors format, writing tensors across ranks without gathering the full model on every rank. * Exported files are balanced by size, with oversized tensors placed in their own files and an index generated for the shards. * Export supports eligible sharded tensors and reverses supported quantization-related weight conversions. * **Bug Fixes** * Exports detect incompatible shards, tensor-name collisions, and write failures across ranks to help prevent incomplete or inconsistent checkpoints. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/hfep-8a-grouped-mm-calibration
branch
from
October 30, 2026 22:13
a2dde55 to
a6758d7
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-8b-nemotron-h-experts
branch
from
October 30, 2026 22:14
bda6d27 to
fbff3f3
Compare
Collaborator
Author
|
/ok to test fbff3f3 |
shengliangxu
added a commit
that referenced
this pull request
Oct 30, 2026
…2a instead of forcing eager (#3074) > [!NOTE] > **[2/2]**: #3122 ([1/2]) is this PR's expert-parallel correctness fix, split out on 2026-10-29 at a reviewer's request so that it lands on its own, with a Bug Fixes changelog entry. This PR keeps the two backward-breaking changes, calibration on `grouped_mm` by default and the new `calibrate()` error, and sat on #3122 until #3122 merged on 2026-10-30. Before that, it was [3/3] of the OMNIML-6073 split, after #3088 (FSDP2's weight-access window) and #2745 (`grouped_mm` quantization), both merged. #3067 (Nemotron-H) sits on this PR. > > **Base:** `main`. #3122 merged on 2026-10-30, so this PR sits on main, its own diff unchanged. > > **In the tree:** #2741's helper base merges #3067, which carries this PR and #2745. <details><summary>Series layout</summary> - `main` - #3088 FSDP2 weight-access window (merged) - #2745 fused experts under grouped_mm (merged) - #3122 expert-parallel experts calibrated and folded on each rank's own experts (merged) - **#3074 calibration on grouped_mm (this PR)** - #3067 Nemotron-H's experts under grouped_mm - #2739 meta-DCP checkpoint load (merged) - #2736 distribution plan, run group, collectives (merged) - #2740 calibration on per-rank data - #2750 keep-fused fold - #2751 distributed-writer pieces - #2752 no-gather export (on #2751, #2750, #2779 and #2740) - #2753 expert-parallel export - #2737 plan decision - #2779 load routing - #2741 expert-parallel loader (on #2779, #3067 and #2740) - #2742 dense sharding - #2743 block-FP8 dequantization - #2780 expert-parallel load routing - #2746 expert parallelism in the model specs - #2747 on-the-fly decompression - #2748 Kimi-K2.5 and Kimi-VL adapters - #2749 Kimi-K3 adapter - #2754 expert-parallel PTQ from hf_ptq.py (on #2746 and #2753) </details> ### What does this PR do? Type of change: New feature, with a backward-incompatible change to MoE calibration (see the changelog entry) #2745 quantizes the `grouped_mm` experts forward, transformers' default, while calibration still runs every model on eager. This PR moves calibration onto `grouped_mm` and the `torch_a2a` expert-parallel forward, and makes everything that reads the experts' inputs or weights during calibration work there. Each gap below otherwise failed silently, by collecting no amax or the wrong one rather than by raising. - **Backend normalization.** `normalize_experts_impl_on_the_fly` replaces `force_eager_experts_impl_on_the_fly`. `eager`, `grouped_mm` and `torch_a2a` are kept; a backend nothing intercepts (`batched_mm`, the opaque `sonicmoe` kernel, hub kernels) is moved to `grouped_mm` before calibration, or to `eager` on a transformers without `_grouped_linear`. A model whose experts quantize only on the eager forward keeps main's forced eager: native FP8 experts, a class with its own wrapper (GPT-OSS's quantizes `down_proj`'s input by intercepting the eager forward's matmuls), and Nemotron-H's non-gated experts until #3067. Such a model is refused on `torch_a2a`, before any config changes, instead of calibrating on raw FP8 bytes or not at all. - **Per-expert captures.** `grouped_mm` and `torch_a2a` feed a projection's input quantizer one expert-sorted batch. The weight-input captures (local Hessian, imatrix, GPTQ) split it back per expert by the cumulative token offsets the grouped interception records, following whichever path ran. The input quantizers calibrate on the rows the grouped GEMM computes only: under transformers' own expert parallelism, the rows routed to other ranks' experts are sorted past the offsets, and before transformers 5.17 the down projection's input holds them uninitialized. - **Expert parallelism.** #3122 makes everything that reads the expert weights outside the forward index each rank's own experts: max calibration's weight pass, `fold_weight`, the weight-input captures and GPTQ. Here the grouped captures split the expert-sorted batch per local expert. The grouped weight quantization needs no indexing: transformers swaps expert-parallel DTensors for their local shards around the experts forward. - **Fail loudly.** `calibrate()` raises, naming the modules, when a fused-experts forward that had to reach a quantizer ran but neither hook fired: a hand-written `bmm` forward on a class without `@use_experts_implementation`, which no config change redirects. Each `calibrate()` judges its own forwards and raises with the model back in its training mode. Under the expert parallelism of transformers 5.16 and later a module's verdict is reduced over its experts' expert-parallel group, so every rank of the group raises, not just the ranks whose experts saw rows. Not reported: experts no calibration token reached (Ernie-4.5-VL's vision experts under text-only calibration, which fall back to weight-derived scales as on main), an eager forward whose rows were all routed to other ranks' experts, and weight-only experts under max or MSE, whose weight quantizers calibrate on the weights alone. **Behavior change.** main forced MoE calibration onto eager; transformers resolves an unrequested backend to `grouped_mm`, so fused-experts models now calibrate there. A single-GPU Qwen3-30B-A3B export, NVFP4 or FP8, matches main in every weight and weight scale, but 4,647 activation scales move, by up to 10.2%. main's own repeat is byte-identical, and loading with `experts_implementation="eager"` reproduces main's checkpoint byte for byte (all 74,835 tensors for NVFP4 and 56,211 for FP8), so the backend alone moves them. The changelog says so. | Moved activation scales, Qwen3-30B-A3B NVFP4, 1 GPU | Tensors | |---|---| | Experts' `down_proj` input (33 layers, 128 experts each) | 4,224 | | Experts' shared `gate_proj` and `up_proj` input (1 layer) | 256 | | Attention input and KV-cache | 167 | | Total, in 45 of 48 layers (none in layers 0, 1 and 5) | 4,647 | The largest moves are 10.2% on the experts (layer 25's `down_proj` input) and 7.8% on attention. *Why they move.* Both forwards feed an input quantizer the same rows: a single GPU has no padding rows (transformers' sentinel rows exist only under its own expert parallelism, where calibration skips them; transformers zeroes them only from 5.17), the routing weights apply after `down_proj` on both, and neither casts. They round the block's output differently. Eager adds a token's 8 expert outputs into a bf16 buffer with `index_add_`, rounding after each add, while `grouped_mm_experts_forward` sums them in fp32 and rounds once, and its GEMMs run other kernels. From the first MoE block's output on, the residual stream differs by rounding, and where a router's 8th and 9th scores nearly tie, the routing flips. A flipped token feeds `down_proj` another expert's activations, and `down_proj`'s input scale is a max over routed (token, expert) rows, so it moves by more than rounding would. The moved scales fit that: the experts' shared `gate_proj`/`up_proj` input, whose rows are every token however it is routed, moved in one layer only, against 33 for `down_proj`. *Why `grouped_mm` off expert parallelism too,* rather than eager there and `grouped_mm` on `torch_a2a` only: it is the forward transformers runs the model on when nothing else is asked for, its fp32 sum is the more precise of the two combines, and expert-parallel calibration runs the grouped path anyway (`torch_a2a` runs the local experts through `_grouped_linear`), so a model calibrates on one forward whether or not it is sharded. `experts_implementation="eager"` keeps earlier releases' checkpoints. ### Usage ```python # To calibrate exactly as earlier releases did: model = AutoModelForCausalLM.from_pretrained(path, experts_implementation="eager") ``` ```bash # The same from hf_ptq.py: python hf_ptq.py --pyt_ckpt_path <model> --qformat nvfp4 --experts_implementation eager # or grouped_mm ``` ### Testing - `tests/unit/torch/quantization/plugins/test_fused_experts.py`, against synthetic and real transformers experts modules: - backend normalization, including the fallback to eager on a transformers without `_grouped_linear`; - GPTQ's Hessians and local Hessian's captures under `grouped_mm` and `torch_a2a` match eager's per expert, and a capture follows whichever path ran; - the interception check raises when the forward ran past both hooks, also after an earlier calibration reached them, and not for unreached experts, an eager forward routed only to other ranks' experts, or weight-only experts beside a calibrated dense input quantizer under max (each expert's weight amax equals its weight's per-channel abs max). Under local Hessian and GPTQ, which read the experts' inputs, those weight-only experts still raise; - GPT-OSS, a real tiny model, calibrating `down_proj`'s input on the eager forward. - Expert parallelism, on CPU under gloo: - the tests that experts sharded the way transformers shards them calibrate, fold, capture and build GPTQ helpers on each rank's own experts moved to #3122 with that fix; - every rank of the experts' expert-parallel group raises when only one rank's experts ran past both hooks; - with fake quantizers inserted, experts split over two ranks train through the `torch_a2a` forward like the whole quantized module: the fake-quantized case of `test_moe_a2a.py`'s whole-module test; - the transport's hop slicing under a graph: `test_moe_a2a.py`'s chunked-hops test also runs forward and backward through hops sliced below the limit, and checks the output and input gradients bit for bit against the unsliced run. - The transport's (#2738, merged) docs say again that ModelOpt quantizes its local experts, now that the `_grouped_linear` override does. Its hop-slicing comments say what the slicing guarantees, every hop within `all_to_all_single`'s int32 element offsets, rather than that it prevents a hang: with the slicing disabled, hops past 2**31 elements completed exactly on NCCL 2.29, 2.30 and 2.31, including at 32 ranks over the network on 2.29 and 2.30. - Against main, on tiny real models: GPT-OSS calibrates on eager with `down_proj`'s input amax equal to main's; Ernie-4.5-VL-MoE quantizes in FP8 and NVFP4 under text-only calibration, as on main. - Byte-for-byte checkpoint comparison against main on one GPU, measured on earlier revisions of #2745: Qwen3-30B-A3B, eager and `grouped_mm`, NVFP4 and FP8; Qwen2.5-0.5B (dense) is byte-identical. On GPU, Nemotron-3-Nano-30B-A3B's NVFP4 calibration with its experts pinned to eager and #3067's with `experts_implementation="eager"` give byte-identical exports, all 18,012 tensors. - The GPU tests of QAT under expert parallelism are in #2742, the first PR with an FSDP2-sharded dense part they also run. - Review fixes, 2026-10-27, carried from #2745: the FSDP2 window's ownership test (unnamed meshes, every shard state); an empty experts forward is no work the interception check expects; the non-gated experts are pinned to eager whichever plugin runs first. - Review fixes, 2026-10-27, last round, each with a test that fails without it: - `grouped_mm` is kept, and redirected to, only where `_grouped_linear` can be intercepted; one capability constant gates the interception, the supported set and the redirect target; - an eager experts forward counts as work for the interception check only when it routes a row to this rank's experts, so a rank whose experts calibration routed nothing to no longer raises alone; - an enabled expert weight quantizer needs the interception only when the calibration captures the experts' inputs to calibrate their weights (GPTQ, local Hessian, imatrix), so weight-only experts under max or MSE no longer raise; the message names what falls back. - Split from #2745 on 2026-10-27, night: #2745's and this PR's commits together give #2745's previous tree exactly, plus this round's fixes. Pre-commit, mypy included, and this PR's unit tests pass at this head on transformers 5.18 and 5.5.4 and on torch 2.8, and the series' PRs merged together add no unit-test failure over main f03651e (tests/unit/torch utils, export, quantization and models). This revision's GPU tests run in CI only. - Review suggestions, 2026-10-28: the weight-input captures (local Hessian, imatrix, GPTQ) read the path that ran, not the one the config named when they were registered, so they follow whichever forward runs and the refusal is gone, pinned by a test that fails on the code before; the interception check's forward count, and the host sync its routing test costs, stop after the first hook call of each calibration; the check's docstring says it is rank-local. Rebased over #2745's fixes, taking back the `torch_a2a` test fixtures #2745 no longer carries. - Rebased over #2745's second-round fixes: the grouped interceptor records the offsets and clears the eager expert index, and #2745's test that captures skip the grouped batch gives way to this PR's per-expert capture tests. - Review follow-up, 2026-10-28, third round: drops the note #2745 now adds, that the grouped captures skip their batch, since here they split it per expert. A new test pins that a calibration forward stops counting, and skips its host-syncing routing test, once a hook has fired; it fails with that gate removed. - Review suggestions, 2026-10-28, fourth round: the experts forward clears `_grouped_offs` when it returns, so a capture reached outside a forward skips instead of splitting its input by a finished forward's offsets; the interception check's docstring says that a rank whose experts saw no rows waits in its next collective until the launcher stops it or the collective times out, so the error is in another rank's log. Rebased over #2745's cleanup. - Rebased on 2026-10-28, evening, onto #2745's head `e7c449fcf87`, after #2745 moved onto main `7092db6ac5b` and took its last review rounds. Conflicts only in the changelog (both sides' entries kept), `_QuantFusedExperts`'s docstring (it keeps #2745's note that the `grouped_mm` fallback does not call `F.linear`) and the fused-experts tests (#2745's transposed-container tests kept). The FP8 experts wrapper's comment from #2745 now names `normalize_experts_impl_on_the_fly`, this PR's name for the function that keeps it on eager. pre-commit, mypy included, passes at this head against main. - Review fix, 2026-10-28, evening: `calibrate()`'s interception check reads the HF plugin from `sys.modules` instead of importing it, so a plugin that failed to load (which `import_plugin` only warns about) no longer fails every calibration with its error. `test_a_plugin_that_failed_to_load_leaves_calibration_alone` breaks the plugin's import and fails without the change. pre-commit, mypy included, passes at this head. - Split, 2026-10-28, evening, at the bot review's request: the FSDP2 window's ownership filter and its ownership test moved to #3088, with a test that the window leaves an ignored expert-parallel DTensor sharded. This PR sits on a helper base merging #2745 and #3088, and its source diff is 334 lines. At this head `test_fused_experts.py`, `test_moe_a2a.py` and `test_utils.py` pass on CPU (200 passed, 6 skipped), and pre-commit, mypy included, passes against main. - Review suggestions, 2026-10-28, evening: `_local_experts` refuses a DTensor whose local shard does not hold the module's experts (one replicated over the ranks, say) instead of indexing the global experts, and the native FP8 wrapper's `unpack_weight` refuses expert weights sharded over the experts instead of filling a buffer for every expert. One two-rank test covers both and fails without either. The suggested all-reduce in the interception check is declined (a collective in every distributed `calibrate()`). Focused tests pass (201 passed, 6 skipped), and pre-commit passes against main. - Review suggestions, 2026-10-28, late evening: the per-expert split of a grouped batch requires exactly one exclusive end per expert, within the rows, as transformers' `grouped_mm` and `a2a_moe` build the offsets, and raises on any other form (a leading zero would credit each expert with the previous one's rows). A test feeds three such forms and fails without the check. The `_grouped_offs` comment names all three captures and the attribute's role as the which-forward-ran signal. Rebased onto the helper base re-merged with #3088's update; focused tests pass (204 passed, 6 skipped), and pre-commit passes against main. - Review suggestion, 2026-10-28, late evening: `_local_experts` checks for a DTensor through the plugin's guarded `torch.distributed.tensor` import, so a torch without distributed does not fail every fused-MoE call. Rebased onto the helper base re-merged with #3088's update; focused tests pass (204 passed, 6 skipped), and pre-commit passes against main. - 2026-10-28, night: `hf_ptq.py` gains `--experts_implementation`, passed on every load path: `get_model` (the compressed-tensors branch included), the low-memory load, and the FSDP2 load, through the config its meta model is built from. The changelog's way back to earlier releases' scales is now reachable from the script. A test loads a tiny Qwen3-MoE with `eager` and fails without the plumbing. The `tests/examples/hf_ptq` unit tests pass (110), and pre-commit passes against main. - Review fixes, 2026-10-28, night: on the FSDP2 path the flag is set on the pre-fetched config's `_experts_implementation`, as `build_meta_causal_lm` sets `attn_implementation`, and a test captures the config the FSDP2 loader receives (it fails without the assignment); `get_model` passes the flag only when given, on every branch. `calibrate()`'s new error for a fused-experts forward that reaches neither hook has its own changelog entry. Rebased onto #2745, itself rebased onto main `bc0b6e03237` after #3088 merged; focused and `tests/examples/hf_ptq` tests pass (316 passed, 6 skipped), and pre-commit passes against main. - Review suggestions and a rebase, 2026-10-28, late night: the backend normalization warns once for each change it makes to a caller's choice, but not for a `None` backend pinned to eager, which already dispatches to the eager forward. Rebased onto #2745 on main `e9f432affbb` (#3091): `DTensor` now comes from the plugin's plain `torch.distributed.tensor` import, since the guarded import is gone with the torch 2.10 floor. Focused and `tests/examples/hf_ptq` tests pass (325 passed, 6 skipped), and pre-commit passes against main. - Review suggestion and a rebase, 2026-10-28, late night: the backend-redirect notice goes through `warn_rank_0`, so an expert-parallel run prints one copy. Rebased onto #2745's comment and warning-text update; `_is_fused_experts_module`'s new docstring names `normalize_experts_impl_on_the_fly`, which keeps GPT-OSS on eager here. `test_fused_experts.py` and `test_utils.py` pass (197), and pre-commit passes against main. - Bot review of `2e253a509d0`, 2026-10-28, late night: the FSDP2 test runs a Qwen3.5-MoE config too and checks the text config its experts dispatch on, which the one assignment reaches through transformers' recursive `_experts_implementation` setter (5.5.4 and 5.18; a tiny Qwen3.5-MoE built on meta dispatches its experts on `eager` with the flag and on `grouped_mm` without it). The test fails when the value reaches the root config only, and the comment at the assignment says the setter carries it. Declined: resolving the routing indices' parameter name from the forward's signature, since every transformers MoE block passes them positionally. The `tests/examples/hf_ptq` unit tests pass on 5.18 and 5.5.4 (120), and pre-commit passes against main. - Bot review of `042e4bf1e55` (approved), 2026-10-28, late night: the interception error names the ways forward: a class no backend change redirects needs its own quantized wrapper, and meanwhile the rest of the model quantizes with these modules' quantizers disabled, which the check honours. Declined: caching the grouped offsets' host copy across the two projections' captures (one small sync per experts module per batch under GPTQ, local Hessian and imatrix, whose per-expert matmuls dominate). pre-commit passes against main. - Rebased on 2026-10-28, late night, onto main `1ca8bd1fc4d` after #2745 merged (squashed with the tree of its last head, so this PR's tree is unchanged); the base is now `main`, so the GPU tests run. pre-commit passes against main. - CI fix, 2026-10-28, late night: the first run on main failed two cases of `tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py` in the trtllm `hf_ptq` lane (2 failed, 267 passed): its stand-in for hf_ptq's parsed arguments lacked the new `experts_implementation`, which `load_model` now reads. The stand-in gains it; the file passes on transformers 5.18 and 5.5.4. - Rebased on 2026-10-28, late night, onto main `bdc45368d10` (#2736 merged), with its own diff unchanged. Its tests with #3067's, #2736's `test_distributed.py` and the `hf_ptq` unit tests pass (376), and pre-commit passes against main. - CI, 2026-10-29: the first full run on main at `df2e3fe4a76`, GPU lanes included, passes. - Split, 2026-10-29, at a reviewer's request: the expert-parallel indexing fix, its tests and a Bug Fixes changelog entry moved to #3122 ([1/2], on main), and this PR now sits on it, its own diff otherwise unchanged. The review's three smaller items are taken: the interception check reduces its verdict over the experts' expert-parallel group (a two-rank test fails with the rank-local check); a test checks that every experts class transformers decorates with `use_experts_implementation` takes the routing indices second (transformers 5.5.4 and 5.18); and `--experts_implementation` offers `eager` and `grouped_mm` as argparse choices. pre-commit, mypy included, passes against main; this PR's tests, `test_moe_a2a.py` and the `hf_ptq` unit tests pass (315). - Review suggestion (Claude), 2026-10-29: the interception check reads the expert-parallel group from the experts' weight, a DTensor sharded over the experts from transformers 5.16 on; from 5.6 to 5.15 the shards are plain tensors that name no group, so the check stays rank-local there, as before. The docstrings say so. Restacked onto #3122's two follow-ups; pre-commit passes against main. - Comment, 2026-10-29: at `calibrate()`'s gate for the interception check, a sentence says that both conditions, the forward loop having run and the HF plugin being loaded, must come out the same on every rank, since under expert parallelism the check all-reduces over the experts' group. pre-commit passes against main. - Test fix, 2026-10-29: `tests/gpu/torch/quantization/test_fsdp2_autoquant.py`'s 2-GPU KV AutoQuantize test builds hf_ptq's arguments by hand and lacked the new `experts_implementation`, which `load_model`'s FSDP2 branch reads. It now passes `None`. The 1-GPU PR runners skip that test, and main's nightly runs it. No other test or example builds hf_ptq's arguments by hand on a path that reads the flag. pre-commit passes against main. - Rebased on 2026-10-30 onto #3122 on main `99c3834f62c`, with its own diff unchanged and no conflicts. pre-commit passes against main. - Review (sychen52), 2026-10-30: with every weight quantizer collecting statistics without quantizing (calibration), the grouped forward returns the stored weight itself instead of building a weight-sized copy slice by slice on each forward. A CPU test checks the identity and the collected per-expert amax, and fails on the previous code. On CPU, one call on a Qwen3-30B-A3B-sized `gate_up_proj` (128 × 1536 × 2048, bf16, 768 MiB) raised peak RSS by 755 MiB before and 0 MiB after; a GPU `max_memory_allocated` comparison is queued. The changelog entry now says the experts stay on `grouped_mm` for every forward after `mtq.quantize` or `mto.restore` (QAT/QAD, AutoQuantize scoring, fake-quant evaluation), with `experts_implementation="eager"` as the way back. `test_fused_experts.py` and `test_moe_a2a.py` pass on transformers 5.18 (220) and 5.5.4 (214, 6 skipped); pre-commit passes against main. - Review (kevalmorabia97), 2026-10-30: the three changes no test covered now have one each, or are gone. The `_local_experts` call in `_quantize_grouped_weight` is dropped as unreachable: the weight arrives as this rank's local shard inside the experts forward. New tests pin that a grouped forward reaching neither hook is counted whatever its routing indices hold, and that a capture after a grouped forward does not split by its offsets. The calibration pass-through calls only quantizers that calibrate (Claude's suggestion), with a constant-amax test case. Each test fails with its change reverted. `test_fused_experts.py`, `test_moe_a2a.py` and `test_hf_ptq_args.py` pass on transformers 5.18 (283) and 5.5.4 (277, 6 skipped); pre-commit passes against main. - GPU calibration peak memory, 2026-10-30 (sychen52's request): `torch.cuda.max_memory_allocated` around the `mtq.quantize(...)` call only, after `reset_peak_memory_stats`. Qwen3-30B-A3B on the default `grouped_mm` backend, recipe `general/ptq/fp8_default-kv_fp8`, `calib_size` 32, `batch_size` 1, one NVIDIA GB200 for both revisions (torch 2.11, transformers 5.13.0, CUDA 13.1). | Revision | Calibration peak (bytes) | MiB | |---|---|---| | Before the pass-through change (`3f3544435d8`) | 864,901,632 | 824.8 | | After (`27c5f77feac`) | 174,759,936 | 166.7 | | Drop | 690,141,696 | 658.2 | - Rebased on 2026-10-30 onto main `55898ee4e3b` after #3122 merged, dropping #3122's pre-squash commits; its own diff is unchanged, with no conflicts. pre-commit passes against main. - CI fix, 2026-10-30: the routing-order test read transformers' modeling sources with the platform's default encoding, cp1252 on Windows, which 16 of them (`modeling_altclip.py`, say) cannot be decoded with; it now reads them as UTF-8, and the windows unit job's only failure goes with it. The `linux` job's one failure on the previous head, a NAS sorting test this PR does not touch, passed on rerun at the same commit. pre-commit passes against main. - `eager` against main, byte for byte, 2026-10-30 (re-verifying at this PR's head the claim that `experts_implementation="eager"` reproduces main): PTQ of Qwen/Qwen3-30B-A3B with recipe `general/ptq/nvfp4_default-kv_fp8`, `calib_size` 512, calibration sequence length 2048, `batch_size` 1, cnn_dailymail, without expert parallelism (one process), on one NVIDIA GB200 (torch 2.11, transformers 5.13.0, CUDA 13.1). This PR at `27c5f77feac` with `--experts_implementation eager`, against main at `99c3834f62c`, both compared at zero tolerance. Later revisions of this PR leave its own diff unchanged (a rebase) or change tests only. | Scale kind | Tensors | Differing | |---|---|---| | `input_scale` | 18,624 | 0 | | `weight_scale` | 18,624 | 0 | | `weight_scale_2` | 18,624 | 0 | | `k_scale` | 48 | 0 | | `v_scale` | 48 | 0 | - Review (juhi10071998), 2026-10-30: under transformers' own expert parallelism the grouped forward sorts the rows routed to other ranks' experts past `offs[-1]`, which the grouped GEMM neither reads nor writes; before transformers 5.17 nothing zeroes them in the down projection's input, so its input quantizer calibrated on uninitialized rows. While an input quantizer collects statistics it now sees only the rows before `offs[-1]` (a host sync, during calibration only). A test hands the grouped projections a NaN or a large finite tail, as pre-5.17 expert parallelism leaves it, and checks the input amaxes match a run without one; both cases fail without the change. `test_fused_experts.py`, `test_moe_a2a.py`, `test_calib.py` and `test_hf_ptq_args.py` pass on transformers 5.18 (307) and 5.5.4 (301, 6 skipped); pre-commit passes against main. - Rebased on 2026-10-30 onto main `059588b89a5` (after #3131, #3124, #2751, #3104, #3128, #3134 and #3149 merged), with its own diff unchanged and no conflicts. pre-commit passes against main. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ❌ Fused-experts MoE calibration now runs on `grouped_mm`, so some activation scales differ from earlier releases; `experts_implementation="eager"` restores the earlier checkpoint exactly. - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ Backward Breaking Changes. - Did you get Claude approval on this PR?: ✅ At `042e4bf1e55`; the error-message change after it is from that review. ### Additional Information Split out of #2745 (part 3 of 21 of expert-parallel PTQ for Hugging Face MoE models), which merged on 2026-10-28. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Behavior Changes** * Hugging Face MoE calibration now uses Transformers’ default `grouped_mm` backend instead of forcing eager execution. Calibrated activation scales may differ; set `experts_implementation="eager"` to use the earlier calibration behavior. * Eager-only expert wrappers continue to use eager execution. Using them with `torch_a2a` raises an error. * **Bug Fixes** * Improved calibration support for grouped and expert-parallel execution, including local expert weights and experts receiving no routed inputs. * Calibration now reports when fused experts run without interception or required inputs are not captured. * Improved FSDP2 parameter handling during weight access and writeback. * **New Features** * Added calibration support for newer Transformers fused-expert implementations, including Qwen3-MoE and native FP8 experts. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Base automatically changed from
shengliangx/hfep-8a-grouped-mm-calibration
to
main
October 30, 2026 23:08
… its model folder Nemotron-H's non-gated experts get their own wrapper, _QuantNemotronHExperts, in modelopt/torch/models/nemotron_h/modeling_ptq.py, registered on transformers' NemotronHExperts (every supported release ships it) in place of the generic plugin's _QuantNonGatedFusedExperts. It is not eager-only: the grouped_mm interception reaches the non-gated layout through _first_proj_attr, so these experts calibrate on the backend transformers selects, as the gated ones do. The generic structural detection now claims only the gated layout. _is_fused_experts_module also counts a class registered with a _QuantFusedExperts subclass, so the backend normalization still finds Nemotron-H's experts, and transformers' non-gated FP8 experts, which stay on eager. The non-gated tests move to tests/unit/torch/models and run the real NemotronHExperts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Nemotron-H's experts register at import here, so no class the on-the-fly registration claims is eager-only, and the stand-in the test needed is gone with the generic non-gated detection. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ed layout Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…near The grouped-vs-eager comparisons would pass just as well if the grouped forward fell back to per-expert F.linear calls, which the eager hook intercepts with the same values. The interception count tells the two hooks apart: one _grouped_linear call per projection, against one F.linear call per projection and expert. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The test this one replaced checked it: for the non-gated layout up_proj is the whole fused weight, so an export that kept it would hold every expert twice, the fused original beside the per-expert copies. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_is_fused_experts_module found Nemotron-H's experts through their class's registration, which the class quantize() converts them to does not carry, so a second quantize() over the converted model skipped the backend normalization for them. Gated experts are found by their shapes either way. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
up_proj's input is the hidden states, which the grouped forward hands _grouped_linear gathered by expert: the same rows eager feeds per expert, so their amax matches exactly. Only down_proj's input passes through a GEMM and keeps the tolerance. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Its one caller, normalize_experts_impl_on_the_fly, runs inside replace_quant_module, which refuses a model that is already quantized, so no converted experts module reaches it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…Experts The non-gated cases of these fused-experts tests went with the stand-in when the non-gated tests moved here. They run on the real class again: the eager forward recovers each up_proj slice's expert, and the quantizer-only iterator yields the weight iterator's quantizers in its order without indexing the fused weight; that test fails if the iterator names the gated projection. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With the non-gated wrapper gone, register_fused_experts_on_the_fly registers only the generic wrapper, which is not eager-only, so looking it up for a class not registered yet cannot change which experts are forced onto eager: those wrappers (FP8's, GPT-OSS's, Nemotron-H's) register when their modules are imported. The transposed check's comment no longer speaks of non-gated layouts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Collaborator
Author
|
/ok to test d452768 |
Fridah-nv
reviewed
Oct 30, 2026
| __all__: list[str] = [] | ||
|
|
||
|
|
||
| class _QuantNemotronHExperts(_QuantFusedExperts): |
Contributor
There was a problem hiding this comment.
Suggestion (non-blocking): the explicit registration skips the is_transposed and shape checks that register_fused_experts_on_the_fly applied, and the docstring says the wrapper relies on transformers' layout. A shape assert can't tell a consistently transposed pair, as you noted. But getattr(module, "is_transposed", False) needs no shape inference. Could setup raise when it is true, so a future layout flip fails loudly instead of slicing the experts wrong?
Collaborator
Author
There was a problem hiding this comment.
Taken in 0da0591. _QuantNemotronHExperts._setup now raises when transformers marks the experts is_transposed, before any quantizer is created, so a future layout flip fails at conversion instead of slicing every expert wrong. The docstring says so. test_experts_marked_transposed_are_refused converts a real NemotronHExperts marked transposed and checks the conversion raises. It fails without the change.
Contributor
|
Question (non-blocking): the Nemotron-3-Nano table shows 1,932 of 5,934 expert input scales moving, up to 4.39%, between |
Registered explicitly, the Nemotron-H experts skip the transposed-layout check register_fused_experts_on_the_fly applies, and the wrapper slices each expert's weight as (out_dim, in_dim). A transposed layout would slice every expert wrong with nothing in the shapes to catch it, so setup now refuses experts transformers marks is_transposed. A test marks real NemotronHExperts transposed and checks the conversion raises; it fails without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Collaborator
Author
|
/ok to test 0da0591 |
Collaborator
Author
|
/claude review |
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.
Note
Split out of #2745 on 2026-10-27, so that the Nemotron-H part is reviewed on its own. #2745, #3122 and #3074 (the calibration this PR builds on) have all merged.
Base:
main.In the tree: #2741's helper base merges this PR instead of #2745 and #3074. The expert-parallel chain quantizes Nemotron-H's experts on the
torch_a2aforward, which refuses an eager-only wrapper.Series layout
main, with [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 (FSDP2 weight-access window,_fsdp2_owned_params) and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 (fused experts under grouped_mm), both merged 2026-10-28What does this PR do?
Type of change: New feature
#2745 quantizes fused experts under
grouped_mm, the experts backend transformers selects by default, in the gated layout (gate_up_proj). Nemotron-H's experts are non-gated,down_proj(act_fn(up_proj(x))), with a single 3-Dup_proj. On #2745 and #3074 they keep main's generic_QuantNonGatedFusedExpertsand calibrate on the eager forward, which main forces and #3074 keeps for them._QuantNemotronHExpertslives inmodelopt/torch/models/nemotron_h/modeling_ptq.py. It registers on transformers'NemotronHExperts, which every supported release (5.5 on) ships, asmodelopt/torch/models/README.mdasks of wrappers only one model needs. The generic plugin's_QuantNonGatedFusedExpertsand itsup_proj-only structural detection are removed: no other transformers experts class has that layout.grouped_mm. The wrapper is not eager-only: [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's grouped interception reaches the non-gated layout through_first_proj_attr, and their quantizer state matches the eager forward's. Under expert parallelism ([OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 on), they run thetorch_a2aforward, which refuses an eager-only wrapper._is_fused_experts_modulealso counts a class registered with a_QuantFusedExpertssubclass. So a backend the quantizers cannot see (batched_mm) is still redirected for Nemotron-H. Transformers' non-gated FP8 experts (an FP8 Nemotron-H's) are still forced onto eager.Behavior change. Nemotron-H's MoE calibration runs on
grouped_mmwhere transformers selects it, so some of its activation scales can differ from main's, as #3074's changelog entry says for fused experts. Loading withexperts_implementation="eager"calibrates as before.Usage
No API change. To calibrate exactly as earlier releases did:
Testing
tests/unit/torch/models/test_nemotron_h_modeling_ptq.py(new) runs transformers' realNemotronHExperts. It covers:up_projquantizers;grouped_mm;grouped_mmcalibration and GPTQ Hessians equal to eager's, for every expert;up_projanddown_proj;main's tests of the non-gated layout ran a hand-written stand-in, which the registration does not reach, so they move here.
test_fused_experts.py's FP8 force-eager test also runs non-gated FP8 experts. It and the backend-redirect test fail with the registry lookup in_is_fused_experts_modulereverted.tests/gpu/torch/models/test_modeling_ptq_registration.pychecks thatNemotronHExpertsis registered whichever module is imported first.Locally at this head: pre-commit, mypy included;
tests/unit/torch/{models,export},test_fused_experts.pyandtest_huggingface.pyon transformers 5.18 and 5.5.4 and on torch 2.8, with the same failures as main in each environment (diffusers). The series' PRs merged together add no unit-test failure over main.GPU, NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 (a BF16 checkpoint, no
quantization_config), one GPU, transformers 5.18.0,general/ptq/nvfp4_default-kv_fp8, 64 calibration samples of 512 tokens, exported tensors compared one by one:grouped_mm, twice (run-to-run floor)grouped_mmvsexperts_implementation="eager"The weights and weight scales do not depend on the backend; activation scales move with it, in the MoE layers and in the layers after them, as [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's changelog entry says for fused experts.
tests/gpu/torch/models/test_modeling_ptq_registration.pypasses on 1 GPU.Re-cut onto [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's last review fixes (2026-10-27): the plugin-order test of the non-gated stand-in is dropped with the stand-in. Pre-commit and this PR's unit tests pass on transformers 5.18 and 5.5.4. A remote-code Nemotron-H keeps its experts as an
nn.ModuleListofnn.LinearMLPs, which the fused detection never claimed, and the remote-code VL model builds transformers'NemotronHForCausalLM, whose experts this PR registers.Review suggestion, 2026-10-27, last round: the wrapper's docstring says the explicit registration skips the transposed-layout check and relies on transformers'
(num_experts, out_dim, in_dim)layout, which every supported release uses.Rebased with [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 onto main f03651e (2026-10-27, night); this PR's own diff is unchanged.
Review suggestions, 2026-10-27, last round: the test pins that
grouped_mmwent through_grouped_linear, two calls rather than oneF.linearcall per expert and projection, and that the export releases the fusedup_projanddown_proj. Re-cut onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074, the second half of [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's split; this PR's own diff is otherwise unchanged.Review suggestions, 2026-10-28: an already-converted Nemotron-H experts module still counts as fused experts for a second
quantize(), pinned by the registration test, which fails without the change;up_proj's input amax is compared exactly, and onlydown_proj's keeps a tolerance.Review suggestion, 2026-10-28, second round: the converted-experts check reuses
_is_quant_fused_experts_module.Review suggestion, 2026-10-28, third round:
_is_fused_experts_moduleno longer counts already-converted experts, undoing the first round's change: its one caller runs insidereplace_quant_module, which refuses a quantized model, so none reaches it.Declined review suggestion, 2026-10-28, third round, raised before: a warning for unregistered non-gated fused experts. transformers'
NemotronHExperts, which is registered, is the only class with the layout, and Nemotron-H's remote code builds per-expert MLPs.Declined review suggestions, 2026-10-28, fourth round: skip guards on the grouped tests, since
NemotronHExpertsand_grouped_linearare in every release the pin allows; and a shape assert in_setup, which cannot tell a consistently transposed pair from the stored one, while the export test built on transformers' ownNemotronHExpertspins the per-expert shapes.Declined review suggestion, 2026-10-28, fifth round (CodeRabbit): resolving a converted experts module's original class in
_is_fused_experts_module; its one caller runs before any module is converted.Rebased on 2026-10-28, late night, onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 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), so the diff shows this PR's own slice only (+292/−362 in 6 files). Conflicts were with [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's later changes:_is_fused_experts_module's docstring keeps [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's GPT-OSS sentence, the FP8 force-eager test runs both [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's backend cases and this PR's gated and non-gated experts, and the FP8 wrapper's comment no longer compares it with the removed non-gated wrapper.Review suggestion (cjluo-nv's bot), same night: the non-gated cases that went with the stand-in run on the real
NemotronHExpertsagain: the eager forward recovers eachup_projslice's expert, and the quantizer-only iterator yields the weight iterator's quantizers in its order without indexing the fused weight (it fails if the iterator names the gated projection). pre-commit passes against main;test_fused_experts.pyand this PR's tests pass on transformers 5.18 and 5.5.4.Rebased on 2026-10-28, late night, onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 on main
bdc45368d10([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 merged), with its own diff unchanged. Tests and pre-commit pass as for [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074.Restacked on 2026-10-29 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 after its expert-parallel fix moved to [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 underneath it, with its own diff unchanged (one import-list conflict with [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's new
assert_fused_experts_interceptedimport). Its tests,test_nemotron_h_modeling_ptq.pyandtest_moe_a2a.pypass (185), and pre-commit passes against main.Rebased on 2026-10-30 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 on main
99c3834f62c, with its own diff unchanged and no conflicts. pre-commit passes against main.Restacked on 2026-10-30 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 on main
55898ee4e3b([OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 merged), with its own diff unchanged and no conflicts. pre-commit passes against main.Restacked on 2026-10-30 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 on main
059588b89a5, with its own diff unchanged and no conflicts. pre-commit passes against main.Rebased on 2026-10-30 onto main
d464d30b120after [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 merged, dropping [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's commits; its own diff is unchanged, with no conflicts. pre-commit passes against main; now main-based, the full CI runs, GPU lanes included.Review suggestion (Fridah-nv), 2026-10-30: the explicitly registered wrapper refuses experts transformers marks
is_transposed, at setup, since nothing in the shapes tells a transposed layout apart. A test marks a realNemotronHExpertstransposed and checks the conversion raises; it fails without the change. The Nemotron-H tests pass on transformers 5.18 and 5.5.4; pre-commit passes against main.Before your PR is "Ready for review"
grouped_mm, so some activation scales differ from earlier releases;experts_implementation="eager"restores the earlier calibration.CONTRIBUTING.md: N/Agrouped_mm.Additional Information
Part of expert-parallel PTQ for Hugging Face MoE models; merges after #2745.
Summary by CodeRabbit