[Dev] Add SiTU-GLU through Transformer Engine

[harryzhou2000]

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

What does this PR do?

Add Kimi K3 SiTU-GLU as a global FFN activation on dev.

Paired main-target PR: #6674

Activation

For gate projection G, up projection U, and K3 defaults beta1=4, beta2=25:

T_g = beta1 * tanh(G / beta1) * sigmoid(G)
T_u = beta2 * tanh(U / beta2)
Y   = T_g * T_u

Reference: Kimi K3 technical report, Figure 4 and Equation 12.

Design

• Add --situ-glu, with --moe-use-situ-glu as an alias, and apply it consistently to dense, routed-expert, and shared-expert FFNs.
• Add situ_glu_beta1=4 and situ_glu_beta2=25 to TransformerConfig, with validation that SiTU-GLU is gated, unclamped, uses zero GLU linear offset, and has finite positive beta values.
• Add situlu as the correct PyTorch pointwise reference and configuration marker. It serves non-TE/custom module paths until PyTorch provides a dedicated torch.nn.functional.situlu-style operation; unary F.silu is not equivalent.
• Use transformer_engine.pytorch.ops.SiTUGLU for ordinary dense, sequential-expert, and ordinary shared-expert paths.
• Use transformer_engine.pytorch.ops.ScaledSiTUGLU for fused grouped dense, routed-expert, and shared-expert paths.
• Fail clearly when a requested TE path does not expose the required SiTU-GLU operator. cuDNN frontend capability selection remains inside TE.
• Extend checkpoint conversion, default FFN sizing, FLOP accounting, and memory estimation to treat SiTU-GLU as a gated GLU.

The required TE interface is available in NVIDIA/TransformerEngine#3402.

With that TE interface, an older cuDNN frontend without SiTU parameters uses TE's unfused GroupedLinear -> ScaledSiTUGLU -> GroupedLinear operation sequence. A SiTU-capable cuDNN frontend enables fused grouped GEMM + SiTU-GLU without an MCore-side backend check.

Test results

Validated on NVIDIA B300 with NVIDIA/TransformerEngine#3402 and cuDNN Frontend develop through NVIDIA/cudnn-frontend#670, using CuTe DSL 4.6.2:

• final focused SiTU-GLU and grouped-MLP suite on the review-fixed tip: 41 passed;
• broader selected dev suites: 114 passed, 11 skipped, 4 deselected; the deselected cases are existing multi-rank shared-expert parameterizations that cannot run on one GPU;
• BF16 real MCore MoE routed forward/backward through ScaledSiTUGLU: passed;
• MXFP8 real MCore MoE forward/backward: routed and fused-shared ScaledSiTUGLU, with the fused cuDNN/CuTe grouped-MLP operation asserted: passed;
• NVFP4 routed-expert forward/backward through ScaledSiTUGLU, with the fused Hadamard forward operation asserted: passed;
• git diff --check, Python compilation, and Ruff on all changed source and test files: passed;
• SSH commit signature and DCO: passed.

Fused-shared NVFP4 currently reaches an existing TE FC1 weight-scale reshape failure that reproduces with an equivalent SwiGLU control, so this PR does not claim that combination. MXFP8 dSiTU validation used the cuDNN frontend correction included in NVIDIA/cudnn-frontend#670; the published cuDNN frontend 1.27.0 package tested here does not yet contain that correction.

Scope and dependencies

This PR depends on NVIDIA/TransformerEngine#3402 for TE-backed SiTU-GLU functionality. The cuDNN frontend support is optional for correctness and required only for fused block-scaled grouped execution. A functional test is deferred until that TE interface is available in Megatron-LM CI; the focused MCore unit and GPU suites cover the current integration boundary.

This PR is independent of Quantile Balancing and #6637.

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: TODO — link a Megatron-LM feature-request issue before marking this draft ready for review.

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"

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.

[yuzhongw-nvidia]

LGTM. Thanks for your outstanding work!

[yuzhongw-nvidia]

/ok to test 4edf750

[yuzhongw-nvidia]

/claude review

[claude]

This gate is broader than SiTU-GLU and looks like an unintended regression for the existing SwiGLU/quick-GEGLU paths.

use_glu_fusion is true for F.silu and quick_gelu as well, so any run that combines --use-transformer-engine-op-fuser with --recompute-modules [... moe_act] now reports unsupported. Since __init__ does assert self._is_fused_impl_supported() when use_transformer_engine_op_fuser is set (experts.py:277-280), that turns into a hard failure for configurations that work on main today.

It also contradicts _make_fused_ops a few lines below, which still computes activation_recompute_in_mlp and forwards it to ScaledSwiGLU / ScaledClampedQGeGLU / ScaledSReLU when TE exposes the parameter — that plumbing becomes dead code.

Only ScaledSiTUGLU is constructed without an activation_recompute_in_mlp kwarg, so scoping the check to SiTU-GLU keeps the new activation safe without changing existing behavior:

Suggested change

	if use_glu_fusion and self.activation_recompute:
	return _unsupported(
	"Transformer Engine scaled GLUs do not support activation recompute"
	)
	if self.config.activation_func is situlu and self.activation_recompute:
	return _unsupported(
	"Transformer Engine ScaledSiTUGLU does not support activation recompute"
	)

[harryzhou2000]

Fixed in 74aafeb52. The recompute guard is now scoped only to situlu/ScaledSiTUGLU. Existing SwiGLU and quick-GEGLU paths retain their prior activation-recompute plumbing, with a parametrized regression test that sets activation_recompute=True for both. The final focused suites pass on dev (41 tests) and main (60 tests).

[claude]

jit_fuser resolves to torch.compile on torch >= 2.2 (megatron/core/jit.py:19-21), so situlu is a compiled callable. Two consequences worth checking:

1. Every other function in this file is unary; situlu is the first one taking extra float args. The three call sites pass betas positionally (situlu(x, beta1, beta2)), which torch.compile handles, but the default values 4.0/25.0 become recompile-triggering guards rather than free constants — worth confirming a config with non-default betas doesn't cause a recompile per call.

2. Identity comparisons (self.activation_func is situlu, config.activation_func == situlu) are load-bearing throughout this PR. They compare against the decorated wrapper, which is consistent as long as nothing re-binds situlu after disable_jit_fuser() / enable_jit_fuser() runs. megatron/core/jit.py rebinds the module-global jit_fuser, not previously-decorated functions, so this holds — just flagging that the whole feature's dispatch rests on it.

[harryzhou2000]

Confirmed. jit_fuser decorates situlu once, and toggling the module-global fuser does not rebind that callable, so the identity-based dispatch remains stable. The beta values are configuration invariants: each call site reuses the same pair for the lifetime of a config/module; a different config may compile a separate specialization, but the values do not vary per invocation. Non-default betas remain covered by the configuration/op-selection tests, so no code change was needed for this observation.

[claude]

Light review. Nice, thorough PR — validation coverage and the config gating are good. Two inline comments (one likely regression), plus the notes below.

1. --situ-glu unconditionally forces use_te_activation_func=True

megatron/training/arguments.py:2248-2252 and the duplicate in megatron/training/argument_utils.py:431-435:

if use_situ_glu:
    kw_args['activation_func'] = situlu
    kw_args['gated_linear_unit'] = True
    kw_args['use_te_activation_func'] = True
    kw_args['bias_activation_fusion'] = False

This makes the PyTorch reference path unreachable from the CLI, even though the PR description says "other paths use the PyTorch reference" and test_transformer_config_accepts_pytorch_situ_glu_fallback explicitly covers use_te_activation_func=False. It also silently overrides use_te_activation_func for anyone who set it via --transformer-impl local style configs, and gpt_layer_specs.py:237-238 raises AssertionError("use_te_activation_func not compatible with using kitchen") — so --situ-glu plus kitchen becomes an unexplained assertion rather than a clear message.

Consider not forcing the flag, and letting the existing TransformerConfig.__post_init__ validation handle the rest:

if use_situ_glu:
    kw_args['activation_func'] = situlu
    kw_args['gated_linear_unit'] = True
    kw_args['bias_activation_fusion'] = False

2. Checkpoint converters in tools/checkpoint/ are not covered

megatron/training/checkpointing.py was updated for the FSDP-DTensor path, but the tools/checkpoint/ converters still key on md.swiglu alone:

• tools/checkpoint/loader_base.py:446 — md.swiglu = self.margs.swiglu
• tools/checkpoint/loader_base.py:321,337 — chunks mlp_l0_weight / mlp_l0_bias into W/V halves only if self.md.swiglu
• tools/checkpoint/saver_base.py:539,573 — the matching re-concat

A SiTU-GLU checkpoint has the same doubled FC1 layout, so converting one through these paths will take the non-gated branch and produce a silently wrong mlp l0 weight. The PR description says "Extend checkpoint conversion ... to treat SiTU-GLU as a gated GLU," which is done for checkpointing.py but not here. Suggested fix in loader_base.py:

md.swiglu = self.margs.swiglu or getattr(self.margs, "situ_glu", False)

That keeps all downstream md.swiglu branches (loader + saver) correct without touching them individually.

3. Pre-existing bug reinforced in TEActivationOp

Not introduced here, but the new elif is now guarded behind it. megatron/core/extensions/transformer_engine.py:481-486:

if config.activation_func == F.silu:
    layer_type = te.pytorch.ops.SwiGLU
elif config.activation_func == F.gelu:
    layer_type = te.pytorch.ops.GEGLU
elif config.activation_func == F.silu:   # <-- duplicate; should be F.relu -> ReGLU
    layer_type = te.pytorch.ops.ReGLU
elif config.activation_func is situlu:

The third branch repeats F.silu so ReGLU is dead (same duplication in the non-gated else block at line 497-498 for ReLU). The situlu branch is fine because it comes after, but the error message this PR edits now advertises ReGLU support that cannot be reached. Fixing the two F.silu typos to F.relu would be a small, contained cleanup if you want it in scope.

4. Test coverage note

Unit coverage for the op selection and config validation is solid. Two gaps:

• No functional test. tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser/ is the natural template (it already sets NVTE_CUTEDSL_FUSED_GROUPED_MLP: 1 and --use-transformer-engine-op-fuser: true); a sibling case swapping --swiglu: true for --situ-glu: true would exercise the routed-expert path end to end. Understandable to defer while the TE dependency (TransformerEngine#3402) is unmerged — worth a note in the PR body so it isn't lost.
• test_situlu_reference_matches_kimi_bf16_precision_and_backward asserts with rtol=0, atol=0 against an expression rebuilt from the same ops in the same order. That pins the FP32-intermediate/BF16-output contract, which is the stated intent, but it will not catch a formula error since the reference mirrors the implementation. A handful of hand-computed expected values at a few points (including large-magnitude inputs where the beta2 soft-cap saturates) would make it a real correctness check against Equation 12.

[harryzhou2000]

Addressed the current review findings on both SiTU-GLU PRs in 74aafeb52 (dev) and 89b265c25 (main):

• Scoped the activation-recompute rejection to ScaledSiTUGLU only, with regressions proving existing SwiGLU and quick-GEGLU recompute support is unchanged.
• Marked checkpoint-converter FC1 metadata as gated for either SwiGLU or SiTU-GLU.
• Added independent fixed-value coverage, including positive and negative saturation.
• Applied the repository's Black 26.3.0/isort/Ruff formatting checks.

I kept the CLI's TE selection intentional: stock SiTU-GLU paths require the TE operators, while situlu remains the numerically correct pointwise reference/configuration marker for direct/custom use. The pre-existing TEActivationOp ReGLU/ReLU typo is unrelated and left out of this focused change. A functional case is noted as deferred until NVIDIA/TransformerEngine#3402 is available in Megatron-LM CI.

Final focused GPU suites on NVIDIA B300: dev 41 passed; main 60 passed.

[FDecaYed]

--use-checkpoint-args does not currently restore situ_glu. load_args_from_checkpoint() restores squared_relu and swiglu, but never calls _set_arg("situ_glu", force=True). As a result, loading a SiTU-GLU checkpoint with --use-checkpoint-args reconstructs it as a non-gated MLP unless the user also repeats --situ-glu, which leads to an FC1 shape mismatch. Could we restore situ_glu from the checkpoint args as well?