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 15, 04:45 UTC
Mark invariant normalizer in Gemma as non-persistent (#19788)
Signed-off-by: Yu-Hang Tang <Tang.Maxin@gmail.com>
Yu-Hang "Maxin" Tang committed
Jun 18, 2025 at 22:56 UTC
83ca9ae47b4476c900738b6db0ff5fcdfce13a7b
4 files changed
+29
-3
tests/models/language/generation/test_gemma.py
new
+20
@@ -0,0 +1,20 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
+import pytest
5
+
6
+MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"]
7
+
8
+
9
+@pytest.mark.parametrize("model", MODELS)
10
+def test_dummy_loader(vllm_runner, model: str) -> None:
11
+ with vllm_runner(
12
+ model,
13
+ load_format="dummy",
14
+ ) as llm:
15
+ normalizers = llm.collective_rpc(lambda self: self.worker.model_runner.
16
+ model.model.normalizer.cpu().item())
17
+ assert np.allclose(
18
+ normalizers,
19
+ llm.llm_engine.model_config.hf_config.hidden_size**0.5,
20
+ rtol=1e-3)
vllm/model_executor/models/gemma.py
+3
-1
@@ -281,7 +281,9 @@ class GemmaModel(nn.Module):
281
# data type such as bfloat16, not float32.
282
# See https://github.com/huggingface/transformers/pull/29402
283
normalizer = self.config.hidden_size**0.5
284
- self.register_buffer("normalizer", torch.tensor(normalizer))
284
+ self.register_buffer("normalizer",
285
+ torch.tensor(normalizer),
286
+ persistent=False)
287
self.make_empty_intermediate_tensors = (
288
make_empty_intermediate_tensors_factory(
289
["hidden_states", "residual"], config.hidden_size))
vllm/model_executor/models/gemma2.py
+3
-1
@@ -267,7 +267,9 @@ class Gemma2Model(nn.Module):
267
# data type such as bfloat16, not float32.
268
# See https://github.com/huggingface/transformers/pull/29402
269
normalizer = self.config.hidden_size**0.5
270
- self.register_buffer("normalizer", torch.tensor(normalizer))
270
+ self.register_buffer("normalizer",
271
+ torch.tensor(normalizer),
272
+ persistent=False)
273
self.make_empty_intermediate_tensors = (
274
make_empty_intermediate_tensors_factory(
275
["hidden_states", "residual"], config.hidden_size))
vllm/model_executor/models/gemma3.py
+3
-1
@@ -371,7 +371,9 @@ class Gemma3Model(nn.Module):
371
# data type such as bfloat16, not float32.
372
# See https://github.com/huggingface/transformers/pull/29402
373
normalizer = self.config.hidden_size**0.5
374
- self.register_buffer("normalizer", torch.tensor(normalizer))
374
+ self.register_buffer("normalizer",
375
+ torch.tensor(normalizer),
376
+ persistent=False)
377
self.make_empty_intermediate_tensors = (
378
make_empty_intermediate_tensors_factory(
379
["hidden_states", "residual"], config.hidden_size))