From df41a57e09cc7590432fbc24c60480fed13ca348 Mon Sep 17 00:00:00 2001 From: Davi Silva Date: Sun, 26 Nov 2023 01:29:04 +0700 Subject: [PATCH] Fix: missing n_kv_heads for smaller models from huggingface (#2438) * fix: missing n_kv_heads for smaller models from huggingface * a lil golfing --- examples/llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llama.py b/examples/llama.py index 519eecd55e..eaa224c7fc 100755 --- a/examples/llama.py +++ b/examples/llama.py @@ -311,7 +311,7 @@ class LLaMa: else: weights = load(str(model_path)) if "model.embed_tokens.weight" in weights: - weights = convert_from_huggingface(weights, model, model_args["n_heads"], model_args["n_kv_heads"]) + weights = convert_from_huggingface(weights, model, model_args["n_heads"], model_args.get("n_kv_heads", model_args["n_heads"])) if quantize: weights = AbsmaxQuantizedLinear.quantize(weights)