diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 1cd6752460..362ffe4b15 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -14,6 +14,7 @@ import torch import torch.multiprocessing as mp from vllm.distributed import get_ep_group +from vllm.platforms import current_platform from vllm.utils.flashinfer import ( has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, @@ -23,6 +24,8 @@ from vllm.utils.network_utils import get_open_port from ..utils import init_test_distributed_environment +DEVICE = current_platform.device_type + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -483,7 +486,7 @@ def _args_dispatch_combine_worker(rank, world_size): from vllm.forward_context import get_forward_context cpu_group = get_ep_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") hidden_size = 64 tokens_per_rank = 16 @@ -624,7 +627,7 @@ def _two_sided_data_worker(rank, world_size): # Use DP group because MnnvlMoe workspace allocation calls get_dp_group() # internally and requires dp_size == ep_size. cpu_group = get_dp_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") num_gpus = torch.accelerator.device_count() hidden_size = 128 @@ -769,7 +772,7 @@ def _one_sided_data_worker(rank, world_size): from vllm.forward_context import get_forward_context cpu_group = get_dp_group().cpu_group - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{DEVICE}:{rank}") hidden_size = 256 tokens_per_rank = 32 diff --git a/tests/utils.py b/tests/utils.py index 81f82d98c8..7ee0e66803 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1427,8 +1427,10 @@ def init_test_distributed_environment( get_current_vllm_config_or_none, set_current_vllm_config, ) + from vllm.platforms import current_platform distributed_init_method = f"tcp://localhost:{distributed_init_port}" + backend = current_platform.dist_backend if data_parallel_size > 1: # For DP we need to set a common DP master port @@ -1449,6 +1451,7 @@ def init_test_distributed_environment( rank=rank % tp_pp_world, distributed_init_method=distributed_init_method, local_rank=local_rank if local_rank >= 0 else rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) return @@ -1460,6 +1463,7 @@ def init_test_distributed_environment( rank=rank, distributed_init_method=distributed_init_method, local_rank=local_rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size) else: @@ -1470,6 +1474,7 @@ def init_test_distributed_environment( rank=rank, distributed_init_method=distributed_init_method, local_rank=local_rank, + backend=backend, ) ensure_model_parallel_initialized(tp_size, pp_size)