Update model/openllama.py

model/openllama.py

@@ -41,7 +41,7 @@ for obj in objs:
         for s in prompted_state:
             for template in prompt_templates:
                 prompted_sentence.append(template.format(s))
-        prompted_sentence = data.load_and_transform_text(prompted_sentence, torch.device('cpu'))#torch.cuda.current_device())
+        prompted_sentence = data.load_and_transform_text(prompted_sentence, torch.cuda.current_device())#torch.cuda.current_device())
         prompt_sentence_obj.append(prompted_sentence)
     prompt_sentences[obj] = prompt_sentence_obj
 
@@ -199,11 +199,11 @@ class OpenLLAMAPEFTModel(nn.Module):
             target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']
         )
 
-        self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16)
+        self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16, 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()
 
-        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False)
+        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.bfloat16, device_map='auto', offload_folder="offload", offload_state_dict = True)
         self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token
         self.llama_tokenizer.padding_side = "right"
         print ('Language decoder initialized.')
@@ -213,7 +213,7 @@ class OpenLLAMAPEFTModel(nn.Module):
         )
 
         self.max_tgt_len = max_tgt_len
-        self.device = torch.device('cpu')#torch.cuda.current_device()
+        self.device = torch.cuda.current_device()
 
 
     def rot90_img(self,x,k):