modify dist_utils & remove child_ips

Feiyang Wu · Feiyang Wu · commit 8ad585af84f4 · 2025-02-21T11:44:03.000+08:00
diff --git a/lightllm/common/basemodel/layer_weights/base_layer_weight.py b/lightllm/common/basemodel/layer_weights/base_layer_weight.py
@@ -2,7 +2,7 @@
 import numpy as np
 import threading
 from lightllm.common.basemodel.layer_weights.meta_weights import BaseWeight
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_current_device_id
 
 
 class BaseLayerWeight:
diff --git a/lightllm/common/basemodel/layer_weights/hf_load_utils.py b/lightllm/common/basemodel/layer_weights/hf_load_utils.py
@@ -3,14 +3,14 @@
 import gc
 from safetensors import safe_open
 import lightllm.utils.petrel_helper as utils
+from lightllm.utils.dist_utils import get_current_device_id
 
 
-def load_func(file_, local_tp_rank, use_safetensors=False, pre_post_layer=None, transformer_layer_list=None, weight_dir=None):
+def load_func(file_, use_safetensors=False, pre_post_layer=None, transformer_layer_list=None, weight_dir=None):
     # fix bug for 多线程加载的时候，每个线程内部的cuda device 会切回 0， 修改后来保证不会出现bug
     import torch.distributed as dist
 
-    # tp_rank = dist.get_rank()
-    torch.cuda.set_device(local_tp_rank)
+    torch.cuda.set_device(get_current_device_id())
 
     if use_safetensors:
         weights = safe_open(os.path.join(weight_dir, file_), "pt", "cpu")
@@ -27,7 +27,7 @@ def load_func(file_, local_tp_rank, use_safetensors=False, pre_post_layer=None,
     gc.collect()
 
 
-def load_hf_weights(data_type, weight_dir, local_tp_rank, pre_post_layer=None, transformer_layer_list=None, weight_dict=None):
+def load_hf_weights(data_type, weight_dir, pre_post_layer=None, transformer_layer_list=None, weight_dict=None):
     if isinstance(data_type, str):
         data_type = torch.float16 if data_type == "fp16" else torch.float32
     if pre_post_layer is not None:
@@ -36,10 +36,10 @@ def load_hf_weights(data_type, weight_dir, local_tp_rank, pre_post_layer=None, t
         assert transformer_layer_list[0].data_type_ == data_type, "type is not right"
     if weight_dict:
         if pre_post_layer is not None:
-            pre_post_layer.load_hf_weights(weight_dict, local_tp_rank)
+            pre_post_layer.load_hf_weights(weight_dict)
         if transformer_layer_list is not None:
             for layer in transformer_layer_list:
-                layer.load_hf_weights(weight_dict, local_tp_rank)
+                layer.load_hf_weights(weight_dict)
         del weight_dict
         return
     use_safetensors = True
@@ -54,7 +54,6 @@ def load_hf_weights(data_type, weight_dir, local_tp_rank, pre_post_layer=None, t
 
     partial_func = partial(
         load_func,
-        local_tp_rank=local_tp_rank,
         use_safetensors=use_safetensors,
         pre_post_layer=pre_post_layer,
         transformer_layer_list=transformer_layer_list,
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py
@@ -1,15 +1,14 @@
 import torch
 from abc import ABC, abstractmethod
-from lightllm.utils.dist_utils import get_world_size, get_rank
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_global_world_size, get_global_rank, get_current_device_id
 
 
 class BaseWeight(ABC):
     def __init__(self):
         pass
 
     @abstractmethod
-    def load_hf_weights(self, weights, local_tp_rank):
+    def load_hf_weights(self, weights):
         pass
 
     @abstractmethod
@@ -19,11 +18,11 @@ def verify_load(self):
 
 class BaseWeightTpl(BaseWeight):
     def __init__(self):
-        self.world_size_ = get_world_size()
-        self.tp_rank_ = get_rank()
+        self.world_size_ = get_global_world_size()
+        self.tp_rank_ = get_global_rank()
         self.device_id_ = get_current_device_id()
 
-    def load_hf_weights(self, weights, local_tp_rank):
+    def load_hf_weights(self, weights):
         pass
 
     def verify_load(self):
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight.py
@@ -5,9 +5,8 @@
 from .base_weight import BaseWeight
 from lightllm.common.quantization import vLLMFP8w8a8QuantizationMethod
 from lightllm.common.quantization.quantize_method import QuantizationMethod
-from lightllm.utils.dist_utils import get_world_size, get_rank
+from lightllm.utils.dist_utils import get_global_world_size, get_global_rank, get_current_device_id
 from lightllm.common.vllm_kernel import _custom_ops as ops
-from lightllm.utils.device_utils import get_current_device_id
 
 
 class FusedMoeWeight(BaseWeight):
@@ -39,7 +38,7 @@ def __init__(
         self.n_routed_experts = n_routed_experts
         self.split_inter_size = split_inter_size
         self.data_type_ = data_type
-        self.tp_rank_ = get_rank()
+        self.tp_rank_ = get_global_rank()
         self.experts_up_projs = [None] * self.n_routed_experts
         self.experts_gate_projs = [None] * self.n_routed_experts
         self.experts_up_proj_scales = [None] * self.n_routed_experts
@@ -159,7 +158,7 @@ def _fuse_weight_scale(self):
                 delattr(self, "experts_gate_proj_scales")
 
     def _load_hf_weights_etp(self, weights):
-        world_size_ = get_world_size()
+        world_size_ = get_global_world_size()
         assert self.n_routed_experts % world_size_ == 0
         n_expert_ep = self.n_routed_experts // world_size_
 
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight.py
@@ -4,7 +4,7 @@
 from typing import Optional, Tuple, List, Dict, Any
 from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
 from lightllm.common.quantization.quantize_method import QuantizationMethod
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_current_device_id
 
 
 def generate_scale_name(name, weight_scale_suffix, act_scale_suffix):
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py
@@ -1,6 +1,6 @@
 import torch
 from .base_weight import BaseWeightTpl
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_current_device_id
 
 
 class NormWeight(BaseWeightTpl):
diff --git a/lightllm/common/mem_manager.py b/lightllm/common/mem_manager.py
@@ -8,7 +8,7 @@
 from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
 from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory
 from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_current_device_id
 
 logger = init_logger(__name__)
 
diff --git a/lightllm/common/quantization/quantize_method.py b/lightllm/common/quantization/quantize_method.py
@@ -1,6 +1,7 @@
 import torch
 from abc import ABC, abstractmethod
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_current_device_id
+
 
 class QuantizationMethod(ABC):
     def __init__(self):
diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py
@@ -76,7 +76,6 @@ def _init_weights(self):
         ]
         load_hf_weights(
             self.data_type,
-            local_tp_rank=self.local_tp_rank,
             weight_dir=self.weight_dir_,
             pre_post_layer=self.pre_post_weight,
             transformer_layer_list=self.trans_layers_weight,
diff --git a/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py
@@ -3,7 +3,7 @@
 import numpy as np
 import torch.nn.functional as F
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_current_device_id
 
 
 class ViTPreAndPostLayerWeight(PreAndPostLayerWeight):
diff --git a/lightllm/models/vit/layer_weights/transformer_layer_weight.py b/lightllm/models/vit/layer_weights/transformer_layer_weight.py
@@ -11,7 +11,7 @@
     MultiROWMMWeight,
     TpNormWeight,
 )
-from lightllm.utils.device_utils import get_current_device_id
+from lightllm.utils.dist_utils import get_current_device_id
 
 
 class ViTTransformerLayerWeight(TransformerLayerWeight):
diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py
@@ -95,12 +95,14 @@ def make_argument_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument("--nnodes", type=int, default=1, help="the number of nodes")
     parser.add_argument("--node_rank", type=int, default=0, help="the rank of the current node")
-    parser.add_argument("--child_ips", type=str, default=None, help="the ips of the child nodes, only provide if this is the parent node")
     parser.add_argument(
-        "--multinode_httpmanager_port", type=int, default=19998, help="the port for multinode http manager, default is 19998"
+        "--multinode_httpmanager_port",
+        type=int,
+        default=40000,
+        help="the port for multinode http manager, default is 40000",
     )
     parser.add_argument(
-        "--multinode_router_port", type=int, default=19999, help="the port for multinode router, default is 19999"
+        "--multinode_router_port", type=int, default=40001, help="the port for multinode router, default is 40001"
     )
     parser.add_argument("--tp", type=int, default=1, help="model tp parral size, the default is 1")
     parser.add_argument(
@@ -115,7 +117,10 @@ def make_argument_parser() -> argparse.ArgumentParser:
         "--max_req_total_len", type=int, default=16384, help="the max value for req_input_len + req_output_len"
     )
     parser.add_argument(
-        "--nccl_host", type=str, default="127.0.0.1", help="the nccl_host to build a distributed environment for PyTorch"
+        "--nccl_host",
+        type=str,
+        default="127.0.0.1",
+        help="the nccl_host to build a distributed environment for PyTorch",
     )
     parser.add_argument(
         "--nccl_port", type=int, default=28765, help="the nccl_port to build a distributed environment for PyTorch"
diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py
@@ -4,6 +4,7 @@
 import uuid
 import subprocess
 import signal
+import socket
 from lightllm.utils.net_utils import alloc_can_use_network_port, PortLocker
 from lightllm.utils.start_utils import process_manager
 from .metrics.manager import start_metric_manager
@@ -164,9 +165,6 @@ def normal_or_p_d_start(args):
 
         args.data_type = get_dtype(args.model_dir)
         assert args.data_type in ["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"]
-    
-    if args.child_ips is not None:
-        assert args.nnodes > 1 and args.node_rank == 0, "child_ips should be set only when nnodes > 1 and node_rank == 0"
 
     already_uesd_ports = args.visual_nccl_ports + [args.nccl_port, args.port]
     if args.run_mode == "decode":
@@ -211,6 +209,30 @@ def normal_or_p_d_start(args):
     if args.run_mode == "decode":
         args.router_max_wait_tokens = 0
 
+    # 传输子node的ip
+    if args.nnodes > 1:
+        import zmq
+
+        if args.node_rank == 0:
+            args.child_ips = None
+            args.child_ips = []
+            for i in range(1, args.nnodes):
+                context = zmq.Context(2)
+                comm_socket = context.socket(zmq.PULL)
+                comm_socket.bind(f"tcp://*:{args.multinode_httpmanager_port + i}")
+                print(f"binding port {args.multinode_httpmanager_port + i}")
+                args.child_ips.append(comm_socket.recv_pyobj())
+                comm_socket.close()
+            print(f"Received child IPs: {args.child_ips}")
+        else:
+            local_ip = socket.gethostbyname(socket.gethostname())
+            context = zmq.Context(2)
+            comm_socket = context.socket(zmq.PUSH)
+            comm_socket.connect(f"tcp://{args.nccl_host}:{args.multinode_httpmanager_port + args.node_rank}")
+            print(f"connecting to {args.nccl_host}:{args.multinode_httpmanager_port + args.node_rank}")
+            comm_socket.send_pyobj(local_ip)
+            comm_socket.close()
+
     set_env(args)
     logger.info(f"all start args:{args}")
 
diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py
@@ -59,9 +59,7 @@ def __init__(
         if args.nnodes > 1:
             if args.node_rank == 0:
                 self.multinode_req_manager = []
-                for child_ip in args.child_ips.split(","):
-                    if ":" in child_ip:
-                        child_ip = child_ip.split(":")[0]
+                for child_ip in args.child_ips:
                     context = zmq.asyncio.Context(2)
                     self.multinode_req_manager.append(context.socket(zmq.PUSH))
                     self.multinode_req_manager[-1].connect(f"tcp://{child_ip}:{args.multinode_httpmanager_port}")
diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py
@@ -50,7 +50,9 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
         self.shm_req_manager = ShmReqManager()
         # 用共享内存进行共享，router 模块读取进行精确的调度估计
         size_per_node = (self.world_size + self.nnodes - 1) // self.nnodes
-        local_world_size = size_per_node if self.node_rank < self.nnodes - 1 else self.world_size - self.node_rank * size_per_node
+        local_world_size = (
+            size_per_node if self.node_rank < self.nnodes - 1 else self.world_size - self.node_rank * size_per_node
+        )
         self.read_only_statics_mem_manager = ReadOnlyStaticsMemoryManager(args.nccl_port, local_world_size)
         # 初始化 radix_cache_client 用于读取 prompt cache 的管理信息
         self.radix_cache_client = None
@@ -77,15 +79,13 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
         self.send_to_detokenization = context.socket(zmq.PUSH)
         self.send_to_detokenization.connect(f"{args.zmq_mode}127.0.0.1:{detokenization_port}")
         self.model_rpc_ports = model_rpc_ports
-        
+
         self.multinode_req_manager = None
         self.multinode_req_queue_lock = asyncio.Lock()
         if args.nnodes > 1:
             if args.node_rank == 0:
                 self.multinode_req_manager = []
-                for child_ip in args.child_ips.split(","):
-                    if ":" in child_ip:
-                        child_ip = child_ip.split(":")[0]
+                for child_ip in args.child_ips:
                     context = zmq.asyncio.Context(2)
                     self.multinode_req_manager.append(context.socket(zmq.PUSH))
                     self.multinode_req_manager[-1].connect(f"tcp://{child_ip}:{args.multinode_router_port}")
@@ -126,8 +126,12 @@ async def wait_to_model_ready(self):
         self.rpc_finished_event = multiprocessing.Event()
 
         size_per_node = (self.world_size + self.nnodes - 1) // self.nnodes
-        local_world_size = size_per_node if self.node_rank < self.nnodes - 1 else self.world_size - self.node_rank * size_per_node
-        for rank_id in range(self.node_rank * size_per_node, min(self.world_size, (self.node_rank + 1) * size_per_node)):
+        local_world_size = (
+            size_per_node if self.node_rank < self.nnodes - 1 else self.world_size - self.node_rank * size_per_node
+        )
+        for rank_id in range(
+            self.node_rank * size_per_node, min(self.world_size, (self.node_rank + 1) * size_per_node)
+        ):
             rpc_model = await start_model_process(
                 args=self.args,
                 tp_rank=rank_id,
@@ -299,7 +303,9 @@ async def get_new_batch():
                 # time.sleep(0.003)  # 这里是为了保证能正确进入推理的流程，保证折叠成功。
                 await asyncio.sleep(0.003)
                 async with self.multinode_req_queue_lock:
-                    new_batch, poped_req_list, remain_req_list = self.req_queue.generate_new_batch(running_batch, current_waiting_list)
+                    new_batch, poped_req_list, remain_req_list = self.req_queue.generate_new_batch(
+                        running_batch, current_waiting_list
+                    )
                     self.req_queue.pop_list(poped_req_list)
                 return new_batch
 
diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py
@@ -3,6 +3,7 @@
 import numpy as np
 import rpyc
 import torch
+import socket
 from datetime import timedelta
 from typing import Dict, List, Tuple
 from transformers.configuration_utils import PretrainedConfig
@@ -37,7 +38,7 @@
 from lightllm.server.router.model_infer.infer_batch import InferReq, InferSamplingParams
 from lightllm.server.router.token_load import TokenLoad
 from lightllm.common.basemodel.infer_lock import g_infer_state_lock, InferStateLock
-from lightllm.utils.device_utils import set_current_device_id
+from lightllm.utils.dist_utils import _init_distributed_env
 from lightllm.server.core.objs import ShmReqManager
 from lightllm.server.router.model_infer.infer_batch import g_infer_context
 import torch.distributed as dist
@@ -84,19 +85,7 @@ def init_model(self, kvargs):
 
         size_per_node = (self.world_size + self.nnodes - 1) // self.nnodes
         self.local_tp_rank = self.tp_rank - size_per_node * self.node_rank
-        torch.cuda.set_device(self.local_tp_rank)
-        set_current_device_id(self.local_tp_rank)
-
-        dist.init_process_group(
-            "nccl",
-            init_method=f'tcp://{kvargs["nccl_host"]}:{kvargs["nccl_port"]}',
-            rank=self.tp_rank,
-            world_size=self.world_size,
-        )
-        # warmup nccl communicator
-        _a = torch.zeros([1]).to(f"cuda:{self.local_tp_rank}")
-        dist.all_reduce(_a)
-        del _a
+        _init_distributed_env(kvargs)
 
         from lightllm.distributed import custom_comm_ops
 
diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/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.device_utils import set_current_device_id
+from lightllm.utils.dist_utils import set_current_device_id
 from lightllm.utils.graceful_utils import graceful_registry
 
 
diff --git a/lightllm/utils/device_utils.py b/lightllm/utils/device_utils.py
@@ -3,23 +3,6 @@
 import subprocess
 
 
-def set_current_device_id(device_id: int):
-    os.environ["CURRENT_DEVICE_ID"] = str(device_id)
-
-
-@lru_cache(maxsize=None)
-def get_current_device_id():
-    import torch
-
-    if torch.cuda.is_available():
-        device_id = os.getenv("CURRENT_DEVICE_ID", None)
-        if device_id is None:
-            raise RuntimeError("set_current_device_id must called first to set current device")
-        return int(device_id)
-    else:
-        raise RuntimeError("Torch CUDA is not avaliable.")
-
-
 @lru_cache(maxsize=None)
 def get_device_sm_count():
     import triton
diff --git a/lightllm/utils/dist_utils.py b/lightllm/utils/dist_utils.py
diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py
diff --git a/test/model/model_infer.py b/test/model/model_infer.py

Original file line number	Diff line number	Diff line change
`@@ -76,7 +76,6 @@ def _init_weights(self):`
`76`	`76`	`]`
`77`	`77`	`load_hf_weights(`
`78`	`78`	`self.data_type,`
`79`		`- local_tp_rank=self.local_tp_rank,`
`80`	`79`	`weight_dir=self.weight_dir_,`
`81`	`80`	`pre_post_layer=self.pre_post_weight,`
`82`	`81`	`transformer_layer_list=self.trans_layers_weight,`