fix: EXAONE when using tie_word_embeddings (#5759)

2025-05-21 11:30:04 -07:00
parent cfe48c5902
commit ada268fd05
1 changed files with 8 additions and 3 deletions
@@ -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()