LLaMA OOM hotfix

2025-04-21 04:25:30 +02:00
parent a355f8b5ea
commit 9d282bcdfd
1 changed files with 5 additions and 1 deletions
@@ -244,8 +244,12 @@ def gguf_clip_loader(path):
    elif arch in {"llama"}:
        temb_key = "token_embd.weight"
        if temb_key in sd and sd[temb_key].shape != (128320, 4096):
-            # This still works. Raise error?
+            # This still works. TODO: Only relevant for llava?
            logging.warning("Warning! token_embd shape may be incorrect for llama 3 model!")
+        if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024):
+            # See note above for T5.
+            logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.")
+            sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16)
        sd = sd_map_replace(sd, LLAMA_SD_MAP)
        sd = llama_permute(sd, 32, 8) # L3
    else: