remove GroupedMLP

[dimapihtar]

What does this PR do ?

Removes GroupedMLM usage since it's going to be removed from MLM: NVIDIA/Megatron-LM#3410

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Resolved an artificial restriction preventing delay_wgrad_compute from being used with the legacy grouped GEMM computation path.

• Refactor

  • Removed unnecessary legacy grouped GEMM configuration parameters to streamline model specification setup.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[dimapihtar]

/ok to test dabcf34

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c8eefcf and dabcf34.

📒 Files selected for processing (4)

• src/megatron/bridge/training/comm_overlap.py
• src/megatron/bridge/training/mlm_compat/model.py
• tests/unit_tests/training/mlm_compat/test_model.py
• tests/unit_tests/training/test_comm_overlap.py

💤 Files with no reviewable changes (4)

• src/megatron/bridge/training/comm_overlap.py
• tests/unit_tests/training/mlm_compat/test_model.py
• src/megatron/bridge/training/mlm_compat/model.py
• tests/unit_tests/training/test_comm_overlap.py

---

📝 Walkthrough

Walkthrough

The changes remove constraints and parameter passing related to legacy grouped GEMM implementation. Specifically, an assertion preventing delay_wgrad_compute with legacy groupedGEMM is eliminated, and the moe_use_legacy_grouped_gemm argument is removed from model layer spec constructors and corresponding test configurations.

Changes

Cohort / File(s)	Summary
Legacy GEMM Constraint Removal src/megatron/bridge/training/comm_overlap.py	Removes assertion preventing delay_wgrad_compute from being enabled with legacy grouped GEMM implementation, allowing both features to be used together.
Model Spec Parameter Updates src/megatron/bridge/training/mlm_compat/model.py	Removes moe_use_legacy_grouped_gemm argument from transformer layer spec calls in both TE-enabled and non-TE paths, eliminating explicit legacy grouped GEMM parameter passing at these call sites.
Test Configuration Updates tests/unit_tests/training/mlm_compat/test_model.py, tests/unit_tests/training/test_comm_overlap.py	Removes moe_use_legacy_grouped_gemm argument from test setup, mock configurations, and assertions to align with updated production code parameter passing.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• #2073: Modifies the same call sites in src/megatron/bridge/training/mlm_compat/model.py by removing the same keyword argument from layer spec constructors.

Suggested labels

cherry-pick, r0.3.0

Suggested reviewers

• oyilmaz-nvidia
• yaoyu-33
• ananthsub

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The PR title 'remove GroupedMLP' does not match the actual changes, which focus on removing moe_use_legacy_grouped_gemm arguments and delay_wgrad_compute assertions, not GroupedMLP.	Update the title to accurately reflect the main changes, such as 'Remove moe_use_legacy_grouped_gemm configuration' or 'Remove legacy grouped GEMM configuration parameters'.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Test Results For Major Changes	✅ Passed	PR removes deprecated internal parameters (moe_use_legacy_grouped_gemm) to align with upstream Megatron-LM changes. Only 15 lines removed, purely internal refactoring with no user-facing API breaking changes. Comprehensive unit tests updated and passing.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch dpykhtar/remove_grouped_mlp

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[dimapihtar]

/ok to test f0a7fa7