A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 06:00 UTC
[ROCm] Fallback pytorch GELU with tanh approximation to GELU() (#29244)
Signed-off-by: Divakar Verma <divakar.verma@amd.com> Signed-off-by: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Divakar Verma committed
Dec 1, 2025 at 20:02 UTC
4b4092499889415c3a3d23dd44f8e616a6c424d1
1 file changed
+28
-2
vllm/model_executor/layers/activation.py
+28
-2
@@ -159,6 +159,13 @@ class GeluAndMulSparse(CustomOp):
159
self.approximate = approximate
160
if approximate not in ("none", "tanh"):
161
raise ValueError(f"Unknown approximate mode: {approximate}")
162
+ if current_platform.is_rocm() and approximate == "tanh":
163
+ # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile
164
+ logger.warning_once(
165
+ "[ROCm] Pytorch's native GELU with tanh approximation is currently "
166
+ "unstable and produces garbage. Fallback to 'none' approximation."
167
+ )
168
+ self.approximate = "none"
169
170
# Sparsity.
171
if activation_sparsity == 0.0:
@@ -209,6 +216,12 @@ class GeluAndMul(CustomOp):
216
self.op = torch.ops._C.gelu_and_mul
217
elif approximate == "tanh":
218
self.op = torch.ops._C.gelu_tanh_and_mul
219
+ if current_platform.is_rocm() and approximate == "tanh":
220
+ logger.warning_once(
221
+ "[ROCm] PyTorch's native GELU with tanh approximation is unstable "
222
+ "with torch.compile. For native implementation, fallback to 'none' "
223
+ "approximation. The custom kernel implementation is unaffected."
224
+ )
225
elif current_platform.is_xpu():
226
from vllm._ipex_ops import ipex_ops
227
@@ -219,8 +232,12 @@ class GeluAndMul(CustomOp):
232
233
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
234
"""PyTorch-native implementation equivalent to forward()."""
235
+ # TODO: [ROCm] PyTorch's native GELU with tanh is unstable with torch.compile
236
+ approximate = self.approximate
237
+ if current_platform.is_rocm() and approximate == "tanh":
238
+ approximate = "none"
239
d = x.shape[-1] // 2
223
- return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:]
240
+ return F.gelu(x[..., :d], approximate=approximate) * x[..., d:]
241
242
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
243
d = x.shape[-1] // 2
@@ -522,7 +539,16 @@ _ACTIVATION_REGISTRY = LazyDict(
539
"gelu": lambda: nn.GELU(),
540
"gelu_fast": lambda: FastGELU(),
541
"gelu_new": lambda: NewGELU(),
525
- "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"),
542
+ "gelu_pytorch_tanh": lambda: (
543
+ # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile
544
+ logger.warning_once(
545
+ "[ROCm] PyTorch's native GELU with tanh approximation is unstable. "
546
+ "Falling back to GELU(approximate='none')."
547
+ ),
548
+ nn.GELU(approximate="none"),
549
+ )[1]
550
+ if current_platform.is_rocm()
551
+ else nn.GELU(approximate="tanh"),
552
"relu": lambda: nn.ReLU(),
553
"relu2": lambda: ReLUSquaredActivation(),
554
"silu": lambda: nn.SiLU(),