fix gpt2 gguf

Signed-off-by: Isotr0py <[email protected]>
vllm-project · Jan 27, 2025 · 4b19360 · 4b19360
1 parent 5204ff5
commit 4b19360
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/vllm/model_executor/models/gpt2.py b/vllm/model_executor/models/gpt2.py
@@ -258,13 +258,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.transformer = GPT2Model(vllm_config=vllm_config,
                                      prefix=maybe_prefix(
                                          prefix, "transformer"))
-        if self.config.tie_word_embeddings:
-            self.lm_head = self.transformer.wte
-        else:
-            self.lm_head = ParallelLMHead(self.config.vocab_size,
-                                          self.config.hidden_size,
-                                          quant_config=quant_config,
-                                          prefix=f"{prefix}.lm_head")
+
+        # GPT-2 ties the weights of the embedding layer and the final
+        # linear layer.
+        self.lm_head = ParallelLMHead(self.config.vocab_size,
+                                      self.config.hidden_size,
+                                      quant_config=quant_config,
+                                      prefix=f"{prefix}.lm_head")
+        self.lm_head = self.lm_head.tie_weights(self.transformer.wte)
+
         self.logits_processor = LogitsProcessor(config.vocab_size)
         self.sampler = get_sampler()
         self.make_empty_intermediate_tensors = (