diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 2df0d9e71c..697be7b407 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -29,9 +29,43 @@ from vllm.distributed.weight_transfer.nccl_engine import ( NCCLWeightTransferInitInfo, NCCLWeightTransferUpdateInfo, ) +from vllm.platforms import current_platform from vllm.utils.network_utils import get_open_port +def _weight_transfer_ray_env_vars() -> dict[str, str]: + if not current_platform.is_rocm(): + return {} + + return { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1", + } + + +def _init_ray_for_weight_transfer() -> None: + if ray.is_initialized(): + return + ray.init( + ignore_reinit_error=True, + runtime_env={"env_vars": _weight_transfer_ray_env_vars()}, + ) + + +def _get_ray_assigned_device() -> torch.device: + gpu_ids = ray.get_gpu_ids() + if not gpu_ids: + return torch.device("cuda:0") + return torch.device(f"cuda:{int(gpu_ids[0])}") + + +def _set_ray_assigned_device() -> torch.device: + device = _get_ray_assigned_device() + torch.accelerator.set_device(device) + return device + + def create_mock_parallel_config( rank: int = 0, world_size: int = 1, @@ -321,6 +355,8 @@ def trainer_broadcast_tensor( """Trainer task that broadcasts a tensor via NCCL.""" import torch + device = _set_ray_assigned_device() + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.utils import StatelessProcessGroup @@ -331,12 +367,11 @@ def trainer_broadcast_tensor( rank=0, world_size=world_size, ) - # Ray sets CUDA_VISIBLE_DEVICES, so device 0 is the assigned GPU - comm = PyNcclCommunicator(pg, device=0) + comm = PyNcclCommunicator(pg, device=device.index) # Create and broadcast the tensor dtype = getattr(torch, tensor_dtype) - tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device="cuda:0") + tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device=device) comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream()) torch.accelerator.synchronize() @@ -356,6 +391,8 @@ def inference_receive_tensor( import torch + _set_ray_assigned_device() + from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.nccl_engine import ( @@ -435,7 +472,7 @@ def test_nccl_weight_transfer_between_processes(): This test verifies that the NCCLWeightTransferEngine can receive tensors broadcast by a trainer process via NCCL. """ - ray.init(ignore_reinit_error=True) + _init_ray_for_weight_transfer() master_address = "127.0.0.1" master_port = get_open_port() @@ -473,6 +510,8 @@ def trainer_broadcast_sparse_tensor( """Trainer task that broadcasts sparse patches via NCCL.""" import torch + device = _set_ray_assigned_device() + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.utils import StatelessProcessGroup from vllm.distributed.weight_transfer.base import SparseWeightPatch @@ -487,12 +526,12 @@ def trainer_broadcast_sparse_tensor( rank=0, world_size=world_size, ) - comm = PyNcclCommunicator(pg, device=0) + comm = PyNcclCommunicator(pg, device=device.index) patch = SparseWeightPatch( name="test.weight", - indices=torch.tensor([1, 7, 25], dtype=torch.int32, device="cuda:0"), - values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device="cuda:0"), + indices=torch.tensor([1, 7, 25], dtype=torch.int32, device=device), + values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device=device), ) NCCLWeightTransferEngine.trainer_send_sparse_weights( iter([patch]), @@ -513,6 +552,8 @@ def inference_receive_sparse_tensor( import torch + device = _set_ray_assigned_device() + from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.nccl_engine import ( @@ -540,7 +581,7 @@ def inference_receive_sparse_tensor( ) ) - target = torch.zeros(30, dtype=torch.float32, device="cuda") + target = torch.zeros(30, dtype=torch.float32, device=device) def apply_sparse_patches(patches: list[SparseWeightPatch]): for patch in patches: @@ -556,9 +597,9 @@ def inference_receive_sparse_tensor( engine.receive_sparse_weights(update_info, apply_sparse_patches) torch.accelerator.synchronize() - expected = torch.zeros(30, dtype=torch.float32, device="cuda") + expected = torch.zeros(30, dtype=torch.float32, device=device) expected[[1, 7, 25]] = torch.tensor( - [10.0, 20.0, 30.0], dtype=torch.float32, device="cuda" + [10.0, 20.0, 30.0], dtype=torch.float32, device=device ) success = torch.equal(target, expected) engine.shutdown() @@ -574,7 +615,7 @@ def inference_receive_sparse_tensor( ) def test_nccl_sparse_weight_transfer_between_processes(): """Test NCCL sparse weight transfer from trainer to inference process.""" - ray.init(ignore_reinit_error=True) + _init_ray_for_weight_transfer() master_address = "127.0.0.1" master_port = get_open_port() @@ -933,16 +974,18 @@ class TrainerActor: """Trainer actor that creates and holds CUDA IPC handles.""" def __init__(self, tensor_shape: list[int], tensor_dtype: str): + device = _set_ray_assigned_device() + # Create tensor on GPU and keep it alive dtype = getattr(torch, tensor_dtype) - self.tensor = torch.ones(tensor_shape, dtype=dtype, device="cuda:0") + self.tensor = torch.ones(tensor_shape, dtype=dtype, device=device) self.tensor.fill_(42.0) # Fill with 42 to verify correct transfer # Create IPC handle (tensor must stay alive for IPC to work) # reduce_tensor returns (rebuild_func, args); we only send args # since the receiver imports rebuild_cuda_tensor directly. _, ipc_args = reduce_tensor(self.tensor) - gpu_uuid = get_physical_gpu_id(0) + gpu_uuid = get_physical_gpu_id(device.index) torch.accelerator.synchronize() @@ -974,6 +1017,8 @@ def inference_receive_ipc_tensor( import torch + _set_ray_assigned_device() + from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.ipc_engine import ( @@ -1072,7 +1117,7 @@ def test_ipc_weight_transfer_between_processes(mode: str): from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy - ray.init(ignore_reinit_error=True) + _init_ray_for_weight_transfer() # Create a placement group to ensure both processes are on the same GPU # Use fractional GPUs so both tasks can share the same GPU bundle