mark invariant normalizer factor as non-persistent to prevent it from being scrambled by the dummy weight loader

Author: yhtang

tests/models/language/generation/test_gemma.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import numpy as np
+
+
+MODELS = [
+    "google/gemma-2b",
+    "google/gemma-2-2b",
+    "google/gemma-3-4b-it"
+]
+
+
+@pytest.mark.parametrize("model", MODELS)
+def test_dummy_loader(vllm_runner, model: str) -> None:
+    with vllm_runner(
+            model,
+            load_format="dummy",
+    ) as llm:
+        normalizers = llm.collective_rpc(
+            lambda self: self.worker.model_runner.model.model.normalizer.cpu().item()
+        )
+        assert np.allclose(
+            normalizers,
+            llm.llm_engine.model_config.hf_config.hidden_size**0.5,
+            rtol=1e-3
+        )

vllm/model_executor/models/gemma.py

@@ -281,7 +281,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         # data type such as bfloat16, not float32.
         # See https://github.com/huggingface/transformers/pull/29402
         normalizer = self.config.hidden_size**0.5
-        self.register_buffer("normalizer", torch.tensor(normalizer))
+        self.register_buffer(
+            "normalizer", torch.tensor(normalizer), persistent=False
+        )
         self.make_empty_intermediate_tensors = (
             make_empty_intermediate_tensors_factory(
                 ["hidden_states", "residual"], config.hidden_size))

vllm/model_executor/models/gemma2.py

@@ -267,7 +267,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         # data type such as bfloat16, not float32.
         # See https://github.com/huggingface/transformers/pull/29402
         normalizer = self.config.hidden_size**0.5
-        self.register_buffer("normalizer", torch.tensor(normalizer))
+        self.register_buffer(
+            "normalizer", torch.tensor(normalizer), persistent=False
+        )
         self.make_empty_intermediate_tensors = (
             make_empty_intermediate_tensors_factory(
                 ["hidden_states", "residual"], config.hidden_size))

vllm/model_executor/models/gemma3.py

@@ -371,7 +371,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         # data type such as bfloat16, not float32.
         # See https://github.com/huggingface/transformers/pull/29402
         normalizer = self.config.hidden_size**0.5
-        self.register_buffer("normalizer", torch.tensor(normalizer))
+        self.register_buffer(
+            "normalizer", torch.tensor(normalizer), persistent=False
+        )
         self.make_empty_intermediate_tensors = (
             make_empty_intermediate_tensors_factory(
                 ["hidden_states", "residual"], config.hidden_size))