Support efficient TEGroupedMLP (moe_grouped_gemm=True) for Minitron Pruning - #1955
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (10)
🚧 Files skipped from review as they are similar to previous changes (10)
📝 WalkthroughWalkthroughAdds grouped-GEMM MoE expert support to dynamic NAS and Minitron pruning, exposes grouped/sequential selection in the pruning example, expands grouped and sequential test coverage, and adds an int32 indexing guard for Mamba calibration prefill. ChangesGrouped GEMM MoE Pruning
Mamba calibration indexing guard
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant PruneExample
participant MegatronBridge
participant DynamicMoE
participant ImportanceRegistry
PruneExample->>MegatronBridge: set moe_grouped_gemm
MegatronBridge->>DynamicMoE: construct TEGroupedMLP experts
DynamicMoE->>ImportanceRegistry: register grouped MLP importance
ImportanceRegistry->>DynamicMoE: collect expert and FFN activations
DynamicMoE-->>PruneExample: export pruned grouped experts
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
kevalmorabia97
changed the title
|
/claude review |
|
@CodeRabbit full review |
✅ Action performedFull review finished. |
|
|
/ok to test 21269a2 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1955 +/- ##
==========================================
- Coverage 78.47% 77.84% -0.63%
==========================================
Files 518 518
Lines 58574 58658 +84
==========================================
- Hits 45967 45664 -303
- Misses 12607 12994 +387
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:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/utils/plugins/mbridge.py (1)
106-107: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winStale docstring contradicts this PR's new grouped-GEMM pruning support.
The
moe_grouped_gemmarg docstring still says "Pruning does not support grouped GEMM yet," but this PR adds exactly that support, andprune_minitron.pynow defaults tomoe_grouped_gemm=True. Please update this line to reflect current behavior.📝 Proposed fix
- moe_grouped_gemm: Whether to use grouped GEMM for MoE. - Pruning does not support grouped GEMM yet. + moe_grouped_gemm: Whether to use grouped GEMM for MoE.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/utils/plugins/mbridge.py` around lines 106 - 107, The docstring for `moe_grouped_gemm` in `MBridge` is stale and still says pruning does not support grouped GEMM; update that description to match the new behavior introduced by this PR, since grouped-GEMM pruning is now supported and `prune_minitron.py` defaults to `moe_grouped_gemm=True`. Keep the wording aligned with the current implementation so the documentation in `mbridge.py` accurately reflects the `moe_grouped_gemm` flag.
🧹 Nitpick comments (1)
modelopt/torch/prune/plugins/mcore_minitron.py (1)
1390-1391: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnclear purpose of the inserted middle dimension.
[:, None, :]inserts a size-1 dim that is immediately collapsed by the subsequent.mean(dim=0)/.sum(dim=0), making it a no-op given the following reductions. If this is intentionally mirroring a[seq, batch, hidden]shape convention from the non-grouped path, a short comment would help; otherwise it can be simplified.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/prune/plugins/mcore_minitron.py` around lines 1390 - 1391, The `acts` computation in `mcore_minitron.py` has an inserted singleton middle dimension that is immediately removed by the later `mean(dim=0)` and `sum(dim=0)`. Either simplify the `gather_from_tensor_model_parallel_region`/`acts` pipeline to remove the no-op `[:, None, :]` shape expansion, or keep it and add a brief clarifying comment in the `acts` assignment block explaining why the `[seq, batch, hidden]`-style dimension is needed for consistency with the non-grouped path.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@modelopt/torch/utils/plugins/mbridge.py`:
- Around line 106-107: The docstring for `moe_grouped_gemm` in `MBridge` is
stale and still says pruning does not support grouped GEMM; update that
description to match the new behavior introduced by this PR, since grouped-GEMM
pruning is now supported and `prune_minitron.py` defaults to
`moe_grouped_gemm=True`. Keep the wording aligned with the current
implementation so the documentation in `mbridge.py` accurately reflects the
`moe_grouped_gemm` flag.
---
Nitpick comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 1390-1391: The `acts` computation in `mcore_minitron.py` has an
inserted singleton middle dimension that is immediately removed by the later
`mean(dim=0)` and `sum(dim=0)`. Either simplify the
`gather_from_tensor_model_parallel_region`/`acts` pipeline to remove the no-op
`[:, None, :]` shape expansion, or keep it and add a brief clarifying comment in
the `acts` assignment block explaining why the `[seq, batch, hidden]`-style
dimension is needed for consistency with the non-grouped path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f6ded485-9fc1-4504-b19d-855a1c305887
📒 Files selected for processing (8)
CHANGELOG.rstexamples/megatron_bridge/prune_minitron.pymodelopt/torch/nas/plugins/megatron.pymodelopt/torch/prune/plugins/mcore_minitron.pymodelopt/torch/utils/plugins/mbridge.pytests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.pytests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.pytests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Summary
Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
Reviewed all 8 changed files. Prioritized the three modelopt/torch/ source files (nas/plugins/megatron.py, prune/plugins/mcore_minitron.py, utils/plugins/mbridge.py), the example driver, and the three GPU test files. (Note: the two-dot diff against the shallow base tip surfaced unrelated main-branch churn in examples/onnx_ptq/* and modelopt/onnx/*; those are not part of this PR's file list and were excluded.)
This is a clean, well-scoped feature that adds a TEGroupedMLP (grouped-GEMM) dynamic-module path alongside the existing SequentialMLP path for Minitron MoE pruning. I traced the algorithm end-to-end:
- Mode/State composability ✅ — The
num_local_expertsTracedHp is the same object shared acrosslinear_fc1,linear_fc2, and (via_DynamicMoELayer) the router, so_slice_order-based expert reorder/drop stays consistent across all three.moe_ffn_hidden_sizeis a single shared hparam (correct for Minitron's homogeneous pruning), and the search-space-size / sortable-hparam-count test assertions were updated accordingly. - Export fidelity ✅ —
_DynamicTEGroupedLinear.export()reads all reordered+slicedweight{i}/bias{i}via the dynamic getters before mutating, clones them, drops the per-expert params, foldsnum_gemms/in_features/out_featuresthroughsuper().export(), then re-registers exactly theactivekept experts. Ordering is correct. - Importance estimator ✅ —
_register_grouped_mlp_importancereuses the SequentialMLP expert-L2 hook fornum_local_expertsand adds an fc2-input activation score formoe_ffn_hidden_sizematching the dense_DynamicMLPreduction. The score is stashed in_activations["ffn_activations"]so it round-trips through the per-rank score checkpoint. - Performance ✅ — No hot-path CPU sync: the
_slice_order.tolist()in the expert getter only runs post-sort, and_DynamicMoELayer.forwardalready blocks forward in the sorted/trimmed state until export.
The single finding is a non-blocking stale docstring in mbridge.py (the inline comment below it was correctly updated, but the Args: entry still says pruning doesn't support grouped GEMM).
Risk: low. Additive change mirroring an established, tested pattern; both paths are exercised by the parametrized GPU tests.
21269a2 to
fc29d46
Compare
|
/ok to test fc29d46 |
1073a71 to
f8b2d51
Compare
|
/ok to test f8b2d51 |
|
/ok to test 1c8cf2b |
db0afc9 to
a1bc8b3
Compare
|
/ok to test a1bc8b3 |
|
/ok to test d83f4ba |
d83f4ba to
1a7dbc2
Compare
|
/ok to test 1a7dbc2 |
1a7dbc2 to
d1dd5f6
Compare
|
/ok to test d1dd5f6 |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/megatron_bridge/prune_minitron.py (1)
513-518: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winValidate the parsed MMLU parameters.
The regex checks only the format, so values such as
mmlu_0pct_bs0andmmlu_101pct_bs1are accepted and passed downstream as invalidfraction/batch_sizevalues. Reject percentages outside1..100and batch sizes below1. Also align the existing help text with the acceptedmmlu_<N>pct_bs<bs>format.Proposed validation
if match: - mmlu_frac = float(match.group(1)) / 100.0 - batch_size = int(match.group(2)) + mmlu_pct = int(match.group(1)) + batch_size = int(match.group(2)) + if not 1 <= mmlu_pct <= 100 or batch_size < 1: + raise ValueError("MMLU percentage must be 1-100 and batch size must be positive") + mmlu_frac = mmlu_pct / 100.0As per coding guidelines, validate external input once at the interface boundary.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/megatron_bridge/prune_minitron.py` around lines 513 - 518, Update the MMLU argument parsing around args.prune_score_func to validate the parsed percentage is between 1 and 100 and batch_size is at least 1 before assigning or passing downstream. Keep the existing mmlu_<N>pct_bs<bs> format and revise the associated help text to describe that format and accepted ranges.Source: Coding guidelines
🧹 Nitpick comments (1)
tests/examples/megatron_bridge/test_prune_minitron.py (1)
63-65: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the new grouped-GEMM default.
Both integration commands pass
--no_moe_grouped_gemm, so they exercise onlyTESequentialMLP. A regression in the new default or inmoe_grouped_gemm=not args.no_moe_grouped_gemmwould go undetected. Keep these compatibility runs if needed, but add or parameterize anemo:26.08+run without the flag.As per path instructions, tests for new behavior must exercise the real implementation.
Also applies to: 130-132
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/examples/megatron_bridge/test_prune_minitron.py` around lines 63 - 65, Add integration-test coverage for the grouped-GEMM default in the prune command setup, while preserving the existing compatibility runs that pass --no_moe_grouped_gemm. Add or parameterize a nemo:26.08+ execution without that flag so the real default path and moe_grouped_gemm behavior are exercised.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/utils/plugins/megatron_generate.py`:
- Around line 103-106: Update the overflow message in the guard around
batch_size, seq_length, and d to calculate the maximum safe batch size with
(2**31 - 1) // (seq_length * d), matching the existing >= 2**31 rejection
condition.
---
Outside diff comments:
In `@examples/megatron_bridge/prune_minitron.py`:
- Around line 513-518: Update the MMLU argument parsing around
args.prune_score_func to validate the parsed percentage is between 1 and 100 and
batch_size is at least 1 before assigning or passing downstream. Keep the
existing mmlu_<N>pct_bs<bs> format and revise the associated help text to
describe that format and accepted ranges.
---
Nitpick comments:
In `@tests/examples/megatron_bridge/test_prune_minitron.py`:
- Around line 63-65: Add integration-test coverage for the grouped-GEMM default
in the prune command setup, while preserving the existing compatibility runs
that pass --no_moe_grouped_gemm. Add or parameterize a nemo:26.08+ execution
without that flag so the real default path and moe_grouped_gemm behavior are
exercised.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 07e15542-5db1-43f7-b695-eaff3c4983f2
📒 Files selected for processing (10)
CHANGELOG.rstexamples/megatron_bridge/prune_minitron.pymodelopt/torch/nas/plugins/megatron.pymodelopt/torch/prune/plugins/mcore_minitron.pymodelopt/torch/utils/plugins/mbridge.pymodelopt/torch/utils/plugins/megatron_generate.pytests/examples/megatron_bridge/test_prune_minitron.pytests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.pytests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.pytests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py
🚧 Files skipped from review as they are similar to previous changes (7)
- CHANGELOG.rst
- tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py
- modelopt/torch/utils/plugins/mbridge.py
- modelopt/torch/prune/plugins/mcore_minitron.py
- tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py
- tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py
- modelopt/torch/nas/plugins/megatron.py
Support pruning MoE models loaded with moe_grouped_gemm=True (fused TEGroupedMLP) in addition to the existing TESequentialMLP path. - Add _DynamicTEGroupedLinear + _DynamicTEGroupedMLP dynamic modules that slice per-expert grouped weights (num_moe_experts / moe_ffn_hidden_size / hidden_size) and reorder/drop experts via num_gemms, mirroring the SequentialMLP path. A single shared moe_ffn_hidden_size is pruned across experts since Minitron prunes homogeneously. - Route ffn/expert importance for grouped experts and store scores in the checkpoint-tracked _activations dict so reprune-from-checkpoint works. - prune_minitron.py uses grouped GEMM by default (--no_moe_grouped_gemm opts out). - Cover both expert impls in tests: parametrize the GPT MoE dynamic-module and pruning tests, and switch the mamba hybrid params NAS test to grouped GEMM (memory NAS test stays SequentialMLP). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
d1dd5f6 to
ab0387f
Compare
2 participants
What does this PR do?
Type of change: New feature
Add Minitron pruning support for MoE models loaded with the efficient fused grouped GEMM experts (
TEGroupedMLP), in addition to the existingTESequentialMLPpath._DynamicTEGroupedLinear+_DynamicTEGroupedMLPdynamic modules slice the per-expert grouped weights (num_moe_experts/moe_ffn_hidden_size/hidden_size) and reorder/drop experts vianum_gemms, mirroring the SequentialMLP path with a minimal DynamicModule diff.moe_ffn_hidden_sizeis pruned across experts (SequentialMLP keeps one per expert). Same pruned width and independent per-expert weights; only the kept-channel index set is shared.examples/megatron_bridge/prune_minitron.pyuses grouped GEMM by default (--no_moe_grouped_gemmto fall back to SequentialMLP).Usage
torchrun --nproc_per_node 2 prune_minitron.py \ --hf_model_name_or_path <moe-model> \ --prune_target_active_params 3e9 \ --output_hf_path /tmp/pruned # add --no_moe_grouped_gemm to use SequentialMLPTesting
Verified in an
nvcr.io/nvidia/nemo:26.06+ MBridge main (as of 23 Jul) mounted so it mimics nemo:26.08 behavior.GPT MoE dynamic-module + pruning + parameter-sorting tests parametrized over both expert impls.
Mamba hybrid NAS metric tests: params-based now covers grouped GEMM, memory-based stays SequentialMLP (exact param counts / top-k / search-space-size assertions hold identically).
NemotronH end-to-end
test_prune_minitronexercises real-forward NAS without grouped GEMM (to be enabled in 26.08 container)Compared Nemotron-3-Nano-30B-A3B pruning: SequentialMLP vs GroupedGEMM on 4x B300 (accuracy + speed)
Top-10 Pruned Candidates — MMLU Scores
Before your PR is "Ready for review"
Summary by CodeRabbit
Summary by CodeRabbit
--no_moe_grouped_gemmto force the sequential expert path when required.--prune_score_funcparsing for MMLU to accept onlymmlu_<N>pct_bs<bs>.