[Bugfix] support tie_word_embeddings for all models (#5724)

2024-08-19 20:00:04 -07:00
parent 0df7ec0b2d
commit f4fc7337bf
30 changed files with 90 additions and 16 deletions
--- a/vllm/model_executor/models/qwen.py
+++ b/vllm/model_executor/models/qwen.py
@@ -252,6 +252,8 @@ class QWenLMHeadModel(nn.Module):
        self.lm_head = ParallelLMHead(config.vocab_size,
                                      config.hidden_size,
                                      quant_config=quant_config)
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.transformer.wte.weight
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()