Feat KoAlpaca Tensor Parallelism

koalpaca.py

@@ -2,13 +2,12 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GenerationConfig
 from peft import PeftModel, PeftConfig
 from model import Model
-from accelerate import Accelerator
+import tensor_parallel as tp
 
 class KoAlpaca(Model):
     def __init__(self):
         peft_model_id = "4n3mone/Komuchat-koalpaca-polyglot-12.8B"
         config = PeftConfig.from_pretrained(peft_model_id)
-        # self.accelerator = Accelerator()
         self.bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_use_double_quant=True,
@@ -17,6 +16,8 @@ class KoAlpaca(Model):
         )
         #self.model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=self.bnb_config, device_map={"":0})
         self.model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=self.bnb_config, device_map='auto')
+        print(self.model.device)
+        self.model = tp.tensor_parallel(self.model, [*next(self.model.parameters()).device])
         self.model = PeftModel.from_pretrained(self.model, peft_model_id)
         self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
         self.gen_config = GenerationConfig.from_pretrained('./models/koalpaca', 'gen_config.json')
@@ -34,4 +35,5 @@ class KoAlpaca(Model):
             generation_config=self.gen_config
         )
         outputs = self.tokenizer.decode(output_ids[0]).split("### 답변: ")[-1]
-        return outputs
+        return outputs
+        

requirements.txt

@@ -3,4 +3,5 @@ gradio==4.1.1
 numpy==1.26.1
 pandas==2.1.2
 torch==2.0.1
+tensor_parallel==2.0.0
 git+https://github.com/huggingface/peft.git