From caa85ea022dc58c89ae5782f1f606d70ae55954d Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:56:57 +0800 Subject: [PATCH] Extract init_torch_distributed and refactor into functions (#31152) --- python/sglang/srt/distributed/bootstrap.py | 291 ++++++++++++++++++ .../sglang/srt/model_executor/model_runner.py | 181 ++--------- 2 files changed, 315 insertions(+), 157 deletions(-) create mode 100644 python/sglang/srt/distributed/bootstrap.py diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py new file mode 100644 index 000000000..f46d9eb11 --- /dev/null +++ b/python/sglang/srt/distributed/bootstrap.py @@ -0,0 +1,291 @@ +import logging +import os +import time +from typing import List, Optional + +import msgspec +import torch +import torch.distributed as dist + +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.distributed import ( + get_default_distributed_backend, + get_pp_group, + get_tp_group, + get_world_group, + init_distributed_environment, + initialize_model_parallel, + set_custom_all_reduce, + set_mscclpp_all_reduce, + set_torch_symm_mem_all_reduce, +) +from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import initialize_dp_attention +from sglang.srt.platforms import current_platform +from sglang.srt.runtime_context import get_parallel +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import ( + cpu_has_amx_support, + get_available_gpu_memory, + is_host_cpu_arm64, + is_npu, + monkey_patch_p2p_access_check, +) +from sglang.srt.utils.network import NetworkAddress +from sglang.srt.utils.patch_torch import register_sgl_tp_rank + +logger = logging.getLogger(__name__) + +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu_arm64 = is_host_cpu_arm64() + + +class TorchDistributedResult(msgspec.Struct, frozen=True, kw_only=True): + tp_group: object + pp_group: object + attention_tp_group: object + pre_model_load_memory: float + + +def init_torch_distributed( + *, + server_args: ServerArgs, + model_config: ModelConfig, + device: str, + gpu_id: int, + tp_rank: int, + tp_size: int, + pp_rank: int, + pp_size: int, + dp_size: int, + attn_cp_size: int, + moe_ep_size: int, + moe_dp_size: int, + dcp_size: int, + dist_port: int, + is_draft_worker: bool, + local_omp_cpuid: Optional[List[int]], +): + tic = time.perf_counter() + logger.info("Init torch distributed begin.") + + try: + torch.get_device_module(device).set_device(gpu_id) + except Exception: + logger.warning( + f"Context: {device=} {gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {tp_rank=} {tp_size=}" + ) + raise + + backend = _resolve_backend(device=device, server_args=server_args, gpu_id=gpu_id) + + before_avail_memory = get_available_gpu_memory(device, gpu_id) + if not server_args.enable_p2p_check: + monkey_patch_p2p_access_check() + + dist_init_method = _resolve_dist_init_method( + server_args=server_args, dist_port=dist_port + ) + _set_all_reduce_flags(server_args=server_args) + + if not is_draft_worker: + if device == "cpu": + _init_cpu_threads_env( + tp_size=tp_size, tp_rank=tp_rank, local_omp_cpuid=local_omp_cpuid + ) + + # Only initialize the distributed environment on the target model worker. + _init_parallel_groups( + backend=backend, + dist_init_method=dist_init_method, + server_args=server_args, + model_config=model_config, + gpu_id=gpu_id, + tp_rank=tp_rank, + tp_size=tp_size, + pp_rank=pp_rank, + pp_size=pp_size, + dp_size=dp_size, + attn_cp_size=attn_cp_size, + moe_ep_size=moe_ep_size, + moe_dp_size=moe_dp_size, + dcp_size=dcp_size, + ) + + # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request + # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs) + if server_args.pre_warm_nccl and ( + tp_size > 1 or pp_size > 1 or moe_ep_size > 1 + ): + _prewarm_nccl(tp_size=tp_size, pp_size=pp_size, moe_ep_size=moe_ep_size) + + pre_model_load_memory = get_available_gpu_memory( + device, + gpu_id, + distributed=get_world_group().world_size > 1, + cpu_group=get_world_group().cpu_group, + ) + tp_group = get_tp_group() + pp_group = get_pp_group() + attention_tp_group = get_parallel().attn_tp_group + + # Check memory for tensor parallelism + local_gpu_memory = get_available_gpu_memory(device, gpu_id) + if tp_size > 1 and not is_draft_worker: + _check_tp_memory_balance( + pre_model_load_memory=pre_model_load_memory, + local_gpu_memory=local_gpu_memory, + ) + + logger.info( + f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, " + f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" + ) + return TorchDistributedResult( + tp_group=tp_group, + pp_group=pp_group, + attention_tp_group=attention_tp_group, + pre_model_load_memory=pre_model_load_memory, + ) + + +def _resolve_backend(*, device: str, server_args: ServerArgs, gpu_id: int) -> str: + backend = get_default_distributed_backend(device) + if device == "cuda" and server_args.elastic_ep_backend == "mooncake": + backend = "mooncake" + if server_args.mooncake_ib_device: + from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ( + get_ib_devices_for_gpu, + ) + + ib_device_for_gpu = get_ib_devices_for_gpu( + server_args.mooncake_ib_device, gpu_id + ) + mooncake_ib_device = ( + ib_device_for_gpu.split(",") if ib_device_for_gpu else [] + ) + try: + from mooncake import ep as mooncake_ep + + mooncake_ep.set_device_filter(mooncake_ib_device) + except: + pass # A warning will be raised in `init_distributed_environment` + return backend + + +def _resolve_dist_init_method(*, server_args: ServerArgs, dist_port: int) -> str: + # Allow external orchestrators (e.g. trainpi) to override the distributed + # init method. When set to "env://", torch uses MASTER_ADDR/MASTER_PORT + # env-vars and an externally-created TCPStore, completely avoiding port + # conflicts with intra-host collocation. + dist_init_method_override = envs.SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE.get() + if dist_init_method_override: + dist_init_method = dist_init_method_override + elif server_args.dist_init_addr: + na = NetworkAddress.parse(server_args.dist_init_addr) + dist_init_method = na.to_tcp() + else: + dist_init_method = NetworkAddress( + server_args.host or "127.0.0.1", dist_port + ).to_tcp() + return dist_init_method + + +def _set_all_reduce_flags(*, server_args: ServerArgs) -> None: + set_custom_all_reduce(not server_args.disable_custom_all_reduce) + set_mscclpp_all_reduce(server_args.enable_mscclpp) + set_torch_symm_mem_all_reduce(server_args.enable_torch_symm_mem) + + +def _init_cpu_threads_env( + *, tp_size: int, tp_rank: int, local_omp_cpuid: Optional[List[int]] +) -> None: + if _is_cpu_amx_available or _is_cpu_arm64: + # Bind OpenMP threads to CPU cores + torch.ops.sgl_kernel.init_cpu_threads_env(local_omp_cpuid) + + # Set local size to hint SGLang to use shared memory based AllReduce + os.environ["LOCAL_SIZE"] = str(tp_size) + torch.ops.sgl_kernel.initialize(tp_size, tp_rank) + + else: + logger.warning( + "init_cpu_threads_env and shared memory based AllReduce is disabled, only intel amx backend and arm64 are supported" + ) + + +def _init_parallel_groups( + *, + backend: str, + dist_init_method: str, + server_args: ServerArgs, + model_config: ModelConfig, + gpu_id: int, + tp_rank: int, + tp_size: int, + pp_rank: int, + pp_size: int, + dp_size: int, + attn_cp_size: int, + moe_ep_size: int, + moe_dp_size: int, + dcp_size: int, +) -> None: + init_distributed_environment( + backend=backend, + world_size=tp_size * pp_size, + rank=tp_size * pp_rank + tp_rank, + local_rank=gpu_id, + distributed_init_method=dist_init_method, + timeout=server_args.dist_timeout, + moe_a2a_backend=server_args.moe_a2a_backend, + recovered_rank=server_args.elastic_ep_rejoin, + ) + initialize_model_parallel( + tensor_model_parallel_size=tp_size, + attention_data_parallel_size=dp_size, + pipeline_model_parallel_size=pp_size, + expert_model_parallel_size=moe_ep_size, + attention_context_model_parallel_size=attn_cp_size, + moe_data_model_parallel_size=moe_dp_size, + decode_context_parallel_size=dcp_size, + duplicate_tp_group=server_args.enable_pdmux, + enable_symm_mem=server_args.enable_symm_mem, + recovered_rank=server_args.elastic_ep_rejoin, + ) + initialize_dp_attention( + server_args=server_args, + model_config=model_config, + ) + if is_npu(): + register_sgl_tp_rank(gpu_id) + + +def _prewarm_nccl(*, tp_size: int, pp_size: int, moe_ep_size: int) -> None: + warmup_start = time.perf_counter() + tp_group_handle = get_tp_group().device_group + + # Single warmup all_reduce to initialize NCCL/RCCL/HCCL communicator + warmup_tensor = torch.zeros(1, device=torch.cuda.current_device()) + dist.all_reduce(warmup_tensor, group=tp_group_handle) + current_platform.synchronize() + + warmup_elapsed = time.perf_counter() - warmup_start + logger.info( + f"NCCL/RCCL/HCCL warmup completed in {warmup_elapsed:.3f}s " + f"(tp_size={tp_size}, pp_size={pp_size}, ep_size={moe_ep_size})" + ) + + +def _check_tp_memory_balance( + *, pre_model_load_memory: float, local_gpu_memory: float +) -> None: + if pre_model_load_memory < local_gpu_memory * 0.9: + msg = "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. " + msg += ( + f"{pre_model_load_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}" + ) + if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get(): + raise RuntimeError(msg) + else: + logger.warning(msg) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 0d55d59dd..e122257c1 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -47,15 +47,9 @@ from sglang.srt.debug_utils.tensor_dump_forward_hook import ( register_forward_hook_for_model, ) from sglang.srt.distributed import ( - get_default_distributed_backend, - get_pp_group, + bootstrap, get_tp_group, get_world_group, - init_distributed_environment, - initialize_model_parallel, - set_custom_all_reduce, - set_mscclpp_all_reduce, - set_torch_symm_mem_all_reduce, ) from sglang.srt.distributed.device_communicators.pynccl_allocator import ( prealloc_symmetric_memory_pool, @@ -105,9 +99,6 @@ from sglang.srt.layers.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.cp.utils import ( get_cp_strategy, ) -from sglang.srt.layers.dp_attention import ( - initialize_dp_attention, -) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.moe.hash_topk import HashTopK from sglang.srt.layers.moe.topk import TopK @@ -164,7 +155,7 @@ from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( ) from sglang.srt.model_loader.utils import resolve_language_model from sglang.srt.platforms import current_platform -from sglang.srt.runtime_context import get_flags, get_parallel, get_server_args +from sglang.srt.runtime_context import get_flags, get_server_args from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.server_args import ( # noqa: F401 (re-export) CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS, @@ -197,7 +188,6 @@ from sglang.srt.utils import ( is_host_cpu_arm64, is_npu, log_info_on_rank0, - monkey_patch_p2p_access_check, numa_utils, require_gathered_buffer, reserve_rope_cache_for_long_sequences, @@ -212,7 +202,6 @@ from sglang.srt.utils.offloader import ( get_offloader, set_offloader, ) -from sglang.srt.utils.patch_torch import register_sgl_tp_rank from sglang.srt.utils.profile_utils import build_step_span_name from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.srt.utils.weight_checker import WeightChecker @@ -475,7 +464,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Get available memory before model loading. # Stored for later use by alloc_memory_pool(). - self.pre_model_load_memory = self.init_torch_distributed() + self.init_torch_distributed() # Initialize MooncakeTransferEngine self.init_shared_mooncake_transfer_engine() @@ -1109,150 +1098,28 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) def init_torch_distributed(self): - tic = time.perf_counter() - logger.info("Init torch distributed begin.") - - try: - torch.get_device_module(self.device).set_device(self.gpu_id) - except Exception: - logger.warning( - f"Context: {self.device=} {self.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {self.tp_rank=} {self.tp_size=}" - ) - raise - - backend = get_default_distributed_backend(self.device) - if self.device == "cuda" and self.server_args.elastic_ep_backend == "mooncake": - backend = "mooncake" - if self.server_args.mooncake_ib_device: - from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ( - get_ib_devices_for_gpu, - ) - - ib_device_for_gpu = get_ib_devices_for_gpu( - self.server_args.mooncake_ib_device, self.gpu_id - ) - mooncake_ib_device = ( - ib_device_for_gpu.split(",") if ib_device_for_gpu else [] - ) - try: - from mooncake import ep as mooncake_ep - - mooncake_ep.set_device_filter(mooncake_ib_device) - except: - pass # A warning will be raised in `init_distributed_environment` - - before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id) - if not self.server_args.enable_p2p_check: - monkey_patch_p2p_access_check() - - # Allow external orchestrators (e.g. trainpi) to override the distributed - # init method. When set to "env://", torch uses MASTER_ADDR/MASTER_PORT - # env-vars and an externally-created TCPStore, completely avoiding port - # conflicts with intra-host collocation. - dist_init_method_override = envs.SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE.get() - if dist_init_method_override: - dist_init_method = dist_init_method_override - elif self.server_args.dist_init_addr: - na = NetworkAddress.parse(self.server_args.dist_init_addr) - dist_init_method = na.to_tcp() - else: - dist_init_method = NetworkAddress( - self.server_args.host or "127.0.0.1", self.dist_port - ).to_tcp() - set_custom_all_reduce(not self.server_args.disable_custom_all_reduce) - set_mscclpp_all_reduce(self.server_args.enable_mscclpp) - set_torch_symm_mem_all_reduce(self.server_args.enable_torch_symm_mem) - - if not self.is_draft_worker: - if self.device == "cpu": - if _is_cpu_amx_available or _is_cpu_arm64: - # Bind OpenMP threads to CPU cores - torch.ops.sgl_kernel.init_cpu_threads_env(self.local_omp_cpuid) - - # Set local size to hint SGLang to use shared memory based AllReduce - os.environ["LOCAL_SIZE"] = str(self.tp_size) - torch.ops.sgl_kernel.initialize(self.tp_size, self.tp_rank) - - else: - logger.warning( - "init_cpu_threads_env and shared memory based AllReduce is disabled, only intel amx backend and arm64 are supported" - ) - - # Only initialize the distributed environment on the target model worker. - init_distributed_environment( - backend=backend, - world_size=self.tp_size * self.pp_size, - rank=self.tp_size * self.pp_rank + self.tp_rank, - local_rank=self.gpu_id, - distributed_init_method=dist_init_method, - timeout=self.server_args.dist_timeout, - moe_a2a_backend=self.server_args.moe_a2a_backend, - recovered_rank=self.server_args.elastic_ep_rejoin, - ) - initialize_model_parallel( - tensor_model_parallel_size=self.tp_size, - attention_data_parallel_size=self.attn_dp_size, - pipeline_model_parallel_size=self.pp_size, - expert_model_parallel_size=self.moe_ep_size, - attention_context_model_parallel_size=self.attn_cp_size, - moe_data_model_parallel_size=self.moe_dp_size, - decode_context_parallel_size=self.dcp_size, - duplicate_tp_group=self.server_args.enable_pdmux, - enable_symm_mem=self.server_args.enable_symm_mem, - recovered_rank=self.server_args.elastic_ep_rejoin, - ) - initialize_dp_attention( - server_args=self.server_args, - model_config=self.model_config, - ) - if is_npu(): - register_sgl_tp_rank(self.gpu_id) - - # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request - # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs) - if self.server_args.pre_warm_nccl and ( - self.tp_size > 1 or self.pp_size > 1 or self.moe_ep_size > 1 - ): - warmup_start = time.perf_counter() - tp_group_handle = get_tp_group().device_group - - # Single warmup all_reduce to initialize NCCL/RCCL/HCCL communicator - warmup_tensor = torch.zeros(1, device=torch.cuda.current_device()) - dist.all_reduce(warmup_tensor, group=tp_group_handle) - current_platform.synchronize() - - warmup_elapsed = time.perf_counter() - warmup_start - logger.info( - f"NCCL/RCCL/HCCL warmup completed in {warmup_elapsed:.3f}s " - f"(tp_size={self.tp_size}, pp_size={self.pp_size}, ep_size={self.moe_ep_size})" - ) - - pre_model_load_memory = get_available_gpu_memory( - self.device, - self.gpu_id, - distributed=get_world_group().world_size > 1, - cpu_group=get_world_group().cpu_group, + result = bootstrap.init_torch_distributed( + server_args=self.server_args, + model_config=self.model_config, + device=self.device, + gpu_id=self.gpu_id, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + pp_rank=self.pp_rank, + pp_size=self.pp_size, + dp_size=self.attn_dp_size, + attn_cp_size=self.attn_cp_size, + moe_ep_size=self.moe_ep_size, + moe_dp_size=self.moe_dp_size, + dcp_size=self.dcp_size, + dist_port=self.dist_port, + is_draft_worker=self.is_draft_worker, + local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None, ) - self.tp_group = get_tp_group() - self.pp_group = get_pp_group() - self.attention_tp_group = get_parallel().attn_tp_group - - # Check memory for tensor parallelism - local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id) - if self.tp_size > 1 and not self.is_draft_worker: - if pre_model_load_memory < local_gpu_memory * 0.9: - msg = "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. " - msg += f"{pre_model_load_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}" - if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get(): - raise RuntimeError(msg) - else: - logger.warning(msg) - - logger.info( - f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, " - f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" - ) - return pre_model_load_memory + self.tp_group = result.tp_group + self.pp_group = result.pp_group + self.attention_tp_group = result.attention_tp_group + self.pre_model_load_memory = result.pre_model_load_memory def init_shared_mooncake_transfer_engine(self): """