From 38806af1c73d7e25ed800af4d78aa9e7705bad07 Mon Sep 17 00:00:00 2001 From: pjdurden Date: Wed, 5 Aug 2026 11:53:30 -0500 Subject: [PATCH 1/3] [None][fix] Fall back to NCCL when CUDA IPC handles cannot be exchanged cudaDeviceCanAccessPeer reporting P2P support does not guarantee that CUDA IPC handles can be exported and imported. On 8x RTX 6000D, and on PCIe-only parts generally, cudaIpcOpenMemHandle fails with cudaErrorInvalidDevice and engine initialization aborted instead of using the NCCL path that already exists. IpcMemory.open_ipc_memory now returns None on such a failure, after agreeing on the outcome across the whole TP group and releasing whatever it managed to allocate, so that either every rank gets IPC buffers or none of them does. IpcMemory keeps null pointers in that case, and lamport_initialize is gated on the buffers actually being open. Fixes #16899 Signed-off-by: pjdurden --- tensorrt_llm/_ipc_utils.py | 49 +++++- .../_torch/distributed/allreduce_helper.py | 4 +- tensorrt_llm/_torch/distributed/ops.py | 28 +++- .../distributed/test_ipc_memory_fallback.py | 157 ++++++++++++++++++ 4 files changed, 226 insertions(+), 12 deletions(-) create mode 100644 tests/unittest/_torch/distributed/test_ipc_memory_fallback.py diff --git a/tensorrt_llm/_ipc_utils.py b/tensorrt_llm/_ipc_utils.py index 3178597b5087..72902fca0a8f 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: @@ -108,7 +108,13 @@ def __init__(self, mapping: Mapping, size: int, open_ipc: bool = True): 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 + else: + self.peer_ptrs, self.local_ptr = ipc_memory def __del__(self): if not sys.is_finalizing() and self.open_ipc: @@ -124,10 +130,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 +158,40 @@ def align_size(size, alignment): _raise_if_error(error) if set_to_zero: _raise_if_error(cudart.cudaMemset(local_ptr, 0, aligned_size)[0]) + error, local_handle = cudart.cudaIpcGetMemHandle(local_ptr) - _raise_if_error(error) + ipc_error = None if error == cudart.cudaError_t.cudaSuccess else error + if ipc_error is not None: + # Exchange a dummy handle to keep the collective below symmetric. + local_handle = cudart.cudaIpcMemHandle_t() handles_reserved = dist.tp_allgather(_ipc_mem_handle_to_bytes(local_handle)) handles = [_ipc_mem_handle_from_bytes(reserved) for reserved in handles_reserved] peer_ptrs = [] + opened_ptrs = [] for node, handle in enumerate(handles): if node == mapping.tp_rank: peer_ptrs.append(local_ptr) - else: + elif ipc_error is None: error, ptr = cudart.cudaIpcOpenMemHandle( handle, cudart.cudaIpcMemLazyEnablePeerAccess ) - _raise_if_error(error) + if error != cudart.cudaError_t.cudaSuccess: + ipc_error = error + continue peer_ptrs.append(ptr) + opened_ptrs.append(ptr) + + if not all(dist.tp_allgather(ipc_error is None)): + if ipc_error is not None: + logger.warning( + f"CUDA IPC is not usable on this system: {ipc_error!r}. " + "Custom all-reduce kernels are disabled, falling back to NCCL." + ) + for ptr in opened_ptrs: + _raise_if_error(cudart.cudaIpcCloseMemHandle(ptr)[0]) + _raise_if_error(cudart.cudaFree(local_ptr)[0]) + 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..9de5ab73da24 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 @@ -115,6 +116,20 @@ def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: return allreduce_workspaces[mapping][1] +def allreduce_workspace_has_ipc(mapping: Mapping) -> bool: + """Whether the workspace returned by get_allreduce_workspace holds real IPC buffers. + + cudaDeviceCanAccessPeer can report P2P support on systems where CUDA IPC handles + still cannot be imported. In that case the workspace pointers are null and only + NCCL can be used. get_allreduce_workspace must have been called first. + """ + allreduce_workspaces = getattr(_thread_local, + f'allreduce_workspaces_{mapping.pp_rank}') + ipc_buffers = allreduce_workspaces[mapping][0] + return all(buffer.open_ipc for buffer in ipc_buffers + if isinstance(buffer, IpcMemory)) + + def allocate_low_presicion_allreduce_workspace(mapping: Mapping) -> None: if not hasattr(_thread_local, 'lowprecision_allreduce_workspaces'): _thread_local.lowprecision_allreduce_workspaces = {} @@ -933,12 +948,21 @@ 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) + # Every custom all-reduce kernel reads the peers' IPC buffers. When + # CUDA IPC is unavailable those pointers are null, so NCCL is the + # only strategy that can run. See allreduce_workspace_has_ipc. + if not allreduce_workspace_has_ipc(self.mapping): + logger.warning( + "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, 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..4209edbf46f1 --- /dev/null +++ b/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py @@ -0,0 +1,157 @@ +# 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. + +`cudaDeviceCanAccessPeer` reporting P2P support does not guarantee that CUDA IPC +handles can be exported/imported. When IPC turns out to be unusable, IpcMemory +must degrade to null pointers (so the runtime falls back to NCCL) instead of +raising, and every rank of the TP group must reach the same decision. + +The CUDA runtime is stubbed, so these tests need neither a GPU nor MPI: + pytest tests/unittest/_torch/distributed/test_ipc_memory_fallback.py -v +""" + +from enum import IntEnum +from unittest.mock import patch + +import pytest + +from tensorrt_llm import _ipc_utils +from tensorrt_llm._ipc_utils import IpcMemory +from tensorrt_llm.mapping import Mapping + +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, open_handle_error=FakeCudaError.cudaSuccess): + self.open_handle_error = open_handle_error + self.next_ptr = 0x1000 + self.allocated = set() + self.opened = set() + + 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): + return FakeCudaError.cudaSuccess, FakeIpcMemHandle() + + def cudaIpcOpenMemHandle(self, handle, flags): + 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_ipc_ok` overrides what the peers report for the boolean agreement + collective, which lets a test exercise "this rank succeeded but another one + did not". + """ + + def __init__(self, tp_size, peer_ipc_ok=True): + self.tp_size = tp_size + 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) + return [obj] * self.tp_size + + +@pytest.fixture +def mapping(): + return Mapping(world_size=TP_SIZE, rank=0, tp_size=TP_SIZE) + + +def _run(mapping, cudart, dist): + from tensorrt_llm._torch.distributed.communicator import Distributed + + with ( + patch.object(_ipc_utils, "cudart", cudart), + patch.object(Distributed, "get", lambda _mapping: dist), + ): + return IpcMemory(mapping, 1 << 20) + + +def test_ipc_memory_is_opened_when_ipc_works(mapping): + cudart = FakeCudart() + + ipc_memory = _run(mapping, cudart, FakeDist(TP_SIZE)) + + assert ipc_memory.open_ipc + 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_ipc_open_failure_falls_back_instead_of_raising(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.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_ipc_failure_disables_ipc_on_this_rank(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.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" From bf1eb79b6e79104505732e97338065e16a081fa0 Mon Sep 17 00:00:00 2001 From: pjdurden Date: Wed, 5 Aug 2026 11:54:00 -0500 Subject: [PATCH 2/3] [None][fix] Narrow the CUDA IPC fallback to real IPC failures Addresses review feedback on #17034. allreduce_workspace_has_ipc() was False for any workspace built without P2P, which includes ordinary inter-node TP, so it rewrote the strategy to NCCL before the MNNVL block and mnnvl_allreduce was never constructed on multi-node NVLink systems. IpcMemory now records ipc_failed separately from open_ipc, and only that failure mode, P2P reported but handles unusable, downgrades the strategy. Configurations that never had P2P keep the behaviour they had. get_allreduce_workspace returns the flag alongside the workspace, so a caller cannot read the thread-local state before it has been populated. The fallback is reported once, from open_ipc_memory where it is detected, and on every rank of the group rather than only on the ones that failed locally. AllReduce is constructed per decoder layer, so it no longer warns per instance. A rank that cannot export its handle contributes None to the allgather instead of an uninitialized cudaIpcMemHandle_t, so no rank imports a garbage handle. The number of collectives is unchanged. MoEAllReduce and MiniMaxAllReduceRMS have no NCCL path, so they now raise an actionable error instead of letting the kernel dereference null peer pointers. Tests cover the strategy downgrade, MNNVL still being selected when the workspace has no IPC buffers by design, and the lamport_initialize gating. Signed-off-by: pjdurden --- tensorrt_llm/_ipc_utils.py | 80 +++++--- tensorrt_llm/_torch/distributed/ops.py | 58 ++++-- .../distributed/test_ipc_memory_fallback.py | 183 ++++++++++++++++-- 3 files changed, 255 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_ipc_utils.py b/tensorrt_llm/_ipc_utils.py index 72902fca0a8f..14244b79c7c2 100644 --- a/tensorrt_llm/_ipc_utils.py +++ b/tensorrt_llm/_ipc_utils.py @@ -104,6 +104,11 @@ 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 @@ -113,6 +118,7 @@ def __init__(self, mapping: Mapping, size: int, open_ipc: bool = True): # 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 @@ -159,38 +165,60 @@ def align_size(size, alignment): 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) - ipc_error = None if error == cudart.cudaError_t.cudaSuccess else error - if ipc_error is not None: - # Exchange a dummy handle to keep the collective below symmetric. - local_handle = cudart.cudaIpcMemHandle_t() - 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 = [] opened_ptrs = [] - for node, handle in enumerate(handles): + open_error = None + for node, reserved in enumerate(handles_reserved): if node == mapping.tp_rank: peer_ptrs.append(local_ptr) - elif ipc_error is None: - error, ptr = cudart.cudaIpcOpenMemHandle( - handle, cudart.cudaIpcMemLazyEnablePeerAccess - ) - if error != cudart.cudaError_t.cudaSuccess: - ipc_error = error - continue - peer_ptrs.append(ptr) - opened_ptrs.append(ptr) - - if not all(dist.tp_allgather(ipc_error is None)): - if ipc_error is not None: - logger.warning( - f"CUDA IPC is not usable on this system: {ipc_error!r}. " - "Custom all-reduce kernels are disabled, falling back to NCCL." - ) - for ptr in opened_ptrs: - _raise_if_error(cudart.cudaIpcCloseMemHandle(ptr)[0]) - _raise_if_error(cudart.cudaFree(local_ptr)[0]) + 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/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 9de5ab73da24..d1958645a508 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -100,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}', {}) @@ -112,22 +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 allreduce_workspace_has_ipc(mapping: Mapping) -> bool: - """Whether the workspace returned by get_allreduce_workspace holds real IPC buffers. +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. - cudaDeviceCanAccessPeer can report P2P support on systems where CUDA IPC handles - still cannot be imported. In that case the workspace pointers are null and only - NCCL can be used. get_allreduce_workspace must have been called first. + 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. """ - allreduce_workspaces = getattr(_thread_local, - f'allreduce_workspaces_{mapping.pp_rank}') - ipc_buffers = allreduce_workspaces[mapping][0] - return all(buffer.open_ipc for buffer in ipc_buffers - if isinstance(buffer, IpcMemory)) + 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: @@ -951,12 +966,14 @@ def __init__(self, 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 - # CUDA IPC is unavailable those pointers are null, so NCCL is the - # only strategy that can run. See allreduce_workspace_has_ipc. - if not allreduce_workspace_has_ipc(self.mapping): - logger.warning( + # 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 @@ -1178,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 @@ -1441,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 index 4209edbf46f1..8960e12977ac 100644 --- a/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py +++ b/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py @@ -13,14 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Tests for the CUDA IPC fallback in IpcMemory. +Tests for the CUDA IPC fallback in IpcMemory and its effect on strategy selection. -`cudaDeviceCanAccessPeer` reporting P2P support does not guarantee that CUDA IPC -handles can be exported/imported. When IPC turns out to be unusable, IpcMemory -must degrade to null pointers (so the runtime falls back to NCCL) instead of -raising, and every rank of the TP group must reach the same decision. +`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. +`test_lamport_skipped` needs one GPU, the workspace helper allocates device tensors. -The CUDA runtime is stubbed, so these tests need neither a GPU nor MPI: pytest tests/unittest/_torch/distributed/test_ipc_memory_fallback.py -v """ @@ -28,9 +33,11 @@ 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 TP_SIZE = 2 @@ -53,11 +60,17 @@ class FakeCudart: cudaIpcMemHandle_t = FakeIpcMemHandle cudaIpcMemLazyEnablePeerAccess = 1 - def __init__(self, open_handle_error=FakeCudaError.cudaSuccess): + 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 @@ -72,9 +85,12 @@ def cudaFree(self, ptr): return (FakeCudaError.cudaSuccess,) def cudaIpcGetMemHandle(self, ptr): - return FakeCudaError.cudaSuccess, FakeIpcMemHandle() + 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 @@ -89,19 +105,20 @@ def cudaIpcCloseMemHandle(self, ptr): class FakeDist: """`tp_allgather` over a fake TP group where the peers mirror this rank. - `peer_ipc_ok` overrides what the peers report for the boolean agreement - collective, which lets a test exercise "this rank succeeded but another one - did not". + `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_ipc_ok=True): + 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) - return [obj] * self.tp_size + peer = obj if self.peer_exports else None + return [obj] + [peer] * (self.tp_size - 1) @pytest.fixture @@ -109,28 +126,34 @@ def mapping(): return Mapping(world_size=TP_SIZE, rank=0, tp_size=TP_SIZE) -def _run(mapping, cudart, dist): +def _patched(cudart, dist): from tensorrt_llm._torch.distributed.communicator import Distributed - with ( + return ( patch.object(_ipc_utils, "cudart", cudart), patch.object(Distributed, "get", lambda _mapping: dist), - ): - return IpcMemory(mapping, 1 << 20) + ) + +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_memory_is_opened_when_ipc_works(mapping): + +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_ipc_open_failure_falls_back_instead_of_raising(mapping): +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) @@ -138,12 +161,13 @@ def test_local_ipc_open_failure_falls_back_instead_of_raising(mapping): 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_ipc_failure_disables_ipc_on_this_rank(mapping): +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() @@ -151,7 +175,126 @@ def test_peer_ipc_failure_disables_ipc_on_this_rank(mapping): 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) + + +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 + + +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 + + +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)) From c484223648ede51bf3c9cb4e73e9d4828b1bc5a5 Mon Sep 17 00:00:00 2001 From: pjdurden Date: Wed, 5 Aug 2026 13:03:41 -0500 Subject: [PATCH 3/3] [None][fix] Keep the IPC fallback tests green on the CPU-only stage #16498 moved tests/unittest/_torch/distributed from l0_dgx_h100 to l0_cpu, so this file now runs with zero GPUs. The three tests that build an AllReduce skip when the trtllm custom ops are not registered, since AllReduce.__init__ resolves torch.ops.trtllm.allreduce before any of the logic under test. test_lamport_skipped still needs a GPU because the workspace helper it drives allocates device tensors. Signed-off-by: pjdurden --- .../distributed/test_ipc_memory_fallback.py | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py b/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py index 8960e12977ac..4642c30ad38e 100644 --- a/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py +++ b/tests/unittest/_torch/distributed/test_ipc_memory_fallback.py @@ -24,7 +24,9 @@ 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. -`test_lamport_skipped` needs one GPU, the workspace helper allocates device tensors. +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 """ @@ -40,6 +42,21 @@ 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 @@ -248,6 +265,7 @@ def _build_allreduce(mapping, strategy, ipc_failed): 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 @@ -258,6 +276,7 @@ def test_mnnvl_kept_without_p2p(mapping): 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) @@ -266,6 +285,7 @@ def test_downgrade_to_nccl(mapping): 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.