[FeSens]

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

What does this PR do?

The Transformer Engine fused MLP classes now apply the SwiGLU clamp (activation_func_clamp_value) or reject SiTU-GLU (activation_func_tanh_clamp_scale). Before this change, TEFusedMLP, TEFusedMLPWithGroupedLinear and FusedSharedExpertMLP silently ran plain SwiGLU when either option was set. TEFusedMLP and TEFusedMLPWithGroupedLinear now use TE's clamped SwiGLU ops. All three classes now raise a ValueError for SiTU-GLU, as the routed experts' fused path already declines it.

What's wrong

MLP, SharedExpertMLP and the routed experts support two optional clamps on the FC1 output:

• activation_func_clamp_value=c (DeepSeek-V4). For SwiGLU it computes silu(min(g, c)) * clamp(u, -c, c).
• activation_func_tanh_clamp_scale=s_g with an optional activation_func_tanh_clamp_scale_linear=s_l (Kimi-K3, TanH Sq-relu and SiTU GLU #7003). With gated silu this is SiTU-GLU: s_g*tanh(g/s_g)*sigmoid(g) * s_l*tanh(u/s_l). With a non-gated activation it is a tanh soft clamp of the input.

Three fused classes build their activation without reading these options:

• TEFusedMLP (a dense MLP with use_transformer_engine_op_fuser=True). _make_te_ops_activation maps (F.silu, gated) to te.pytorch.ops.SwiGLU whatever the clamp settings.
• TEFusedMLPWithGroupedLinear (use_grouped_gemm_for_dense_mlp=True). It always appends ScaledSwiGLU(glu_interleave_size=32).
• FusedSharedExpertMLP (use_grouped_gemm_for_shared_expert=True). It applies activation_func_clamp_value through ScaledClampedQGeGLU, but it never checks activation_func_tanh_clamp_scale, so a SiTU-GLU config runs plain SwiGLU.

None of them raises an error or a warning. The routed experts already handle both options. TEGroupedMLP._is_fused_impl_supported uses ScaledClampedQGeGLU for the clamp. For SiTU-GLU it returns False ("A clamped gated activation is SiTU-GLU, which the fused GLU path does not implement"), so those layers fall back to the unfused path.

Reproduction

# torchrun --nproc-per-node 1 repro_short.py   (from a Megatron-LM checkout, one GPU, TE >= 2.17)
import torch
import torch.nn.functional as F

from megatron.core.extensions.transformer_engine import (
    TEFusedMLP,
    TELayerNormColumnParallelLinear,
    TERowParallelLinear,
)
from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
from megatron.core.transformer.mlp import MLP, MLPSubmodules
from megatron.core.transformer.transformer_config import TransformerConfig
from tests.unit_tests.test_utilities import Utils

Utils.initialize_model_parallel(1, 1)
model_parallel_cuda_manual_seed(123)
submodules = MLPSubmodules(linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear)
for act in (
    {"activation_func_clamp_value": 10.0},  # DeepSeek-V4 SwiGLU clamp
    {"activation_func_tanh_clamp_scale": 4.0, "activation_func_tanh_clamp_scale_linear": 25.0},  # K3
):
    config = TransformerConfig(
        num_layers=1, hidden_size=64, ffn_hidden_size=112, num_attention_heads=4, num_moe_experts=4,
        activation_func=F.silu, gated_linear_unit=True, add_bias_linear=False,
        bias_activation_fusion=False, **act,
    )
    mlp = MLP(config, submodules, ffn_hidden_size=112).cuda()
    with torch.no_grad():
        mlp.linear_fc1.weight.normal_(0.0, 1.0)  # FC1 outputs with std ~8, so the clamps are active
    try:
        fused = TEFusedMLP(config, submodules, ffn_hidden_size=112).cuda()
    except ValueError as e:
        print(act, "ValueError:", e)
        continue
    fused.load_state_dict(mlp.state_dict())
    x = torch.randn(32, 2, 64, device="cuda")
    print(act, "max |TEFusedMLP - MLP|:", (fused(x)[0] - mlp(x)[0]).abs().max().item())
Utils.destroy_model_parallel()

Results:

• On main (8e13411), this prints max |TEFusedMLP - MLP|: 19.77 for the clamp and 19.50 for SiTU-GLU.
• With this PR, it prints 0.0 for the clamp, and SiTU-GLU raises ValueError: TEFusedMLP does not support activation_func_tanh_clamp_scale: ....

A longer script also compares gradients and covers the other two classes. Each fused module and its reference have the same weights. The relative error is ||a - b|| / ||b||.

Path and config	main	this PR
TEFusedMLP, SwiGLU (control), fp32	op SwiGLU; output, dx and dW_fc1 rel err 0	same
TEFusedMLP, clamp 10, fp32	op SwiGLU; output rel err 0.616 (max abs 14.45); dx 0.679; dW_fc1 0.696	op ClampedSwiGLU; output 0; dx 1.1e-6; dW_fc1 1.1e-6
TEFusedMLP, SiTU-GLU 4/25, fp32	op SwiGLU; output rel err 1.92 (max abs 20.98); dx 1.68; dW_fc1 1.74	ValueError
FusedSharedExpertMLP, SiTU-GLU 4/25, bf16 + NVFP4	builds ScaledSwiGLU; vs SharedExpertMLP with SiTU-GLU: 1.219; vs SharedExpertMLP with SwiGLU: 0.248	ValueError
TEFusedMLPWithGroupedLinear, clamp 10, bf16 + NVFP4	builds ScaledSwiGLU; vs MLP with the clamp: 1.456; vs MLP without it: 0.247	builds ScaledClampedQGeGLU; vs MLP with the clamp: 0.277; vs MLP without it: 0.663

In the NVFP4 rows, the fused path's error against the reference it matches (0.25 to 0.28) is mostly quantization error. For comparison, the clamped and unclamped MLP references differ by 1.45 on these inputs, and the SiTU-GLU and SwiGLU references differ by 1.18. The FC1 weights are drawn so that FC1 outputs have a std of about 6 to 11, which keeps the clamps active on many elements.

Impact

• A config is affected when it sets activation_func_clamp_value or activation_func_tanh_clamp_scale and also enables one of these fused paths:

  • use_transformer_engine_op_fuser=True for dense MLP layers;
  • use_grouped_gemm_for_dense_mlp=True;
  • use_grouped_gemm_for_shared_expert=True.

  Routed experts are not affected.

• These layers train and evaluate with a different activation from the one configured, and nothing reports it. On the inputs above, the outputs differ from the configured activation by 62% to 192% (relative norm), and the gradients by 68% to 174%.

• Weights trained through a fused path behave differently when loaded into the unfused classes, for example MLP, which apply the configured activation.

• I have not measured the effect on a real training run.

Fix

• _make_te_ops_activation takes the clamp value. For a clamped SwiGLU it returns te.pytorch.ops.ClampedSwiGLU(limit=c, alpha=1.0, glu_linear_offset=0.0), which computes silu(min(g, c)) * clamp(u, -c, c), the same as MLP. The config already requires glu_linear_offset == 0 with the SwiGLU clamp. TEFusedMLP._make_fused_impl passes activation_func_clamp_value.

• When the clamp is set, TEFusedMLPWithGroupedLinear._make_fused_impl builds ScaledClampedQGeGLU(glu_interleave_size=32, alpha=1.0, limit=c, glu_linear_offset=0.0). These are the same arguments FusedSharedExpertMLP and TEGroupedMLP use.

• TEFusedMLP.__init__, which TEFusedMLPWithGroupedLinear inherits, raises two errors:

  • a ValueError when activation_func_tanh_clamp_scale is set;
  • a RuntimeError when the clamp is set with TE < 2.17.0.dev0. That is the bound FusedSharedExpertMLP and TEGroupedMLP use for glu_linear_offset.

  These checks run at construction, because the fused ops are only built on the first forward.

• FusedSharedExpertMLP._validate_fused_grouped_swiglu raises a ValueError when activation_func_tanh_clamp_scale is set, next to its other config checks.

I chose to reject SiTU-GLU rather than use TE's SiTUGLU and ScaledSiTUGLU ops (TE 2.20) for three reasons:

• The routed experts do not use those ops yet.
• Megatron allows the gate scale without the linear scale. The TE ops require a finite, positive beta2 for the linear half, so they cannot express that case.
• [MCore] Add SiTU-GLU through Transformer Engine #6674 is the open PR that adds SiTU-GLU through TE.

This PR stops the silent fallback until #6674 or similar work lands. That work would replace these ValueErrors.

Tests

• New test_te_fused_mlp_activation.py:
  • TEFusedMLP must reject SiTU-GLU, with and without the linear scale.
  • With the clamp at 10, TEFusedMLP must match MLP in fp32 on the same weights: the output, the input gradient and all parameter gradients, within 1e-4.
  • On main, the two rejection tests fail with DID NOT RAISE. The match test fails on 2048 of 2048 elements, with a max abs difference of 19.77.
• test_te_fused_mlp_with_grouped_linear_spec.py:
  • With the fake TE ops, _make_fused_impl must build ScaledClampedQGeGLU with the arguments above when the clamp is set.
  • With the real TE class, SiTU-GLU must raise.
• test_shared_experts.py:
  • One more case in test_validate_fused_grouped_swiglu_rejects_unsupported_configs.
  • test_fused_shared_expert_rejects_situ_glu, which constructs the real FusedSharedExpertMLP with MXFP8 and SiTU-GLU. It only checks construction: my GPU has no MXFP8, so I could not test a forward pass there.
• determinism/kernels/test_te_wrappers.py:
  • test_te_fused_mlp_clamped_swiglu_replays, a bit-exact replay of TEFusedMLP with the clamp (bf16, clamp 1.0).
  • It passes before and after the fix; on main it runs the unclamped op.

Results on one GPU (RTX 5070 Ti, sm_120, torch 2.14.1+cu130, TE 2.20.2). I ran them on 8e13411. The branch is now based on 6463d38, where none of the touched files has changed:

Test set	Before	After
New and changed cases	7 failed, 6 passed	13 passed
Five related files	11 failed, 99 passed, 22 skipped	4 failed, 106 passed, 22 skipped

The five related files are test_te_fused_mlp_activation.py, test_te_fused_mlp_with_grouped_linear_spec.py, moe/test_shared_experts.py, fusions/test_te_fused_ops.py and determinism/kernels/test_te_wrappers.py.

The 4 remaining failures are the same before and after:

• Three test_shared_expert_forward_backward parametrizations need more than one GPU (world_size (1) is not divisible by 4 / 2).
• test_te_lm_head_two_uses_replays uses MXFP8, which TE does not support on sm_120.

The skipped tests need SM100, a world size divisible by EP=2, or a TE GroupedLinear with split_alignment.

black 26.3.0, isort 5.13.2, ruff 0.9 and pylint 3.2.6 pass on the changed files with the repo's settings.

The RTX 5070 Ti limits these results in four ways:

• The new tests use ffn_hidden_size=112, which is not a multiple of 32. On sm_120, TE 2.20.2 sends GLU activations whose column count is divisible by 32 to cast_gated_tma, which fails with CUDA Error: invalid argument. This happens for plain SwiGLU too, so it is unrelated to this PR. TE main checks the shared-memory limit before it selects that kernel.
• sm_120 has no MXFP8, so the NVFP4 rows above stand in for it. The B300 run below covers MXFP8 and the cuDNN fused grouped-MLP kernel.
• On main, the repro's FusedSharedExpertMLP forward needed a workaround that only applies to the repro. TE 2.20.2's unfused ScaledSwiGLU asserts scales_tensor.numel() == rows, but FusedSharedExpertMLP caches a 1-element unit scale, so I preset that scale to one entry per row.
• It has one GPU. The 8x H100 run below covers the multi-GPU cases.

Tested on B300 (MXFP8) and 8x H100. Same software as above (torch 2.14.1+cu130, TE 2.20.2, cudnn-frontend 1.32.0), main 6463d38 with this PR's tests against this PR's head:

• B300 (sm_103), MXFP8, NVTE_CUTEDSL_FUSED_GROUPED_MLP=1. Every case below ran on the cuDNN fused kernel (GroupedMLP_CuTeGEMMGLU, forward and backward).
  • TEFusedMLPWithGroupedLinear with the clamp at 10, relative error against MLP with the clamp (output / input gradient): 1.448 / 1.327 on main, where it matches the unclamped MLP (0.066 / 0.083), and 0.076 / 0.183 with this PR. The SwiGLU control is 0.066 / 0.083 on both.
  • FusedSharedExpertMLP with SiTU-GLU 4/25: on main the forward runs and is 1.179 off SharedExpertMLP with SiTU-GLU, 0.066 off plain SwiGLU. With this PR it raises ValueError.
  • Five related files: 10 failed, 100 passed, 22 skipped on main; 3 failed, 107 passed, 22 skipped with this PR. The 7 differences are the new tests, and the 3 failures are the multi-GPU test_shared_expert_forward_backward cases. test_te_lm_head_two_uses_replays passes on both. The skips are the same as on sm_120: the batch-invariant grouped-GEMM cases require exactly sm_100, and the others need EP=2 or TE features newer than 2.20.2.
• 8x H100: all four test_shared_expert_forward_backward parametrizations pass on all 8 ranks, before and after.

The input-gradient error with the clamp (0.183 against 0.083 for the control) comes from MXFP8. The gradient is zero on one side of the clamp, and rounding the FC1 operands moves some pre-activations across it. With the clamp at 10 and FC1 outputs of std ~11, 1522 of 131072 clamp decisions change, which alone changes dx by 0.161 (relative). Against an fp32 reference that takes its clamp decisions from MXFP8-quantized operands, the kernel's dx error is 0.106. With the clamp at 30 it is 0.084 to 0.089, the control level. An MXFP8 test of a clamped activation therefore needs a looser dx tolerance than an unclamped one, or a reference whose clamp decisions come from the quantized operands. None of the tests in this PR compares a clamped gradient under MXFP8.

tools/check_kernel_determinism_coverage.py reports one issue. transformer_engine.py is also a source of the paged_stash_marker_propagation entry, and that entry's test, test_fused_activations.py, is unchanged. The new test_te_wrappers.py case covers the file's other two entries, transformer_engine_wrappers and te_fused_cross_entropy. This change does not touch paged stashing. Could a maintainer add the determinism-exempt label? I can't add labels myself. If you prefer, I can add a case to test_fused_activations.py instead.

Not covered:

• GEGLU or ReGLU with activation_func_clamp_value. The unfused MLP clamps any gated activation, but TEFusedMLP still ignores the clamp for these two. The config documents the clamp only for SwiGLU and quick-GELU, which the TE op path does not support. MLP's own fused bias_geglu_impl also ignores the clamp, so I left this as is.
• The activation_func_tanh_clamp_scale docstring still mentions only the routed-expert fused path.

Related open PRs:

• [MCore] Add SiTU-GLU through Transformer Engine #6674 (changes requested, last updated 2026-09-16) adds SiTU-GLU through TE with its own config (situlu, situ_glu_beta1, situ_glu_beta2). That config predates activation_func_tanh_clamp_scale from TanH Sq-relu and SiTU GLU #7003. [MCore] Add SiTU-GLU through Transformer Engine #6674 edits the same functions in transformer_engine.py and shared_experts.py, so whichever lands second needs a rebase.
• Use the MLP's TP group in TE fused MLPs and require a CP group for CP attention #8038 changes TEFusedMLP's TP group, and Use the shared expert's TP group in the fused SwiGLU path #8024 and Fix the FC1 size of the fused shared expert with tensor parallel #8145 change FusedSharedExpertMLP. The source changes in this PR apply to each of their heads with git apply --3way without conflicts. One or two of the test files conflict on each head. Those heads are older than this PR's base, and I have not checked which conflicts come from the PRs themselves.

The code, the tests and this description were drafted with an AI coding agent (Claude Code).

⚠️ 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: none. I found no existing issue or PR for this bug. The searches are listed below.

Duplicate search

I ran read-only searches of issues and PRs in all states on 2026-10-28, with gh search issues --include-prs --repo NVIDIA/Megatron-LM "<query>" --limit 15:

Query	Hits	Relevant results
TEFusedMLP clamp	0	none
TEFusedMLP activation	0	none
fused MLP clamp ignored	0	none
FusedSharedExpertMLP	7	#8024 and #8145 (open; other FusedSharedExpertMLP bugs), #5604 (merged; adds the class). None reports this bug.
activation_func_tanh_clamp_scale	5	#7003 (merged; adds the option), #7496 (merged; fused tanh-SReLU for routed experts). None touches these classes.
activation_func_clamp_value	12	#5940 (merged; clamped SwiGLU across the expert paths, not the dense fused MLP), #6918 (open issue; a different clamp bug under Inductor), #8566 (open; inference experts)
SiTU	15	#6674 (open; see above), #6673 (merged on dev), #8497 (open on dev; MoK kernels, does not touch these files), #6635 and #6636 (closed)
ClampedSwiGLU	6	#5130 (merged on dev; the clamp in the routed-expert op fuser, experts.py only), #5941 (closed)
use_transformer_engine_op_fuser	15	none about the activation
TEFusedMLPWithGroupedLinear	8	#8038 (open; TP group)

I also listed the files of #6674, #8497, #8024, #8145, #8038, #8566, #5130, #5940, #7003, #7769 and #6810 with gh pr view <n> --json files. Only #6674 changes the activation in these classes.

Code provenance

You should always make a best effort to avoid copying code from other repositories. In exceptional cases in which this is unavoidable and you copied or adapted code from an external source, provide the following details and copy to the file header:

No code was copied or adapted from an external source.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• 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"

1. When your PR is ready, click Ready for Review.
2. 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]

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.