From 08b19224054088e36aa25762a99b3e53adaadb21 Mon Sep 17 00:00:00 2001 From: Yash Agarwal Date: Wed, 16 Sep 2026 03:46:39 +0530 Subject: [PATCH] fix(gemma4): set lm_head_is_tied for Gemma4UnifiedForConditionalGeneration (#35809) Co-authored-by: kpham-sgl Co-authored-by: Claude Fable 5.1 --- python/sglang/srt/models/gemma4_unified.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/gemma4_unified.py b/python/sglang/srt/models/gemma4_unified.py index 67d0771d4..610f0c20c 100644 --- a/python/sglang/srt/models/gemma4_unified.py +++ b/python/sglang/srt/models/gemma4_unified.py @@ -190,7 +190,8 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration): ) text_tie = getattr(text_config, "tie_word_embeddings", True) - if self.pp_group.world_size == 1 and text_tie: + self.lm_head_is_tied = self.pp_group.world_size == 1 and text_tie + if self.lm_head_is_tied: self.lm_head = self.language_model.embed_tokens elif self.pp_group.is_last_rank: self.lm_head = ParallelLMHead(