Repository navigation
Skip long-lived objects in the per-step gc.collect() of QAT and QAD - #3104
Merged
Merged
Conversation
Contributor
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Collection and freeze lifecycle modelopt/torch/opt/plugins/transformers.py, examples/llm_qat/ARGUMENTS.md, tests/unit/torch/opt/plugins/test_trainer_gc.py, CHANGELOG.rst |
Manual GC collects during training and prediction. Training collection freezes the GC generation after the first optimizer step. train() unfreezes it on exit if this trainer froze it. The documentation describes this behavior, and the test checks that a model in a reference cycle can be collected after training. The changelog states that Hugging Face QAT and QAD can skip model and optimizer state during pre-step collection. |
QAT collection integration modelopt/torch/quantization/plugins/transformers_trainer.py |
QAT training calls collection with freezing enabled. Prediction calls collection without freezing. |
Priority:
Estimated code review effort: 2 (Simple) | ~12 minutes
Change: Bug fix
Sequence Diagram(s)
sequenceDiagram
participant QATTrainer
participant ModelOptHFTrainer
participant PythonGC
QATTrainer->>ModelOptHFTrainer: training_step()
ModelOptHFTrainer->>PythonGC: collect and freeze after first optimizer step
QATTrainer->>ModelOptHFTrainer: train() exits
ModelOptHFTrainer->>PythonGC: unfreeze if this trainer froze the generation
Suggested reviewers: cjluo-nv, sugunav14
Merge Risk: 
5539b
-
No identified issue prevents merging 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: skipping long-lived objects during per-step garbage collection for QAT and QAD. Docstring Coverage Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The PR changes only two Python files under the reviewed scopes: modelopt/torch/opt/plugins/transformers.pyandmodelopt/torch/quantization/plugins/transformers_trainer.py. The additions conc…Full details: Docstring Coverage
-
Explanation
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 8 functions across 3 files. (1 skipped: 1 unsupported.)
Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only two Python files under the reviewed scopes:
modelopt/torch/opt/plugins/transformers.pyandmodelopt/torch/quantization/plugins/transformers_trainer.py. The additions concerngc.collect(),gc.freeze(), andgc.unfreeze()only. They add notorch.load(..., weights_only=False), pickle-enabled NumPy load, hardcodedtrust_remote_code=True,eval(),exec(), or# nosec. The PR also changes nopyproject.tomlor requirements file, so it adds no dependency requiring license review.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- 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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🔇 Additional comments: 1
modelopt/torch/opt/plugins/transformers.py-578-578: 1
578-578:Major · Stability & Availability · Heavy lift
⚠️ Unverified finding
Verification ran but could not confirm this finding. It is shown for review, not as a verified issue.Do not freeze transient objects for the full training run.
If a callback retains a cyclic GPU buffer through the first optimizer step and releases it after the next step,
gc.freeze()moves that buffer into the permanent generation. Latergc.collect()calls cannot release its GPU memory untiltrain()returns. A long QAT run can then run out of memory. Provide a way to keep per-step collection without freezing for workloads that release cyclic objects during training. Python documents thatgc.freeze()freezes all tracked objects, not only the model and optimizer. (docs.python.org)
-
-
- Source mode: Auto
- Repository: NVIDIA/Model-Optimizer
- Branch: main
- File: .coderabbit.yaml
- Commit: bdc4536
- Layers: Repository YAML
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
92aa0978-ec93-4671-8ca8-1315ec8eb232
-
CHANGELOG.rstexamples/llm_qat/ARGUMENTS.mdmodelopt/torch/opt/plugins/transformers.pymodelopt/torch/quantization/plugins/transformers_trainer.pytests/unit/torch/opt/plugins/test_trainer_gc.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
-
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3104 +/- ##
==========================================
+ Coverage 71.63% 78.95% +7.31%
==========================================
Files 669 669
Lines 77033 77045 +12
==========================================
+ Hits 55182 60830 +5648
+ Misses 21851 16215 -5636
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:
|
kevalmorabia97
left a comment
kevalmorabia97
left a comment
Collaborator
There was a problem hiding this comment.
The change looks correct and the measurements are convincing; one required fix.
Needs action:
- Add a test that the freeze happens, not only that
train()unfreezes (inline comment).
No action needed:
gc.unfreeze()intrain()also undoes any freeze the application did beforetrain(). The code comment says so; it only costs later scans.- The CodeRabbit finding about transient cyclic GPU objects being frozen for the whole run is a real risk in principle. The 120-step FSDP2 run shows flat memory, and
ARGUMENTS.mddocuments the behaviour, so I wouldn't block on it.
| train_dataset=dataset, | ||
| ) | ||
| try: | ||
| trainer.train() |
Collaborator
There was a problem hiding this comment.
Required: this test only guards gc.unfreeze(). I ran it at this head: it fails with the unfreeze removed, but it still passes with gc.freeze() removed, so nothing guards the speed-up and a refactor could turn the freeze into a no-op.
Please assert the freeze too. For example, add a TrainerCallback whose on_step_end records gc.get_freeze_count(), and check after train() that the count was above its starting value from step 2 on, and that it is back at the starting value once train() returns. Python starts with a few hundred frozen objects, so compare against the count taken before train().
Contributor
Author
There was a problem hiding this comment.
Done in 36dbabd: a callback records gc.get_freeze_count() after each step, and the test checks it is above the pre-train() count from step 2 on, and 0 once train() returns. It is 0 rather than the starting count because gc.unfreeze() also returns the objects Python freezes at startup (the comment at the unfreeze says so). On B200 (NeMo 26.08, transformers 5.18) it passes in 1.2 s, and fails with either gc.freeze() or gc.unfreeze() removed.
shengliangxu
left a comment
shengliangxu
left a comment
Collaborator
There was a problem hiding this comment.
Pensieve comment review using shengliangx's profile (revision 1).
Comment: the required regression assertion that freezing actually occurs remains missing at this unchanged head.
Needs action:
- Add the freeze assertion requested in
4409536438totest_trainer_gc.py; the current test still passes ifgc.freeze()is removed. Verify freezing during training and cleanup afterward. - Confirm the new test's NVIDIA license header matches the repository's canonical
LICENSE_HEADER; that file is unavailable in this snapshot.
No action needed:
- The existing test covers collection of a cyclic model after training; no existing tests were modified or deleted.
- No duplicate findings added. Duplicate checking is limited to the supplied discussion and bounded review history.
- Repository navigation, canonical-header verification, and test execution were unavailable; reported test results were not independently verified.
Review context was incomplete.
Review run: 29a0f5f744f34bc1a312fc5b5862353d
Contributor
Author
|
QATTrainer runs gc.collect() before every training step, and manual_gc runs another one. Each walks every object the cycle collector tracks: about 1.9M on Qwen3-30B-A3B QAD (models, quantizers, optimizer state), 0.5 s per collection and 1.1 s of a 3.6 s step. From the second training step on, the trainer freezes what is alive after the collection once (gc.freeze()), so later collections skip it, and unfreezes it when train() returns so the model can be freed. Python 3.12 starts with a few hundred frozen objects, so the trainer tracks its own freeze instead of checking gc.get_freeze_count(). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
The test passed with gc.freeze() removed, so nothing guarded the speed-up. A callback now records gc.get_freeze_count() after each step: it must be above its value before train() from the second step on, and zero once train() returns, since gc.unfreeze() also returns what Python froze at startup. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
sychen52
force-pushed
the
qat-gc-freeze
branch
from
October 30, 2026 18:01
36dbabd to
5539b11
Compare
kevalmorabia97
approved these changes
Oct 30, 2026
kevalmorabia97
left a comment
kevalmorabia97
left a comment
Collaborator
There was a problem hiding this comment.
The freeze test now guards both calls: it fails with gc.freeze() removed and with gc.unfreeze() removed.
kevalmorabia97
enabled auto-merge (squash)
October 30, 2026 18:01
This was referenced Oct 30, 2026
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: bug fix (performance)
QATTrainer(and soQADTrainer) runsgc.collect()before every training step, and--manual_gcruns another one. Each full collection walks every object the cycle collector tracks. On Qwen3-30B-A3B QAD that is about 1.9 M objects, mostly the student, the teacher, their quantizers and the optimizer state. That takes about 0.5 s per collection, twice per step: 1.1 s of every 3.6 s step went to collections that free a few hundred objects.Now, from the second training step on, after the first collection the trainer calls
gc.freeze()once. Objects that live through training move to the permanent generation, so later collections skip them. They still free per-step garbage, which is what the per-step collection is there for.train()callsgc.unfreeze()when it returns, so the model and the trainer can be freed afterwards.Python 3.12 starts with a few hundred frozen objects (
python -I -S -c "import gc; print(gc.get_freeze_count())"prints 375), so the trainer tracks its own freeze instead of checkinggc.get_freeze_count().Usage
No new API. QAT and QAD with
QATTrainer,QADTrainer, ormanual_gc=Trueget it automatically.Testing
New test:
tests/unit/torch/opt/plugins/test_trainer_gc.py::test_model_is_freed_after_train_with_manual_gctrains a tiny CPU model held in a reference cycle for 3 steps, then checks it is freed aftertrain()returns (1 s). Without thegc.unfreeze()intrain(), it fails.Full Qwen3-30B-A3B NVFP4 QAD with a BF16 teacher on 8x B200: FSDP2 + activation checkpointing as in
examples/llm_qat, 16K tokens per GPU, Liger on,manual_gcon as inqad_nvfp4.yaml. Same node, median of steps 4-8:A 120-step run with this PR stays at 2.39 s/step, with allocated memory flat at 49.9 GiB and peak at 72.9 GiB from step 4 to step 120. Losses match main at steps 1-2 and stay within run-to-run noise after.
Also ran (NeMo 26.08 container, transformers 5.18):
tests/unit/torch/opt/pluginsandtests/unit/torch/distill(all pass except 9 tests that also fail on main because the container's/opt/Megatron-Bridgeis not readable), andtests/examples/llm_qat/test_llm_qat.py -k "qat_nvfp4 and fsdp2".Not run: multi-node, DeepSpeed or DDP, and a full-length training run.
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Independent of the context-parallelism PR (#3103). Both edit
QATTrainer.training_stepandCHANGELOG.rstnext to each other, so whichever lands second needs a small rebase.🤖 Generated with Claude Code
Summary by CodeRabbit