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 11, 07:30 UTC
[Model] Gemma3: Support untied word embeddings (#30827)
Signed-off-by: www-spam <panmahm@naver.com>
KimHyemin committed
Dec 18, 2025 at 00:11 UTC
196cdc3224112df7f68c901fe4c5314875a65be8
1 file changed
+15
-4
vllm/model_executor/models/gemma3.py
+15
-4
@@ -39,7 +39,10 @@ from vllm.model_executor.layers.linear import (
39
from vllm.model_executor.layers.logits_processor import LogitsProcessor
40
from vllm.model_executor.layers.quantization import QuantizationConfig
41
from vllm.model_executor.layers.rotary_embedding import get_rope
42
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
42
+from vllm.model_executor.layers.vocab_parallel_embedding import (
43
+ ParallelLMHead,
44
+ VocabParallelEmbedding,
45
+)
46
from vllm.model_executor.model_loader.weight_utils import (
47
default_weight_loader,
48
maybe_remap_kv_scale_name,
@@ -532,12 +535,20 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
535
536
super().__init__()
537
self.config = config
535
- # currently all existing Gemma models have `tie_word_embeddings` enabled
536
- assert config.tie_word_embeddings
538
self.quant_config = quant_config
539
self.model = Gemma3Model(
540
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
541
)
542
+
543
+ self.lm_head = ParallelLMHead(
544
+ config.vocab_size,
545
+ config.hidden_size,
546
+ quant_config=quant_config,
547
+ prefix=maybe_prefix(prefix, "lm_head"),
548
+ )
549
+ if config.tie_word_embeddings:
550
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
551
+
552
self.logits_processor = LogitsProcessor(
553
config.vocab_size, soft_cap=config.final_logit_softcapping
554
)
@@ -565,7 +576,7 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
576
self,
577
hidden_states: torch.Tensor,
578
) -> torch.Tensor | None:
568
- logits = self.logits_processor(self.model.embed_tokens, hidden_states)
579
+ logits = self.logits_processor(self.lm_head, hidden_states)
580
return logits
581
582
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: