Update model/openllama.py

model/openllama.py

@@ -1,7 +1,6 @@
 from header import *
 import os
 import torch.nn.functional as F
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 from .ImageBind import *
 from .ImageBind import data
 from .modeling_llama import LlamaForCausalLM
@@ -103,13 +102,7 @@ class OpenLLAMAPEFTModel(nn.Module):
             target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']
         )
 
-        with init_empty_weights():
-            config = LlamaConfig.from_pretrained(vicuna_ckpt_path, use_auth_token=os.environ['API_TOKEN'])
-            self.llama_model = LlamaForCausalLM(config)
-
-        self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map='sequential')
-
-        # self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, use_auth_token=os.environ['API_TOKEN'])
+        self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, use_auth_token=os.environ['API_TOKEN'])
         self.llama_model = get_peft_model(self.llama_model, peft_config)
         self.llama_model.print_trainable_parameters()