ChatGLM Support (#1261)

2023-11-07 08:09:33 +08:00
parent e7f579eb97
commit 1a2bbc9301
7 changed files with 490 additions and 4 deletions
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -166,6 +166,10 @@ class ModelConfig:
        if getattr(self.hf_config, "num_key_value_heads", None) is not None:
            return (self.hf_config.num_key_value_heads //
                    parallel_config.tensor_parallel_size)
+        # For ChatGLM-2:
+        if getattr(self.hf_config, "multi_query_group_num", None) is not None:
+            return (self.hf_config.multi_query_group_num //
+                    parallel_config.tensor_parallel_size)
        total_num_attention_heads = self.hf_config.num_attention_heads
        return total_num_attention_heads // parallel_config.tensor_parallel_size