TanH Sq-relu and SiTU GLU by abhinav-khattar · Pull Request #7003 · NVIDIA/Megatron-LM

abhinav-khattar · 2026-08-31T22:48:46Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

TanH Sq-relu and SiTU GLU activations

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/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"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
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, the Final Review label 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 Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-08-31T22:48:49Z

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.

github-actions · 2026-08-31T22:48:58Z

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:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

wdykas · 2026-09-01T18:34:10Z

/claude review

zhongbozhu · 2026-09-01T18:44:42Z

@@ -996,26 +998,41 @@
                    bias_parallel is None
                ), "Bias is not supported with fused weighted squared relu."
                intermediate_parallel = weighted_squared_relu_impl(


what's the plan for fused kernel support when enabling use_transformer_engine_op_fuser?

I think we may be silently dropping the tanh part here when going through this path?

I think with use_transformer_engine_op_fuser are constructing ScaledSwiGLU or ScaledSReLU without reading either tanh scale.

edited to disable use_transformer_engine_op_fuser if tanh_clamp is set. we can update once we have the cudnn fe kernel ready

sidsingh-nvidia · 2026-09-01T18:48:34Z


-def _get_activation_func(activation_type: ActivationType, fused_quant: bool = False) -> Callable:
+def _get_activation_func(
+    activation_type: ActivationType, fused_quant: bool = False, clamp_scale: Optional[float] = None


can we have **activation_kwargs as a generic interface to pass activation function specific args like clamp_scale?

claude · 2026-09-01T18:56:31Z

Light review. Two findings.

1. `activation_func_tanh_clamp_scale` is silently ignored on the TE op-fuser paths

The PR is careful to hard-fail on the inference backends that can't express the clamp (_flashinfer_forward assert, vllm_fused_moe assert, _get_activation_func NotImplementedError), but the two training TE op-fuser paths have no such guard and will silently train the wrong activation:

TEGroupedMLP._make_fused_ops (megatron/core/transformer/moe/experts.py:551-627) builds ScaledSwiGLU / ScaledClampedQGeGLU / ScaledSReLU purely from activation_func + activation_func_clamp_value. None of those TE ops take a tanh clamp scale, so with use_transformer_engine_op_fuser=True the clamp / SiTU-GLU is dropped and the forward runs plain SwiGLU or squared-ReLU.
FusedSharedExpertMLP._make_fused_grouped_swiglu_ops (megatron/core/transformer/moe/shared_experts.py:509-516) has the same problem under use_grouped_gemm_for_shared_expert=True.

The new __post_init__ validation in transformer_config.py rejects use_te_activation_func and unsupported bias_activation_fusion, but doesn't cover either of these flags. Suggest extending the validation block (~transformer_config.py:2634):

if self.use_transformer_engine_op_fuser:
    raise ValueError(
        "activation_func_tanh_clamp_scale is not supported with "
        "use_transformer_engine_op_fuser; the TE ScaledSwiGLU / ScaledClampedQGeGLU / "
        "ScaledSReLU fused ops have no tanh clamp."
    )
if self.use_grouped_gemm_for_shared_expert:
    raise ValueError(
        "activation_func_tanh_clamp_scale is not supported with "
        "use_grouped_gemm_for_shared_expert; FusedSharedExpertMLP builds TE "
        "ScaledSwiGLU / ScaledClampedQGeGLU, which cannot apply the clamp."
    )

(Or add the checks to TEGroupedMLP._is_fused_impl_supported / _validate_fused_grouped_swiglu if you'd rather keep the failure local to the modules.)

2. No test coverage for the training-side kernels

All five new tests exercise inference Triton kernels (test_moe_permute.py, test_vllm_fused_moe.py). The training-side additions are untested:

fused_bias_swiglu.py: situ_glu, bias_situ_glu, weighted_situ_glu and their hand-derived backwards (_situ_glu_grads, situ_glu_back, bias_situ_glu_back, weighted_situ_glu_back).
fused_weighted_squared_relu.py: weighted_clamped_squared_relu + weighted_clamped_squared_relu_back.
The new __post_init__ validation errors.

The hand-written gradients are the highest-risk part of the diff — dy_1 = (sech2 + c * (1 - sig)) * sig * lin * g is easy to get subtly wrong and nothing here would catch it. tests/unit_tests/fusions/test_swiglu_fusion.py and tests/unit_tests/fusions/test_weighted_squared_relu_fusion.py already have the right shape for this: compare the fused impl (fwd + bwd) against an autograd reference built from megatron.core.activations.situ_glu / squared_relu(tanh_soft_clamp(x, s)), plus a clamp_scale=None case asserting bit-identity with the existing unclamped path. test_swiglu_fusion.py::test_clamped_swiglu_config_* is likewise the pattern to follow for the new config validation.

zhongbozhu · 2026-09-01T18:56:38Z

            .to(intermediate_parallel.dtype)
        )

    def _is_fused_impl_supported(self) -> bool:


Please change this function accordingly if needed

wdykas

I left a few comments but dont want my review to block anything so will put the accept here assuming everything is addressed

wdykas · 2026-09-01T18:47:53Z

+    out_dtype = x.dtype
+    t = _tanh_relu_over_scale(x, clamp_scale)
+
+    c = (clamp_scale * t).to(out_dtype)


This is odd down casting I am assuming like other things is for training and inference parity? we might want to note it

its to ensure the fused/unfused paths are identical, since the unfused path does a cast to out_dtype after the tanh activation which is a separate activation function

removed the intermediate downcasts

wdykas · 2026-09-01T18:55:10Z

@@ -996,26 +998,41 @@
                    bias_parallel is None
                ), "Bias is not supported with fused weighted squared relu."
                intermediate_parallel = weighted_squared_relu_impl(


I think we may be silently dropping the tanh part here when going through this path?

wdykas · 2026-09-01T18:56:03Z

@@ -996,26 +998,41 @@
                    bias_parallel is None
                ), "Bias is not supported with fused weighted squared relu."
                intermediate_parallel = weighted_squared_relu_impl(


I think with use_transformer_engine_op_fuser are constructing ScaledSwiGLU or ScaledSReLU without reading either tanh scale.

santhnm2 · 2026-09-01T22:00:10Z

/ok to test 49cd794

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

abhinav-khattar · 2026-09-02T16:50:17Z

/ok to test fa69213

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

abhinav-khattar · 2026-09-02T23:00:35Z

/ok to test d40cecb

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 · 2026-09-03T04:14:00Z

/ok to test 68e6005

nemo-automation-bot · 2026-09-03T09:44:25Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/33740494110

abhinav-khattar requested review from a team as code owners August 31, 2026 22:48

svcnvidia-nemo-ci marked this pull request as draft August 31, 2026 22:48

abhinav-khattar marked this pull request as ready for review September 1, 2026 18:20

abhinav-khattar requested a review from a team as a code owner September 1, 2026 18:20

github-actions Bot added the complexity: high label Sep 1, 2026

zhongbozhu requested changes Sep 1, 2026

View reviewed changes

sidsingh-nvidia requested changes Sep 1, 2026

View reviewed changes

zhongbozhu requested changes Sep 1, 2026

View reviewed changes

wdykas approved these changes Sep 1, 2026

View reviewed changes

sidsingh-nvidia approved these changes Sep 1, 2026

View reviewed changes

copy-pr-bot Bot had a problem deploying to test September 1, 2026 22:01 Error

abhinav-khattar force-pushed the tanh-activations branch from 49cd794 to 0ffd278 Compare September 1, 2026 22:02

abhinav-khattar requested a review from a team as a code owner September 1, 2026 22:02

abhinav-khattar requested a review from a team September 1, 2026 22:02

abhinav-khattar force-pushed the tanh-activations branch from 0ffd278 to 8e5af92 Compare September 1, 2026 22:12

abhinav-khattar and others added 6 commits September 1, 2026 15:15

Add Tanh soft clamping + Fused path

0ef9558

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

SITU GLU

ea53c9d

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

Add inference support for tanh

e6a023a

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Fix unit tests

aadb65f

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

remove cast in fused kernel, add allclose tests with unfused path

4aa04c4

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

Make inference rounding behavior match fused training kernel

d74d902

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

copy-pr-bot Bot temporarily deployed to public September 2, 2026 07:18 Inactive

abhinav-khattar enabled auto-merge September 2, 2026 08:12

fix arg name

fa69213

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

copy-pr-bot Bot temporarily deployed to test September 2, 2026 16:51 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 16:57 Inactive

fix wgrad tol for older swiglu fusion test

d40cecb

Signed-off-by: Abhinav Khattar <akhattar@nvidia.com>

copy-pr-bot Bot temporarily deployed to test September 2, 2026 23:01 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 23:04 Inactive

Fix test_flashinfer_mxfp8

68e6005

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

copy-pr-bot Bot temporarily deployed to test September 3, 2026 04:15 Inactive

copy-pr-bot Bot temporarily deployed to public September 3, 2026 04:17 Inactive

cuichenx approved these changes Sep 3, 2026

View reviewed changes

abhinav-khattar added this pull request to the merge queue Sep 3, 2026

Merged via the queue into NVIDIA:main with commit 9d8183b Sep 3, 2026
102 of 103 checks passed

abhinav-khattar deleted the tanh-activations branch September 3, 2026 11:18

lhb8125 mentioned this pull request Sep 7, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

FeSens mentioned this pull request Oct 28, 2026

fix(MoE): apply the SwiGLU clamp and reject SiTU-GLU in TE fused MLPs #8640

Draft

3 of 7 tasks

Conversation

abhinav-khattar commented Aug 31, 2026

What does this PR do?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

wdykas commented Sep 1, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude Bot commented Sep 1, 2026

1. activation_func_tanh_clamp_scale is silently ignored on the TE op-fuser paths

2. No test coverage for the training-side kernels

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wdykas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

abhinav-khattar Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

santhnm2 commented Sep 1, 2026

Uh oh!

abhinav-khattar commented Sep 2, 2026

Uh oh!

abhinav-khattar commented Sep 2, 2026

Uh oh!

santhnm2 commented Sep 3, 2026

Uh oh!

nemo-automation-bot Bot commented Sep 3, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants

1. `activation_func_tanh_clamp_scale` is silently ignored on the TE op-fuser paths

abhinav-khattar Sep 1, 2026 •

edited

Loading