diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index edddb6ec84..b6769e4d51 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -9,6 +9,7 @@ import torch import torch.distributed as dist from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa +from vllm.distributed.device_communicators import custom_all_reduce as car from vllm.distributed.parallel_state import get_tp_group, graph_capture from ..utils import ( @@ -23,6 +24,67 @@ for i, v in enumerate(test_sizes): test_sizes[i] -= v % 8 +@pytest.mark.parametrize( + ("major", "local_multicast", "expected"), + [ + (8, True, False), + (9, True, False), + (10, False, False), + (10, True, True), + ], +) +def test_cross_node_mnnvl_gate_checks_generation_and_multicast( + monkeypatch, + major, + local_multicast, + expected, +): + def has_device_capability(capability, device_id): + assert capability == 100 + assert device_id == 3 + return major >= 10 + + monkeypatch.setattr( + car.current_platform, + "has_device_capability", + has_device_capability, + ) + monkeypatch.setattr( + car, + "_has_local_multicast_support", + lambda _device: local_multicast, + ) + monkeypatch.setattr(car.dist, "all_reduce", lambda *_args, **_kwargs: None) + + assert car._group_can_attempt_mnnvl(object(), torch.device("cuda:3")) is expected + + +def test_cross_node_mnnvl_gate_requires_support_on_every_rank(monkeypatch): + monkeypatch.setattr( + car.current_platform, + "has_device_capability", + lambda *_args: True, + ) + monkeypatch.setattr( + car, + "_has_local_multicast_support", + lambda _device: True, + ) + + def report_unsupported_peer(support, **_kwargs): + support.zero_() + + monkeypatch.setattr(car.dist, "all_reduce", report_unsupported_peer) + + assert not car._group_can_attempt_mnnvl(object(), torch.device("cuda:0")) + + +def test_local_multicast_support_rejects_non_cuda(monkeypatch): + monkeypatch.setattr(car.current_platform, "is_cuda", lambda: False) + + assert not car._has_local_multicast_support(torch.device("cuda:0")) + + @ray.remote(num_gpus=1, max_calls=1) def graph_allreduce( monkeypatch: pytest.MonkeyPatch, diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 86180a51a7..d1a841e259 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -33,6 +33,57 @@ except ImportError: logger = init_logger(__name__) +def _has_local_multicast_support(device: torch.device) -> bool: + """Return whether this CUDA device can allocate multicast memory.""" + if torch_symm_mem is None or not current_platform.is_cuda(): + return False + if device.index is None: + return False + + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + return _SymmetricMemory.has_multicast_support( + DeviceType.CUDA, + device.index, + ) + except Exception as error: + logger.debug("MNNVL capability probe failed: %s", error) + return False + + +def _group_can_attempt_mnnvl( + group: ProcessGroup, + device: torch.device, +) -> bool: + """Return whether every rank can enter the cross-node MNNVL path. + + MNNVL is available only on Blackwell-class GPUs. Local multicast support + is necessary but does not establish that the process group spans an MNNVL + domain; the symmetric-memory rendezvous below performs that group-level + check. The CPU all-reduce keeps every rank on the same control-flow path + when a heterogeneous or partially configured group is encountered. + """ + device_index = device.index + local_support = ( + device_index is not None + and current_platform.has_device_capability(100, device_index) + and _has_local_multicast_support(device) + ) + group_support = torch.tensor( + int(local_support), + dtype=torch.int32, + device="cpu", + ) + dist.all_reduce( + group_support, + op=dist.ReduceOp.MIN, + group=group, + ) + return bool(group_support.item()) + + def _can_p2p(rank: int, world_size: int) -> bool: for i in range(world_size): if i == rank: @@ -146,6 +197,12 @@ class CustomAllreduce: # now `device` is a `torch.device` object assert isinstance(device, torch.device) self.device = device + if not same_node and not _group_can_attempt_mnnvl(group, device): + logger.warning( + "Custom collectives are disabled because this multi-node " + "group does not have MNNVL-capable GPUs on every rank." + ) + return device_capability = current_platform.get_device_capability() if ( current_platform.is_cuda()