ba22cb6765 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -29,9 +29,43 @@ from vllm.distributed.weight_transfer.nccl_engine import (

29 NCCLWeightTransferInitInfo,

30 NCCLWeightTransferUpdateInfo,

31 )

32 +from vllm.platforms import current_platform

33 from vllm.utils.network_utils import get_open_port

34

35

36 +def _weight_transfer_ray_env_vars() -> dict[str, str]:

37 + if not current_platform.is_rocm():

38 + return {}

39 +

40 + return {

41 + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",

42 + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1",

43 + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1",

44 + }

45 +

46 +

47 +def _init_ray_for_weight_transfer() -> None:

48 + if ray.is_initialized():

49 + return

50 + ray.init(

51 + ignore_reinit_error=True,

52 + runtime_env={"env_vars": _weight_transfer_ray_env_vars()},

53 + )

54 +

55 +

56 +def _get_ray_assigned_device() -> torch.device:

57 + gpu_ids = ray.get_gpu_ids()

58 + if not gpu_ids:

59 + return torch.device("cuda:0")

60 + return torch.device(f"cuda:{int(gpu_ids[0])}")

61 +

62 +

63 +def _set_ray_assigned_device() -> torch.device:

64 + device = _get_ray_assigned_device()

65 + torch.accelerator.set_device(device)

66 + return device

67 +

68 +

69 def create_mock_parallel_config(

70 rank: int = 0,

71 world_size: int = 1,

@@ -321,6 +355,8 @@ def trainer_broadcast_tensor(

355 """Trainer task that broadcasts a tensor via NCCL."""

356 import torch

357

358 + device = _set_ray_assigned_device()

359 +

360 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator

361 from vllm.distributed.utils import StatelessProcessGroup

362

@@ -331,12 +367,11 @@ def trainer_broadcast_tensor(

367 rank=0,

368 world_size=world_size,

369 )

334 - # Ray sets CUDA_VISIBLE_DEVICES, so device 0 is the assigned GPU

335 - comm = PyNcclCommunicator(pg, device=0)

370 + comm = PyNcclCommunicator(pg, device=device.index)

371

372 # Create and broadcast the tensor

373 dtype = getattr(torch, tensor_dtype)

339 - tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device="cuda:0")

374 + tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device=device)

375 comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream())

376 torch.accelerator.synchronize()

377

@@ -356,6 +391,8 @@ def inference_receive_tensor(

391

392 import torch

393

394 + _set_ray_assigned_device()

395 +

396 from vllm.config.parallel import ParallelConfig

397 from vllm.config.weight_transfer import WeightTransferConfig

398 from vllm.distributed.weight_transfer.nccl_engine import (

@@ -435,7 +472,7 @@ def test_nccl_weight_transfer_between_processes():

472 This test verifies that the NCCLWeightTransferEngine can receive

473 tensors broadcast by a trainer process via NCCL.

474 """

438 - ray.init(ignore_reinit_error=True)

475 + _init_ray_for_weight_transfer()

476

477 master_address = "127.0.0.1"

478 master_port = get_open_port()

@@ -473,6 +510,8 @@ def trainer_broadcast_sparse_tensor(

510 """Trainer task that broadcasts sparse patches via NCCL."""

511 import torch

512

513 + device = _set_ray_assigned_device()

514 +

515 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator

516 from vllm.distributed.utils import StatelessProcessGroup

517 from vllm.distributed.weight_transfer.base import SparseWeightPatch

@@ -487,12 +526,12 @@ def trainer_broadcast_sparse_tensor(

526 rank=0,

527 world_size=world_size,

528 )

490 - comm = PyNcclCommunicator(pg, device=0)

529 + comm = PyNcclCommunicator(pg, device=device.index)

530

531 patch = SparseWeightPatch(

532 name="test.weight",

494 - indices=torch.tensor([1, 7, 25], dtype=torch.int32, device="cuda:0"),

495 - values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device="cuda:0"),

533 + indices=torch.tensor([1, 7, 25], dtype=torch.int32, device=device),

534 + values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device=device),

535 )

536 NCCLWeightTransferEngine.trainer_send_sparse_weights(

537 iter([patch]),

@@ -513,6 +552,8 @@ def inference_receive_sparse_tensor(

552

553 import torch

554

555 + device = _set_ray_assigned_device()

556 +

557 from vllm.config.parallel import ParallelConfig

558 from vllm.config.weight_transfer import WeightTransferConfig

559 from vllm.distributed.weight_transfer.nccl_engine import (

@@ -540,7 +581,7 @@ def inference_receive_sparse_tensor(

581 )

582 )

583

543 - target = torch.zeros(30, dtype=torch.float32, device="cuda")

584 + target = torch.zeros(30, dtype=torch.float32, device=device)

585

586 def apply_sparse_patches(patches: list[SparseWeightPatch]):

587 for patch in patches:

@@ -556,9 +597,9 @@ def inference_receive_sparse_tensor(

597 engine.receive_sparse_weights(update_info, apply_sparse_patches)

598 torch.accelerator.synchronize()

599

559 - expected = torch.zeros(30, dtype=torch.float32, device="cuda")

600 + expected = torch.zeros(30, dtype=torch.float32, device=device)

601 expected[[1, 7, 25]] = torch.tensor(

561 - [10.0, 20.0, 30.0], dtype=torch.float32, device="cuda"

602 + [10.0, 20.0, 30.0], dtype=torch.float32, device=device

603 )

604 success = torch.equal(target, expected)

605 engine.shutdown()

@@ -574,7 +615,7 @@ def inference_receive_sparse_tensor(

615 )

616 def test_nccl_sparse_weight_transfer_between_processes():

617 """Test NCCL sparse weight transfer from trainer to inference process."""

577 - ray.init(ignore_reinit_error=True)

618 + _init_ray_for_weight_transfer()

619

620 master_address = "127.0.0.1"

621 master_port = get_open_port()

@@ -933,16 +974,18 @@ class TrainerActor:

974 """Trainer actor that creates and holds CUDA IPC handles."""

975

976 def __init__(self, tensor_shape: list[int], tensor_dtype: str):

977 + device = _set_ray_assigned_device()

978 +

979 # Create tensor on GPU and keep it alive

980 dtype = getattr(torch, tensor_dtype)

938 - self.tensor = torch.ones(tensor_shape, dtype=dtype, device="cuda:0")

981 + self.tensor = torch.ones(tensor_shape, dtype=dtype, device=device)

982 self.tensor.fill_(42.0) # Fill with 42 to verify correct transfer

983

984 # Create IPC handle (tensor must stay alive for IPC to work)

985 # reduce_tensor returns (rebuild_func, args); we only send args

986 # since the receiver imports rebuild_cuda_tensor directly.

987 _, ipc_args = reduce_tensor(self.tensor)

945 - gpu_uuid = get_physical_gpu_id(0)

988 + gpu_uuid = get_physical_gpu_id(device.index)

989

990 torch.accelerator.synchronize()

991

@@ -974,6 +1017,8 @@ def inference_receive_ipc_tensor(

1017

1018 import torch

1019

1020 + _set_ray_assigned_device()

1021 +

1022 from vllm.config.parallel import ParallelConfig

1023 from vllm.config.weight_transfer import WeightTransferConfig

1024 from vllm.distributed.weight_transfer.ipc_engine import (

@@ -1072,7 +1117,7 @@ def test_ipc_weight_transfer_between_processes(mode: str):

1117 from ray.util.placement_group import placement_group

1118 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

1119

1075 - ray.init(ignore_reinit_error=True)

1120 + _init_ray_for_weight_transfer()

1121

1122 # Create a placement group to ensure both processes are on the same GPU

1123 # Use fractional GPUs so both tasks can share the same GPU bundle

[ROCm][Ray][CI] Keep assigned GPU visible for weight transfer (#47000)