export

InternLM · Dec 9, 2024 · 5e967c0 · 5e967c0
1 parent 00c435a
commit 5e967c0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lmdeploy/turbomind/deploy/module.py b/lmdeploy/turbomind/deploy/module.py
@@ -229,6 +229,8 @@ def _export(self, idx: int, qkvo, kind: str, pack_fn, **kwargs):
         if is_lora_a:
             qkv, o = map(transpose, qkvo)
         else:
+            if self.model.repeat_kv:
+                qkvo = self._repeat_kv(qkvo, kind)
             qkv, o = self._reorder_and_merge(qkvo)
         self.model.save_split(pack_fn(qkv),
                               self._attn.format(idx, 'w_qkv', kind),