add tp

app.py

@@ -257,7 +257,7 @@ def predict(
     repair_time = repair_end_time - repair_start_time
 
     encoded_answer = tokenizer(partial_output_repair, padding=True, truncation=True, return_tensors='pt')
-    print("loc token/s:",len(encoded_answer['input_ids'][0])/repair_time)
+    print("repair token/s:",len(encoded_answer['input_ids'][0])/repair_time)
 
     # yield response, "null test", "Generate: Success"
     yield [[prompt,highlight_response],[repair_prompt,highlight_response_repair]], [["null test","null test2"]], "Generate: Success"

kimi_dev/serve/inference.py

@@ -5,6 +5,7 @@ from transformers import (
     AutoConfig,
     AutoTokenizer
 )
+import tensor_parallel as tp
 
 logger = logging.getLogger(__name__)
 
@@ -20,6 +21,8 @@ def load_model(model_path: str = "moonshotai/Kimi-Dev-72B"):
         device_map="auto",
         trust_remote_code=True,
     )
+    model = tp.tensor_parallel(model, ["cuda:0", "cuda:1", "cuda:2", "cuda:3"])  # <- each GPU has half the weights
+
 
     tokenizer = AutoTokenizer.from_pretrained(model_path)
 

requirements.txt

@@ -18,3 +18,4 @@ colorama
 Pygments
 markdown
 SentencePiece
+tensor_parallel