support falcon40b

intel · Mar 13, 2024 · 8fe6c74 · 8fe6c74
1 parent 1c018ec
commit 8fe6c74
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/neural_speed/convert/convert_quantized_falcon.py b/neural_speed/convert/convert_quantized_falcon.py
@@ -54,10 +54,7 @@ def main(args_in: Optional[List[str]] = None) -> None:
         ftype = 1
 
     # 1. write hparams
-    #
-    # "n_head": 71,
-    # "n_layer": 32,
-    n_head_kv = hparams.get("num_kv_heads", 1)
+    n_head_kv = hparams.get("n_head_kv", 1)
     n_head = hparams["n_head"]
     head_dim = hparams["hidden_size"] // n_head