From c7a03fa7d840af6c90362d40ef5fc1a1d60cb696 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Thu, 20 Aug 2026 19:45:14 +0800 Subject: [PATCH] [Core] Add SiTU-GLU activation support Signed-off-by: Harry Zhou --- megatron/core/activations.py | 17 +++ .../core/extensions/transformer_engine.py | 61 +++++++- megatron/core/transformer/mlp.py | 10 +- megatron/core/transformer/moe/experts.py | 38 ++++- .../core/transformer/moe/shared_experts.py | 46 ++++-- .../core/transformer/transformer_config.py | 24 ++- megatron/training/argument_utils.py | 14 +- megatron/training/arguments.py | 19 ++- megatron/training/checkpointing.py | 5 +- megatron/training/theoretical_memory_usage.py | 2 +- megatron/training/training.py | 8 +- .../models/test_hybrid_moe_model.py | 2 + tests/unit_tests/test_checkpointing.py | 33 ++++ .../test_num_floating_point_operations.py | 10 ++ .../test_weight_and_optimizer_memory.py | 9 ++ .../transformer/moe/test_grouped_mlp.py | 59 ++++++- .../transformer/moe/test_shared_experts.py | 32 +++- tests/unit_tests/transformer/test_situ_glu.py | 144 ++++++++++++++++++ ...t_te_fused_mlp_with_grouped_linear_spec.py | 34 ++++- tools/checkpoint/loader_base.py | 3 +- 20 files changed, 516 insertions(+), 54 deletions(-) create mode 100644 tests/unit_tests/transformer/test_situ_glu.py diff --git a/megatron/core/activations.py b/megatron/core/activations.py index 8b422d73a35..02f109557cf 100644 --- a/megatron/core/activations.py +++ b/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) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 48f8ff2498b..722eab4af40 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/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,6 +608,13 @@ 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 @@ -614,10 +622,14 @@ def __new__(cls, config: TransformerConfig): 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 diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index 66a69ff3f37..6f72a1a1c74 100644 --- a/megatron/core/transformer/mlp.py +++ b/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) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 6b2ccae8466..22838946a29 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/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) diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 6eef5dee9ff..b7d1f0279c9 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/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. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 245a164fb63..fb254562ada 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -29,6 +29,7 @@ from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from .._rank_utils import log_single_rank +from ..activations import situlu from ..fusions.fused_bias_geglu import quick_gelu from ..model_parallel_config import ModelParallelConfig from ..utils import ( @@ -1219,6 +1220,12 @@ class TransformerConfig(ModelParallelConfig): use_te_activation_func: bool = False """Whether to use ffn activation functions implemented by TransformerEngine""" + situ_glu_beta1: float = 4.0 + """SiTU-GLU gate tanh soft-cap.""" + + situ_glu_beta2: float = 25.0 + """SiTU-GLU up-branch tanh soft-cap.""" + use_te_rng_tracker: bool = False """ Whether to use the TE or MCore version of the RNG tracker. """ @@ -2518,13 +2525,26 @@ def __post_init__(self): ) if self.use_te_activation_func: - if self.activation_func not in (F.gelu, F.silu, F.relu): + if self.activation_func not in (F.gelu, F.silu, F.relu, situlu): raise ValueError( - "TransformerEngine only support gelu, geglu, silu, swiglu, relu, reglu. " + "TransformerEngine only supports gelu, geglu, silu, swiglu, relu, reglu, " + "and situ-glu. " "If you don't want to use TransformerEngine activation function, set " "use_te_activation_func to False" ) + if self.activation_func == situlu: + if not self.gated_linear_unit: + raise ValueError("SiTU-GLU requires gated_linear_unit=True.") + if self.activation_func_clamp_value is not None: + raise ValueError("SiTU-GLU does not use activation_func_clamp_value.") + if self.glu_linear_offset != 0.0: + raise ValueError("SiTU-GLU requires glu_linear_offset=0.0.") + if not math.isfinite(self.situ_glu_beta1) or self.situ_glu_beta1 <= 0: + raise ValueError("situ_glu_beta1 must be finite and positive.") + if not math.isfinite(self.situ_glu_beta2) or self.situ_glu_beta2 <= 0: + raise ValueError("situ_glu_beta2 must be finite and positive.") + if self.activation_func_fp8_input_store: if self.activation_func != F.silu or not self.gated_linear_unit: raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.") diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 124083ead83..5c0ff1a59a0 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -275,7 +275,7 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: def core_transformer_config_from_args(args, config_class=None): - from megatron.core.activations import squared_relu + from megatron.core.activations import situlu, squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu from megatron.core.quantization.utils import ( kitchen_quantization_recipe_config, @@ -311,17 +311,23 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['num_layers_in_last_pipeline_stage']= args.decoder_last_pipeline_num_layers kw_args['fp8_param'] = args.fp8_param_gather kw_args['fp4_param'] = args.fp4_param_gather - if args.swiglu: + use_situ_glu = getattr(args, 'situ_glu', False) + 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 + elif args.swiglu: kw_args['activation_func'] = F.silu kw_args['gated_linear_unit'] = True kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion else: kw_args['bias_activation_fusion'] = args.bias_gelu_fusion if args.squared_relu: - assert not args.swiglu + assert not args.swiglu and not use_situ_glu kw_args['activation_func'] = squared_relu elif args.quick_geglu: - assert not args.swiglu + assert not args.swiglu and not use_situ_glu kw_args['gated_linear_unit'] = True kw_args['activation_func'] = quick_gelu if args.init_method_xavier_uniform: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 12a9067be7e..a777b381a76 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1320,8 +1320,16 @@ def validate_args(args, defaults={}): _check_arg_is_not_none(args, req_arg) # Checks. + use_situ_glu = getattr(args, 'situ_glu', False) + if use_situ_glu: + if args.swiglu or args.quick_geglu or args.squared_relu: + raise ValueError( + "--situ-glu is mutually exclusive with --swiglu, --quick-geglu, " + "and --squared-relu." + ) + if args.ffn_hidden_size is None: - if args.swiglu: + if args.swiglu or use_situ_glu: # reduce the dimnesion for MLP since projections happens on # two linear layers. this keeps the number of paramters in # the same ballpark as the counterpart with 4*h size @@ -2452,6 +2460,15 @@ def _add_network_size_args(parser): help='Use squared relu activation instead of default gelu') group.add_argument('--swiglu', action='store_true', help='Use gated linear units and SiLU activation instead of default gelu') + group.add_argument( + '--situ-glu', + '--moe-use-situ-glu', + action='store_true', + help=( + 'Use SiTU-GLU in all dense and MoE FFNs. TE-backed paths select SiTUGLU ' + 'or ScaledSiTUGLU; other paths use the PyTorch reference.' + ), + ) group.add_argument('--quick-geglu', action='store_true', help='Use quick geglu activation instead of default gelu') group.add_argument('--onnx-safe', type=bool, required=False, diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 9bbc2a7d8ad..4daea787dc9 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1740,7 +1740,7 @@ def apply(handler): if 'optimizer' in state_dict: state_dict['optimizer'] = optimizer_state_dict - if args.swiglu: + if args.swiglu or getattr(args, "situ_glu", False): apply(handle_swiglu_in_state_dict) # Split a fused MLA q/kv down-projection (mla_down_proj_fusion) back into the unfused # layout used on disk. No-op for unfused models. @@ -2288,6 +2288,9 @@ def _set_arg(arg_name, old_arg_name=None, force=False): _set_arg('add_qkv_bias', force=True) _set_arg('squared_relu', force=True) _set_arg('swiglu', force=True) + _set_arg('situ_glu', force=True) + _set_arg('situ_glu_beta1', force=True) + _set_arg('situ_glu_beta2', force=True) _set_arg('untie_embeddings_and_output_weights', force=True) _set_arg('apply_layernorm_1p', force=True) _set_arg('normalization', force=True) diff --git a/megatron/training/theoretical_memory_usage.py b/megatron/training/theoretical_memory_usage.py index ee398d3bf66..25def866b64 100644 --- a/megatron/training/theoretical_memory_usage.py +++ b/megatron/training/theoretical_memory_usage.py @@ -18,7 +18,7 @@ def compute_weight_and_optimizer_memory(args, verbose=False): args.num_query_groups = args.num_attention_heads # MoE. num_experts = 1 if args.num_experts is None else args.num_experts - gated_linear_multiplier = 3 / 2 if args.swiglu else 1 + gated_linear_multiplier = 3 / 2 if (args.swiglu or getattr(args, "situ_glu", False)) else 1 shared_expert_ffn_hidden_size = ( 0 diff --git a/megatron/training/training.py b/megatron/training/training.py index e4b0afb56f3..01fa92ea2ee 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1065,9 +1065,9 @@ def transformer_flops(): forward_backward_expansion_factor = 3 # - 2x: A GEMM of a m*n tensor with a n*k tensor requires 2mnk floating-point operations. fma_expansion_factor = 2 - # - 3x (SwiGLU enabled): h->2*ffn_h GEMM and ffn_h->h GEMM are stacked. - # - 2x (SwiGLU disabled): h->ffn_h GEMM and ffn_h->h GEMM are stacked. - ffn_expansion_factor = 3 if args.swiglu else 2 + # - 3x (gated GLU): h->2*ffn_h GEMM and ffn_h->h GEMM are stacked. + # - 2x (non-gated): h->ffn_h GEMM and ffn_h->h GEMM are stacked. + ffn_expansion_factor = 3 if (args.swiglu or getattr(args, "situ_glu", False)) else 2 # self_attn is split into a token-linear part (projections, multiplied by # ``batch_size * args.seq_length`` like all other token-linear work) and a @@ -1366,7 +1366,7 @@ def _split_spec_part(part): gqa_groups=args.num_query_groups, kv_channels=args.kv_channels, mlp_expansion=args.ffn_hidden_size / args.hidden_size, - swiglu=args.swiglu, + swiglu=(args.swiglu or getattr(args, "situ_glu", False)), use_gated_delta_product=_uses_gated_delta_product_spec(args), moe_latent_size=args.moe_latent_size, moe_ffn_hidden_size=(args.moe_ffn_hidden_size if args.moe_ffn_hidden_size is not None diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 51a0780e65d..1639e805ad7 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -298,6 +298,8 @@ "recompute_num_layers": None, "rotary_interleaved": False, "sequence_parallel": True, + "situ_glu_beta1": 4.0, + "situ_glu_beta2": 25.0, "softmax_scale": None, "softmax_type": "vanilla", "symmetric_ar_type": None, diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index d5c625ca8ca..c59c2f5214b 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -200,6 +200,39 @@ def test_load_args_restores_gdp_num_householder_from_checkpoint( assert restored_args.gdp_num_householder == expected_num_householder +@pytest.mark.parametrize( + "checkpoint_args", + [ + SimpleNamespace(swiglu=True, situ_glu=False, situ_glu_beta1=4.0, situ_glu_beta2=25.0), + SimpleNamespace(swiglu=False, situ_glu=True, situ_glu_beta1=3.0, situ_glu_beta2=20.0), + ], +) +def test_load_args_restores_glu_activation_from_checkpoint(checkpoint_args): + """Checkpoint arguments select the saved gated activation and SiTU scaling.""" + args = SimpleNamespace( + load="checkpoint", + iteration=0, + swiglu=not checkpoint_args.swiglu, + situ_glu=not checkpoint_args.situ_glu, + situ_glu_beta1=1.0, + situ_glu_beta2=1.0, + use_tokenizer_model_from_checkpoint_args=False, + use_mp_args_from_checkpoint_args=False, + ) + state_dict = {"args": checkpoint_args, "iteration": 12} + + with mock.patch( + "megatron.training.checkpointing._load_base_checkpoint", + return_value=(state_dict, "checkpoint", False, CheckpointType.LEGACY), + ): + restored_args, _ = load_args_from_checkpoint(args) + + assert restored_args.swiglu is checkpoint_args.swiglu + assert restored_args.situ_glu is checkpoint_args.situ_glu + assert restored_args.situ_glu_beta1 == checkpoint_args.situ_glu_beta1 + assert restored_args.situ_glu_beta2 == checkpoint_args.situ_glu_beta2 + + def create_checkpoint(load_path, ckpt_format): """Setup a dummy checkpoint directory.""" iteration = 123 diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index df5e4191843..f15953447e6 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -100,6 +100,16 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s return args +def test_situ_glu_counts_the_same_ffn_gemms_as_swiglu(): + swiglu_args = _make_gpt_args(swiglu=True) + situ_glu_args = _make_gpt_args(swiglu=False) + situ_glu_args.situ_glu = True + + assert num_floating_point_operations( + situ_glu_args, batch_size=8 + ) == num_floating_point_operations(swiglu_args, batch_size=8) + + class TestBSHDBackwardCompat: """For unpacked BSHD, the new optional arg must not change the result.""" diff --git a/tests/unit_tests/training/test_weight_and_optimizer_memory.py b/tests/unit_tests/training/test_weight_and_optimizer_memory.py index 19a532c35c2..f7de1a02d75 100644 --- a/tests/unit_tests/training/test_weight_and_optimizer_memory.py +++ b/tests/unit_tests/training/test_weight_and_optimizer_memory.py @@ -41,6 +41,15 @@ def _make_args(**overrides): return args +def test_situ_glu_counts_the_same_ffn_parameters_as_swiglu(): + swiglu_args = _make_args(swiglu=True) + situ_glu_args = _make_args(situ_glu=True) + + assert compute_weight_and_optimizer_memory( + situ_glu_args + ) == compute_weight_and_optimizer_memory(swiglu_args) + + def test_weight_and_optimizer_memory_accounts_for_expert_parallelism(): args = _make_args(pipeline_model_parallel_size=2, world_size=64) diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 8cc1d96db2a..0f9c81ea395 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -9,7 +9,7 @@ import torch.nn.functional as F import megatron.core.transformer.moe.experts as experts_module -from megatron.core.activations import squared_relu +from megatron.core.activations import situlu, squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_submodules, @@ -516,6 +516,13 @@ def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): self.glu_interleave_size = glu_interleave_size self.activation_recompute_in_mlp = activation_recompute_in_mlp + class FakeScaledSiTUGLU(torch.nn.Module): + def __init__(self, glu_interleave_size, *, beta1, beta2): + super().__init__() + self.glu_interleave_size = glu_interleave_size + self.beta1 = beta1 + self.beta2 = beta2 + class FakeScaledClampedQGeGLU(torch.nn.Module): def __init__( self, @@ -549,6 +556,7 @@ def register_forward_pre_hook(self, hook): ops=SimpleNamespace( GroupedLinear=FakeGroupedLinear, ScaledSwiGLU=FakeScaledSwiGLU, + ScaledSiTUGLU=FakeScaledSiTUGLU, ScaledClampedQGeGLU=FakeScaledClampedQGeGLU, ScaledSReLU=FakeScaledSReLU, Sequential=FakeSequential, @@ -670,7 +678,7 @@ def test_make_fused_ops_rejects_scaled_srelu_with_gated_linear_unit(monkeypatch) module.linear_fc2.weight0 = torch.nn.Parameter(torch.ones(4, 8)) module.linear_fc2.weight1 = torch.nn.Parameter(torch.ones(4, 8)) - with pytest.raises(RuntimeError, match="expected SwiGLU, quick_gelu"): + with pytest.raises(RuntimeError, match="expected SwiGLU, SiTU-GLU, quick_gelu"): module._make_fused_ops() @@ -709,11 +717,16 @@ def _make_fused_impl_support_module( gated_linear_unit=gated_linear_unit, use_fused_weighted_squared_relu=use_fused_weighted_squared_relu, moe_apply_probs_on_input=False, + moe_mlp_glu_interleave_size=32, + delay_wgrad_compute=False, + situ_glu_beta1=4.0, + situ_glu_beta2=25.0, ) module.activation_func = object() module.tp_group = SimpleNamespace(size=lambda: 1) module.offload_expert_fc1 = False module.offload_moe_act = False + module.activation_recompute = False common = dict( device="cuda", dtype=torch.bfloat16, @@ -725,7 +738,44 @@ def _make_fused_impl_support_module( return module -def test_is_fused_impl_supported_uses_config_activation_for_swiglu(monkeypatch): +def test_make_fused_ops_selects_scaled_situ_glu(monkeypatch): + """The routed-expert op-fuser passes SiTU betas to TE.""" + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + module = _make_fused_impl_support_module( + FakeGroupedLinear, activation_func=situlu, gated_linear_unit=True + ) + for linear in (module.linear_fc1, module.linear_fc2): + linear.weight0 = torch.nn.Parameter(torch.ones(linear.out_features, linear.in_features)) + linear.weight1 = torch.nn.Parameter(torch.ones(linear.out_features, linear.in_features)) + + ops = module._make_fused_ops() + + activation = ops[1] + assert type(activation).__name__ == "FakeScaledSiTUGLU" + assert activation.glu_interleave_size == 32 + assert activation.beta1 == 4.0 + assert activation.beta2 == 25.0 + + +def test_is_fused_impl_supported_rejects_situ_glu_activation_recompute(monkeypatch): + """ScaledSiTUGLU cannot honor selective MoE activation recompute.""" + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + monkeypatch.setattr(experts_module, "HAVE_TE", True) + monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True) + _install_fake_te_ops_modules(monkeypatch, fake_te) + module = _make_fused_impl_support_module( + FakeGroupedLinear, activation_func=situlu, gated_linear_unit=True + ) + module.activation_recompute = True + + assert module._is_fused_impl_supported() is False + + +@pytest.mark.parametrize("activation_func", [F.silu, quick_gelu], ids=("swiglu", "quick-geglu")) +def test_is_fused_impl_supported_preserves_existing_glu_recompute(monkeypatch, activation_func): + """Existing scaled GLUs continue to forward activation recompute to TE.""" fake_te, FakeGroupedLinear = _make_fake_te_namespace() monkeypatch.setattr(experts_module, "te", fake_te) monkeypatch.setattr(experts_module, "HAVE_TE", True) @@ -733,8 +783,9 @@ def test_is_fused_impl_supported_uses_config_activation_for_swiglu(monkeypatch): _install_fake_te_ops_modules(monkeypatch, fake_te) module = _make_fused_impl_support_module( - FakeGroupedLinear, activation_func=F.silu, gated_linear_unit=True + FakeGroupedLinear, activation_func=activation_func, gated_linear_unit=True ) + module.activation_recompute = True assert module._is_fused_impl_supported() is True diff --git a/tests/unit_tests/transformer/moe/test_shared_experts.py b/tests/unit_tests/transformer/moe/test_shared_experts.py index c71f665259f..cf22c7f44d2 100644 --- a/tests/unit_tests/transformer/moe/test_shared_experts.py +++ b/tests/unit_tests/transformer/moe/test_shared_experts.py @@ -7,6 +7,7 @@ import torch import torch.nn.functional as F +from megatron.core.activations import situlu from megatron.core.models.gpt import moe_module_specs from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules from megatron.core.parallel_state import get_tensor_model_parallel_world_size @@ -52,6 +53,14 @@ def __init__(self, glu_interleave_size): self.glu_interleave_size = glu_interleave_size +class _FakeTEScaledSiTUGLU(torch.nn.Module): + def __init__(self, glu_interleave_size, *, beta1, beta2): + super().__init__() + self.glu_interleave_size = glu_interleave_size + self.beta1 = beta1 + self.beta2 = beta2 + + class _FakeTEScaledClampedQGeGLU(torch.nn.Module): def __init__(self, glu_interleave_size, *, alpha, limit, glu_linear_offset): super().__init__() @@ -99,6 +108,7 @@ def _fake_te_module(linear_cls=_FakeTELinear): ops=SimpleNamespace( GroupedLinear=_FakeTEGroupedLinear, ScaledSwiGLU=_FakeTEScaledSwiGLU, + ScaledSiTUGLU=_FakeTEScaledSiTUGLU, ScaledClampedQGeGLU=_FakeTEScaledClampedQGeGLU, Sequential=_FakeTESequential, ), @@ -134,6 +144,9 @@ def _fake_shared_expert(**config_kwargs): gated_linear_unit=True, activation_func=F.silu, activation_func_clamp_value=None, + situ_glu_beta1=4.0, + situ_glu_beta2=25.0, + use_fused_weighted_squared_relu=False, moe_shared_expert_glu_interleave_size=32, delay_wgrad_compute=False, sequence_parallel=False, @@ -215,8 +228,8 @@ def test_validate_fused_grouped_swiglu_requires_clamped_te_support(monkeypatch, ("config_kwargs", "bad_linear", "match"), [ ({"add_bias_linear": True}, None, "add_bias_linear"), - ({"activation_func": F.gelu}, None, "SwiGLU activation"), - ({"gated_linear_unit": False}, None, "SwiGLU activation"), + ({"activation_func": F.gelu}, None, "SwiGLU or SiTU-GLU activation"), + ({"gated_linear_unit": False}, None, "SwiGLU or SiTU-GLU activation"), ({"moe_shared_expert_glu_interleave_size": None}, None, "glu_interleave_size"), ({}, "linear_fc1", "FC1"), ({}, "linear_fc2", "FC2"), @@ -267,6 +280,21 @@ def test_make_fused_grouped_swiglu_ops_builds_grouped_pipeline(monkeypatch): assert fc2_op.weight0 is shared_expert.linear_fc2.weight +def test_make_fused_grouped_swiglu_ops_selects_situ_glu(monkeypatch): + _patch_fake_shared_expert_te(monkeypatch) + shared_expert = _fake_shared_expert(activation_func=situlu) + + shared_expert._validate_fused_grouped_swiglu() + ops = shared_expert._make_fused_grouped_swiglu_ops() + + activation_op = list(ops.children())[1] + assert isinstance(activation_op, _FakeTEScaledSiTUGLU) + assert activation_op.glu_interleave_size == 32 + assert activation_op.beta1 == 4.0 + assert activation_op.beta2 == 25.0 + assert activation_op._grouped_mlp_unit_activation_scale is True + + def test_make_fused_grouped_swiglu_ops_builds_clamped_activation(monkeypatch): _patch_fake_shared_expert_te(monkeypatch) shared_expert = _fake_shared_expert(activation_func_clamp_value=7.0) diff --git a/tests/unit_tests/transformer/test_situ_glu.py b/tests/unit_tests/transformer/test_situ_glu.py new file mode 100644 index 00000000000..d9e4603597e --- /dev/null +++ b/tests/unit_tests/transformer/test_situ_glu.py @@ -0,0 +1,144 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +import megatron.core.extensions.transformer_engine as te_extension +from megatron.core.activations import situlu +from megatron.core.transformer.transformer_config import TransformerConfig + + +class _FakeActivation: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +def _patch_te_activation_ops(monkeypatch, *, include_situ=True): + ops = SimpleNamespace() + if include_situ: + ops.SiTUGLU = type("SiTUGLU", (_FakeActivation,), {}) + + monkeypatch.setattr( + te_extension, "te", SimpleNamespace(pytorch=SimpleNamespace(ops=ops)), raising=False + ) + monkeypatch.setattr(te_extension, "HAVE_TE", True) + monkeypatch.setattr(te_extension, "is_te_min_version", lambda *_args, **_kwargs: True) + return ops + + +def _config(activation_func, gated_linear_unit, **kwargs): + defaults = dict( + activation_func=activation_func, + gated_linear_unit=gated_linear_unit, + activation_func_fp8_input_store=False, + activation_func_clamp_value=None, + use_fused_weighted_squared_relu=False, + situ_glu_beta1=4.0, + situ_glu_beta2=25.0, + ) + defaults.update(kwargs) + return SimpleNamespace(**defaults) + + +def test_ordinary_te_situ_glu_activation_selection(monkeypatch): + """The ordinary TE path selects SiTUGLU with the configured betas.""" + _patch_te_activation_ops(monkeypatch) + op = te_extension.TEActivationOp(_config(situlu, True)) + + assert type(op).__name__ == "SiTUGLU" + assert op.kwargs == {"beta1": 4.0, "beta2": 25.0} + + +def test_situ_glu_requires_te_operations(monkeypatch): + """An older TE cannot cause the SiTU marker to fall through to SwiGLU.""" + _patch_te_activation_ops(monkeypatch, include_situ=False) + + with pytest.raises(RuntimeError, match="pytorch.ops.SiTUGLU"): + te_extension.TEActivationOp(_config(situlu, True)) + + +def test_situlu_reference_matches_kimi_bf16_precision_and_backward(): + """The fallback evaluates both branches in FP32 and returns the original dtype.""" + x = torch.linspace(-20, 20, 30, device="cuda", dtype=torch.bfloat16).reshape(3, 10) + x = x.detach().requires_grad_(True) + reference_x = x.detach().clone().requires_grad_(True) + gate, up = reference_x.chunk(2, dim=-1) + gate = gate.float() + up = up.float() + expected = 3.0 * torch.tanh(gate / 3.0) * torch.sigmoid(gate) + expected = (expected * (7.0 * torch.tanh(up / 7.0))).to(reference_x.dtype) + + actual = situlu(x, 3.0, 7.0) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + grad = torch.randn_like(actual) + actual.backward(grad) + expected.backward(grad) + torch.testing.assert_close(x.grad, reference_x.grad, rtol=0, atol=0) + + +def test_situlu_reference_matches_fixed_values(): + """Check independent values, including positive and negative saturation.""" + x = torch.tensor( + [[0.0, 1.0, 0.0, 2.0], [-1.0, 4.0, -2.0, 25.0], [100.0, -100.0, 100.0, -100.0]], + device="cuda", + ) + expected = torch.tensor( + [[0.0, 1.4293511], [0.5258289, 56.959320], [99.932930, 3.717581e-42]], device="cuda" + ) + + torch.testing.assert_close(situlu(x), expected, rtol=1e-6, atol=1e-5) + + +def test_transformer_config_accepts_situ_glu_defaults(): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + activation_func=situlu, + gated_linear_unit=True, + use_te_activation_func=True, + ) + + assert config.situ_glu_beta1 == 4.0 + assert config.situ_glu_beta2 == 25.0 + + +def test_transformer_config_accepts_pytorch_situ_glu_fallback(): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + activation_func=situlu, + gated_linear_unit=True, + use_te_activation_func=False, + ) + + assert config.activation_func is situlu + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"gated_linear_unit": False}, "gated_linear_unit=True"), + ({"activation_func_clamp_value": 1.0}, "does not use activation_func_clamp_value"), + ({"glu_linear_offset": 1.0}, "requires glu_linear_offset=0.0"), + ({"situ_glu_beta1": 0.0}, "situ_glu_beta1 must be finite and positive"), + ({"situ_glu_beta2": float("inf")}, "situ_glu_beta2 must be finite and positive"), + ], +) +def test_transformer_config_rejects_unsupported_situ_glu(overrides, match): + kwargs = dict( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + activation_func=situlu, + gated_linear_unit=True, + use_te_activation_func=True, + ) + kwargs.update(overrides) + + with pytest.raises(ValueError, match=match): + TransformerConfig(**kwargs) diff --git a/tests/unit_tests/transformer/test_te_fused_mlp_with_grouped_linear_spec.py b/tests/unit_tests/transformer/test_te_fused_mlp_with_grouped_linear_spec.py index 2713502c329..a614a876b5b 100644 --- a/tests/unit_tests/transformer/test_te_fused_mlp_with_grouped_linear_spec.py +++ b/tests/unit_tests/transformer/test_te_fused_mlp_with_grouped_linear_spec.py @@ -7,6 +7,7 @@ import torch import torch.nn.functional as F +from megatron.core.activations import situlu from megatron.core.extensions.transformer_engine import ( HAVE_TE, TEFusedMLP, @@ -74,13 +75,18 @@ class FakeScaledSwiGLU: def __init__(self, **kwargs): self.kwargs = kwargs + class FakeScaledSiTUGLU(FakeScaledSwiGLU): + pass + fake_ops = SimpleNamespace( Sequential=FakeSequential, LayerNorm=FakeNorm, RMSNorm=FakeNorm, GroupedLinear=FakeGroupedLinear, ScaledSwiGLU=FakeScaledSwiGLU, + ScaledSiTUGLU=FakeScaledSiTUGLU, ) + monkeypatch.setattr(te_ext, "HAVE_TE", True) monkeypatch.setattr(te_ext.te.pytorch, "LayerNormLinear", FakeLayerNormLinear, raising=False) monkeypatch.setattr(te_ext.te.pytorch, "Linear", FakeLinear, raising=False) monkeypatch.setattr(te_ext.te.pytorch, "ops", fake_ops, raising=False) @@ -95,11 +101,20 @@ def __init__(self, **kwargs): Linear=FakeLinear, GroupedLinear=FakeGroupedLinear, ScaledSwiGLU=FakeScaledSwiGLU, + ScaledSiTUGLU=FakeScaledSiTUGLU, ) -def _make_fake_grouped_mlp(fake_te, normalization="LayerNorm"): +def _make_fake_grouped_mlp(fake_te, normalization="LayerNorm", activation_func=F.silu): module = TEFusedMLPWithGroupedLinear.__new__(TEFusedMLPWithGroupedLinear) + module.config = SimpleNamespace( + activation_func=activation_func, + activation_func_clamp_value=None, + gated_linear_unit=True, + situ_glu_beta1=4.0, + situ_glu_beta2=25.0, + use_fused_weighted_squared_relu=False, + ) fc1 = fake_te.LayerNormLinear() fc1.normalization = normalization @@ -147,8 +162,8 @@ def fake_is_te_min_version(version, *args, **kwargs): ("config_overrides", "match"), [ ({"add_bias_linear": True}, "add_bias_linear"), - ({"activation_func": F.gelu}, "SwiGLU activation"), - ({"gated_linear_unit": False}, "SwiGLU activation"), + ({"activation_func": F.gelu}, "SwiGLU or SiTU-GLU activation"), + ({"gated_linear_unit": False}, "SwiGLU or SiTU-GLU activation"), ], ) def test_init_validates_supported_dense_swiglu_config( @@ -208,6 +223,15 @@ def test_make_fused_impl_builds_grouped_linear_pipeline(self, monkeypatch, norma assert fc2_op.kwargs["accumulate_into_main_grad"] is False assert fc2_op.weight0 is module.linear_fc2.weight + def test_make_fused_impl_selects_scaled_situ_glu(self, monkeypatch): + fake_te = _patch_fake_te_ops(monkeypatch) + module = _make_fake_grouped_mlp(fake_te, activation_func=situlu) + + fused_impl = TEFusedMLPWithGroupedLinear._make_fused_impl(module) + + assert isinstance(fused_impl[1], fake_te.ScaledSiTUGLU) + assert fused_impl[1].kwargs == {"glu_interleave_size": 32, "beta1": 4.0, "beta2": 25.0} + @pytest.mark.parametrize(("bad_attr", "match"), [("linear_fc1", "FC1"), ("linear_fc2", "FC2")]) def test_make_fused_impl_validates_te_linear_types(self, monkeypatch, bad_attr, match): fake_te = _patch_fake_te_ops(monkeypatch) @@ -333,12 +357,12 @@ def test_instantiation(self): def test_wrong_activation_raises(self): config = _make_config(activation_func=F.gelu, gated_linear_unit=False) - with pytest.raises(ValueError, match="SwiGLU activation"): + with pytest.raises(ValueError, match="SwiGLU or SiTU-GLU activation"): TEFusedMLPWithGroupedLinear(config, _make_submodules()) def test_gated_linear_unit_false_raises(self): config = _make_config(gated_linear_unit=False) - with pytest.raises(ValueError, match="SwiGLU activation"): + with pytest.raises(ValueError, match="SwiGLU or SiTU-GLU activation"): TEFusedMLPWithGroupedLinear(config, _make_submodules()) def test_add_bias_linear_raises(self): diff --git a/tools/checkpoint/loader_base.py b/tools/checkpoint/loader_base.py index 3cf717fd486..a8c97d38a74 100644 --- a/tools/checkpoint/loader_base.py +++ b/tools/checkpoint/loader_base.py @@ -447,7 +447,7 @@ def build_checkpoint_metadata(self, true_vocab_size): md.linear_bias = self.margs.add_bias_linear md.qkv_bias = self.margs.add_qkv_bias md.norm_has_bias = norm_has_bias - md.swiglu = self.margs.swiglu + md.swiglu = self.margs.swiglu or getattr(self.margs, "situ_glu", False) md.previous_tensor_parallel_size = self.margs.tensor_model_parallel_size md.previous_pipeline_parallel_size = self.margs.pipeline_model_parallel_size md.true_vocab_size = true_vocab_size @@ -487,4 +487,3 @@ def import_model_provider(self): def send_model_over_queue(self): """Creates model schema and sends the model over the queue""" raise NotImplementedError -