回退单核

handler.py

@@ -4,20 +4,13 @@ import os
 from threading import Thread
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from torch.nn.parallel import DistributedDataParallel as DDP
-import torch.distributed as dist
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 512
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))
 
-def setup_distributed():
-    dist.init_process_group(backend='nccl')
-    torch.cuda.set_device(dist.get_rank())
-
 class EndpointHandler:
     def __init__(self, path=""):
-        setup_distributed()
         local_config_path = "./config.json"
         remote_model_name = "threadshare/Peach-9B-8k-Roleplay"
         
@@ -28,11 +21,9 @@ class EndpointHandler:
             self.model_name_or_path = remote_model_name
             
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True, flash_atten=True)
-        model = AutoModelForCausalLM.from_pretrained(
+        self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name_or_path, torch_dtype=torch.bfloat16,
             trust_remote_code=True, device_map="auto")
-        
-        self.model = DDP(model.to(dist.get_rank()), device_ids=[dist.get_rank()])
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         # print json data
@@ -75,7 +66,7 @@ class EndpointHandler:
         if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
             input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
 
-        input_ids = input_ids.to(dist.get_rank())
+        input_ids = input_ids.to("cuda")
         streamer = TextIteratorStreamer(self.tokenizer, timeout=50.0, skip_prompt=True, skip_special_tokens=True)
         generate_kwargs = dict(
             input_ids=input_ids,
@@ -89,7 +80,7 @@ class EndpointHandler:
             no_repeat_ngram_size=8,
             repetition_penalty=repetition_penalty
         )
-        t = Thread(target=self.model.module.generate, kwargs=generate_kwargs)
+        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
         t.start()
         outputs = []
         for text in streamer: