NPU needs to be initialized when starting a new process (#2843)

lm-sys · Dec 28, 2023 · 1368f3f · 1368f3f
1 parent 3623d9b
commit 1368f3f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/fastchat/serve/model_worker.py b/fastchat/serve/model_worker.py
@@ -31,7 +31,6 @@
     str_to_torch_dtype,
 )
 
-
 worker_id = str(uuid.uuid4())[:8]
 logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
 
@@ -103,6 +102,10 @@ def __init__(
             self.init_heart_beat()
 
     def generate_stream_gate(self, params):
+        if self.device == "npu":
+            import torch_npu
+
+            torch_npu.npu.set_device("npu:0")
         self.call_ct += 1
 
         try: