diff --git a/tensorrt_llm/_ipc_utils.py b/tensorrt_llm/_ipc_utils.py index 3178597b5087..14244b79c7c2 100644 --- a/tensorrt_llm/_ipc_utils.py +++ b/tensorrt_llm/_ipc_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -16,7 +16,7 @@ import ctypes import struct import sys -from typing import List, Tuple +from typing import List, Optional, Tuple try: from cuda.bindings import driver as cuda @@ -34,10 +34,10 @@ def _raise_if_error(error: cudart.cudaError_t | cuda.CUresult): if isinstance(error, cudart.cudaError_t): if error != cudart.cudaError_t.cudaSuccess: - raise RuntimeError(f"CUDA Runtime API error: {repr(error)}") + raise RuntimeError(f"CUDA Runtime API error: {error!r}") if isinstance(error, cuda.CUresult): if error != cuda.CUresult.CUDA_SUCCESS: - raise RuntimeError(f"CUDA Driver API error: {repr(error)}") + raise RuntimeError(f"CUDA Driver API error: {error!r}") def _ipc_mem_handle_to_bytes(handle) -> bytes: @@ -104,11 +104,23 @@ class IpcMemory: def __init__(self, mapping: Mapping, size: int, open_ipc: bool = True): self.mapping = mapping self.open_ipc = open_ipc and mapping.tp_size <= mapping.gpus_per_node + # Set only when IPC was expected to work (P2P reported and the TP group fits in + # one node) but the handle exchange failed anyway. Buffers that never requested + # IPC in the first place -- inter-node TP for instance -- keep `open_ipc` False + # with this flag unset, so that callers can tell the two cases apart. + self.ipc_failed = False self.peer_ptrs = [0] * mapping.tp_size self.local_ptr = 0 if self.open_ipc: - self.peer_ptrs, self.local_ptr = IpcMemory.open_ipc_memory(self.mapping, size, True) + ipc_memory = IpcMemory.open_ipc_memory(self.mapping, size, True) + if ipc_memory is None: + # CUDA IPC is unusable on this system. Keep the pointers null so that + # callers fall back to the non-IPC (NCCL) path instead of failing hard. + self.open_ipc = False + self.ipc_failed = True + else: + self.peer_ptrs, self.local_ptr = ipc_memory def __del__(self): if not sys.is_finalizing() and self.open_ipc: @@ -124,10 +136,16 @@ def serialize(self) -> List[int]: @staticmethod def open_ipc_memory( mapping: Mapping, size: int, set_to_zero: bool = False - ) -> Tuple[List[int], int]: + ) -> Optional[Tuple[List[int], int]]: """Allocates a buffer with the given *size* on each GPU. Then, enables IPC communication between TP groups. Returns a list of buffer pointers, buffers[i] is a handle to the corresponding buffer residing on GPU #i. Call close_ipc_handle with the *buffer*. + + Returns None when CUDA IPC is not usable. `cudaDeviceCanAccessPeer` only reports that the + GPUs can address each other's memory, it does not guarantee that IPC handles can be + exported/imported: that additionally fails on GPUs without CUDA IPC support and when the + ranks do not share an IPC namespace. The outcome is agreed upon by the whole TP group, so + that either every rank gets IPC buffers or none of them does. """ def align_size(size, alignment): @@ -146,21 +164,62 @@ def align_size(size, alignment): _raise_if_error(error) if set_to_zero: _raise_if_error(cudart.cudaMemset(local_ptr, 0, aligned_size)[0]) + + def disable_ipc(reason: str, opened_ptrs: List[int]) -> None: + # Every rank of the group takes this path, including the ones whose own + # handles were fine, so log unconditionally: otherwise a rank that is + # falling back because of a peer would degrade silently. + logger.warning_once( + f"CUDA IPC is not usable on this system: {reason} " + "Custom all-reduce kernels are disabled, falling back to NCCL.", + key="cuda-ipc-unavailable", + ) + for ptr in opened_ptrs: + _raise_if_error(cudart.cudaIpcCloseMemHandle(ptr)[0]) + _raise_if_error(cudart.cudaFree(local_ptr)[0]) + + # A rank that cannot export its handle contributes None instead of an + # uninitialized `cudaIpcMemHandle_t`, so that no rank ever hands garbage to + # `cudaIpcOpenMemHandle`. This single collective carries both the payload and + # the agreement on whether exporting worked everywhere. error, local_handle = cudart.cudaIpcGetMemHandle(local_ptr) - _raise_if_error(error) - handles_reserved = dist.tp_allgather(_ipc_mem_handle_to_bytes(local_handle)) - handles = [_ipc_mem_handle_from_bytes(reserved) for reserved in handles_reserved] + get_error = None if error == cudart.cudaError_t.cudaSuccess else error + handles_reserved = dist.tp_allgather( + _ipc_mem_handle_to_bytes(local_handle) if get_error is None else None + ) + + if any(reserved is None for reserved in handles_reserved): + disable_ipc( + f"cudaIpcGetMemHandle failed with {get_error!r}." + if get_error is not None + else "cudaIpcGetMemHandle failed on a peer rank of the TP group.", + [], + ) + return None peer_ptrs = [] - for node, handle in enumerate(handles): + opened_ptrs = [] + open_error = None + for node, reserved in enumerate(handles_reserved): if node == mapping.tp_rank: peer_ptrs.append(local_ptr) - else: - error, ptr = cudart.cudaIpcOpenMemHandle( - handle, cudart.cudaIpcMemLazyEnablePeerAccess - ) - _raise_if_error(error) - peer_ptrs.append(ptr) + continue + handle = _ipc_mem_handle_from_bytes(reserved) + error, ptr = cudart.cudaIpcOpenMemHandle(handle, cudart.cudaIpcMemLazyEnablePeerAccess) + if error != cudart.cudaError_t.cudaSuccess: + open_error = error + break + peer_ptrs.append(ptr) + opened_ptrs.append(ptr) + + if not all(dist.tp_allgather(open_error is None)): + disable_ipc( + f"cudaIpcOpenMemHandle failed with {open_error!r}." + if open_error is not None + else "cudaIpcOpenMemHandle failed on a peer rank of the TP group.", + opened_ptrs, + ) + return None return peer_ptrs, local_ptr diff --git a/tensorrt_llm/_torch/distributed/allreduce_helper.py b/tensorrt_llm/_torch/distributed/allreduce_helper.py index ef555be727f6..25f3b7afaeda 100644 --- a/tensorrt_llm/_torch/distributed/allreduce_helper.py +++ b/tensorrt_llm/_torch/distributed/allreduce_helper.py @@ -129,7 +129,9 @@ def allocate_allreduce_fusion_workspace( ipc_barriers = IpcMemory(mapping, 256 * mapping.tp_size, is_p2p_supported) lamport_buffers_size = size * mapping.tp_size lamport_buffers = IpcMemory(mapping, 3 * lamport_buffers_size, is_p2p_supported) - if is_p2p_supported: + # `lamport_buffers.open_ipc` is False when IPC turned out to be unusable even though + # `can_access_peer` reported P2P support, in which case `local_ptr` is null. + if lamport_buffers.open_ipc: lamport_initialize( lamport_buffers.local_ptr, 3 * lamport_buffers_size, diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index b3ea554da186..d1958645a508 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -22,6 +22,7 @@ import torch from torch import nn +from tensorrt_llm._ipc_utils import IpcMemory from tensorrt_llm._mnnvl_utils import HelixCpMnnvlMemory, MnnvlMemory from tensorrt_llm._torch.distributed.allreduce_helper import \ CustomAllReduceHelper @@ -99,7 +100,18 @@ class _MnnvlWorkspace(TypedDict): mpi_comm: Optional[_MpiCommProtocol] -def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: +def get_allreduce_workspace(mapping: Mapping) -> Tuple[torch.LongTensor, bool]: + """Returns the all-reduce workspace and whether CUDA IPC failed while building it. + + cudaDeviceCanAccessPeer can report P2P support on systems where CUDA IPC handles + still cannot be imported. The workspace pointers are null in that case and only + NCCL can run, which the second element of the returned tuple signals. + + It stays False when the workspace holds no IPC buffers by design, i.e. when P2P + was never available in the first place (inter-node TP for instance). Those + configurations behave as they always did and must keep their strategy, MNNVL in + particular. + """ if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): setattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}', {}) @@ -111,8 +123,26 @@ def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: CustomAllReduceHelper.max_workspace_size_auto( mapping.tp_size, support_deterministic=False), ) - allreduce_workspaces[mapping] = (ipc_buffers, workspace) - return allreduce_workspaces[mapping][1] + ipc_failed = any(buffer.ipc_failed for buffer in ipc_buffers + if isinstance(buffer, IpcMemory)) + allreduce_workspaces[mapping] = (ipc_buffers, workspace, ipc_failed) + _, workspace, ipc_failed = allreduce_workspaces[mapping] + return workspace, ipc_failed + + +def _require_ipc_workspace(mapping: Mapping, op_name: str) -> torch.LongTensor: + """Workspace for fused ops that reinterpret it as `void**` and have no NCCL path. + + AllReduce degrades to NCCL when CUDA IPC turns out to be unusable, but these + kernels cannot: they would dereference the null peer pointers. Fail here with an + actionable message rather than in the kernel with an illegal memory access. + """ + workspace, ipc_failed = get_allreduce_workspace(mapping) + if ipc_failed: + raise RuntimeError( + f"{op_name} requires CUDA IPC, which is unavailable on this system, " + "and it has no NCCL fallback.") + return workspace def allocate_low_presicion_allreduce_workspace(mapping: Mapping) -> None: @@ -933,12 +963,23 @@ def __init__(self, # Note: SYMM_MEM now also needs workspace for fallback scenarios (fused ops, etc.) # Only UB doesn't need workspace if self.strategy != AllReduceStrategy.UB: - if self.strategy == AllReduceStrategy.LOWPRECISION: - allocate_low_presicion_allreduce_workspace(self.mapping) if self.strategy not in (AllReduceStrategy.UB, AllReduceStrategy.NCCL, AllReduceStrategy.NCCL_SYMMETRIC): - self.workspace = get_allreduce_workspace(self.mapping) + self.workspace, ipc_failed = get_allreduce_workspace( + self.mapping) + # Every custom all-reduce kernel reads the peers' IPC buffers. When + # P2P was reported but the IPC handles could not be exchanged those + # pointers are null, so NCCL is the only strategy left. The user + # already got a warning from open_ipc_memory at that point. + if ipc_failed: + logger.debug( + "CUDA IPC is unavailable, falling back from " + f"{self.strategy.name} to NCCL allreduce.") + self.strategy = AllReduceStrategy.NCCL + self.workspace = None + if self.strategy == AllReduceStrategy.LOWPRECISION: + allocate_low_presicion_allreduce_workspace(self.mapping) # Initialize MNNVL if using AUTO or MNNVL strategy if self.strategy in (AllReduceStrategy.AUTO, @@ -1154,7 +1195,7 @@ def __init__(self, mapping: Mapping): """ super().__init__() self.mapping = mapping - self.workspace = get_allreduce_workspace(self.mapping) + self.workspace = _require_ipc_workspace(self.mapping, "MoEAllReduce") # Pls keep this value in sync with the kOneShotMaxToken in moeAllReduceFusionKernels.h self.max_token = 128 @@ -1417,7 +1458,8 @@ class MiniMaxAllReduceRMS(nn.Module): def __init__(self, mapping: Mapping): super().__init__() self.mapping = mapping - self.workspace = get_allreduce_workspace(self.mapping) + self.workspace = _require_ipc_workspace(self.mapping, + "MiniMaxAllReduceRMS") def forward(self, input: torch.Tensor, rms_weights: torch.Tensor, eps: float): diff --git a/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py b/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py new file mode 100644 index 000000000000..4642c30ad38e --- /dev/null +++ b/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the CUDA IPC fallback in IpcMemory and its effect on strategy selection. + +`cudaDeviceCanAccessPeer` reporting P2P does not guarantee that IPC handles can be +exported/imported. When IPC turns out to be unusable, IpcMemory keeps null pointers +so the runtime falls back to NCCL instead of raising, and the whole TP group agrees +on that. + +A workspace with no IPC buffers because P2P was never there (inter-node TP) is not a +failure though, and must not change the strategy, or MNNVL never gets selected. + +The CUDA runtime and the collectives are stubbed, so no multi-GPU and no MPI needed. +This directory runs in the CPU-only l0_cpu stage since #16498, so everything here has +to pass with zero GPUs. `test_lamport_skipped` is the exception and skips there: the +workspace helper it drives allocates device tensors. + + pytest tests/unittest/_torch/distributed/test_ipc_memory_fallback.py -v +""" + +from enum import IntEnum +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm import _ipc_utils +from tensorrt_llm._ipc_utils import IpcMemory +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping + + +def _has_allreduce_op() -> bool: + """Whether the trtllm custom ops are registered in this build.""" + try: + return torch.ops.trtllm.allreduce is not None + except Exception: + return False + + +# AllReduce.__init__ resolves torch.ops.trtllm.allreduce before any of the logic +# under test, so these skip rather than error on a build without the extension. +requires_allreduce_op = pytest.mark.skipif( + not _has_allreduce_op(), reason="trtllm allreduce custom op not registered" +) + +TP_SIZE = 2 + + +class FakeCudaError(IntEnum): + cudaSuccess = 0 + cudaErrorInvalidDevice = 101 + + +class FakeIpcMemHandle: + def __init__(self): + self.reserved = b"\x00" * 64 + + +class FakeCudart: + """Minimal stand-in for `cuda.bindings.runtime` that tracks allocations.""" + + cudaError_t = FakeCudaError + cudaIpcMemHandle_t = FakeIpcMemHandle + cudaIpcMemLazyEnablePeerAccess = 1 + + def __init__( + self, + get_handle_error=FakeCudaError.cudaSuccess, + open_handle_error=FakeCudaError.cudaSuccess, + ): + self.get_handle_error = get_handle_error + self.open_handle_error = open_handle_error + self.next_ptr = 0x1000 + self.allocated = set() + self.opened = set() + self.open_calls = [] + + def cudaMalloc(self, size): + self.next_ptr += size + self.allocated.add(self.next_ptr) + return FakeCudaError.cudaSuccess, self.next_ptr + + def cudaMemset(self, ptr, value, size): + return (FakeCudaError.cudaSuccess,) + + def cudaFree(self, ptr): + self.allocated.discard(ptr) + return (FakeCudaError.cudaSuccess,) + + def cudaIpcGetMemHandle(self, ptr): + handle = FakeIpcMemHandle() + handle.reserved = b"\xab" * 64 + return self.get_handle_error, handle + + def cudaIpcOpenMemHandle(self, handle, flags): + self.open_calls.append(bytes(handle.reserved)) + if self.open_handle_error != FakeCudaError.cudaSuccess: + return self.open_handle_error, 0 + self.next_ptr += 0x1000 + self.opened.add(self.next_ptr) + return FakeCudaError.cudaSuccess, self.next_ptr + + def cudaIpcCloseMemHandle(self, ptr): + self.opened.discard(ptr) + return (FakeCudaError.cudaSuccess,) + + +class FakeDist: + """`tp_allgather` over a fake TP group where the peers mirror this rank. + + `peer_exports` and `peer_ipc_ok` override what the peers contribute, which + lets a test exercise "this rank succeeded but another one did not". + """ + + def __init__(self, tp_size, peer_exports=True, peer_ipc_ok=True): + self.tp_size = tp_size + self.peer_exports = peer_exports + self.peer_ipc_ok = peer_ipc_ok + + def tp_allgather(self, obj): + if isinstance(obj, bool): + return [obj] + [self.peer_ipc_ok] * (self.tp_size - 1) + peer = obj if self.peer_exports else None + return [obj] + [peer] * (self.tp_size - 1) + + +@pytest.fixture +def mapping(): + return Mapping(world_size=TP_SIZE, rank=0, tp_size=TP_SIZE) + + +def _patched(cudart, dist): + from tensorrt_llm._torch.distributed.communicator import Distributed + + return ( + patch.object(_ipc_utils, "cudart", cudart), + patch.object(Distributed, "get", lambda _mapping: dist), + ) + + +def _run(mapping, cudart, dist, open_ipc=True): + cudart_patch, dist_patch = _patched(cudart, dist) + with cudart_patch, dist_patch: + return IpcMemory(mapping, 1 << 20, open_ipc) + + +def test_ipc_ok(mapping): + cudart = FakeCudart() + + ipc_memory = _run(mapping, cudart, FakeDist(TP_SIZE)) + + assert ipc_memory.open_ipc + assert not ipc_memory.ipc_failed + assert ipc_memory.local_ptr != 0 + assert all(ptr != 0 for ptr in ipc_memory.peer_ptrs) + assert len(cudart.opened) == TP_SIZE - 1 + + +def test_local_open_fails(mapping): + # Reproduces github.com/NVIDIA/TensorRT-LLM/issues/16899: cudaIpcOpenMemHandle + # fails with cudaErrorInvalidDevice even though cudaDeviceCanAccessPeer passed. + cudart = FakeCudart(open_handle_error=FakeCudaError.cudaErrorInvalidDevice) + + ipc_memory = _run(mapping, cudart, FakeDist(TP_SIZE, peer_ipc_ok=False)) + + assert not ipc_memory.open_ipc + assert ipc_memory.ipc_failed + assert ipc_memory.local_ptr == 0 + assert ipc_memory.peer_ptrs == [0] * TP_SIZE + assert not cudart.allocated, "the local buffer must be released on fallback" + + +def test_peer_open_fails(mapping): + # IPC works locally but another rank of the TP group failed: this rank must + # fall back as well, otherwise it would use buffers the peers never mapped. + cudart = FakeCudart() + + ipc_memory = _run(mapping, cudart, FakeDist(TP_SIZE, peer_ipc_ok=False)) + + assert not ipc_memory.open_ipc + assert ipc_memory.ipc_failed + assert ipc_memory.local_ptr == 0 + assert ipc_memory.peer_ptrs == [0] * TP_SIZE + assert not cudart.allocated + assert not cudart.opened, "handles opened before the fallback must be closed" + + +def test_local_export_fails(mapping): + cudart = FakeCudart(get_handle_error=FakeCudaError.cudaErrorInvalidDevice) + + ipc_memory = _run(mapping, cudart, FakeDist(TP_SIZE)) + + assert not ipc_memory.open_ipc + assert ipc_memory.ipc_failed + assert not cudart.allocated + + +def test_peer_export_fails(mapping): + # A rank that cannot export contributes None rather than an uninitialized + # cudaIpcMemHandle_t, so the peers never import a garbage handle. + cudart = FakeCudart() + + ipc_memory = _run(mapping, cudart, FakeDist(TP_SIZE, peer_exports=False)) + + assert not ipc_memory.open_ipc + assert ipc_memory.ipc_failed + assert cudart.open_calls == [] + assert not cudart.allocated + + +def test_ipc_not_requested(mapping): + # `can_access_peer` returning False (inter-node TP, no P2P) is not a failure: + # the workspace is null by design and strategy selection must not react to it. + cudart = FakeCudart() + + ipc_memory = _run(mapping, cudart, FakeDist(TP_SIZE), open_ipc=False) + + assert not ipc_memory.open_ipc + assert not ipc_memory.ipc_failed + assert not cudart.allocated + + +def test_inter_node_tp(): + inter_node = Mapping(world_size=16, rank=0, tp_size=16, gpus_per_node=8) + cudart = FakeCudart() + + ipc_memory = _run(inter_node, cudart, FakeDist(16)) + + assert not ipc_memory.open_ipc + assert not ipc_memory.ipc_failed + + +class FakeMNNVLAllReduce: + @staticmethod + def is_mnnvl(mapping, dtype): + return True + + def __init__(self, mapping, dtype): + self.mapping = mapping + + +def _build_allreduce(mapping, strategy, ipc_failed): + """Builds an AllReduce with the workspace and MNNVL support both stubbed out.""" + from tensorrt_llm._torch.distributed import ops + + workspace = torch.zeros(1, dtype=torch.int64) + with ( + patch.object(ops, "get_allreduce_workspace", lambda _mapping: (workspace, ipc_failed)), + patch.object(ops, "MNNVLAllReduce", FakeMNNVLAllReduce), + ): + return ops.AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + + +@requires_allreduce_op +def test_mnnvl_kept_without_p2p(mapping): + # Regression guard: an inter-node workspace holds no IPC buffers, but that must + # not rewrite the strategy, otherwise MNNVL is never constructed on NVLink + # multi-node systems. + allreduce = _build_allreduce(mapping, AllReduceStrategy.AUTO, ipc_failed=False) + + assert allreduce.strategy == AllReduceStrategy.AUTO + assert allreduce.mnnvl_allreduce is not None + + +@requires_allreduce_op +def test_downgrade_to_nccl(mapping): + allreduce = _build_allreduce(mapping, AllReduceStrategy.ONESHOT, ipc_failed=True) + + assert allreduce.strategy == AllReduceStrategy.NCCL + assert allreduce.workspace is None + assert allreduce.mnnvl_allreduce is None + + +@requires_allreduce_op +def test_moe_allreduce_raises(mapping): + # MoEAllReduce has no NCCL path: it must fail with an actionable message rather + # than dereference the null peer pointers inside the kernel. + from tensorrt_llm._torch.distributed import ops + + workspace = torch.zeros(1, dtype=torch.int64) + with patch.object(ops, "get_allreduce_workspace", lambda _mapping: (workspace, True)): + with pytest.raises(RuntimeError, match="requires CUDA IPC"): + ops.MoEAllReduce(mapping) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_lamport_skipped(mapping): + # lamport_initialize writes through lamport_buffers.local_ptr, which is null + # once IPC has been disabled. + from tensorrt_llm._torch.distributed import allreduce_helper + from tensorrt_llm._torch.distributed.allreduce_helper import CustomAllReduceHelper + + cudart = FakeCudart(open_handle_error=FakeCudaError.cudaErrorInvalidDevice) + cudart_patch, dist_patch = _patched(cudart, FakeDist(TP_SIZE, peer_ipc_ok=False)) + calls = [] + + with ( + cudart_patch, + dist_patch, + patch.object(allreduce_helper, "can_access_peer", lambda _mapping: True), + patch.object(allreduce_helper, "lamport_initialize", lambda *args: calls.append(args)), + ): + buffers, _ = CustomAllReduceHelper.allocate_allreduce_fusion_workspace(mapping, 1 << 20) + + assert calls == [], "lamport_initialize must not run on a null buffer" + assert all(buffer.ipc_failed for buffer in buffers if isinstance(buffer, IpcMemory))