fix dist_utils

Author: shihaobai

lightllm/utils/dist_utils.py

@@ -22,11 +22,11 @@ def _init_distributed_env(kvargs):
     set_current_dp_rank(get_global_rank() // get_dp_world_size())
     set_current_dp_inner_rank(get_global_rank() % get_dp_world_size())
 
-    size_per_node = (kvargs["world_size"] + kvargs["args"].nnodes - 1) // kvargs["args"].nnodes
-    local_tp_rank = kvargs["rank_id"] - size_per_node * kvargs["args"].node_rank
-    set_current_device_id(local_tp_rank)
-    torch.cuda.set_device(local_tp_rank)
-    print(local_tp_rank)
+    assert kvargs["world_size"] % kvargs["args"].nnodes == 0, "world_size should be divided by nnodes"
+    size_per_node = kvargs["world_size"] // kvargs["args"].nnodes
+    device_id = kvargs["rank_id"] % size_per_node
+    set_current_device_id(device_id)
+    torch.cuda.set_device(device_id)
     if kvargs["world_size"] > 1:
         dist.init_process_group(
             "nccl",
@@ -35,25 +35,25 @@ def _init_distributed_env(kvargs):
             world_size=kvargs["world_size"],
         )
         # warmup nccl communicator
-        _a = torch.zeros([1]).to(f"cuda:{local_tp_rank}")
+        _a = torch.zeros([1]).to(f"cuda:{device_id}")
         dist.all_reduce(_a)
         del _a
 
 
 def set_global_rank(global_rank: int):
-    set_environ("GLOBAL_RANK", global_rank)
+    set_environ("LIGHTLLM_GLOBAL_RANK", global_rank)
 
 
 def get_global_rank():
-    return int(get_environ("GLOBAL_RANK"))
+    return int(get_environ("LIGHTLLM_GLOBAL_RANK"))
 
 
 def set_global_world_size(world_size: int):
-    set_environ("GLOBAL_WORLD_SIZE", world_size)
+    set_environ("LIGHTLLM_GLOBAL_WORLD_SIZE", world_size)
 
 
 def get_global_world_size():
-    return int(get_environ("GLOBAL_WORLD_SIZE"))
+    return int(get_environ("LIGHTLLM_GLOBAL_WORLD_SIZE"))
 
 
 def set_dp_size(dp_size: int):
@@ -68,32 +68,32 @@ def get_dp_size():
 
 
 def set_dp_world_size(world_size: int):
-    set_environ("DP_WORLD_SIZE", world_size)
+    set_environ("LIGHTLLM_DP_WORLD_SIZE", world_size)
 
 
 def get_dp_world_size():
-    return int(get_environ("DP_WORLD_SIZE"))
+    return int(get_environ("LIGHTLLM_DP_WORLD_SIZE"))
 
 
 def set_current_dp_rank(rank: int):
-    set_environ("CURRENT_DP_RANK", rank)
+    set_environ("LIGHTLLM_CURRENT_DP_RANK", rank)
 
 
 def get_current_dp_rank():
-    return int(get_environ("CURRENT_DP_RANK"))
+    return int(get_environ("LIGHTLLM_CURRENT_DP_RANK"))
 
 
 def set_current_dp_inner_rank(rank: int):
-    set_environ("CURRENT_DP_INNER_RANK", rank)
+    set_environ("LIGHTLLM_CURRENT_DP_INNER_RANK", rank)
 
 
 def get_current_dp_inner_rank():
-    return get_environ("CURRENT_DP_INNER_RANK")
+    return get_environ("LIGHTLLM_CURRENT_DP_INNER_RANK")
 
 
 def set_current_device_id(device_id: int):
-    set_environ("CURRENT_DEVICE_ID", device_id)
+    set_environ("LIGHTLLM_CURRENT_DEVICE_ID", device_id)
 
 
 def get_current_device_id():
-    return int(get_environ("CURRENT_DEVICE_ID"))
+    return int(get_environ("LIGHTLLM_CURRENT_DEVICE_ID"))