minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Extract init_torch_distributed and refactor into functions (#31152)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-14 15:56:57 +08:00
committed by GitHub
parent 205a2f2de4
commit caa85ea022
2 changed files with 315 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/bootstrap.py
+291
View File
@@ -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)
python/sglang/srt/model_executor/model_runner.py
+24 -157
View File
@@ -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):
"""