[Bugfix] Load untied Gemma LM head weights (#53170)
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com>
Kevin H. Luu committed
Aug 20, 2026 at 17:47 UTC
d29f7f5c9294be8e489dac34d45a939b95a06336
2 files changed
+58
-2
tests/models/language/generation/test_gemma.py
+45
@@ -1,11 +1,56 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+from typing import cast
5
+
6
import numpy as np
7
import pytest
8
+import torch
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
12
+from vllm.model_executor.models import gemma
13
14
MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"]
15
16
17
+@pytest.mark.cpu_test
18
+@pytest.mark.usefixtures("dist_init")
19
+def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None:
20
+ """A physical LM head must load after checkpoint-driven untying."""
21
+
22
+ class StubGemmaModel(torch.nn.Module):
23
+ def __init__(self, *, vllm_config, prefix):
24
+ super().__init__()
25
+ self.embed_tokens = VocabParallelEmbedding(4, 2)
26
+ self.make_empty_intermediate_tensors = None
27
+
28
+ monkeypatch.setattr(gemma, "GemmaModel", StubGemmaModel)
29
+ config = SimpleNamespace(
30
+ vocab_size=4,
31
+ hidden_size=2,
32
+ tie_word_embeddings=False,
33
+ )
34
+ vllm_config = SimpleNamespace(
35
+ model_config=SimpleNamespace(hf_config=config),
36
+ quant_config=None,
37
+ )
38
+ model = gemma.GemmaForCausalLM(vllm_config=cast(VllmConfig, vllm_config))
39
+ embedding_weight = torch.full((4, 2), 1.0)
40
+ lm_head_weight = torch.full((4, 2), 2.0)
41
+
42
+ loaded = model.load_weights(
43
+ [
44
+ ("model.embed_tokens.weight", embedding_weight),
45
+ ("lm_head.weight", lm_head_weight),
46
+ ]
47
+ )
48
+
49
+ assert loaded == {"model.embed_tokens.weight", "lm_head.weight"}
50
+ assert torch.equal(model.model.embed_tokens.weight[:4], embedding_weight)
51
+ assert torch.equal(model.lm_head.weight[:4], lm_head_weight)
52
+
53
+
54
@pytest.mark.parametrize("model", MODELS)
55
def test_dummy_loader(vllm_runner, monkeypatch, model: str) -> None:
56
with monkeypatch.context() as m:
vllm/model_executor/models/gemma.py
+13
-2
@@ -41,7 +41,10 @@ from vllm.model_executor.layers.linear import (
41
from vllm.model_executor.layers.logits_processor import LogitsProcessor
42
from vllm.model_executor.layers.quantization import QuantizationConfig
43
from vllm.model_executor.layers.rotary_embedding import get_rope
44
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
44
+from vllm.model_executor.layers.vocab_parallel_embedding import (
45
+ ParallelLMHead,
46
+ VocabParallelEmbedding,
47
+)
48
from vllm.sequence import IntermediateTensors
49
50
from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant
@@ -351,6 +354,14 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
354
self.model = GemmaModel(
355
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
356
)
357
+ self.lm_head = ParallelLMHead(
358
+ config.vocab_size,
359
+ config.hidden_size,
360
+ quant_config=quant_config,
361
+ prefix=maybe_prefix(prefix, "lm_head"),
362
+ )
363
+ if config.tie_word_embeddings:
364
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
365
self.logits_processor = LogitsProcessor(config.vocab_size)
366
self.make_empty_intermediate_tensors = (
367
self.model.make_empty_intermediate_tensors
@@ -375,7 +386,7 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
386
self,
387
hidden_states: torch.Tensor,
388
) -> torch.Tensor | None:
378
- logits = self.logits_processor(self.model.embed_tokens, hidden_states)
389
+ logits = self.logits_processor(self.lm_head, hidden_states)
390
return logits
391
392
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: