| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test the communication operators. |
| 4 | |
| 5 | Run `pytest tests/distributed/test_comm_ops.py`. |
| 6 | """ |
| 7 | |
| 8 | from collections import deque |
| 9 | from collections.abc import Callable |
| 10 | from typing import Any |
| 11 | from unittest.mock import Mock |
| 12 | |
| 13 | import pytest |
| 14 | import ray |
| 15 | import torch |
| 16 | |
| 17 | from vllm.distributed import ( |
| 18 | broadcast_tensor_dict, |
| 19 | get_pp_group, |
| 20 | tensor_model_parallel_all_gather, |
| 21 | tensor_model_parallel_all_reduce, |
| 22 | tensor_model_parallel_reduce_scatter, |
| 23 | ) |
| 24 | from vllm.distributed.device_communicators import flashinfer_all_reduce |
| 25 | from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator |
| 26 | from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata |
| 27 | from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors |
| 28 | |
| 29 | from ..utils import ( |
| 30 | init_test_distributed_environment, |
| 31 | multi_gpu_test, |
| 32 | multi_process_parallel, |
| 33 | ) |
| 34 | |
| 35 | |
| 36 | @ray.remote(num_gpus=1, max_calls=1) |
| 37 | def all_reduce_test_worker( |
| 38 | monkeypatch: pytest.MonkeyPatch, |
| 39 | tp_size: int, |
| 40 | pp_size: int, |
| 41 | rank: int, |
| 42 | distributed_init_port: str, |
| 43 | ): |
| 44 | # it is important to delete the CUDA_VISIBLE_DEVICES environment variable |
| 45 | # so that each worker can see all the GPUs |
| 46 | # they will be able to set the device to the correct GPU |
| 47 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 48 | |
| 49 | device = torch.device(f"cuda:{rank}") |
| 50 | torch.accelerator.set_device_index(device) |
| 51 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 52 | num_elements = 8 |
| 53 | all_tensors = [ |
| 54 | torch.arange(num_elements, dtype=torch.float32, device="cuda") * (r + 1) |
| 55 | for r in range(tp_size) |
| 56 | ] |
| 57 | expected = torch.sum(torch.stack(all_tensors, dim=0), dim=0) |
| 58 | t = all_tensors[rank % tp_size] |
| 59 | t = tensor_model_parallel_all_reduce(t) |
| 60 | torch.testing.assert_close(t, expected) |
| 61 | |
| 62 | |
| 63 | @ray.remote(num_gpus=1, max_calls=1) |
| 64 | def reduce_scatter_test_worker( |
| 65 | monkeypatch: pytest.MonkeyPatch, |
| 66 | tp_size: int, |
| 67 | pp_size: int, |
| 68 | rank: int, |
| 69 | distributed_init_port: str, |
| 70 | ): |
| 71 | # it is important to delete the CUDA_VISIBLE_DEVICES environment variable |
| 72 | # so that each worker can see all the GPUs |
| 73 | # they will be able to set the device to the correct GPU |
| 74 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 75 | device = torch.device(f"cuda:{rank}") |
| 76 | torch.accelerator.set_device_index(device) |
| 77 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 78 | |
| 79 | num_elements = 8 |
| 80 | all_tensors = [ |
| 81 | torch.arange(num_elements, dtype=torch.float32, device="cuda") * (r + 1) |
| 82 | for r in range(tp_size) |
| 83 | ] |
| 84 | |
| 85 | index = rank % tp_size |
| 86 | partition_size = num_elements // tp_size |
| 87 | all_reduce = torch.sum(torch.stack(all_tensors, dim=0), dim=0) |
| 88 | expected = all_reduce[index * partition_size : (index + 1) * partition_size] |
| 89 | t = all_tensors[index] |
| 90 | t = tensor_model_parallel_reduce_scatter(t, 0) |
| 91 | torch.testing.assert_close(t, expected) |
| 92 | |
| 93 | |
| 94 | @ray.remote(num_gpus=1, max_calls=1) |
| 95 | def all_gather_test_worker( |
| 96 | monkeypatch: pytest.MonkeyPatch, |
| 97 | tp_size: int, |
| 98 | pp_size: int, |
| 99 | rank: int, |
| 100 | distributed_init_port: str, |
| 101 | ): |
| 102 | # it is important to delete the CUDA_VISIBLE_DEVICES environment variable |
| 103 | # so that each worker can see all the GPUs |
| 104 | # they will be able to set the device to the correct GPU |
| 105 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 106 | device = torch.device(f"cuda:{rank}") |
| 107 | torch.accelerator.set_device_index(device) |
| 108 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 109 | num_dimensions = 3 |
| 110 | tensor_size = list(range(2, num_dimensions + 2)) |
| 111 | total_size = 1 |
| 112 | for s in tensor_size: |
| 113 | total_size *= s |
| 114 | for all_gather_dimension in range(num_dimensions): |
| 115 | all_tensors = [ |
| 116 | torch.arange(total_size, dtype=torch.float32, device="cuda").reshape( |
| 117 | tensor_size |
| 118 | ) |
| 119 | * (r + 1) |
| 120 | for r in range(tp_size) |
| 121 | ] |
| 122 | expected = torch.cat(all_tensors, dim=all_gather_dimension) |
| 123 | t = all_tensors[rank % tp_size] |
| 124 | t = tensor_model_parallel_all_gather(t, all_gather_dimension) |
| 125 | torch.testing.assert_close(t, expected) |
| 126 | |
| 127 | |
| 128 | @ray.remote(num_gpus=1, max_calls=1) |
| 129 | def broadcast_tensor_dict_test_worker( |
| 130 | monkeypatch: pytest.MonkeyPatch, |
| 131 | tp_size: int, |
| 132 | pp_size: int, |
| 133 | rank: int, |
| 134 | distributed_init_port: str, |
| 135 | ): |
| 136 | # it is important to delete the CUDA_VISIBLE_DEVICES environment variable |
| 137 | # so that each worker can see all the GPUs |
| 138 | # they will be able to set the device to the correct GPU |
| 139 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 140 | device = torch.device(f"cuda:{rank}") |
| 141 | torch.accelerator.set_device_index(device) |
| 142 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 143 | test_dict = { |
| 144 | # device tensor |
| 145 | "a": torch.arange(8, dtype=torch.float32, device="cuda"), |
| 146 | # CPU tensor |
| 147 | "b": torch.arange(16, dtype=torch.int8, device="cpu"), |
| 148 | "c": "test", |
| 149 | "d": [1, 2, 3], |
| 150 | "e": {"a": 1, "b": 2}, |
| 151 | # empty tensor |
| 152 | "f": torch.tensor([], dtype=torch.float32, device="cuda"), |
| 153 | } |
| 154 | |
| 155 | if (rank % tp_size) == 0: |
| 156 | broadcast_tensor_dict(test_dict, src=0) |
| 157 | else: |
| 158 | recv_dict = broadcast_tensor_dict(src=0) |
| 159 | assert len(recv_dict) == len(test_dict) |
| 160 | torch.testing.assert_close(recv_dict["a"], test_dict["a"]) |
| 161 | torch.testing.assert_close(recv_dict["b"], test_dict["b"]) |
| 162 | assert recv_dict["c"] == test_dict["c"] |
| 163 | assert recv_dict["d"] == test_dict["d"] |
| 164 | assert recv_dict["e"] == test_dict["e"] |
| 165 | torch.testing.assert_close(recv_dict["f"], test_dict["f"]) |
| 166 | |
| 167 | |
| 168 | @ray.remote(num_gpus=1, max_calls=1) |
| 169 | def send_recv_tensor_dict_test_worker( |
| 170 | monkeypatch: pytest.MonkeyPatch, |
| 171 | tp_size: int, |
| 172 | pp_size: int, |
| 173 | rank: int, |
| 174 | distributed_init_port: str, |
| 175 | ): |
| 176 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 177 | device = torch.device(f"cuda:{rank}") |
| 178 | torch.accelerator.set_device_index(device) |
| 179 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 180 | |
| 181 | test_dict = { |
| 182 | # device tensor |
| 183 | "a": torch.arange(8, dtype=torch.float32, device="cuda"), |
| 184 | # CPU tensor |
| 185 | "b": torch.arange(16, dtype=torch.int8, device="cpu"), |
| 186 | "c": "test", |
| 187 | "d": [1, 2, 3], |
| 188 | "e": {"a": 1, "b": 2}, |
| 189 | # empty tensor |
| 190 | "f": torch.tensor([], dtype=torch.float32, device="cuda"), |
| 191 | } |
| 192 | |
| 193 | if not get_pp_group().is_first_rank: |
| 194 | recv_dict = get_pp_group().recv_tensor_dict() |
| 195 | |
| 196 | if not get_pp_group().is_last_rank: |
| 197 | get_pp_group().send_tensor_dict(test_dict) |
| 198 | |
| 199 | if not get_pp_group().is_first_rank: |
| 200 | assert len(recv_dict) == len(test_dict) |
| 201 | torch.testing.assert_close(recv_dict["a"], test_dict["a"]) |
| 202 | torch.testing.assert_close(recv_dict["b"], test_dict["b"]) |
| 203 | assert recv_dict["c"] == test_dict["c"] |
| 204 | assert recv_dict["d"] == test_dict["d"] |
| 205 | assert recv_dict["e"] == test_dict["e"] |
| 206 | torch.testing.assert_close(recv_dict["f"], test_dict["f"]) |
| 207 | |
| 208 | |
| 209 | class _DummyWork: |
| 210 | def __init__(self) -> None: |
| 211 | self.wait_calls = 0 |
| 212 | self.completed = False |
| 213 | |
| 214 | def wait(self) -> None: |
| 215 | self.wait_calls += 1 |
| 216 | self.completed = True |
| 217 | |
| 218 | def is_completed(self) -> bool: |
| 219 | return self.completed |
| 220 | |
| 221 | |
| 222 | class _DummyAllGatherGroup: |
| 223 | def __init__(self, world_size: int, rank_in_group: int) -> None: |
| 224 | self.world_size = world_size |
| 225 | self.rank_in_group = rank_in_group |
| 226 | |
| 227 | def all_gather(self, t: torch.Tensor, dim: int = 0) -> torch.Tensor: |
| 228 | # duplicate local slice across ranks. |
| 229 | assert dim == 0 |
| 230 | return torch.cat([t for _ in range(self.world_size)], dim=0) |
| 231 | |
| 232 | |
| 233 | def _make_group_for_unit_test( |
| 234 | rank_in_group: int = 0, world_size: int = 2 |
| 235 | ) -> GroupCoordinator: |
| 236 | # avoid running GroupCoordinator.__init__ (it wires up real process groups). |
| 237 | g = GroupCoordinator.__new__(GroupCoordinator) |
| 238 | g.world_size = world_size |
| 239 | g.rank_in_group = rank_in_group |
| 240 | g.ranks = list(range(world_size)) |
| 241 | g.use_cpu_custom_send_recv = False |
| 242 | g.device_group = None |
| 243 | g.cpu_group = None |
| 244 | g._pending_isends = deque() |
| 245 | return g |
| 246 | |
| 247 | |
| 248 | def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys( |
| 249 | monkeypatch: pytest.MonkeyPatch, |
| 250 | ) -> None: |
| 251 | def fake_irecv(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork: |
| 252 | t.fill_(1) |
| 253 | return _DummyWork() |
| 254 | |
| 255 | monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) |
| 256 | monkeypatch.setattr(torch.distributed, "irecv", fake_irecv) |
| 257 | |
| 258 | g = _make_group_for_unit_test(rank_in_group=0, world_size=2) |
| 259 | # 2 tensors so we can catch late-binding bugs in postprocess closures. |
| 260 | metadata_list = [ |
| 261 | ("a", TensorMetadata("cpu", torch.int32, torch.Size([4]))), |
| 262 | ("b", TensorMetadata("cpu", torch.int32, torch.Size([4]))), |
| 263 | ] |
| 264 | g.recv_object = lambda src=None: metadata_list # type: ignore[method-assign] |
| 265 | |
| 266 | ag = _DummyAllGatherGroup(world_size=2, rank_in_group=0) |
| 267 | td, handles, postprocess = g.irecv_tensor_dict(all_gather_group=ag) |
| 268 | |
| 269 | assert td is not None |
| 270 | assert len(handles) == 2 |
| 271 | assert len(postprocess) == 2 |
| 272 | |
| 273 | # before postprocess, dict holds the TP slice (shape 2). |
| 274 | assert td["a"].shape == torch.Size([2]) |
| 275 | assert td["b"].shape == torch.Size([2]) |
| 276 | |
| 277 | # simulate worker-side "defer wait": wait + postprocess later. |
| 278 | for handle in handles: |
| 279 | handle.wait() |
| 280 | for fn in postprocess: |
| 281 | fn() |
| 282 | |
| 283 | # after postprocess, dict values are reconstructed to full shape (shape 4), |
| 284 | # and each key should be updated independently |
| 285 | assert td["a"].shape == torch.Size([4]) |
| 286 | assert td["b"].shape == torch.Size([4]) |
| 287 | torch.testing.assert_close(td["a"], torch.ones(4, dtype=torch.int32)) |
| 288 | torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32)) |
| 289 | |
| 290 | |
| 291 | @pytest.mark.parametrize("aliased", [False, True]) |
| 292 | def test_cuda_communicator_checkpoints_flashinfer_workspaces( |
| 293 | monkeypatch: pytest.MonkeyPatch, |
| 294 | aliased: bool, |
| 295 | ) -> None: |
| 296 | group = object() |
| 297 | normal_workspace = Mock() |
| 298 | quant_workspace = normal_workspace if aliased else Mock() |
| 299 | unique_workspaces = ( |
| 300 | [normal_workspace] if aliased else [normal_workspace, quant_workspace] |
| 301 | ) |
| 302 | |
| 303 | monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", normal_workspace) |
| 304 | monkeypatch.setattr( |
| 305 | flashinfer_all_reduce, "_fi_ar_quant_workspace", quant_workspace |
| 306 | ) |
| 307 | monkeypatch.setattr( |
| 308 | flashinfer_all_reduce, |
| 309 | "_fi_ar_workspace_groups", |
| 310 | {id(workspace): group for workspace in unique_workspaces}, |
| 311 | ) |
| 312 | monkeypatch.setattr( |
| 313 | flashinfer_all_reduce, "TorchDistBackend", lambda group: group, raising=False |
| 314 | ) |
| 315 | |
| 316 | communicator = CudaCommunicator.__new__(CudaCommunicator) |
| 317 | communicator.cpu_group = group |
| 318 | communicator.fi_ar_comm = None |
| 319 | communicator.all2all_manager = None |
| 320 | communicator.checkpoint_prepare() |
| 321 | communicator.checkpoint_restore() |
| 322 | |
| 323 | for workspace in unique_workspaces: |
| 324 | workspace.checkpoint_prepare.assert_called_once_with() |
| 325 | workspace.checkpoint_restore.assert_called_once_with(group) |
| 326 | |
| 327 | |
| 328 | @pytest.mark.parametrize( |
| 329 | ("backend", "capability", "world_size", "nodes", "expected"), |
| 330 | [ |
| 331 | ("mnnvl", 103, 4, 1, 80 * flashinfer_all_reduce.MiB - 1), |
| 332 | ("mnnvl", 103, 8, 2, 64 * flashinfer_all_reduce.MiB - 1), |
| 333 | ("mnnvl", 103, 16, 4, 8 * flashinfer_all_reduce.MiB - 1), |
| 334 | ("mnnvl", 103, 2, 1, None), |
| 335 | ("mnnvl", 103, 12, 3, None), |
| 336 | ("mnnvl", 103, 8, 1, 64 * flashinfer_all_reduce.MiB - 1), |
| 337 | ("trtllm", 103, 8, 2, None), |
| 338 | ("mnnvl", 90, 8, 2, None), |
| 339 | ], |
| 340 | ) |
| 341 | def test_flashinfer_standalone_size_tuning( |
| 342 | monkeypatch: pytest.MonkeyPatch, |
| 343 | backend: str, |
| 344 | capability: int, |
| 345 | world_size: int, |
| 346 | nodes: int, |
| 347 | expected: int | None, |
| 348 | ) -> None: |
| 349 | monkeypatch.setattr( |
| 350 | flashinfer_all_reduce, |
| 351 | "current_platform", |
| 352 | Mock(get_device_capability=lambda: Mock(to_int=lambda: capability)), |
| 353 | ) |
| 354 | monkeypatch.setattr(flashinfer_all_reduce, "_node_count", lambda _: nodes) |
| 355 | |
| 356 | assert ( |
| 357 | flashinfer_all_reduce._get_tuned_standalone_max_size( |
| 358 | world_size, backend, Mock() |
| 359 | ) |
| 360 | == expected |
| 361 | ) |
| 362 | |
| 363 | |
| 364 | @pytest.mark.parametrize(("enabled", "expected"), [(True, 4681), (False, 128)]) |
| 365 | def test_flashinfer_standalone_workspace_size( |
| 366 | monkeypatch: pytest.MonkeyPatch, enabled: bool, expected: int |
| 367 | ) -> None: |
| 368 | create_workspace = Mock(return_value=Mock(backend="mnnvl")) |
| 369 | monkeypatch.setattr( |
| 370 | flashinfer_all_reduce.envs, "VLLM_ALLREDUCE_USE_FLASHINFER", enabled |
| 371 | ) |
| 372 | monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", None) |
| 373 | monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_quant_workspace", None) |
| 374 | monkeypatch.setattr( |
| 375 | flashinfer_all_reduce, |
| 376 | "_resolve_fi_ar_backend", |
| 377 | Mock(return_value=("mnnvl", False)), |
| 378 | ) |
| 379 | monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 2) |
| 380 | monkeypatch.setattr( |
| 381 | flashinfer_all_reduce, |
| 382 | "_get_tuned_standalone_max_size", |
| 383 | Mock(return_value=64 * flashinfer_all_reduce.MiB - 1), |
| 384 | ) |
| 385 | monkeypatch.setattr(flashinfer_all_reduce, "_create_workspace", create_workspace) |
| 386 | |
| 387 | flashinfer_all_reduce.get_fi_ar_workspace(8, 0, 128, 7168, torch.bfloat16, Mock()) |
| 388 | |
| 389 | assert create_workspace.call_args.args[3] == expected |
| 390 | |
| 391 | |
| 392 | def test_flashinfer_workspace_failure_is_not_retried( |
| 393 | monkeypatch: pytest.MonkeyPatch, |
| 394 | ) -> None: |
| 395 | """Creation is collective; retrying it per all-reduce desyncs ranks.""" |
| 396 | create_workspace = Mock(return_value=None) |
| 397 | monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", None) |
| 398 | monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_quant_workspace", None) |
| 399 | monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace_failed", False) |
| 400 | monkeypatch.setattr( |
| 401 | flashinfer_all_reduce, |
| 402 | "_resolve_fi_ar_backend", |
| 403 | Mock(return_value=("mnnvl", True)), |
| 404 | ) |
| 405 | monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 1) |
| 406 | monkeypatch.setattr( |
| 407 | flashinfer_all_reduce, "_get_tuned_standalone_max_size", Mock(return_value=None) |
| 408 | ) |
| 409 | monkeypatch.setattr(flashinfer_all_reduce, "_create_workspace", create_workspace) |
| 410 | args = (2, 0, 128, 4096, torch.bfloat16, Mock()) |
| 411 | |
| 412 | assert flashinfer_all_reduce.get_fi_ar_workspace(*args) is None |
| 413 | assert flashinfer_all_reduce.get_fi_ar_workspace(*args) is None |
| 414 | # One mnnvl attempt plus one trtllm fallback, never repeated. |
| 415 | assert create_workspace.call_count == 2 |
| 416 | |
| 417 | flashinfer_all_reduce.destroy_fi_ar_workspace() |
| 418 | flashinfer_all_reduce.get_fi_ar_workspace(*args) |
| 419 | assert create_workspace.call_count == 4 |
| 420 | |
| 421 | |
| 422 | def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) -> None: |
| 423 | output = torch.empty(2) |
| 424 | fi_ar_comm = Mock(disabled=False) |
| 425 | fi_ar_comm.should_use_fi_ar.return_value = True |
| 426 | fi_ar_comm.all_reduce.return_value = output |
| 427 | communicator = CudaCommunicator.__new__(CudaCommunicator) |
| 428 | communicator.fi_ar_comm = fi_ar_comm |
| 429 | communicator.fi_pcie_ipc_ar_comm = None |
| 430 | communicator.pynccl_comm = Mock(world_size=8) |
| 431 | communicator.qr_comm = None |
| 432 | nccl_selector = Mock(return_value=True) |
| 433 | monkeypatch.setattr( |
| 434 | "vllm.distributed.device_communicators.cuda_communicator." |
| 435 | "should_nccl_symm_mem_allreduce", |
| 436 | nccl_selector, |
| 437 | ) |
| 438 | |
| 439 | assert communicator.all_reduce(torch.empty(1)) is output |
| 440 | nccl_selector.assert_not_called() |
| 441 | |
| 442 | |
| 443 | def test_aiter_all_gather_precedes_pynccl(monkeypatch: pytest.MonkeyPatch) -> None: |
| 444 | """Ensure can use aiter all_gather if available even if pynccl is not available.""" |
| 445 | output = torch.empty(2) |
| 446 | aiter_comm = Mock() |
| 447 | aiter_comm.should_custom_ag.return_value = True |
| 448 | aiter_comm.custom_all_gather.return_value = output |
| 449 | communicator = CudaCommunicator.__new__(CudaCommunicator) |
| 450 | communicator.world_size = 2 |
| 451 | communicator.aiter_ar_comm = aiter_comm |
| 452 | communicator.pynccl_comm = Mock(disabled=True) |
| 453 | monkeypatch.setattr(communicator, "_can_use_aiter_ag_rs", Mock(return_value=True)) |
| 454 | |
| 455 | assert communicator.all_gatherv(torch.empty(1)) is output |
| 456 | |
| 457 | |
| 458 | def test_isend_object_posts_size_then_object_and_releases_on_wait( |
| 459 | monkeypatch: pytest.MonkeyPatch, |
| 460 | ) -> None: |
| 461 | posted: list[tuple[torch.Tensor, _DummyWork]] = [] |
| 462 | |
| 463 | def fake_isend(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork: |
| 464 | w = _DummyWork() |
| 465 | posted.append((t, w)) |
| 466 | return w |
| 467 | |
| 468 | monkeypatch.setattr(torch.distributed, "isend", fake_isend) |
| 469 | |
| 470 | g = _make_group_for_unit_test(rank_in_group=0, world_size=2) |
| 471 | handle = g.isend_object({"k": [1, 2, 3]}, dst=1) |
| 472 | |
| 473 | # two sends, in size-then-object order (preserves gloo FIFO). |
| 474 | assert len(posted) == 2 |
| 475 | assert posted[0][0].dtype == torch.long |
| 476 | assert posted[0][0].shape == torch.Size([1]) |
| 477 | assert posted[0][0].item() == posted[1][0].numel() |
| 478 | assert posted[1][0].dtype == torch.uint8 |
| 479 | |
| 480 | # retain holds both source tensors until wait. |
| 481 | assert handle._retained == (posted[0][0], posted[1][0]) |
| 482 | |
| 483 | handle.wait() |
| 484 | |
| 485 | # both underlying works drained, retain dropped. |
| 486 | assert all(w.wait_calls == 1 for _, w in posted) |
| 487 | assert handle._retained == () |
| 488 | |
| 489 | # wait is idempotent: gloo Work.wait() is single-shot for p2p sends (a |
| 490 | # second wait blocks forever), and both the lazy FIFO reap in |
| 491 | # ``_reap_completed_isends`` and any explicit isend caller may wait the |
| 492 | # same metadata handle. |
| 493 | handle.wait() |
| 494 | assert all(w.wait_calls == 1 for _, w in posted) |
| 495 | assert handle._retained == () |
| 496 | |
| 497 | |
| 498 | def test_isend_tensor_dict_includes_metadata_handle( |
| 499 | monkeypatch: pytest.MonkeyPatch, |
| 500 | ) -> None: |
| 501 | posted: list[torch.Tensor] = [] |
| 502 | |
| 503 | def fake_isend(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork: |
| 504 | posted.append(t) |
| 505 | return _DummyWork() |
| 506 | |
| 507 | monkeypatch.setattr(torch.distributed, "isend", fake_isend) |
| 508 | |
| 509 | g = _make_group_for_unit_test(rank_in_group=0, world_size=2) |
| 510 | td = {"a": torch.arange(4, dtype=torch.float32, device="cpu")} |
| 511 | handles = g.isend_tensor_dict(td, dst=1) |
| 512 | |
| 513 | # size + object (metadata) + one tensor send. |
| 514 | assert len(posted) == 3 |
| 515 | assert posted[0].dtype == torch.long |
| 516 | assert posted[1].dtype == torch.uint8 |
| 517 | assert posted[2].dtype == torch.float32 |
| 518 | |
| 519 | # composite metadata handle + one tensor handle. |
| 520 | assert len(handles) == 2 |
| 521 | |
| 522 | for handle in handles: |
| 523 | handle.wait() |
| 524 | assert handles[0]._retained == () |
| 525 | |
| 526 | |
| 527 | def test_isend_tensor_dict_self_retains_for_fire_and_forget( |
| 528 | monkeypatch: pytest.MonkeyPatch, |
| 529 | ) -> None: |
| 530 | works: list[_DummyWork] = [] |
| 531 | |
| 532 | def fake_isend(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork: |
| 533 | w = _DummyWork() |
| 534 | works.append(w) |
| 535 | return w |
| 536 | |
| 537 | monkeypatch.setattr(torch.distributed, "isend", fake_isend) |
| 538 | |
| 539 | g = _make_group_for_unit_test(rank_in_group=0, world_size=2) |
| 540 | td = {"a": torch.arange(4, dtype=torch.float32, device="cpu")} |
| 541 | |
| 542 | # fire-and-forget: the returned handles are dropped by the caller, so the |
| 543 | # group must self-retain the handles and the source tensor. |
| 544 | g.isend_tensor_dict(td, dst=1) |
| 545 | assert len(g._pending_isends) == 1 |
| 546 | handles0, tensors0 = g._pending_isends[0] |
| 547 | assert len(handles0) == 2 # metadata composite + one tensor send |
| 548 | assert tensors0 == [td["a"]] |
| 549 | |
| 550 | # the first send is still in flight, so the next call's reap keeps it. |
| 551 | g.isend_tensor_dict(td, dst=1) |
| 552 | assert len(g._pending_isends) == 2 |
| 553 | |
| 554 | # metadata-only send (no tensors) is dropped best-effort on the next |
| 555 | # reap, since it has no reliable completion signal. |
| 556 | g.isend_tensor_dict({"meta": "no-tensors"}, dst=1) |
| 557 | assert len(g._pending_isends) == 3 |
| 558 | |
| 559 | # complete every posted work; the next call reaps all three old entries |
| 560 | # (the metadata-only one unconditionally) and keeps only its own. |
| 561 | for w in works: |
| 562 | w.completed = True |
| 563 | g.isend_tensor_dict(td, dst=1) |
| 564 | assert len(g._pending_isends) == 1 |
| 565 | # reaped entries had their metadata handles waited as a backstop. |
| 566 | assert handles0[0]._retained == () |
| 567 | |
| 568 | |
| 569 | def test_send_tensor_dict_sync_path_does_not_self_retain( |
| 570 | monkeypatch: pytest.MonkeyPatch, |
| 571 | ) -> None: |
| 572 | def fake_isend(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork: |
| 573 | return _DummyWork() |
| 574 | |
| 575 | monkeypatch.setattr(torch.distributed, "isend", fake_isend) |
| 576 | monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) |
| 577 | monkeypatch.setattr(GroupCoordinator, "send_object", lambda self, obj, dst: None) |
| 578 | |
| 579 | g = _make_group_for_unit_test(rank_in_group=0, world_size=2) |
| 580 | td = {"a": torch.arange(4, dtype=torch.float32, device="cpu")} |
| 581 | g.send_tensor_dict(td, dst=1) |
| 582 | |
| 583 | # the sync path waited every handle, so no retention entry may leak. |
| 584 | assert len(g._pending_isends) == 0 |
| 585 | |
| 586 | |
| 587 | def test_async_intermediate_tensors_lazy_wait() -> None: |
| 588 | work = _DummyWork() |
| 589 | post_calls = {"n": 0} |
| 590 | |
| 591 | def post() -> None: |
| 592 | post_calls["n"] += 1 |
| 593 | |
| 594 | it = AsyncIntermediateTensors( |
| 595 | {"x": torch.tensor([1])}, |
| 596 | comm_handles=[work], |
| 597 | comm_postprocess=[post], |
| 598 | ) |
| 599 | |
| 600 | # accessing non-tensor attributes should not trigger wait. |
| 601 | assert it._comm_handles is not None |
| 602 | assert work.wait_calls == 0 |
| 603 | assert post_calls["n"] == 0 |
| 604 | |
| 605 | # first access of `.tensors` triggers wait + postprocess. |
| 606 | _ = it.tensors |
| 607 | assert work.wait_calls == 1 |
| 608 | assert post_calls["n"] == 1 |
| 609 | |
| 610 | # subsequent access should not re-wait. |
| 611 | _ = it.tensors |
| 612 | assert work.wait_calls == 1 |
| 613 | assert post_calls["n"] == 1 |
| 614 | |
| 615 | |
| 616 | @ray.remote(num_gpus=1, max_calls=1) |
| 617 | def send_recv_test_worker( |
| 618 | monkeypatch: pytest.MonkeyPatch, |
| 619 | tp_size: int, |
| 620 | pp_size: int, |
| 621 | rank: int, |
| 622 | distributed_init_port: str, |
| 623 | ): |
| 624 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 625 | device = torch.device(f"cuda:{rank}") |
| 626 | torch.accelerator.set_device_index(device) |
| 627 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 628 | |
| 629 | size = 64 |
| 630 | test_tensor = torch.arange(64, dtype=torch.float32, device="cuda") |
| 631 | |
| 632 | if not get_pp_group().is_first_rank: |
| 633 | recv_tensor = get_pp_group().recv(size, dtype=torch.float32) |
| 634 | |
| 635 | if not get_pp_group().is_last_rank: |
| 636 | get_pp_group().send(test_tensor) |
| 637 | |
| 638 | if not get_pp_group().is_first_rank: |
| 639 | torch.testing.assert_close(test_tensor, recv_tensor) |
| 640 | |
| 641 | |
| 642 | @multi_gpu_test(num_gpus=2) |
| 643 | @pytest.mark.parametrize("tp_size", [2]) |
| 644 | @pytest.mark.parametrize( |
| 645 | "test_target", |
| 646 | [all_reduce_test_worker, all_gather_test_worker, broadcast_tensor_dict_test_worker], |
| 647 | ) |
| 648 | def test_multi_process_tensor_parallel( |
| 649 | monkeypatch: pytest.MonkeyPatch, |
| 650 | tp_size: int, |
| 651 | test_target: Callable[..., Any], |
| 652 | ): |
| 653 | multi_process_parallel(monkeypatch, tp_size, 1, test_target) |
| 654 | |
| 655 | |
| 656 | @multi_gpu_test(num_gpus=2) |
| 657 | @pytest.mark.parametrize("pp_size", [2]) |
| 658 | @pytest.mark.parametrize( |
| 659 | "test_target", [send_recv_test_worker, send_recv_tensor_dict_test_worker] |
| 660 | ) |
| 661 | def test_multi_process_pipeline_parallel( |
| 662 | monkeypatch: pytest.MonkeyPatch, |
| 663 | pp_size: int, |
| 664 | test_target: Callable[..., Any], |
| 665 | ): |
| 666 | multi_process_parallel(monkeypatch, 1, pp_size, test_target) |
| 667 | |
| 668 | |
| 669 | @multi_gpu_test(num_gpus=4) |
| 670 | @pytest.mark.parametrize("tp_size", [2]) |
| 671 | @pytest.mark.parametrize("pp_size", [2]) |
| 672 | @pytest.mark.parametrize( |
| 673 | "test_target", |
| 674 | [ |
| 675 | send_recv_test_worker, |
| 676 | send_recv_tensor_dict_test_worker, |
| 677 | all_reduce_test_worker, |
| 678 | all_gather_test_worker, |
| 679 | broadcast_tensor_dict_test_worker, |
| 680 | ], |
| 681 | ) |
| 682 | def test_multi_process_tensor_parallel_pipeline_parallel( |
| 683 | tp_size: int, |
| 684 | pp_size: int, |
| 685 | test_target: Callable[..., Any], |
| 686 | monkeypatch: pytest.MonkeyPatch, |
| 687 | ): |
| 688 | multi_process_parallel(monkeypatch, tp_size, pp_size, test_target) |