hotfix: Fix number of KV heads (#2202)

Fix number of KV heads
huggingface · Jul 26, 2024 · b2431ca · b2431ca
1 parent 361aae7
commit b2431ca
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py
@@ -906,8 +906,8 @@ def __init__(
         # Validation is done in the model itself
         if num_kv_heads is None:
             # Order is important here.
-            for attr in ["num_key_value_heads", "num_key_value_heads", "n_head"]:
-                num_kv_heads = getattr(config, "num_attention_heads", None)
+            for attr in ["num_key_value_heads", "num_attention_heads", "n_head"]:
+                num_kv_heads = getattr(config, attr, None)
                 if num_kv_heads is not None:
                     break
             if num_kv_heads is None: