Fix internlm after vllm-project#2860 (vllm-project#2861)

xjpang · Feb 22, 2024 · 51ed1d5 · 51ed1d5
1 parent 824c2c0
commit 51ed1d5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py
@@ -175,7 +175,8 @@ def __init__(
         self.self_attn = LlamaAttention(
             hidden_size=self.hidden_size,
             num_heads=config.num_attention_heads,
-            num_kv_heads=config.num_key_value_heads,
+            num_kv_heads=getattr(config, "num_key_value_heads",
+                                 config.num_attention_heads),
             rope_theta=rope_theta,
             rope_scaling=rope_scaling,
             max_position_embeddings=max_position_embeddings,