Update model/openllama.py

model/openllama.py

@@ -215,17 +215,17 @@ 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.float16, device_map='auto', load_in_8bit=True, offload_folder="offload1")
+        self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map='auto', load_in_8bit=True)
         # # except:
         #     pass
         # finally:
         #     print(self.llama_model.hf_device_map)
         self.llama_model = get_peft_model(self.llama_model, peft_config)
-        delta_ckpt = torch.load(args['delta_ckpt_path'], map_location=torch.device('cpu'))
-        self.llama_model.load_state_dict(delta_ckpt, strict=False)
+        # delta_ckpt = torch.load(args['delta_ckpt_path'], map_location=torch.device('cpu'))
+        # self.llama_model.load_state_dict(delta_ckpt, strict=False)
         self.llama_model.print_trainable_parameters()
 
-        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.float16, device_map='auto', offload_folder="offload2")
+        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.float16)
         self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token
         self.llama_tokenizer.padding_side = "right"
         print ('Language decoder initialized.')