Add MoE TEGroupedMLP support for Minitron Pruning

[kevalmorabia97]

What does this PR do?

Add support for MoE TEGroupedMLP (moe_grouped_gemm=True) which is supposed to be much more efficient than current SequentialMLP

Since kernels are different, pruned model may differ slightly

Testing

• Unit tests updated and passing
• Compare Qwen3-30B-A3B -> 24B pruned model MMLU
• Compare Nemotron3-Nano-30B-A3B -> 24B pruned model MMLU

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.).

• Is this change backward compatible?: ✅ Users can continue using SequentialMLP for bw compatibility or use efficient TEGroupedMLP
• 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?: ❌

Additional Information

[copy-pr-bot]

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.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ce7ff4c5-f7f8-4de8-8178-6957d736f7c7

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch kmorabia/minitron-grouped-gemm

📝 Coding Plan

• Generate coding plan for human review comments

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.08%. Comparing base (8225445) to head (50973d4).

Additional details and impacted files

@@                        Coverage Diff                         @@
##           kmorabia/minitron-full-te-spec    #1038      +/-   ##
==================================================================
- Coverage                           70.11%   70.08%   -0.03%     
==================================================================
  Files                                 221      221              
  Lines                               25459    25499      +40     
==================================================================
+ Hits                                17851    17872      +21     
- Misses                               7608     7627      +19     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.