Update model/openllama.py

model/openllama.py

@@ -208,7 +208,9 @@ class OpenLLAMAPEFTModel(nn.Module):
         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)
+
+        device_map = infer_auto_device_map(self.llama_model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")
+        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', offload_folder="offload", offload_state_dict = True)