Support ScaledSReLU in TE grouped MLP fuser - #4859
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review Summary
CRITICAL: 0 | IMPORTANT: 1 | SUGGESTION: 2
Findings
-
[IMPORTANT Compatibility] — The
self.activation_func→self.config.activation_funcchange in_is_fused_impl_supportedand_make_fused_opsis broader than just ScaledSReLU support. It also enables the fused TE op path for SwiGLU and quick_gelu whenuse_te_activation_func=True, which was previously blocked (assertion failure). This is a correct fix — the fused TE ops handle activation themselves — but it's a behavior change for existing configurations that deserves mention in the PR description. -
[SUGGESTION Simplification] —
use_srelu_fusiondoesn't guard againstgated_linear_unit=True. Addingand not self.config.gated_linear_unitwould be a cheap defensive check against a misconfiguration whereScaledSReLU(non-GLU) would be used with doubled FC1 output. Not a regression (the same gap exists in the non-fused path at line 661), but easy to fix while the block is being refactored. -
[SUGGESTION] — The test at line 526 (
module.activation_func = object()) is a good pattern that verifies_is_fused_impl_supportedreads fromconfig.activation_funcrather thanself.activation_func.
Overall Assessment
Low risk. The core implementation is correct — ScaledSReLU is properly gated behind use_fused_weighted_squared_relu, the TE import check (try/except ImportError) follows the existing pattern for ScaledClampedQGeGLU, and the tests cover both the _is_fused_impl_supported gate (parametrized True/False) and _make_fused_ops construction. The self.config.activation_func switch is a sensible fix that makes the fused path work correctly regardless of use_te_activation_func. No dtype, parallelism, or checkpoint compatibility concerns.
|
/ok to test 68349b0 |
|
/ok to test b079324 |
|
/ok to test e15f4eb |
|
/ok to test d4dad68 |
|
/ok to test 1292dcc |
|
/ok to test 092f812 |
@sraman-rgb, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
|
/ok to test dcf3a7e |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26656553469 |
10 participants
What does this PR do ?
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.