Update model/openllama.py

model/openllama.py

@@ -205,16 +205,17 @@ class OpenLLAMAPEFTModel(nn.Module):
             target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']
         )
 
-        # config = AutoConfig.from_pretrained(vicuna_ckpt_path)
-        # with init_empty_weights():
-        #     self.llama_model = AutoModelForCausalLM.from_config(config)
-        # self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map="auto", no_split_module_classes=["OPTDecoderLayer"], offload_folder="offload", offload_state_dict = True)
-        try:
-            self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map='auto', offload_folder="offload", offload_state_dict = True)
-        except:
-            pass
-        finally:
-            print(self.llama_model.hf_device_map)
+        config = AutoConfig.from_pretrained(vicuna_ckpt_path)
+        with init_empty_weights():
+            self.llama_model = AutoModelForCausalLM.from_config(config)
+        self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map="auto", no_split_module_classes=["OPTDecoderLayer"], 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', offload_folder="offload", offload_state_dict = True)
+        # except:
+        #     pass
+        # finally:
+        #     print(self.llama_model.hf_device_map)
         self.llama_model = get_peft_model(self.llama_model, peft_config)
         self.llama_model.print_trainable_parameters()