Support efficient TEGroupedMLP (moe_grouped_gemm=True) for Minitron Pruning#1955
Support efficient TEGroupedMLP (moe_grouped_gemm=True) for Minitron Pruning#1955kevalmorabia97 wants to merge 1 commit into
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds grouped-GEMM MoE expert support (TEGroupedMLP) to the dynamic module and pruning frameworks, introducing new dynamic classes, importance estimation for grouped weights, a CLI flag in the pruning example script, and corresponding test coverage for both grouped and sequential MoE paths. ChangesGrouped GEMM MoE Pruning
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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 76.33% 76.20% -0.13%
==========================================
Files 522 522
Lines 58452 58542 +90
==========================================
- Hits 44618 44611 -7
- Misses 13834 13931 +97
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.
kevalmorabia97
force-pushed
the
kmorabia/minitron-moe-grouped-gemm
branch
from
21269a2 to
fc29d46
Compare
|
/ok to test fc29d46 |
kevalmorabia97
force-pushed
the
kmorabia/minitron-moe-grouped-gemm
branch
from
fc29d46 to
5db72ad
Compare
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>
kevalmorabia97
force-pushed
the
kmorabia/minitron-moe-grouped-gemm
branch
from
5db72ad to
1073a71
Compare
1 participant
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.06container (TE 2.16, 2x RTX 6000 Ada):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)Compare Nemotron-3-Nano-30B-A3B pruning: SequentialMLP vs GroupedGEMM (accuracy + speed)
Before your PR is "Ready for review"
Summary by CodeRabbit
New Features
Documentation