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 348847e7399..20cc9d96b75 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -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 @@ -483,6 +484,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 @@ -490,10 +498,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) @@ -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: @@ -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) @@ -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.""" @@ -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 diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index 0a107a0b0cf..1ec912bc4f0 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, @@ -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 @@ -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) @@ -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) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index ec9ba66e809..4254c8c9c34 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 @@ -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 = ( @@ -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") @@ -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 = { @@ -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) @@ -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 @@ -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) diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 50e2ef6c0ce..ed7cd9de965 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 @@ -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: @@ -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) @@ -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 " @@ -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( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 275bb1720c3..cad4faf38be 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -31,6 +31,7 @@ from megatron.core.utils import experimental_api 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 ( @@ -1317,6 +1318,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. """ @@ -2756,13 +2763,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 db927573166..2a6b65c209f 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -388,7 +388,7 @@ def _resolve_dsa_kernel_backend_cli_default(args, kw_args): 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, @@ -427,17 +427,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 9c379511b20..ad279bcdfeb 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -14,7 +14,7 @@ import torch.nn.functional as F from packaging.version import Version as PkgVersion -from megatron.core.activations import squared_relu +from megatron.core.activations import situlu, squared_relu from megatron.core.dist_checkpointing.validation import StrictHandling from megatron.core.fusions.fused_bias_geglu import quick_gelu from megatron.core.model_parallel_config import _parse_pad_packed_seq_alignment @@ -1412,8 +1412,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 @@ -2243,17 +2251,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: @@ -2994,6 +3008,15 @@ def _add_network_size_args(parser): 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', diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 395d5a57eca..24a92ac54f8 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1369,7 +1369,7 @@ def generate_state_dict( def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model): state_dict = raw_state_dict.copy() handle_fp8_extra_state_case(state_dict["model"]) - if args.swiglu: + if args.swiglu or getattr(args, "situ_glu", False): if "optimizer" in state_dict: model_state_dict, optimizer_state_dict = handle_swiglu_in_state_dict( model, state_dict["model"], state_dict["optimizer"] @@ -1927,6 +1927,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 4c55abfd22c..b1c1f451de8 100644 --- a/megatron/training/theoretical_memory_usage.py +++ b/megatron/training/theoretical_memory_usage.py @@ -19,7 +19,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 85fc717d901..b6c428233f0 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -966,9 +966,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 @@ -1337,7 +1337,7 @@ def transformer_flops(): 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)), moe_latent_size=args.moe_latent_size, moe_ffn_hidden_size=( args.moe_ffn_hidden_size diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index a9c015b4011..f6f3d203ca8 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -301,6 +301,8 @@ "rotary_base_per_layer": 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 5e35d3bda9f..918c0bd51fc 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -24,6 +24,7 @@ _build_sharded_state_dict_metadata, _load_base_checkpoint, get_checkpoint_tracker_filename, + load_args_from_checkpoint, load_checkpoint, read_metadata, save_checkpoint, @@ -74,6 +75,39 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() +@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 f1872e3f744..153edae6d92 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -126,6 +126,16 @@ def _make_mla_hybrid_args(): 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 09e26605a98..e04cdcd6094 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, @@ -435,6 +435,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, @@ -471,6 +478,7 @@ def register_forward_hook(self, hook): ops=SimpleNamespace( GroupedLinear=FakeGroupedLinear, ScaledSwiGLU=FakeScaledSwiGLU, + ScaledSiTUGLU=FakeScaledSiTUGLU, ScaledClampedQGeGLU=FakeScaledClampedQGeGLU, ScaledSReLU=FakeScaledSReLU, Sequential=FakeSequential, @@ -583,7 +591,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() @@ -624,11 +632,15 @@ def _make_fused_impl_support_module( use_fused_weighted_squared_relu=use_fused_weighted_squared_relu, moe_mlp_glu_interleave_size=moe_mlp_glu_interleave_size, moe_apply_probs_on_input=False, + 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, @@ -640,7 +652,45 @@ 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) + monkeypatch.setenv("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "1") + _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) @@ -649,8 +699,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 08da6c1ed0e..63538e0acd7 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 _FakeTESequential(torch.nn.Module): def append(self, module): self.add_module(str(len(self._modules)), module) @@ -90,6 +99,7 @@ def _fake_te_module(linear_cls=_FakeTELinear): ops=SimpleNamespace( GroupedLinear=_FakeTEGroupedLinear, ScaledSwiGLU=_FakeTEScaledSwiGLU, + ScaledSiTUGLU=_FakeTEScaledSiTUGLU, Sequential=_FakeTESequential, ), fp8_autocast=_FakeFP8Autocast, @@ -123,6 +133,10 @@ def _fake_shared_expert(**config_kwargs): add_bias_linear=False, 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, @@ -187,8 +201,8 @@ def test_validate_fused_grouped_swiglu_requires_te(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"), @@ -238,6 +252,20 @@ 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 + + def test_fused_grouped_swiglu_ops_replay_linear_pre_forward_hooks(monkeypatch): _patch_fake_shared_expert_te(monkeypatch) shared_expert = _fake_shared_expert() @@ -476,6 +504,9 @@ def test_shared_expert_clamped_swiglu(self, bias_activation_fusion): moe_shared_expert_overlap=False, moe_token_dispatcher_type="alltoall", activation_func_clamp_value=None, + situ_glu_beta1=4.0, + situ_glu_beta2=25.0, + use_fused_weighted_squared_relu=False, bias_activation_fusion=bias_activation_fusion, ).to(dtype=torch.bfloat16) moe_layer_unclamped.load_state_dict(moe_layer_overlap.state_dict()) 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_dense_mlp_spec.py b/tests/unit_tests/transformer/test_te_fused_dense_mlp_spec.py index d83cca760be..7fa13b52bf8 100644 --- a/tests/unit_tests/transformer/test_te_fused_dense_mlp_spec.py +++ b/tests/unit_tests/transformer/test_te_fused_dense_mlp_spec.py @@ -54,12 +54,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"): TEFusedDenseMLP(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"): TEFusedDenseMLP(config, _make_submodules()) def test_add_bias_linear_raises(self): 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 de31391d076..84e65ce3d47 100644 --- a/tools/checkpoint/loader_base.py +++ b/tools/checkpoint/loader_base.py @@ -443,7 +443,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