Update model/openllama.py

model/openllama.py

@@ -10,6 +10,8 @@ import kornia as K
 
 import torch
 from torch.nn.utils import rnn
+from transformers import AutoConfig, AutoModelForCausalLM
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 CLASS_NAMES = ['bottle', 'cable', 'capsule', 'carpet', 'grid', 'hazelnut', 'leather', 'metal nut', 'pill', 'screw', 'tile', 'toothbrush', 'transistor', 'wood', 'zipper', 'object',
                'candle', 'cashew', 'chewinggum', 'fryum', 'macaroni', 'pcb', 'pipe fryum']
@@ -203,7 +205,11 @@ class OpenLLAMAPEFTModel(nn.Module):
             target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']
         )
 
-        self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map='auto', offload_folder="offload", offload_state_dict = True)
+        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"])
+        # self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map='auto', offload_folder="offload", offload_state_dict = True)
         self.llama_model = get_peft_model(self.llama_model, peft_config)
         self.llama_model.print_trainable_parameters()