diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 8ef49579354..fad6966d0a8 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations +import inspect import logging from collections.abc import Callable from contextlib import nullcontext @@ -339,16 +340,34 @@ def _is_fused_impl_supported(self) -> bool: if not isinstance(self.linear_fc2, te.pytorch.GroupedLinear): return False - # Check activation: SwiGLU (ScaledSwiGLU) or quick GEGLU (ScaledClampedQGeGLU, TE >= 2.15) - if self.config.gated_linear_unit: - if self.activation_func == F.silu: - return True - if self.activation_func == quick_gelu: - try: - from transformer_engine.pytorch.ops import ScaledClampedQGeGLU # noqa: F401 - except ImportError: - return False - return True + # Check activation: SwiGLU, 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, + quick_gelu, + ) + use_srelu_fusion = ( + self.config.activation_func == squared_relu + and self.config.use_fused_weighted_squared_relu + and not self.config.gated_linear_unit + ) + if not (use_glu_fusion or use_srelu_fusion): + return False + if self.config.activation_func == F.silu: + return True + if self.config.activation_func == quick_gelu: + try: + from transformer_engine.pytorch.ops import ScaledClampedQGeGLU # noqa: F401 + except ImportError: + return False + return True + if self.config.activation_func == squared_relu: + try: + from transformer_engine.pytorch.ops import ScaledSReLU # noqa: F401 + except ImportError: + return False + return True return False @@ -410,21 +429,64 @@ def _make_fused_ops(self) -> torch.nn.Module: setattr(op, "bias", getattr(self.linear_fc1, "bias")) ops.append(op) - # Activation and post-multiply probs (SwiGLU or clamped quick-GeGLU) + # Activation and post-multiply probs (SwiGLU, clamped quick-GeGLU, or SReLU) glu_interleave = self.config.moe_mlp_glu_interleave_size - if self.activation_func == F.silu and self.config.gated_linear_unit: - op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave) - elif self.activation_func == quick_gelu and self.config.gated_linear_unit: + activation_recompute_in_mlp = bool(getattr(self, "activation_recompute", False)) + if self.config.activation_func == F.silu and self.config.gated_linear_unit: + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledSwiGLU).parameters + ): + op = te.pytorch.ops.ScaledSwiGLU( + glu_interleave_size=glu_interleave, + activation_recompute_in_mlp=activation_recompute_in_mlp, + ) + else: + op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave) + elif self.config.activation_func == quick_gelu and self.config.gated_linear_unit: clamp = self.config.activation_func_clamp_value if clamp is not None: - op = te.pytorch.ops.ScaledClampedQGeGLU( - glu_interleave_size=glu_interleave, limit=clamp + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledClampedQGeGLU).parameters + ): + op = te.pytorch.ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave, + activation_recompute_in_mlp=activation_recompute_in_mlp, + limit=clamp, + ) + else: + op = te.pytorch.ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave, limit=clamp + ) + else: + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledClampedQGeGLU).parameters + ): + op = te.pytorch.ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave, + activation_recompute_in_mlp=activation_recompute_in_mlp, + ) + else: + op = te.pytorch.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave) + elif ( + self.config.activation_func == squared_relu + and self.config.use_fused_weighted_squared_relu + and not self.config.gated_linear_unit + ): + if ( + "activation_recompute_in_mlp" + in inspect.signature(te.pytorch.ops.ScaledSReLU).parameters + ): + op = te.pytorch.ops.ScaledSReLU( + activation_recompute_in_mlp=activation_recompute_in_mlp ) else: - op = te.pytorch.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave) + op = te.pytorch.ops.ScaledSReLU() else: raise RuntimeError( - "_make_fused_ops expected SwiGLU or quick_gelu with gated_linear_unit; " + "_make_fused_ops expected SwiGLU, quick_gelu, or weighted squared_relu; " "call _is_fused_impl_supported() before constructing fused ops." ) ops.append(op) diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index 6926f34940d..f6cef08b1b6 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -1,13 +1,15 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. import argparse -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace import pytest import torch 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.fusions.fused_bias_geglu import quick_gelu from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_submodules, @@ -77,9 +79,10 @@ def need_backward_dw(self): return False class FakeScaledSwiGLU(torch.nn.Module): - def __init__(self, glu_interleave_size): + def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): super().__init__() self.glu_interleave_size = glu_interleave_size + self.activation_recompute_in_mlp = activation_recompute_in_mlp class FakeSequential(list): def register_forward_pre_hook(self, hook): @@ -100,7 +103,10 @@ def register_forward_pre_hook(self, hook): module = TEGroupedMLP.__new__(TEGroupedMLP) torch.nn.Module.__init__(module) module.config = SimpleNamespace( - moe_mlp_glu_interleave_size=16, delay_wgrad_compute=False, activation_func_clamp_value=None + moe_mlp_glu_interleave_size=16, + delay_wgrad_compute=False, + activation_func_clamp_value=None, + activation_func=F.silu, ) module.activation_func = F.silu module.config.gated_linear_unit = True @@ -279,9 +285,10 @@ def need_backward_dw(self): return False class FakeScaledSwiGLU(torch.nn.Module): - def __init__(self, glu_interleave_size): + def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): super().__init__() self.glu_interleave_size = glu_interleave_size + self.activation_recompute_in_mlp = activation_recompute_in_mlp class FakeSequential(list): def register_forward_pre_hook(self, hook): @@ -305,9 +312,11 @@ def register_forward_pre_hook(self, hook): moe_mlp_glu_interleave_size=8, delay_wgrad_compute=False, activation_func_clamp_value=None, + activation_func=F.silu, gated_linear_unit=True, ) module.activation_func = F.silu + module.activation_recompute = True module.linear_fc1 = FakeGroupedLinear( 2, 4, @@ -338,6 +347,7 @@ def register_forward_pre_hook(self, hook): assert ops[0].weight is module.linear_fc1.weight assert ops[1].glu_interleave_size == 8 + assert ops[1].activation_recompute_in_mlp is True assert ops[2].weight0 is module.linear_fc2.weight0 assert ops[2].weight1 is module.linear_fc2.weight1 assert ops[2].bias0 is module.linear_fc2.bias0 @@ -378,16 +388,23 @@ def need_backward_dw(self): return False class FakeScaledSwiGLU(torch.nn.Module): - def __init__(self, glu_interleave_size): + def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False): super().__init__() self.glu_interleave_size = glu_interleave_size + self.activation_recompute_in_mlp = activation_recompute_in_mlp class FakeScaledClampedQGeGLU(torch.nn.Module): - def __init__(self, glu_interleave_size, limit=None): + def __init__(self, glu_interleave_size, *, activation_recompute_in_mlp=False, limit=None): super().__init__() self.glu_interleave_size = glu_interleave_size + self.activation_recompute_in_mlp = activation_recompute_in_mlp self.limit = limit + class FakeScaledSReLU(torch.nn.Module): + def __init__(self, *, activation_recompute_in_mlp=False): + super().__init__() + self.activation_recompute_in_mlp = activation_recompute_in_mlp + class FakeSequential(list): def register_forward_pre_hook(self, hook): self.forward_pre_hook = hook @@ -400,6 +417,7 @@ def register_forward_pre_hook(self, hook): GroupedLinear=FakeGroupedLinear, ScaledSwiGLU=FakeScaledSwiGLU, ScaledClampedQGeGLU=FakeScaledClampedQGeGLU, + ScaledSReLU=FakeScaledSReLU, Sequential=FakeSequential, ), ) @@ -419,9 +437,11 @@ def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): moe_mlp_glu_interleave_size=4, delay_wgrad_compute=False, activation_func_clamp_value=7.0, + activation_func=quick_gelu, gated_linear_unit=True, ) module.activation_func = quick_gelu + module.activation_recompute = True common = dict( device="cuda", dtype=torch.bfloat16, @@ -440,9 +460,177 @@ def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): activation = ops[1] assert type(activation).__name__ == "FakeScaledClampedQGeGLU" assert activation.glu_interleave_size == 4 + assert activation.activation_recompute_in_mlp is True assert activation.limit == 7.0 +def test_make_fused_ops_uses_scaled_srelu_for_weighted_squared_relu(monkeypatch): + """squared_relu + weighted squared-ReLU fusion -> ScaledSReLU.""" + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + moe_mlp_glu_interleave_size=None, + delay_wgrad_compute=False, + activation_func_clamp_value=None, + activation_func=squared_relu, + gated_linear_unit=False, + use_fused_weighted_squared_relu=True, + ) + module.activation_func = squared_relu + module.activation_recompute = True + common = dict( + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=False, + single_grouped_weight=False, + ) + module.linear_fc1 = FakeGroupedLinear(2, 4, 8, bias=False, **common) + module.linear_fc2 = FakeGroupedLinear(2, 8, 4, bias=False, **common) + module.linear_fc1.weight0 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc1.weight1 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc2.weight0 = torch.nn.Parameter(torch.ones(4, 8)) + module.linear_fc2.weight1 = torch.nn.Parameter(torch.ones(4, 8)) + + ops = module._make_fused_ops() + + assert type(ops[1]).__name__ == "FakeScaledSReLU" + assert ops[1].activation_recompute_in_mlp is True + + +def test_make_fused_ops_rejects_scaled_srelu_with_gated_linear_unit(monkeypatch): + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + moe_mlp_glu_interleave_size=None, + delay_wgrad_compute=False, + activation_func_clamp_value=None, + activation_func=squared_relu, + gated_linear_unit=True, + use_fused_weighted_squared_relu=True, + ) + module.activation_func = squared_relu + common = dict( + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=False, + single_grouped_weight=False, + ) + module.linear_fc1 = FakeGroupedLinear(2, 4, 8, bias=False, **common) + module.linear_fc2 = FakeGroupedLinear(2, 8, 4, bias=False, **common) + module.linear_fc1.weight0 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc1.weight1 = torch.nn.Parameter(torch.ones(8, 4)) + 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"): + module._make_fused_ops() + + +def _install_fake_te_ops_modules( + monkeypatch, fake_te, *, include_clamped_qgeglu=True, include_scaled_srelu=True +): + transformer_engine = ModuleType("transformer_engine") + pytorch = ModuleType("transformer_engine.pytorch") + ops = ModuleType("transformer_engine.pytorch.ops") + ops.GroupedLinear = fake_te.pytorch.GroupedLinear + ops.ScaledSwiGLU = fake_te.pytorch.ops.ScaledSwiGLU + if include_clamped_qgeglu: + ops.ScaledClampedQGeGLU = fake_te.pytorch.ops.ScaledClampedQGeGLU + if include_scaled_srelu: + ops.ScaledSReLU = fake_te.pytorch.ops.ScaledSReLU + pytorch.ops = ops + transformer_engine.pytorch = pytorch + monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine) + monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", pytorch) + monkeypatch.setitem(sys.modules, "transformer_engine.pytorch.ops", ops) + + +def _make_fused_impl_support_module( + FakeGroupedLinear, *, activation_func, gated_linear_unit, use_fused_weighted_squared_relu=False +): + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + activation_func=activation_func, + gated_linear_unit=gated_linear_unit, + use_fused_weighted_squared_relu=use_fused_weighted_squared_relu, + moe_apply_probs_on_input=False, + ) + module.activation_func = object() + module.tp_group = SimpleNamespace(size=lambda: 1) + module.offload_expert_fc1 = False + module.offload_moe_act = False + common = dict( + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=False, + single_grouped_weight=False, + ) + module.linear_fc1 = FakeGroupedLinear(2, 4, 8, bias=False, **common) + module.linear_fc2 = FakeGroupedLinear(2, 8, 4, bias=False, **common) + return module + + +def test_is_fused_impl_supported_uses_config_activation_for_swiglu(monkeypatch): + 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=F.silu, gated_linear_unit=True + ) + + assert module._is_fused_impl_supported() is True + + +@pytest.mark.parametrize( + ("use_fused_weighted_squared_relu", "gated_linear_unit", "expected"), + [(True, False, True), (False, False, False), (True, True, False)], +) +def test_is_fused_impl_supported_gates_scaled_srelu_on_weighted_flag_and_non_glu( + monkeypatch, use_fused_weighted_squared_relu, gated_linear_unit, expected +): + 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=squared_relu, + gated_linear_unit=gated_linear_unit, + use_fused_weighted_squared_relu=use_fused_weighted_squared_relu, + ) + + assert module._is_fused_impl_supported() is expected + + +def test_is_fused_impl_supported_requires_scaled_srelu_op(monkeypatch): + 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, include_scaled_srelu=False) + + module = _make_fused_impl_support_module( + FakeGroupedLinear, + activation_func=squared_relu, + gated_linear_unit=False, + use_fused_weighted_squared_relu=True, + ) + + assert module._is_fused_impl_supported() is False + + def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): """single_grouped_bias=True → bias attached as `bias` (not `bias{idx}`).""" fake_te, FakeGroupedLinear = _make_fake_te_namespace() @@ -454,6 +642,7 @@ def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): moe_mlp_glu_interleave_size=2, delay_wgrad_compute=False, activation_func_clamp_value=None, + activation_func=F.silu, gated_linear_unit=True, ) module.activation_func = F.silu