[MCore] Add SiTU-GLU through Transformer Engine

megatron/core/activations.py

@@ -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:
+    """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)

megatron/core/extensions/transformer_engine.py

@@ -20,6 +20,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 ShardedObject, ShardedStateDict
 from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
 from megatron.core.enums import Fp4Recipe, Fp8Recipe
@@ -607,17 +608,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)
@@ -956,7 +968,11 @@ def _make_te_ops_bias_from_tensor(
         return op
 
     def _make_te_ops_activation(
-        activation_func: Callable, gated_linear_unit: bool, cache_quantized_input: bool
+        activation_func: Callable,
+        gated_linear_unit: bool,
+        cache_quantized_input: bool,
+        situ_glu_beta1: float = 4.0,
+        situ_glu_beta2: float = 25.0,
     ) -> te.pytorch.ops.FusibleOperation:
         """Construct a TE activation op."""
         op_type = None
@@ -974,6 +990,12 @@ def _make_te_ops_activation(
             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."
+                )
 
         if op_type is None:
             raise NotImplementedError(
@@ -983,6 +1005,8 @@ def _make_te_ops_activation(
             )
 
         kwargs = {}
+        if activation_func is situlu:
+            kwargs.update(beta1=situ_glu_beta1, beta2=situ_glu_beta2)
         if is_te_min_version("2.3"):
             kwargs["cache_quantized_input"] = cache_quantized_input
         return op_type(**kwargs)
@@ -3077,6 +3101,8 @@ def _make_fused_impl(self) -> te.pytorch.ops.Sequential:
                 self.activation_func,
                 self.config.gated_linear_unit,
                 self.config.activation_func_fp8_input_store,
+                self.config.situ_glu_beta1,
+                self.config.situ_glu_beta2,
             )
             fused_impl.append(op)
 
@@ -3166,14 +3192,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."""
@@ -3250,9 +3285,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

megatron/core/transformer/mlp.py

@@ -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,
@@ -267,7 +268,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)
@@ -323,6 +329,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)

megatron/core/transformer/moe/experts.py

@@ -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
@@ -412,11 +412,12 @@ def _is_fused_impl_supported(self) -> bool:
         ):
             return False  # Older TE op-fuser versions cannot scale FC2 bias by router probabilities
 
-        # Check activation: SwiGLU, quick GEGLU, or weighted squared ReLU.
+        # Check activation: SwiGLU, SiTU-GLU, quick GEGLU, or weighted squared ReLU.
         # 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 = (
@@ -426,7 +427,14 @@ def _is_fused_impl_supported(self) -> bool:
         )
         if not (use_glu_fusion or use_srelu_fusion):
             return False
-        if self.config.activation_func == F.silu:
+        if self.config.activation_func is situlu and self.activation_recompute:
+            return False
+        if self.config.activation_func == situlu:
+            try:
+                from transformer_engine.pytorch.ops import ScaledSiTUGLU  # noqa: F401
+            except ImportError:
+                return False
+        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 False
@@ -542,10 +550,16 @@ def register_grouped_linear_params(
             op.ep_mxfp8_carrier_input = True
         ops.append(op)
 
-        # Activation and post-multiply probs (SwiGLU, clamped GLU, or SReLU).
+        # Activation and post-multiply probs (SwiGLU, SiTU-GLU, clamped GLU, or SReLU).
         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_value = self.config.activation_func_clamp_value
             if clamp_value is not None:
                 clamped_glu_kwargs = {
@@ -614,7 +628,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)
@@ -955,7 +970,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
@@ -1000,6 +1022,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)

megatron/core/transformer/moe/shared_experts.py

@@ -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
@@ -272,7 +273,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:
@@ -297,6 +305,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 (clamp_value := self.config.activation_func_clamp_value) is not None:
                             x_glu = x_glu.clamp(min=None, max=clamp_value)
@@ -415,13 +425,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.activation_func_clamp_value is not None and (
             not is_te_min_version("2.17.0.dev0")
             or not hasattr(te.pytorch.ops, "ScaledClampedQGeGLU")
@@ -490,16 +505,23 @@ def _make_fused_grouped_swiglu_ops(self) -> torch.nn.Module:
         op._glu_interleave_size = glu_interleave_size
         ops.append(op)
 
-        clamp_value = self.config.activation_func_clamp_value
-        if clamp_value is None:
-            activation_op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size)
-        else:
-            activation_op = te.pytorch.ops.ScaledClampedQGeGLU(
+        if self.config.activation_func is situlu:
+            activation_op = te.pytorch.ops.ScaledSiTUGLU(
                 glu_interleave_size=glu_interleave_size,
-                alpha=1.0,
-                limit=clamp_value,
-                glu_linear_offset=0.0,
+                beta1=self.config.situ_glu_beta1,
+                beta2=self.config.situ_glu_beta2,
             )
+        else:
+            clamp_value = self.config.activation_func_clamp_value
+            if clamp_value is None:
+                activation_op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size)
+            else:
+                activation_op = te.pytorch.ops.ScaledClampedQGeGLU(
+                    glu_interleave_size=glu_interleave_size,
+                    alpha=1.0,
+                    limit=clamp_value,
+                    glu_linear_offset=0.0,
+                )
         # Shared experts are not router-gated. Mark this fused-op instance so
         # TE can omit the optional forward cuDNN probability tensor without
         # changing the semantics of routed single-group MLPs.