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

xjpang · Mar 4, 2024 · ed47c4e · ed47c4e
1 parent 92d45f8
commit ed47c4e
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,