Update llama.py

higgsfield-ai · Feb 13, 2024 · d12a36e · d12a36e
1 parent f7a4fbf
commit d12a36e
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/higgsfield/llama/llama.py b/higgsfield/llama/llama.py
@@ -59,14 +59,14 @@ def __init__(
         if not checkpoint_path:
             if cpu_init_rank0: 
                 if rank == 0:
-                    model = LlamaForCausalLM.from_pretrained(model_name)
+                    model = LlamaForCausalLM.from_pretrained(model_name, use_cache=False)
                 else:
-                    llama_config = LlamaConfig.from_pretrained(model_name)
+                    llama_config = LlamaConfig.from_pretrained(model_name, use_cache=False)
 
                     with torch.device('meta'):
                         model = LlamaForCausalLM(llama_config)
             else:
-                model = LlamaForCausalLM.from_pretrained(model_name)
+                model = LlamaForCausalLM.from_pretrained(model_name, use_cache=False)
         else:
             if not cpu_init_rank0:
                 print("Ignoring cpu_init_rank0=False while loading model from checkpoint path")
@@ -298,4 +298,4 @@ def __init__(
             precision,
             cpu_init_rank0,
             cpu_offload,
-        )
+        )