refactor: extract device-to-backend mapping into get_default_distributed_backend (#19202)
Signed-off-by: Wang, Yi <yi.a.wang@intel.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Liangsheng Yin <lsyincs@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
@@ -25,6 +25,7 @@ from sglang.srt.configs.load_config import LoadConfig
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.disaggregation.encode_receiver import EmbeddingData
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
get_default_distributed_backend,
|
||||
get_mooncake_transfer_engine,
|
||||
get_tp_group,
|
||||
init_distributed_environment,
|
||||
@@ -176,6 +177,7 @@ class MMEncoder:
|
||||
self.use_image_processor_gpu = use_image_processor_gpu
|
||||
|
||||
init_distributed_environment(
|
||||
backend=get_default_distributed_backend(self.device),
|
||||
world_size=server_args.tp_size,
|
||||
rank=rank,
|
||||
distributed_init_method=dist_init_method,
|
||||
|
||||
@@ -1578,6 +1578,20 @@ def set_torch_symm_mem_all_reduce(enable: bool):
|
||||
_ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = enable
|
||||
|
||||
|
||||
_DEVICE_TO_DISTRIBUTED_BACKEND = {
|
||||
"cuda": "nccl",
|
||||
"xpu": "xccl",
|
||||
"hpu": "hccl",
|
||||
"cpu": "gloo",
|
||||
"npu": "hccl",
|
||||
"musa": "mccl",
|
||||
}
|
||||
|
||||
|
||||
def get_default_distributed_backend(device: str) -> str:
|
||||
return _DEVICE_TO_DISTRIBUTED_BACKEND.get(device, "gloo")
|
||||
|
||||
|
||||
def init_distributed_environment(
|
||||
world_size: int = -1,
|
||||
rank: int = -1,
|
||||
|
||||
@@ -55,6 +55,7 @@ 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,
|
||||
get_tp_group,
|
||||
get_world_group,
|
||||
@@ -734,29 +735,17 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
)
|
||||
raise
|
||||
|
||||
if self.device == "cuda":
|
||||
if self.server_args.elastic_ep_backend == "mooncake":
|
||||
backend = "mooncake"
|
||||
if self.server_args.mooncake_ib_device:
|
||||
mooncake_ib_device = self.server_args.mooncake_ib_device.split(",")
|
||||
try:
|
||||
from mooncake import ep as mooncake_ep
|
||||
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:
|
||||
mooncake_ib_device = self.server_args.mooncake_ib_device.split(",")
|
||||
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`
|
||||
else:
|
||||
backend = "nccl"
|
||||
elif self.device == "xpu":
|
||||
backend = "xccl"
|
||||
elif self.device == "hpu":
|
||||
backend = "hccl"
|
||||
elif self.device == "cpu":
|
||||
backend = "gloo"
|
||||
elif self.device == "npu":
|
||||
backend = "hccl"
|
||||
elif self.device == "musa":
|
||||
backend = "mccl"
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user