Update model/openllama.py

model/openllama.py

@@ -215,7 +215,7 @@ class OpenLLAMAPEFTModel(nn.Module):
         # # self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map=device_map, offload_folder="offload", offload_state_dict = True)
         # # self.llama_model.to(torch.float16)
         # # try:
-        self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.bfloat16, device_map='auto', offload_folder="offload1", offload_state_dict = True)
+        self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.bfloat16, device_map='auto', load_in_4bit=True, offload_folder="offload1", offload_state_dict = True)
         # # except:
         #     pass
         # finally: