[None][fix] Fall back to NCCL when CUDA IPC handles cannot be exchanged

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
 

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,

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):