| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import random |
| 5 | |
| 6 | import pytest |
| 7 | import ray |
| 8 | import torch |
| 9 | import torch.distributed as dist |
| 10 | |
| 11 | from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa |
| 12 | from vllm.distributed.device_communicators import custom_all_reduce as car |
| 13 | from vllm.distributed.parallel_state import get_tp_group, graph_capture |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | from ..utils import ( |
| 17 | ensure_model_parallel_initialized, |
| 18 | init_test_distributed_environment, |
| 19 | multi_process_parallel, |
| 20 | ) |
| 21 | |
| 22 | random.seed(42) |
| 23 | test_sizes = [random.randint(1024, 2048 * 1024) for _ in range(8)] |
| 24 | for i, v in enumerate(test_sizes): |
| 25 | test_sizes[i] -= v % 8 |
| 26 | |
| 27 | |
| 28 | def _bf16_ulps(a: torch.Tensor, b: torch.Tensor) -> int: |
| 29 | """Largest distance between two BF16 tensors in units in the last place.""" |
| 30 | |
| 31 | def ordered(t: torch.Tensor) -> torch.Tensor: |
| 32 | bits = t.view(torch.int16).int() & 0xFFFF |
| 33 | return torch.where(bits >= 0x8000, 0x8000 - bits, bits) |
| 34 | |
| 35 | return int((ordered(a) - ordered(b)).abs().max()) |
| 36 | |
| 37 | |
| 38 | def _unfused_all_reduce_mhc(peers, residual, post, comb, pre, weight, eps): |
| 39 | """All-reduce, mHC post, collapse and RMSNorm in the fused kernel's FP32 |
| 40 | order, rounding to BF16 where the unfused path returns BF16.""" |
| 41 | reduced = peers[0].float() |
| 42 | for peer in peers[1:]: |
| 43 | reduced = reduced + peer.float() |
| 44 | reduced = reduced.bfloat16().float() |
| 45 | output = torch.empty_like(residual) |
| 46 | collapse = torch.zeros_like(reduced) |
| 47 | for target in range(4): |
| 48 | mixed = reduced * post[:, target : target + 1] |
| 49 | for source in range(4): |
| 50 | mixed = torch.addcmul( |
| 51 | mixed, residual[:, source].float(), comb[:, source, target, None] |
| 52 | ) |
| 53 | output[:, target] = mixed.bfloat16() |
| 54 | collapse = torch.addcmul( |
| 55 | collapse, output[:, target].float(), pre[:, target : target + 1] |
| 56 | ) |
| 57 | prenorm = collapse.bfloat16().float() |
| 58 | inv_rms = torch.rsqrt(prenorm.square().mean(-1, keepdim=True) + eps) |
| 59 | return output, (prenorm * inv_rms * weight.float()).bfloat16() |
| 60 | |
| 61 | |
| 62 | @ray.remote(num_gpus=1, max_calls=1) |
| 63 | def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port): |
| 64 | from vllm.models.deepseek_v41.nvidia.ops.cute_dsl import AllReduceMHC |
| 65 | |
| 66 | with monkeypatch.context() as m: |
| 67 | m.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 68 | device = torch.device(f"cuda:{rank}") |
| 69 | torch.accelerator.set_device_index(device) |
| 70 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 71 | ensure_model_parallel_initialized(tp_size, pp_size) |
| 72 | op = AllReduceMHC( |
| 73 | hidden_size=5120, hc_mult=4, max_num_tokens=16, top_k=6, device=device |
| 74 | ) |
| 75 | |
| 76 | def check_eager_and_replayed(fused, check, halved, doubled): |
| 77 | check(fused()) |
| 78 | graph = torch.cuda.CUDAGraph() |
| 79 | with torch.cuda.graph(graph): |
| 80 | for _ in range(5): |
| 81 | fused() |
| 82 | captured = fused() |
| 83 | for _ in range(20): |
| 84 | graph.replay() |
| 85 | check(captured) |
| 86 | # Replays must pick up in-place input changes. |
| 87 | halved.mul_(0.5) |
| 88 | doubled.mul_(2) |
| 89 | graph.replay() |
| 90 | check(captured) |
| 91 | |
| 92 | def mhc_inputs(n): |
| 93 | torch.manual_seed(123) |
| 94 | residual = torch.randn(n, 4, 5120, device=device, dtype=torch.bfloat16) |
| 95 | post = torch.rand(n, 4, device=device) |
| 96 | comb = torch.randn(n, 4, 4, device=device) * 0.1 |
| 97 | pre = torch.rand(n, 4, device=device) |
| 98 | weight = torch.randn(5120, device=device, dtype=torch.bfloat16) |
| 99 | return residual, post, comb, pre, weight |
| 100 | |
| 101 | def run(n): |
| 102 | torch.manual_seed(42 + rank) |
| 103 | x = torch.randn(n, 5120, device=device, dtype=torch.bfloat16) |
| 104 | # Packed +0/-0 pairs collide with the Lamport sentinel. |
| 105 | x[:, :16] = 0 |
| 106 | x[:, 9:16:2] = -0.0 |
| 107 | residual, post, comb, pre, weight = mhc_inputs(n) |
| 108 | |
| 109 | def fused(): |
| 110 | return op(x, residual, post, comb, pre, weight, 1e-6) |
| 111 | |
| 112 | def check(outputs): |
| 113 | peers = get_tp_group().all_gather(x, dim=0).view(tp_size, n, 5120) |
| 114 | output, normalized = _unfused_all_reduce_mhc( |
| 115 | peers, residual, post, comb, pre, weight, 1e-6 |
| 116 | ) |
| 117 | # The mixed hc streams match bit for bit. The RMSNorm sums the |
| 118 | # squares in another order and uses an approximate rsqrt. |
| 119 | assert torch.equal(outputs[0], output) |
| 120 | assert _bf16_ulps(outputs[1], normalized) <= 1 |
| 121 | |
| 122 | check_eager_and_replayed(fused, check, x, residual) |
| 123 | |
| 124 | def run_finalize(n): |
| 125 | torch.manual_seed(7 + rank) |
| 126 | # A padded permuted GEMM2 buffer, like the MoE's. |
| 127 | rows = n * 6 + 5 |
| 128 | gemm2 = torch.randn(rows, 5120, device=device, dtype=torch.bfloat16) |
| 129 | permuted = torch.randperm(rows, device=device)[: n * 6].view(n, 6).int() |
| 130 | # A route to an expert this rank does not hold. |
| 131 | permuted[0, -1] = -1 |
| 132 | weights = torch.rand(n, 6, device=device) |
| 133 | shared = torch.randn(n, 5120, device=device, dtype=torch.bfloat16) |
| 134 | residual, post, comb, pre, weight = mhc_inputs(n) |
| 135 | |
| 136 | def fused(): |
| 137 | return op.finalize( |
| 138 | gemm2, |
| 139 | weights, |
| 140 | permuted, |
| 141 | shared, |
| 142 | residual, |
| 143 | post, |
| 144 | comb, |
| 145 | pre, |
| 146 | weight, |
| 147 | 1e-6, |
| 148 | ) |
| 149 | |
| 150 | def check(outputs): |
| 151 | # One FP32 FMA per route in route order, the shared add, one |
| 152 | # rounding, then the plain path: only the finalize differs. |
| 153 | acc = torch.zeros(n, 5120, device=device) |
| 154 | for k in range(6): |
| 155 | valid = (permuted[:, k] >= 0).unsqueeze(-1) |
| 156 | rows_k = gemm2[permuted[:, k].clamp_min(0).long()].float() |
| 157 | torch.addcmul( |
| 158 | acc, |
| 159 | torch.where(valid, rows_k, 0.0), |
| 160 | torch.where(valid, weights[:, k : k + 1], 0.0), |
| 161 | out=acc, |
| 162 | ) |
| 163 | x = (acc + shared.float()).bfloat16() |
| 164 | expected = op(x, residual, post, comb, pre, weight, 1e-6) |
| 165 | assert torch.equal(outputs[0], expected[0]) |
| 166 | assert torch.equal(outputs[1], expected[1]) |
| 167 | |
| 168 | check_eager_and_replayed(fused, check, gemm2, shared) |
| 169 | |
| 170 | # Changing shapes cover shrinking and growing batches in one mailbox. |
| 171 | for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1): |
| 172 | run(n) |
| 173 | run_finalize(n) |
| 174 | |
| 175 | |
| 176 | @pytest.mark.skipif( |
| 177 | not current_platform.is_device_capability_family(100), reason="Requires SM100" |
| 178 | ) |
| 179 | def test_all_reduce_mhc_matches_unfused_path(monkeypatch): |
| 180 | if torch.accelerator.device_count() < 4: |
| 181 | pytest.skip("Requires four GPUs with NVLink multicast") |
| 182 | multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc) |
| 183 | |
| 184 | |
| 185 | @pytest.mark.parametrize( |
| 186 | ("dtype", "expected"), |
| 187 | [ |
| 188 | (torch.float32, True), |
| 189 | (torch.float16, True), |
| 190 | (torch.bfloat16, True), |
| 191 | (torch.int8, False), |
| 192 | (torch.float8_e4m3fn, False), |
| 193 | ], |
| 194 | ) |
| 195 | def test_custom_allreduce_filters_dtype( |
| 196 | dtype: torch.dtype, |
| 197 | expected: bool, |
| 198 | ) -> None: |
| 199 | communicator = car.CustomAllreduce.__new__(car.CustomAllreduce) |
| 200 | communicator.disabled = False |
| 201 | communicator._ptr = 0 |
| 202 | communicator.world_size = 2 |
| 203 | communicator.max_size = 1024 |
| 204 | |
| 205 | assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) is expected |
| 206 | |
| 207 | |
| 208 | @pytest.mark.parametrize("batch_invariant", [False, True]) |
| 209 | def test_custom_allreduce_size_gate_ignored_under_batch_invariance( |
| 210 | batch_invariant: bool, |
| 211 | ) -> None: |
| 212 | """Batch invariance must not switch backends based on the input size.""" |
| 213 | communicator = car.CustomAllreduce.__new__(car.CustomAllreduce) |
| 214 | communicator.disabled = False |
| 215 | communicator.world_size = 2 |
| 216 | communicator.max_size = 1024 |
| 217 | communicator.batch_invariant = batch_invariant |
| 218 | communicator._ptr = 0 |
| 219 | |
| 220 | oversized = torch.empty(1024, dtype=torch.float16) |
| 221 | assert communicator.should_custom_ar(oversized) is batch_invariant |
| 222 | |
| 223 | |
| 224 | @pytest.mark.parametrize("batch_invariant", [False, True]) |
| 225 | def test_custom_reduce_scatter_disabled_under_batch_invariance( |
| 226 | monkeypatch, batch_invariant: bool |
| 227 | ) -> None: |
| 228 | """Reduce-scatter stays on one backend under batch invariance.""" |
| 229 | monkeypatch.setattr(car.current_platform, "is_cuda", lambda: True) |
| 230 | communicator = car.CustomAllreduce.__new__(car.CustomAllreduce) |
| 231 | communicator.disabled = False |
| 232 | communicator.world_size = 2 |
| 233 | communicator.fully_connected = True |
| 234 | communicator.mnnvl_only = False |
| 235 | communicator.mnnvl_multicast_ptr = 0 |
| 236 | communicator.max_reduce_scatter_size = 1024 |
| 237 | communicator.max_mnnvl_reduce_scatter_size = 1024 |
| 238 | communicator.batch_invariant = batch_invariant |
| 239 | |
| 240 | in_range = torch.empty(64, dtype=torch.float16) |
| 241 | assert communicator.should_custom_reduce_scatter(in_range) is not batch_invariant |
| 242 | |
| 243 | |
| 244 | @pytest.mark.parametrize( |
| 245 | ("major", "local_multicast", "expected"), |
| 246 | [ |
| 247 | (8, True, False), |
| 248 | (9, True, False), |
| 249 | (10, False, False), |
| 250 | (10, True, True), |
| 251 | ], |
| 252 | ) |
| 253 | def test_cross_node_mnnvl_gate_checks_generation_and_multicast( |
| 254 | monkeypatch, |
| 255 | major, |
| 256 | local_multicast, |
| 257 | expected, |
| 258 | ): |
| 259 | def has_device_capability(capability, device_id): |
| 260 | assert capability == 100 |
| 261 | assert device_id == 3 |
| 262 | return major >= 10 |
| 263 | |
| 264 | monkeypatch.setattr( |
| 265 | car.current_platform, |
| 266 | "has_device_capability", |
| 267 | has_device_capability, |
| 268 | ) |
| 269 | monkeypatch.setattr( |
| 270 | car, |
| 271 | "_has_local_multicast_support", |
| 272 | lambda _device: local_multicast, |
| 273 | ) |
| 274 | monkeypatch.setattr(car.dist, "all_reduce", lambda *_args, **_kwargs: None) |
| 275 | |
| 276 | assert car._group_can_attempt_mnnvl(object(), torch.device("cuda:3")) is expected |
| 277 | |
| 278 | |
| 279 | def test_cross_node_mnnvl_gate_requires_support_on_every_rank(monkeypatch): |
| 280 | monkeypatch.setattr( |
| 281 | car.current_platform, |
| 282 | "has_device_capability", |
| 283 | lambda *_args: True, |
| 284 | ) |
| 285 | monkeypatch.setattr( |
| 286 | car, |
| 287 | "_has_local_multicast_support", |
| 288 | lambda _device: True, |
| 289 | ) |
| 290 | |
| 291 | def report_unsupported_peer(support, **_kwargs): |
| 292 | support.zero_() |
| 293 | |
| 294 | monkeypatch.setattr(car.dist, "all_reduce", report_unsupported_peer) |
| 295 | |
| 296 | assert not car._group_can_attempt_mnnvl(object(), torch.device("cuda:0")) |
| 297 | |
| 298 | |
| 299 | def test_local_multicast_support_rejects_non_cuda(monkeypatch): |
| 300 | monkeypatch.setattr(car.current_platform, "is_cuda", lambda: False) |
| 301 | |
| 302 | assert not car._has_local_multicast_support(torch.device("cuda:0")) |
| 303 | |
| 304 | |
| 305 | @pytest.mark.parametrize( |
| 306 | ("world_size", "device_capability", "local_multicast", "expected"), |
| 307 | [ |
| 308 | (2, (10, 0), True, True), |
| 309 | (4, (10, 3), True, True), |
| 310 | (8, (10, 0), True, True), |
| 311 | (8, (10, 3), True, True), |
| 312 | (6, (10, 3), True, False), |
| 313 | (8, (10, 1), True, False), |
| 314 | (8, (9, 0), True, False), |
| 315 | (8, (10, 3), False, False), |
| 316 | ], |
| 317 | ) |
| 318 | def test_mnnvl_multimem_reduce_scatter_platform_gate( |
| 319 | monkeypatch, |
| 320 | world_size, |
| 321 | device_capability, |
| 322 | local_multicast, |
| 323 | expected, |
| 324 | ): |
| 325 | def is_device_capability(capability, device_id): |
| 326 | assert capability in ((10, 0), (10, 3)) |
| 327 | assert device_id == 3 |
| 328 | return device_capability == capability |
| 329 | |
| 330 | monkeypatch.setattr( |
| 331 | car.current_platform, |
| 332 | "is_device_capability", |
| 333 | is_device_capability, |
| 334 | ) |
| 335 | monkeypatch.setattr( |
| 336 | car, |
| 337 | "_has_local_multicast_support", |
| 338 | lambda _device: local_multicast, |
| 339 | ) |
| 340 | |
| 341 | supported = car._supports_mnnvl_multimem_reduce_scatter( |
| 342 | torch.device("cuda:3"), world_size |
| 343 | ) |
| 344 | assert supported is expected |
| 345 | |
| 346 | |
| 347 | @pytest.mark.parametrize( |
| 348 | ( |
| 349 | "message_bytes", |
| 350 | "multimem_ptr", |
| 351 | "multimem_initialized", |
| 352 | "batch_invariant", |
| 353 | "expected", |
| 354 | ), |
| 355 | [ |
| 356 | (16 * 1024 * 1024, 1, True, False, "mnnvl_lamport"), |
| 357 | (16 * 1024 * 1024 + 128, 1, True, False, "mnnvl_multimem"), |
| 358 | (64 * 1024 * 1024, 1, True, False, "mnnvl_multimem"), |
| 359 | (64 * 1024 * 1024 + 128, 1, True, False, None), |
| 360 | (32 * 1024 * 1024, 0, True, False, None), |
| 361 | (32 * 1024 * 1024, 0, False, False, "mnnvl_multimem"), |
| 362 | (8 * 1024 * 1024, 1, True, True, "mnnvl_lamport"), |
| 363 | (32 * 1024 * 1024, 1, True, True, None), |
| 364 | ], |
| 365 | ) |
| 366 | @pytest.mark.parametrize("world_size", [2, 4, 8]) |
| 367 | def test_mnnvl_reduce_scatter_backend_gate( |
| 368 | monkeypatch, |
| 369 | world_size, |
| 370 | message_bytes, |
| 371 | multimem_ptr, |
| 372 | multimem_initialized, |
| 373 | batch_invariant, |
| 374 | expected, |
| 375 | ): |
| 376 | monkeypatch.setattr(car.current_platform, "is_cuda", lambda: True) |
| 377 | monkeypatch.setattr(car.envs, "VLLM_BATCH_INVARIANT", batch_invariant) |
| 378 | communicator = car.CustomAllreduce.__new__(car.CustomAllreduce) |
| 379 | communicator.disabled = False |
| 380 | communicator._ptr = 0 |
| 381 | communicator.world_size = world_size |
| 382 | communicator.mnnvl_only = False |
| 383 | communicator.fully_connected = True |
| 384 | communicator.mnnvl_multicast_ptr = 1 |
| 385 | communicator.mnnvl_multimem_rs_supported = True |
| 386 | communicator.mnnvl_multimem_rs_initialized = multimem_initialized |
| 387 | communicator.mnnvl_multimem_rs_multicast_ptr = multimem_ptr |
| 388 | communicator.max_mnnvl_reduce_scatter_size = 16 * 1024 * 1024 |
| 389 | communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024 |
| 390 | communicator.max_reduce_scatter_size = 16 * 1024 * 1024 |
| 391 | inp = torch.empty( |
| 392 | (world_size, message_bytes // torch.bfloat16.itemsize // world_size), |
| 393 | dtype=torch.bfloat16, |
| 394 | ) |
| 395 | |
| 396 | assert inp.nbytes == message_bytes |
| 397 | assert communicator._select_reduce_scatter_backend(inp) == expected |
| 398 | assert communicator.should_custom_reduce_scatter(inp) is (expected is not None) |
| 399 | assert communicator.should_mnnvl_multimem_reduce_scatter(inp) is ( |
| 400 | expected == "mnnvl_multimem" |
| 401 | ) |
| 402 | |
| 403 | |
| 404 | def test_mnnvl_multimem_reduce_scatter_skips_rendezvous_after_peer_alloc_failure( |
| 405 | monkeypatch, |
| 406 | ): |
| 407 | events = [] |
| 408 | |
| 409 | class FakeSymmMem: |
| 410 | @staticmethod |
| 411 | def empty(*_args, **_kwargs): |
| 412 | events.append("empty") |
| 413 | return torch.empty(1, dtype=torch.uint8) |
| 414 | |
| 415 | @staticmethod |
| 416 | def rendezvous(*_args, **_kwargs): |
| 417 | events.append("rendezvous") |
| 418 | return None |
| 419 | |
| 420 | def report_peer_allocation_failure(group_value, **_kwargs): |
| 421 | events.append("all_reduce") |
| 422 | assert group_value.item() == 1 |
| 423 | group_value.zero_() |
| 424 | |
| 425 | monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem) |
| 426 | monkeypatch.setattr(car.ops, "meta_size", lambda: 128) |
| 427 | monkeypatch.setattr(car.dist, "all_reduce", report_peer_allocation_failure) |
| 428 | warnings = [] |
| 429 | monkeypatch.setattr( |
| 430 | car.logger, |
| 431 | "warning_once", |
| 432 | lambda message, *_args, **_kwargs: warnings.append(message), |
| 433 | ) |
| 434 | |
| 435 | communicator = car.CustomAllreduce.__new__(car.CustomAllreduce) |
| 436 | communicator.disabled = True |
| 437 | communicator._ptr = 0 |
| 438 | communicator.group = object() |
| 439 | communicator.device = torch.device("cuda:0") |
| 440 | communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024 |
| 441 | communicator.mnnvl_multimem_rs_supported = True |
| 442 | communicator.mnnvl_multimem_rs_initialized = False |
| 443 | communicator.mnnvl_multimem_rs_buffer = None |
| 444 | communicator.mnnvl_multimem_rs_multicast_ptr = 0 |
| 445 | |
| 446 | communicator._init_mnnvl_multimem_reduce_scatter_buffer() |
| 447 | |
| 448 | assert events == ["empty", "all_reduce"] |
| 449 | assert communicator.mnnvl_multimem_rs_initialized |
| 450 | assert communicator.mnnvl_multimem_rs_buffer is None |
| 451 | assert communicator.mnnvl_multimem_rs_multicast_ptr == 0 |
| 452 | assert warnings == [ |
| 453 | "MNNVL multimem reduce-scatter symmetric-memory allocation " |
| 454 | "failed on at least one rank; falling back to NCCL." |
| 455 | ] |
| 456 | |
| 457 | |
| 458 | def test_mnnvl_multimem_reduce_scatter_warns_on_rendezvous_failure(monkeypatch): |
| 459 | events = [] |
| 460 | |
| 461 | class FakeSymmMem: |
| 462 | @staticmethod |
| 463 | def empty(*_args, **_kwargs): |
| 464 | events.append("empty") |
| 465 | return torch.empty(1, dtype=torch.uint8) |
| 466 | |
| 467 | @staticmethod |
| 468 | def rendezvous(*_args, **_kwargs): |
| 469 | events.append("rendezvous") |
| 470 | raise RuntimeError("rendezvous failed") |
| 471 | |
| 472 | def preserve_local_result(_group_value, **_kwargs): |
| 473 | events.append("all_reduce") |
| 474 | |
| 475 | warnings = [] |
| 476 | monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem) |
| 477 | monkeypatch.setattr(car.ops, "meta_size", lambda: 128) |
| 478 | monkeypatch.setattr(car.dist, "all_reduce", preserve_local_result) |
| 479 | monkeypatch.setattr( |
| 480 | car.logger, |
| 481 | "warning_once", |
| 482 | lambda message, *_args, **_kwargs: warnings.append(message), |
| 483 | ) |
| 484 | |
| 485 | communicator = car.CustomAllreduce.__new__(car.CustomAllreduce) |
| 486 | communicator.disabled = True |
| 487 | communicator._ptr = 0 |
| 488 | communicator.group = type("Group", (), {"group_name": "test"})() |
| 489 | communicator.device = torch.device("cuda:0") |
| 490 | communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024 |
| 491 | communicator.mnnvl_multimem_rs_supported = True |
| 492 | communicator.mnnvl_multimem_rs_initialized = False |
| 493 | communicator.mnnvl_multimem_rs_buffer = None |
| 494 | communicator.mnnvl_multimem_rs_multicast_ptr = 0 |
| 495 | |
| 496 | communicator._init_mnnvl_multimem_reduce_scatter_buffer() |
| 497 | |
| 498 | assert events == ["empty", "all_reduce", "rendezvous", "all_reduce"] |
| 499 | assert communicator.mnnvl_multimem_rs_initialized |
| 500 | assert communicator.mnnvl_multimem_rs_buffer is None |
| 501 | assert communicator.mnnvl_multimem_rs_multicast_ptr == 0 |
| 502 | assert warnings == [ |
| 503 | "MNNVL multimem reduce-scatter symmetric-memory rendezvous " |
| 504 | "failed on at least one rank; falling back to NCCL." |
| 505 | ] |
| 506 | |
| 507 | |
| 508 | def test_mnnvl_multimem_reduce_scatter_initializes_signals(monkeypatch): |
| 509 | events = [] |
| 510 | buffers = [] |
| 511 | |
| 512 | class FakeHandle: |
| 513 | multicast_ptr = 0x3000 |
| 514 | |
| 515 | class FakeSymmMem: |
| 516 | @staticmethod |
| 517 | def empty(size, **_kwargs): |
| 518 | events.append(("empty", size)) |
| 519 | buffer = torch.ones(size, dtype=torch.uint8) |
| 520 | buffers.append(buffer) |
| 521 | return buffer |
| 522 | |
| 523 | @staticmethod |
| 524 | def rendezvous(*_args, **_kwargs): |
| 525 | events.append(("rendezvous", None)) |
| 526 | return FakeHandle() |
| 527 | |
| 528 | def preserve_local_result(*_args, **_kwargs): |
| 529 | events.append(("all_reduce", None)) |
| 530 | |
| 531 | monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem) |
| 532 | monkeypatch.setattr(car.ops, "meta_size", lambda: 128) |
| 533 | monkeypatch.setattr( |
| 534 | car.torch.accelerator, |
| 535 | "synchronize", |
| 536 | lambda: events.append(("synchronize", None)), |
| 537 | ) |
| 538 | monkeypatch.setattr(car.dist, "all_reduce", preserve_local_result) |
| 539 | |
| 540 | communicator = car.CustomAllreduce.__new__(car.CustomAllreduce) |
| 541 | communicator.disabled = True |
| 542 | communicator._ptr = 0 |
| 543 | communicator.group = type("Group", (), {"group_name": "test"})() |
| 544 | communicator.device = torch.device("cpu") |
| 545 | communicator.max_mnnvl_multimem_reduce_scatter_size = 129 |
| 546 | communicator.mnnvl_multimem_rs_supported = True |
| 547 | communicator.mnnvl_multimem_rs_initialized = False |
| 548 | communicator.mnnvl_multimem_rs_buffer = None |
| 549 | communicator.mnnvl_multimem_rs_multicast_ptr = 0 |
| 550 | |
| 551 | communicator._init_mnnvl_multimem_reduce_scatter_buffer() |
| 552 | |
| 553 | assert events == [ |
| 554 | ("empty", 257), |
| 555 | ("all_reduce", None), |
| 556 | ("rendezvous", None), |
| 557 | ("synchronize", None), |
| 558 | ("all_reduce", None), |
| 559 | ] |
| 560 | assert torch.all(buffers[0][:128] == 0) |
| 561 | assert torch.all(buffers[0][128:] == 1) |
| 562 | assert communicator.mnnvl_multimem_rs_buffer_size == 129 |
| 563 | assert communicator.mnnvl_multimem_rs_local_ptr == buffers[0].data_ptr() + 128 |
| 564 | assert communicator.mnnvl_multimem_rs_multicast_ptr == 0x3080 |
| 565 | |
| 566 | |
| 567 | @ray.remote(num_gpus=1, max_calls=1) |
| 568 | def graph_allreduce( |
| 569 | monkeypatch: pytest.MonkeyPatch, |
| 570 | tp_size, |
| 571 | pp_size, |
| 572 | rank, |
| 573 | distributed_init_port, |
| 574 | ): |
| 575 | with monkeypatch.context() as m: |
| 576 | m.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 577 | m.delenv("HIP_VISIBLE_DEVICES", raising=False) |
| 578 | device = torch.device(f"cuda:{rank}") |
| 579 | torch.accelerator.set_device_index(device) |
| 580 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 581 | ensure_model_parallel_initialized(tp_size, pp_size) |
| 582 | group = get_tp_group().device_group |
| 583 | |
| 584 | # A small all_reduce for warmup. |
| 585 | # this is needed because device communicators might be created lazily |
| 586 | # (e.g. NCCL). This will ensure that the communicator is initialized |
| 587 | # before any communication happens, so that this group can be used for |
| 588 | # graph capture immediately. |
| 589 | data = torch.zeros(1) |
| 590 | data = data.to(device=device) |
| 591 | torch.distributed.all_reduce(data, group=group) |
| 592 | torch.accelerator.synchronize() |
| 593 | del data |
| 594 | |
| 595 | # we use the first group to communicate once |
| 596 | # and the second group to communicate twice |
| 597 | # and so on |
| 598 | # this is used to demonstrate that each group can |
| 599 | # communicate independently |
| 600 | num_communication = rank // tp_size + 1 |
| 601 | |
| 602 | for sz in test_sizes: |
| 603 | for dtype in [torch.float32, torch.float16, torch.bfloat16]: |
| 604 | with graph_capture(device=device) as graph_capture_context: |
| 605 | # use integers so result matches NCCL exactly |
| 606 | device_idx = torch.accelerator.current_device_index() |
| 607 | inp1 = torch.randint(1, 16, (sz,), dtype=dtype, device=device_idx) |
| 608 | inp2 = torch.randint(1, 16, (sz,), dtype=dtype, device=device_idx) |
| 609 | |
| 610 | torch.accelerator.synchronize() |
| 611 | graph = torch.cuda.CUDAGraph() |
| 612 | with torch.cuda.graph(graph, stream=graph_capture_context.stream): |
| 613 | for i in range(num_communication): |
| 614 | out1 = tensor_model_parallel_all_reduce(inp1) |
| 615 | # the input buffer is immediately modified to test |
| 616 | # synchronization |
| 617 | dist.all_reduce(inp1, group=group) |
| 618 | out2 = tensor_model_parallel_all_reduce(inp2) |
| 619 | dist.all_reduce(inp2, group=group) |
| 620 | graph.replay() |
| 621 | torch.testing.assert_close(out1, inp1) |
| 622 | torch.testing.assert_close(out2, inp2) |
| 623 | |
| 624 | |
| 625 | @ray.remote(num_gpus=1, max_calls=1) |
| 626 | def eager_allreduce( |
| 627 | monkeypatch: pytest.MonkeyPatch, |
| 628 | tp_size, |
| 629 | pp_size, |
| 630 | rank, |
| 631 | distributed_init_port, |
| 632 | ): |
| 633 | with monkeypatch.context() as m: |
| 634 | m.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 635 | m.delenv("HIP_VISIBLE_DEVICES", raising=False) |
| 636 | device = torch.device(f"cuda:{rank}") |
| 637 | torch.accelerator.set_device_index(device) |
| 638 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 639 | |
| 640 | # we use the first group to communicate once |
| 641 | # and the second group to communicate twice |
| 642 | # and so on |
| 643 | # this is used to demonstrate that each group can |
| 644 | # communicate independently |
| 645 | num_communication = rank // tp_size + 1 |
| 646 | sz = 1024 |
| 647 | fa = get_tp_group().device_communicator.ca_comm |
| 648 | inp = torch.ones(sz, dtype=torch.float32, device=device) |
| 649 | out = inp |
| 650 | for _ in range(num_communication): |
| 651 | out = fa.all_reduce(out, registered=False) |
| 652 | torch.testing.assert_close(out, inp * (tp_size**num_communication)) |
| 653 | |
| 654 | inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device) |
| 655 | out = inp |
| 656 | for _ in range(num_communication): |
| 657 | out = fa.all_reduce(out, registered=False) |
| 658 | torch.testing.assert_close(out, inp * (tp_size**num_communication)) |
| 659 | |
| 660 | |
| 661 | @ray.remote(num_gpus=1, max_calls=1) |
| 662 | def chunked_allreduce( |
| 663 | monkeypatch: pytest.MonkeyPatch, |
| 664 | tp_size, |
| 665 | pp_size, |
| 666 | rank, |
| 667 | distributed_init_port, |
| 668 | ): |
| 669 | """Inputs above max_size are reduced in chunks and match the fixed-order |
| 670 | reference bitwise, both eagerly and inside a CUDA graph.""" |
| 671 | with monkeypatch.context() as m: |
| 672 | m.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 673 | m.delenv("HIP_VISIBLE_DEVICES", raising=False) |
| 674 | m.setenv("VLLM_CUSTOM_ALLREDUCE_ALGO", "1stage") |
| 675 | device = torch.device(f"cuda:{rank}") |
| 676 | torch.accelerator.set_device_index(device) |
| 677 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 678 | ensure_model_parallel_initialized(tp_size, pp_size) |
| 679 | group = get_tp_group().device_group |
| 680 | fa = get_tp_group().device_communicator.ca_comm |
| 681 | |
| 682 | # Two full chunks plus a partial one. |
| 683 | dtype = torch.bfloat16 |
| 684 | chunk_numel = fa.max_size // dtype.itemsize |
| 685 | inp = torch.randn(2 * chunk_numel + 4096, device=device).to(dtype) |
| 686 | gathered = [torch.empty_like(inp) for _ in range(tp_size)] |
| 687 | dist.all_gather(gathered, inp, group=group) |
| 688 | ref = gathered[0].float() |
| 689 | for peer in gathered[1:]: |
| 690 | ref = ref + peer.float() |
| 691 | ref = ref.to(torch.bfloat16) |
| 692 | |
| 693 | out = fa.all_reduce(inp, registered=False) |
| 694 | assert torch.equal(out, ref) |
| 695 | |
| 696 | # Weak-contiguous but not C-contiguous: a transposed matrix. |
| 697 | inp_t, ref_t = inp.view(-1, 4096).t(), ref.view(-1, 4096).t() |
| 698 | assert torch.equal(fa.all_reduce(inp_t, registered=False), ref_t) |
| 699 | |
| 700 | with fa.capture(): |
| 701 | graph = torch.cuda.CUDAGraph() |
| 702 | with torch.cuda.graph(graph): |
| 703 | out = fa.all_reduce(inp, registered=True) |
| 704 | graph.replay() |
| 705 | assert torch.equal(out, ref) |
| 706 | |
| 707 | |
| 708 | @pytest.mark.parametrize("tp_size", [2, 4]) |
| 709 | def test_custom_allreduce_chunked(monkeypatch: pytest.MonkeyPatch, tp_size): |
| 710 | if tp_size > torch.accelerator.device_count(): |
| 711 | pytest.skip("Not enough GPUs to run the test.") |
| 712 | multi_process_parallel(monkeypatch, tp_size, 1, chunked_allreduce) |
| 713 | |
| 714 | |
| 715 | @pytest.mark.parametrize("tp_size", [2]) |
| 716 | @pytest.mark.parametrize("pipeline_parallel_size", [1, 2]) |
| 717 | @pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) |
| 718 | def test_custom_allreduce( |
| 719 | monkeypatch: pytest.MonkeyPatch, |
| 720 | tp_size, |
| 721 | pipeline_parallel_size, |
| 722 | test_target, |
| 723 | ): |
| 724 | world_size = tp_size * pipeline_parallel_size |
| 725 | if world_size > torch.accelerator.device_count(): |
| 726 | pytest.skip("Not enough GPUs to run the test.") |
| 727 | multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) |