[CPU] Fix shm allreduce collision and sglang-router import (#37179)
This commit is contained in:
@@ -93,7 +93,10 @@ def init_torch_distributed(
|
|||||||
if not is_draft_worker:
|
if not is_draft_worker:
|
||||||
if device == "cpu":
|
if device == "cpu":
|
||||||
_init_cpu_threads_env(
|
_init_cpu_threads_env(
|
||||||
tp_size=ps.tp_size, tp_rank=ps.tp_rank, local_omp_cpuid=local_omp_cpuid
|
tp_size=ps.tp_size,
|
||||||
|
tp_rank=ps.tp_rank,
|
||||||
|
local_omp_cpuid=local_omp_cpuid,
|
||||||
|
dist_init_method=dist_init_method,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Only initialize the distributed environment on the target model worker.
|
# Only initialize the distributed environment on the target model worker.
|
||||||
@@ -204,8 +207,25 @@ def _set_all_reduce_flags(*, server_args: ServerArgs) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _set_shm_master_env(dist_init_method: Optional[str]) -> None:
|
||||||
|
# setdefault so an explicit user-provided MASTER_ADDR/MASTER_PORT wins.
|
||||||
|
prefix = "tcp://"
|
||||||
|
if (
|
||||||
|
dist_init_method
|
||||||
|
and dist_init_method.startswith(prefix)
|
||||||
|
and ":" in dist_init_method[len(prefix) :]
|
||||||
|
):
|
||||||
|
host, port = dist_init_method[len(prefix) :].rsplit(":", 1)
|
||||||
|
os.environ.setdefault("MASTER_ADDR", host)
|
||||||
|
os.environ.setdefault("MASTER_PORT", port)
|
||||||
|
|
||||||
|
|
||||||
def _init_cpu_threads_env(
|
def _init_cpu_threads_env(
|
||||||
*, tp_size: int, tp_rank: int, local_omp_cpuid: Optional[List[int]]
|
*,
|
||||||
|
tp_size: int,
|
||||||
|
tp_rank: int,
|
||||||
|
local_omp_cpuid: Optional[List[int]],
|
||||||
|
dist_init_method: Optional[str] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
if _is_cpu_amx_available or _is_cpu_arm64:
|
if _is_cpu_amx_available or _is_cpu_arm64:
|
||||||
# Bind OpenMP threads to CPU cores
|
# Bind OpenMP threads to CPU cores
|
||||||
@@ -213,6 +233,13 @@ def _init_cpu_threads_env(
|
|||||||
|
|
||||||
# Set local size to hint SGLang to use shared memory based AllReduce
|
# Set local size to hint SGLang to use shared memory based AllReduce
|
||||||
os.environ["LOCAL_SIZE"] = str(tp_size)
|
os.environ["LOCAL_SIZE"] = str(tp_size)
|
||||||
|
|
||||||
|
# shm.cpp names its /dev/shm segments from MASTER_ADDR/MASTER_PORT.
|
||||||
|
# Feed each engine's unique dist_init_method (tcp://host:port) into
|
||||||
|
# these env vars so co-located engines get distinct segment names and
|
||||||
|
# don't collide.
|
||||||
|
_set_shm_master_env(dist_init_method)
|
||||||
|
|
||||||
torch.ops.sgl_kernel.initialize(tp_size, tp_rank)
|
torch.ops.sgl_kernel.initialize(tp_size, tp_rank)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -1,2 +1,3 @@
|
|||||||
# Temporarily do this to avoid changing all imports in the repo
|
# Temporarily do this to avoid changing all imports in the repo
|
||||||
from sglang.srt.utils.common import *
|
from sglang.srt.utils.common import *
|
||||||
|
from sglang.srt.utils.network import is_port_available # noqa: F401
|
||||||
|
|||||||
@@ -481,4 +481,8 @@ def init_threads_binding(
|
|||||||
f"in this case the available memory amount of each rank cannot be determined in prior. "
|
f"in this case the available memory amount of each rank cannot be determined in prior. "
|
||||||
f"Please set proper `--max-total-tokens` to avoid the out-of-memory error."
|
f"Please set proper `--max-total-tokens` to avoid the out-of-memory error."
|
||||||
)
|
)
|
||||||
|
logger.info(
|
||||||
|
f"init_threads_binding: numa_index={numa_index}, world_size={world_size}, "
|
||||||
|
f"local_omp_cpuid={local_omp_cpuid}"
|
||||||
|
)
|
||||||
return local_omp_cpuid
|
return local_omp_cpuid
|
||||||
|
|||||||
Reference in New Issue
Block a user