Fix llama3.2-1b inference error by handling tie_word_embedding (#2568)

InternLM · Oct 9, 2024 · 231e5bb · 231e5bb
1 parent c722ff5
commit 231e5bb
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/lmdeploy/pytorch/models/llama.py b/lmdeploy/pytorch/models/llama.py
@@ -371,6 +371,11 @@ def forward(
         )
         return hidden_states
 
+    def update_weights(self):
+        """update weights."""
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
     def get_logits(self, hidden_states: torch.Tensor):
         """compute logits of the model output."""
         return self.lm_head(hidden_states)