Skip to content

[Dev] Add SiTU-GLU through Transformer Engine #6673

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
harryzhou2000 wants to merge 2 commits into
base: dev
Choose a base branch
from
Open

[Dev] Add SiTU-GLU through Transformer Engine #6673

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5c1d0f6
[Dev] Add SiTU-GLU activation support
harryzhou2000 Aug 20, 2026
79f681c
Fix SiTU-GLU error assertions
FDecaYed Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions megatron/core/activations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,20 @@ def quick_gelu(x: torch.Tensor) -> torch.Tensor:
def fast_gelu(x: torch.Tensor) -> torch.Tensor:
"""Fast GELU activation"""
return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))


@jit_fuser
def situlu(x: torch.Tensor, beta1: float = 4.0, beta2: float = 25.0) -> torch.Tensor:
Comment on lines +26 to +27

@claude claude Bot Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 harryzhou2000 Aug 25, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

"""Apply SiTU-GLU to contiguous gate/up halves of an FC1 output.

This is the slow PyTorch reference and config marker until PyTorch provides
a dedicated ``torch.nn.functional.situlu``-style operation. Unary
``F.silu`` is not equivalent because SiTU-GLU transforms both branches.
"""
input_dtype = x.dtype
gate, up = torch.chunk(x, 2, dim=-1)
gate = gate.float()
up = up.float()
gate = beta1 * torch.tanh(gate / beta1) * torch.sigmoid(gate)
up = beta2 * torch.tanh(up / beta2)
return (gate * up).to(input_dtype)
53 changes: 46 additions & 7 deletions megatron/core/extensions/transformer_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from torch.nn.parameter import Parameter
from typing_extensions import override

from megatron.core.activations import situlu
from megatron.core.dist_checkpointing.mapping import ShardedStateDict
from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
from megatron.core.enums import Fp4Recipe, Fp8Recipe
Expand Down Expand Up @@ -483,17 +484,28 @@ def __new__(cls, config: TransformerConfig):
layer_type = te.pytorch.ops.GEGLU
elif config.activation_func == F.silu:
layer_type = te.pytorch.ops.ReGLU
elif config.activation_func is situlu:
layer_type = getattr(te.pytorch.ops, "SiTUGLU", None)
if layer_type is None:
raise RuntimeError(
"SiTU-GLU requires Transformer Engine with "
"pytorch.ops.SiTUGLU support."
)
else:
if config.activation_func == F.gelu:
layer_type = te.pytorch.ops.GELU
elif config.activation_func == F.silu:
layer_type = te.pytorch.ops.ReLU
if layer_type is None:
raise Exception(
'Only SwiGLU, GEGLU, ReGLU, GELU, ReLU are supported by '
'Only SwiGLU, SiTU-GLU, GEGLU, ReGLU, GELU, ReLU are supported by '
'transformer engine. Please set use_te_activation_func=False'
)
activation_func_kwargs = {}
if config.activation_func is situlu:
activation_func_kwargs.update(
beta1=config.situ_glu_beta1, beta2=config.situ_glu_beta2
)
if config.activation_func_fp8_input_store:
activation_func_kwargs["cache_quantized_input"] = True
layer = layer_type(**activation_func_kwargs)
Expand Down Expand Up @@ -2729,6 +2741,12 @@ def _make_activation_op(
op_type = te.pytorch.ops.ReLU
elif (activation_func, gated_linear_unit) == (F.relu, True):
op_type = te.pytorch.ops.ReGLU
elif (activation_func, gated_linear_unit) == (situlu, True):
op_type = getattr(te.pytorch.ops, "SiTUGLU", None)
if op_type is None:
raise RuntimeError(
"SiTU-GLU requires Transformer Engine with " "pytorch.ops.SiTUGLU support."
)

# Could not find corresponding activation op
if op_type is None:
Expand All @@ -2740,6 +2758,8 @@ def _make_activation_op(

# Construct op
kwargs = {}
if activation_func is situlu:
kwargs.update(beta1=self.config.situ_glu_beta1, beta2=self.config.situ_glu_beta2)
if is_te_min_version("2.3"):
kwargs["cache_quantized_input"] = cache_quantized_input
return op_type(**kwargs)
Expand Down Expand Up @@ -2908,14 +2928,23 @@ def __init__(self, *args, **kwargs):
f"{self.__class__.__name__} does not support add_bias_linear=True; "
"the CuTeGEMM fused kernel requires bias-free linear layers."
)
if self.config.activation_func != F.silu or not self.config.gated_linear_unit:
if not self.config.gated_linear_unit or self.config.activation_func not in (
F.silu,
situlu,
):
raise ValueError(
f"{self.__class__.__name__} requires SwiGLU activation "
"(activation_func=F.silu, gated_linear_unit=True) "
f"{self.__class__.__name__} requires SwiGLU or SiTU-GLU activation "
"with gated_linear_unit=True "
"for the CuTeGEMM fused kernel, but got "
f"activation_func={self.config.activation_func}, "
f"gated_linear_unit={self.config.gated_linear_unit}."
)
if self.config.activation_func is situlu:
if not hasattr(te.pytorch.ops, "ScaledSiTUGLU"):
raise RuntimeError(
"SiTU-GLU requires Transformer Engine with "
"pytorch.ops.ScaledSiTUGLU support."
)

def _make_fused_impl(self) -> te.pytorch.ops.Sequential:
"""Construct fused module with GroupedLinear(num_groups=1) + ScaledSwiGLU."""
Expand Down Expand Up @@ -2992,9 +3021,19 @@ def _make_fused_impl(self) -> te.pytorch.ops.Sequential:
op._glu_interleave_size = _GLU_INTERLEAVE_SIZE # signals fuser_forward to interleave
fused_impl.append(op)

# ScaledSwiGLU with glu_interleave_size=32
# Required by ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8
fused_impl.append(te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=32))
if self.config.activation_func is situlu:
# ScaledSiTUGLU with glu_interleave_size=32.
fused_impl.append(
te.pytorch.ops.ScaledSiTUGLU(
glu_interleave_size=32,
beta1=self.config.situ_glu_beta1,
beta2=self.config.situ_glu_beta2,
)
)
else:
# ScaledSwiGLU with glu_interleave_size=32
# Required by ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8
fused_impl.append(te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=32))

# FC2: GroupedLinear(num_groups=1) instead of BasicLinear
weight = self.linear_fc2.weight
Expand Down
18 changes: 11 additions & 7 deletions megatron/core/transformer/mlp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import torch
import torch.nn.functional as F

from megatron.core.activations import situlu
from megatron.core.dist_checkpointing import ShardedTensor
from megatron.core.dist_checkpointing.mapping import (
ReplicaId,
Expand Down Expand Up @@ -187,12 +188,8 @@ def __init__(
if ffn_hidden_size is None:
if is_expert:
raise ValueError("MoE MLP requires `ffn_hidden_size`, but it was not provided.")
warnings.warn(
"MLP requires ffn_hidden_size, but it was not provided. Using \
config.ffn_hidden_size by default.",
DeprecationWarning,
stacklevel=2,
)
warnings.warn("MLP requires ffn_hidden_size, but it was not provided. Using \
config.ffn_hidden_size by default.", DeprecationWarning, stacklevel=2)
ffn_hidden_size = not_none(self.config.ffn_hidden_size)

# If this is a gated linear unit we double the output width
Expand Down Expand Up @@ -262,7 +259,12 @@ def forward(
if self.config.use_te_activation_func:
if bias_parallel is not None:
intermediate_parallel = intermediate_parallel + bias_parallel
intermediate_parallel = self.activation_func(intermediate_parallel)
if self.activation_func is situlu:
intermediate_parallel = situlu(
intermediate_parallel, self.config.situ_glu_beta1, self.config.situ_glu_beta2
)
else:
intermediate_parallel = self.activation_func(intermediate_parallel)
if per_token_scale is not None:
original_dtype = intermediate_parallel.dtype
intermediate_parallel = intermediate_parallel * per_token_scale.unsqueeze(-1)
Expand Down Expand Up @@ -317,6 +319,8 @@ def forward(
if self.config.gated_linear_unit:

def glu(x):
if self.config.activation_func is situlu:
return situlu(x, self.config.situ_glu_beta1, self.config.situ_glu_beta2)
x_glu, x_linear = torch.chunk(x, 2, dim=-1)
if (val := self.config.activation_func_clamp_value) is not None:
x_glu = x_glu.clamp(min=None, max=val)
Expand Down
38 changes: 31 additions & 7 deletions megatron/core/transformer/moe/experts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import torch.nn.functional as F

from megatron.core import tensor_parallel
from megatron.core.activations import squared_relu
from megatron.core.activations import situlu, squared_relu
from megatron.core.dist_checkpointing.mapping import ShardedStateDict
from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
from megatron.core.enums import Fp4Recipe, Fp8Recipe
Expand Down Expand Up @@ -364,13 +364,14 @@ def _activation_name():
if not isinstance(self.linear_fc2, te.pytorch.GroupedLinear):
return _unsupported(f"linear_fc2 is {type(self.linear_fc2).__name__}")

# Check activation: SwiGLU, quick GEGLU, or weighted squared ReLU.
# Check activation: SwiGLU, SiTU-GLU, quick GEGLU, or weighted squared ReLU.
# Clamped SwiGLU (e.g. DSv4) routes through ScaledClampedQGeGLU with
# alpha=1.0, since the cuDNN geglu kernel is a superset of swiglu.
# Use config.activation_func instead of self.activation_func because when
# use_te_activation_func is True, self.activation_func is a TE module, not the raw function.
use_glu_fusion = self.config.gated_linear_unit and self.config.activation_func in (
F.silu,
situlu,
quick_gelu,
)
use_srelu_fusion = (
Expand All @@ -384,7 +385,14 @@ def _activation_name():
f"(gated_linear_unit={self.config.gated_linear_unit}, "
f"use_fused_weighted_squared_relu={self.config.use_fused_weighted_squared_relu})"
)
if self.config.activation_func == F.silu:
if self.config.activation_func is situlu and self.activation_recompute:
return _unsupported(
"Transformer Engine ScaledSiTUGLU does not support activation recompute"
)
if self.config.activation_func == situlu:
if not hasattr(te_ops, "ScaledSiTUGLU"):
return _unsupported("SiTU-GLU needs ScaledSiTUGLU")
elif self.config.activation_func == F.silu:
if self.config.activation_func_clamp_value is not None:
if not is_te_min_version("2.17.0.dev0"):
return _unsupported("clamped SwiGLU needs TE >= 2.17.0.dev0")
Expand Down Expand Up @@ -487,14 +495,20 @@ def register_grouped_linear_params(
)
ops.append(op)

# Activation and post-multiply probs (SwiGLU, clamped GeGLU, or SReLU).
# Activation and post-multiply probs (SwiGLU, SiTU-GLU, clamped GeGLU, or SReLU).
# TE's ScaledClampedQGeGLU computes sigmoid(alpha * x) * x, so
# alpha=1.702 gives quick_gelu and alpha=1.0 gives silu/swiglu.
# With cuDNN FE >= 1.24.0 the alpha, limit and offset are
# forwarded as runtime params to the cuDNN kernel.
glu_interleave = self.config.moe_mlp_glu_interleave_size
activation_recompute_in_mlp = bool(getattr(self, "activation_recompute", False))
if self.config.activation_func == F.silu and self.config.gated_linear_unit:
if self.config.activation_func is situlu and self.config.gated_linear_unit:
op = te.pytorch.ops.ScaledSiTUGLU(
glu_interleave_size=glu_interleave,
beta1=self.config.situ_glu_beta1,
beta2=self.config.situ_glu_beta2,
)
elif self.config.activation_func == F.silu and self.config.gated_linear_unit:
clamp = self.config.activation_func_clamp_value
if clamp is not None:
qgeglu_kwargs = {
Expand Down Expand Up @@ -562,7 +576,8 @@ def register_grouped_linear_params(
op = te.pytorch.ops.ScaledSReLU()
else:
raise RuntimeError(
"_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; "
"_make_fused_ops expected SwiGLU, SiTU-GLU, quick_gelu, or weighted "
"squared_relu; "
"call _is_fused_impl_supported() before constructing fused ops."
)
ops.append(op)
Expand Down Expand Up @@ -814,7 +829,14 @@ def bias_act_func(intermediate_parallel, bias_parallel, permuted_probs):
intermediate_parallel = self._remove_glu_interleaving(
intermediate_parallel, self.config.moe_mlp_glu_interleave_size
)
intermediate_parallel = self.activation_func(intermediate_parallel)
if self.activation_func is situlu:
intermediate_parallel = situlu(
intermediate_parallel,
self.config.situ_glu_beta1,
self.config.situ_glu_beta2,
)
else:
intermediate_parallel = self.activation_func(intermediate_parallel)
if permuted_probs is not None:
original_dtype = intermediate_parallel.dtype
intermediate_parallel = intermediate_parallel * permuted_probs
Expand Down Expand Up @@ -859,6 +881,8 @@ def glu(x):
x = self._remove_glu_interleaving(
x, self.config.moe_mlp_glu_interleave_size
)
if self.config.activation_func is situlu:
return situlu(x, self.config.situ_glu_beta1, self.config.situ_glu_beta2)
x_glu, x_linear = torch.chunk(x, 2, dim=-1)
if (val := self.config.activation_func_clamp_value) is not None:
x_glu = x_glu.clamp(min=None, max=val)
Expand Down
34 changes: 29 additions & 5 deletions megatron/core/transformer/moe/shared_experts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import torch
import torch.nn.functional as F

from megatron.core.activations import situlu
from megatron.core.dist_checkpointing.mapping import ShardedStateDict
from megatron.core.extensions.transformer_engine import HAVE_TE
from megatron.core.fusions.fused_bias_geglu import bias_geglu_impl
Expand Down Expand Up @@ -276,7 +277,14 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None):
if self.config.use_te_activation_func:
if bias_parallel is not None:
intermediate_parallel = intermediate_parallel + bias_parallel
intermediate_parallel = self.activation_func(intermediate_parallel)
if self.activation_func is situlu:
intermediate_parallel = situlu(
intermediate_parallel,
self.config.situ_glu_beta1,
self.config.situ_glu_beta2,
)
else:
intermediate_parallel = self.activation_func(intermediate_parallel)
elif self.config.bias_activation_fusion:
if self.activation_func == F.gelu:
if self.config.gated_linear_unit:
Expand All @@ -301,6 +309,8 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None):
if self.config.gated_linear_unit:

def glu(x):
if self.config.activation_func is situlu:
return situlu(x, self.config.situ_glu_beta1, self.config.situ_glu_beta2)
x_glu, x_linear = torch.chunk(x, 2, dim=-1)
if (val := self.config.activation_func_clamp_value) is not None:
x_glu = x_glu.clamp(min=None, max=val)
Expand Down Expand Up @@ -417,13 +427,18 @@ def _validate_fused_grouped_swiglu(self) -> None:
f"{self.__class__.__name__} does not support add_bias_linear=True; "
"the CuTeGEMM fused kernel requires bias-free linear layers."
)
if not self.config.gated_linear_unit or self.config.activation_func != F.silu:
if not self.config.gated_linear_unit or self.config.activation_func not in (F.silu, situlu):
raise ValueError(
f"{self.__class__.__name__} requires SwiGLU activation "
"(activation_func=F.silu, gated_linear_unit=True) for the CuTeGEMM "
f"{self.__class__.__name__} requires SwiGLU or SiTU-GLU activation "
"with gated_linear_unit=True for the CuTeGEMM "
f"fused kernel, but got activation_func={self.config.activation_func}, "
f"gated_linear_unit={self.config.gated_linear_unit}."
)
if self.config.activation_func is situlu and not hasattr(te.pytorch.ops, "ScaledSiTUGLU"):
raise RuntimeError(
f"{self.__class__.__name__} requires Transformer Engine with "
"pytorch.ops.ScaledSiTUGLU for SiTU-GLU."
)
if self.config.moe_shared_expert_glu_interleave_size is None:
raise ValueError(
f"{self.__class__.__name__} requires "
Expand Down Expand Up @@ -483,7 +498,16 @@ def _make_fused_grouped_swiglu_ops(self) -> torch.nn.Module:
op._glu_interleave_size = glu_interleave_size
ops.append(op)

ops.append(te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size))
if self.config.activation_func is situlu:
ops.append(
te.pytorch.ops.ScaledSiTUGLU(
glu_interleave_size=glu_interleave_size,
beta1=self.config.situ_glu_beta1,
beta2=self.config.situ_glu_beta2,
)
)
else:
ops.append(te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size))

fc2_weight = self.linear_fc2.weight
op = te.pytorch.ops.GroupedLinear(
Expand Down
Loading
Loading