fix visualserver

Author: shihaobai

lightllm/server/visualserver/model_infer/model_rpc.py

@@ -16,7 +16,7 @@
 from lightllm.server.embed_cache.utils import tensor2bytes, read_shm, create_shm, get_shm_name_data, get_shm_name_embed
 from lightllm.utils.infer_utils import set_random_seed
 from lightllm.utils.infer_utils import calculate_time, mark_start, mark_end
-from lightllm.utils.dist_utils import set_current_device_id
+from lightllm.utils.dist_utils import _init_vision_distributed_env
 from lightllm.utils.graceful_utils import graceful_registry
 
 
@@ -31,20 +31,11 @@ def exposed_init_model(self, kvargs):
         self.tp_rank_id = kvargs["tp_rank_id"]
         self.cache_port = kvargs["cache_port"]
         weight_dir = kvargs["weight_dir"]
-        visual_gpu_ids = kvargs["visual_gpu_ids"]
-        visual_nccl_port = kvargs["visual_nccl_port"]
         self.vit_rank_id = kvargs["vit_rank_id"]
         self.cache_client = rpyc.connect("localhost", self.cache_port)
         self.data_type = kvargs["data_type"]
 
-        torch.cuda.set_device(visual_gpu_ids[self.vit_rank_id])
-        set_current_device_id(visual_gpu_ids[self.vit_rank_id])
-        dist.init_process_group(
-            backend="nccl",
-            init_method=f"tcp://127.0.0.1:{visual_nccl_port}",
-            rank=self.tp_rank_id,
-            world_size=self.vit_tp,
-        )
+        _init_vision_distributed_env(kvargs)
         model_cfg, _ = PretrainedConfig.get_config_dict(weight_dir)
 
         try:

lightllm/utils/dist_utils.py

@@ -25,6 +25,26 @@ def get_environ(environ_name):
     return value
 
 
+def _init_vision_distributed_env(kvargs):
+    world_size = kvargs["vit_tp"]
+    set_global_rank(kvargs["tp_rank_id"])
+    set_global_world_size(world_size)
+    visual_gpu_ids = kvargs["visual_gpu_ids"]
+    device_id = visual_gpu_ids[kvargs["vit_rank_id"]]
+    set_current_device_id(device_id)
+    torch.cuda.set_device(device_id)
+    dist.init_process_group(
+        "nccl",
+        init_method=f'tcp://127.0.0.1:{kvargs["visual_nccl_port"]}',
+        rank=kvargs["tp_rank_id"],
+        world_size=world_size,
+    )
+    # warmup nccl communicator
+    _a = torch.zeros([1]).to(f"cuda:{device_id}")
+    dist.all_reduce(_a)
+    del _a
+
+
 def _init_distributed_env(kvargs):
     assert kvargs["world_size"] % kvargs["args"].nnodes == 0, "world_size should be divided by nnodes"
     node_world_size = kvargs["world_size"] // kvargs["args"].nnodes
@@ -47,7 +67,6 @@ def _init_distributed_env(kvargs):
         rank=kvargs["rank_id"],
         world_size=kvargs["world_size"],
     )
-    # if kvargs["world_size"] > 1:
     # warmup nccl communicator
     _a = torch.zeros([1]).to(f"cuda:{device_id}")
     dist.all_reduce(_a)