From ada268fd05b214ca8c674517d6d46958ee476af8 Mon Sep 17 00:00:00 2001 From: Kyungmin Lee <30465912+lkm2835@users.noreply.github.com> Date: Thu, 22 May 2025 03:30:04 +0900 Subject: [PATCH] fix: EXAONE when using tie_word_embeddings (#5759) --- python/sglang/srt/models/exaone.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/models/exaone.py b/python/sglang/srt/models/exaone.py index 430c1d58b..1e4dfb3df 100644 --- a/python/sglang/srt/models/exaone.py +++ b/python/sglang/srt/models/exaone.py @@ -307,9 +307,14 @@ class ExaoneForCausalLM(nn.Module): self.transformer = ExaoneModel( config, quant_config=quant_config, prefix=add_prefix("transformer", prefix) ) - self.lm_head = ParallelLMHead( - config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix) - ) + if self.config.tie_word_embeddings: + self.lm_head = self.transformer.wte + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("lm_head", prefix), + ) self.logits_processor = LogitsProcessor(config) @torch.no_grad()