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 6, 20:30 UTC
Make Gemma and Gemma 2 accept `inputs_embeds` like Gemma 3 (#36787)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 11, 2026 at 18:12 UTC
65986db6ba71abf4cf0639c5fd1477b0d8df8f5e
4 files changed
+25
-4
tests/basic_correctness/test_basic_correctness.py
+11
@@ -11,6 +11,8 @@ from unittest.mock import Mock
11
12
import pytest
13
import torch
14
+from packaging.version import Version
15
+from transformers import __version__ as TRANSFORMERS_VERSION
16
17
from vllm import LLM
18
from vllm.platforms import current_platform
@@ -91,6 +93,15 @@ def test_models(
93
if enable_prompt_embeds:
94
with torch.no_grad():
95
prompt_embeds = hf_model.get_prompt_embeddings(example_prompts)
96
+ if model == "hmellor/tiny-random-Gemma2ForCausalLM" and (
97
+ Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0")
98
+ ):
99
+ # For Gemma 1/2 models with Transformers 5.4.0+, the prompt embeddings
100
+ # are normalised in `get_prompt_embeddings`, like Gemma 3.
101
+ # For older versions, we need to manually normalise.
102
+ embed_scale = hf_model.config.hidden_size**0.5
103
+ normalizer = torch.tensor(embed_scale, dtype=prompt_embeds[0].dtype)
104
+ prompt_embeds = [p_e * normalizer for p_e in prompt_embeds]
105
106
with VllmRunner(
107
model,
tests/models/language/generation/test_common.py
+12
@@ -3,6 +3,8 @@
3
4
import pytest
5
import torch
6
+from packaging.version import Version
7
+from transformers import __version__ as TRANSFORMERS_VERSION
8
9
from vllm.platforms import current_platform
10
@@ -151,6 +153,16 @@ def test_models(
153
if prompt_embeds is not None:
154
embed = hf_model.model.get_input_embeddings()(token_ids)
155
156
+ if "gemma" in model.lower() and (
157
+ Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0")
158
+ ):
159
+ # For Gemma 1/2 models with Transformers 5.4.0+, the prompt
160
+ # embeddings are normalised in `get_prompt_embeddings`,
161
+ # like Gemma 3. For older versions, we need to manually normalise.
162
+ embed_scale = hf_model.config.hidden_size**0.5
163
+ normalizer = torch.tensor(embed_scale, dtype=embed.dtype)
164
+ embed *= normalizer
165
+
166
# MiniCPM models apply scale_emb to embeddings internally.
167
# vLLM expects pre-scaled embeddings when using inputs_embeds.
168
if model in EMBED_SCALING_MODELS:
vllm/model_executor/models/gemma.py
+1
-2
@@ -293,7 +293,7 @@ class GemmaModel(nn.Module):
293
)
294
295
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
296
- return self.embed_tokens(input_ids)
296
+ return self.embed_tokens(input_ids) * self.normalizer
297
298
def forward(
299
self,
@@ -307,7 +307,6 @@ class GemmaModel(nn.Module):
307
hidden_states = inputs_embeds
308
else:
309
hidden_states = self.embed_input_ids(input_ids)
310
- hidden_states *= self.normalizer
310
residual = None
311
else:
312
hidden_states = intermediate_tensors["hidden_states"]
vllm/model_executor/models/gemma2.py
+1
-2
@@ -284,7 +284,7 @@ class Gemma2Model(nn.Module):
284
)
285
286
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
287
- return self.embed_tokens(input_ids)
287
+ return self.embed_tokens(input_ids) * self.normalizer
288
289
def forward(
290
self,
@@ -298,7 +298,6 @@ class Gemma2Model(nn.Module):
298
hidden_states = inputs_embeds
299
else:
300
hidden_states = self.embed_input_ids(input_ids)
301
- hidden_states *= self.normalizer
301
residual = None
302
else:
303
assert intermediate_tensors is not None