[Models][Gemma3/Gemma4] Support hidden_act variants in gated MLP (#40588)

Signed-off-by: Hiroaki Mikami <[email protected]>

tests/model_executor/test_gemma_hidden_act.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.activation import (
+    GeluAndMul,
+    SiluAndMul,
+    get_act_and_mul_fn,
+    get_act_fn,
+)
+from vllm.model_executor.models.gemma3 import Gemma3MLP
+from vllm.model_executor.models.gemma4 import Gemma4MLP
+
+
+@pytest.mark.parametrize(
+    ("activation_name", "expected_type"),
+    [
+        ("gelu_pytorch_tanh", GeluAndMul),
+        ("silu", SiluAndMul),
+        ("swish", SiluAndMul),
+    ],
+)
+def test_get_act_and_mul_fn_supports_gemma_hidden_act_aliases(
+    activation_name: str,
+    expected_type: type[torch.nn.Module],
+    default_vllm_config,
+) -> None:
+    assert isinstance(get_act_and_mul_fn(activation_name), expected_type)
+
+
+def test_get_act_fn_supports_swish_alias() -> None:
+    assert isinstance(get_act_fn("swish"), torch.nn.SiLU)
+
+
+@pytest.mark.parametrize("mlp_cls", [Gemma3MLP, Gemma4MLP])
+@pytest.mark.parametrize(
+    ("activation_name", "expected_type"),
+    [
+        ("gelu_pytorch_tanh", GeluAndMul),
+        ("silu", SiluAndMul),
+        ("swish", SiluAndMul),
+    ],
+)
+def test_gemma_mlp_supports_hidden_act_variants(
+    mlp_cls: type[torch.nn.Module],
+    activation_name: str,
+    expected_type: type[torch.nn.Module],
+    default_vllm_config,
+    dist_init,
+) -> None:
+    mlp = mlp_cls(
+        hidden_size=16,
+        intermediate_size=32,
+        hidden_activation=activation_name,
+    )
+
+    assert isinstance(mlp.act_fn, expected_type)
+    assert mlp(torch.randn(3, 16)).shape == (3, 16)

vllm/model_executor/layers/activation.py

@@ -712,6 +712,7 @@ _ACTIVATION_REGISTRY = LazyDict(
         "relu": lambda: nn.ReLU(),
         "relu2": lambda: ReLUSquaredActivation(),
         "silu": lambda: nn.SiLU(),
+        "swish": lambda: nn.SiLU(),
         "quick_gelu": lambda: QuickGELU(),
         "tanh": lambda: nn.Tanh(),
         "sigmoid": lambda: nn.Sigmoid(),
@@ -751,7 +752,9 @@ def get_act_fn(act_fn_name: str) -> nn.Module:
 _ACTIVATION_AND_MUL_REGISTRY: LazyDict[nn.Module] = LazyDict(
     {
         "gelu": lambda: GeluAndMul(),
+        "gelu_pytorch_tanh": lambda: GeluAndMul(approximate="tanh"),
         "silu": lambda: SiluAndMul(),
+        "swish": lambda: SiluAndMul(),
         "geglu": lambda: GeluAndMul(),
         "swigluoai": lambda: SwigluOAIAndMul(),
     }

vllm/model_executor/models/gemma3.py

@@ -26,7 +26,7 @@ from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
-from vllm.model_executor.layers.activation import GeluAndMul
+from vllm.model_executor.layers.activation import get_act_and_mul_fn
 from vllm.model_executor.layers.attention import (
     Attention,
     EncoderOnlyAttention,
@@ -88,13 +88,7 @@ class Gemma3MLP(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
         )
-        if hidden_activation != "gelu_pytorch_tanh":
-            raise ValueError(
-                "Gemma3 uses `gelu_pytorch_tanh` as the hidden activation "
-                "function. Please set `hidden_act` and `hidden_activation` to "
-                "`gelu_pytorch_tanh`."
-            )
-        self.act_fn = GeluAndMul(approximate="tanh")
+        self.act_fn = get_act_and_mul_fn(hidden_activation)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         gate_up, _ = self.gate_up_proj(x)

vllm/model_executor/models/gemma4.py

@@ -35,7 +35,7 @@ from vllm.distributed import (
 )
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
-from vllm.model_executor.layers.activation import GeluAndMul
+from vllm.model_executor.layers.activation import get_act_and_mul_fn
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
@@ -238,13 +238,7 @@ class Gemma4MLP(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
         )
-        if hidden_activation != "gelu_pytorch_tanh":
-            raise ValueError(
-                "Gemma4 uses `gelu_pytorch_tanh` as the hidden activation "
-                "function. Please set `hidden_act` and `hidden_activation` to "
-                "`gelu_pytorch_tanh`."
-            )
-        self.act_fn = GeluAndMul(approximate="tanh")
+        self.act_fn = get_act_and_mul_fn(hidden_activation)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         gate_up, _ = self.gate_up_proj(x)