fix: Llama AWS with GQA (#114)

predibase · Dec 8, 2023 · 6b65aac · 6b65aac
1 parent 41c657e
commit 6b65aac
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/lorax_server/models/custom_modeling/flash_llama_modeling.py b/server/lorax_server/models/custom_modeling/flash_llama_modeling.py
@@ -184,7 +184,7 @@ def _load_gqa(config, prefix: str, weights):
         dim=0,
     )
 
-    if config.quantize not in ["gptq"]:
+    if config.quantize not in ["gptq", "awq"]:
         weight = weight.to(dtype=weights.dtype).to(device=weights.device)
 
         head_size = config.hidden_size // config.num_attention_heads