From 90f145aaf724194ccffeb3ea6a68e9457ff00169 Mon Sep 17 00:00:00 2001 From: Hiroaki Mikami Date: Sat, 9 May 2026 03:29:11 +0900 Subject: [PATCH] [Models][Gemma3/Gemma4] Support hidden_act variants in gated MLP (#40588) Signed-off-by: Hiroaki Mikami --- tests/model_executor/test_gemma_hidden_act.py | 60 +++++++++++++++++++ vllm/model_executor/layers/activation.py | 3 + vllm/model_executor/models/gemma3.py | 10 +--- vllm/model_executor/models/gemma4.py | 10 +--- 4 files changed, 67 insertions(+), 16 deletions(-) create mode 100644 tests/model_executor/test_gemma_hidden_act.py diff --git a/tests/model_executor/test_gemma_hidden_act.py b/tests/model_executor/test_gemma_hidden_act.py new file mode 100644 index 00000000000..d34851a6f10 --- /dev/null +++ b/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) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index df9459012ae..d5b67ef04e2 100644 --- a/vllm/model_executor/layers/activation.py +++ b/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(), } diff --git a/vllm/model_executor/models/gemma3.py b/vllm/model_executor/models/gemma3.py index b2352a3c926..f61f7c6f780 100644 --- a/vllm/model_executor/models/gemma3.py +++ b/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) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 31f2d6a28dd..62b7f557667 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/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)