Skip to content

[https://nvbugs/5923949][fix] Improve NCCL library load stability #12015

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Tabrizian merged 5 commits into from
Mar 11, 2026
Merged

[https://nvbugs/5923949][fix] Improve NCCL library load stability #12015

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fd6f020
pass shared_ptr<ncclComm_t> into ScopedNCCLWindowBuffer and createNCC…
nv-lschneider Feb 25, 2026
52e01b9
replace NCCLHelper with nccl version macros and runtime checks
nv-lschneider Mar 9, 2026
b0a827a
guard against unchecked shared ptr dereference
nv-lschneider Mar 9, 2026
4fe79de
Merge branch 'main' into lschneider/nccl-window-shared-ptr-lifetime
nv-lschneider Mar 9, 2026
ead6bba
Merge branch 'main' into lschneider/nccl-window-shared-ptr-lifetime
nv-lschneider Mar 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 28 additions & 150 deletions cpp/tensorrt_llm/common/ncclUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,145 +211,43 @@ size_t NcclCommResourceManager::getResourceCount(ncclComm_t comm) const noexcept
}

//==============================================================================
// NCCLHelper Implementation
// NCCLWindowAllocator Implementation
//==============================================================================

NCCLHelper& NCCLHelper::getInstance()
{
static NCCLHelper instance;
return instance;
}

NCCLHelper::NCCLHelper()
: mLibraryHandle(nullptr)
, mNCCLCommWindowRegister(nullptr)
, mNCCLMemAlloc(nullptr)
, mIsLoaded(false)
{
loadNCCLLibrary();
}
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)

NCCLHelper::~NCCLHelper()
NCCLWindowAllocator& NCCLWindowAllocator::getInstance()
{
if (mLibraryHandle)
{
#ifdef _WIN32
FreeLibrary(mLibraryHandle);
#else
dlclose(mLibraryHandle);
#endif
mLibraryHandle = nullptr;
}
static NCCLWindowAllocator instance;
return instance;
}

void NCCLHelper::loadNCCLLibrary()
NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size)
{
try
{
#ifdef _WIN32
char const* libraryNames[] = {"nccl.dll"};
#else
char const* libraryNames[] = {"libnccl.so"};
#endif

for (auto const* name : libraryNames)
// One-time runtime version check: the runtime NCCL library must also support window buffers.
static std::once_flag versionCheckFlag;
static bool runtimeVersionOk = false;
std::call_once(versionCheckFlag,
[]()
{
mLibraryHandle = loadLibraryHandle(name);
if (mLibraryHandle)
int version = 0;
if (ncclGetVersion(&version) == ncclSuccess && version >= NCCL_VERSION(2, 28, 0))
{
TLLM_LOG_INFO("Successfully loaded NCCL library: %s", name);
break;
runtimeVersionOk = true;
}
}

if (!mLibraryHandle)
{
TLLM_LOG_WARNING("Failed to load NCCL library");
return;
}

// Load the required symbols
mNCCLCommWindowRegister
= reinterpret_cast<ncclCommWindowRegisterFunc>(getSymbolAddress(mLibraryHandle, "ncclCommWindowRegister"));

mNCCLMemAlloc = reinterpret_cast<ncclMemAllocFunc>(getSymbolAddress(mLibraryHandle, "ncclMemAlloc"));

if (mNCCLCommWindowRegister == nullptr)
{
TLLM_LOG_WARNING("Failed to load ncclCommWindowRegister symbol, NCCL symmetric will not be supported.");
}

if (mNCCLMemAlloc == nullptr)
{
TLLM_LOG_WARNING("Failed to load ncclMemAlloc symbol, NCCL symmetric will not be supported.");
}

if (mNCCLCommWindowRegister != nullptr && mNCCLMemAlloc != nullptr)
{
mIsLoaded = true;
}
else
{
TLLM_LOG_WARNING(
"Failed to load required NCCL symbols (both ncclCommWindowRegister and ncclMemAlloc are required)");
}
}
catch (std::exception const& e)
{
TLLM_LOG_WARNING("Exception while loading NCCL library: %s", e.what());
}
}

void* NCCLHelper::loadLibraryHandle(char const* libName)
{
#ifdef _WIN32
return LoadLibraryA(libName);
#else
return dlopen(libName, RTLD_LAZY | RTLD_GLOBAL);
#endif
}

void* NCCLHelper::getSymbolAddress(void* handle, char const* symbolName)
{
if (!handle)
else
{
TLLM_LOG_WARNING(
"[NCCLUtil] NCCL runtime version %d.%d.%d does not support window buffers; "
"falling back to regular tensors.",
version / 10000, (version % 10000) / 100, version % 100);
}
});
Comment thread
nv-lschneider marked this conversation as resolved.
if (!runtimeVersionOk)
{
return nullptr;
return NCCLWindowBuffer();
}

#ifdef _WIN32
return GetProcAddress(static_cast<HMODULE>(handle), symbolName);
#else
return dlsym(handle, symbolName);
#endif
}

NCCLHelper::ncclCommWindowRegisterFunc NCCLHelper::getNCCLCommWindowRegister()
{
return mNCCLCommWindowRegister;
}

NCCLHelper::ncclMemAllocFunc NCCLHelper::getNCCLMemAlloc()
{
return mNCCLMemAlloc;
}

bool NCCLHelper::isLoaded() const
{
return mIsLoaded;
}

//==============================================================================
// NCCLWindowAllocator Implementation
//==============================================================================

NCCLWindowAllocator& NCCLWindowAllocator::getInstance()
{
static NCCLWindowAllocator instance;
return instance;
}

NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size)
{
TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null");
TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0");

Expand Down Expand Up @@ -508,38 +406,16 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm,
NCCLWindowBuffer buffer;
buffer.handle = handle;

// Get NCCL helper for dynamic symbol loading
auto& ncclHelper = NCCLHelper::getInstance();
if (!ncclHelper.isLoaded())
{
TLLM_THROW("NCCL library could not be loaded for dynamic symbol access");
}

auto ncclMemAllocFunc = ncclHelper.getNCCLMemAlloc();
auto ncclCommWindowRegisterFunc = ncclHelper.getNCCLCommWindowRegister();

// Defensive checks: both function pointers must be non-null
if (ncclMemAllocFunc == nullptr)
{
TLLM_THROW("ncclMemAlloc function pointer is null, cannot allocate NCCL window buffer");
}

if (ncclCommWindowRegisterFunc == nullptr)
{
TLLM_THROW("ncclCommWindowRegister function pointer is null, cannot register NCCL window buffer");
}

// Allocate device memory using ncclMemAlloc
ncclResult_t allocResult = ncclMemAllocFunc(&buffer.ptr, size);
ncclResult_t allocResult = ncclMemAlloc(&buffer.ptr, size);
if (allocResult != ncclSuccess)
{
TLLM_THROW("ncclMemAlloc failed with error: %d", allocResult);
}
buffer.size = size;

// Register the buffer with NCCL as a window
ncclResult_t regResult
= ncclCommWindowRegisterFunc(comm, buffer.ptr, size, &buffer.window, NCCL_WIN_COLL_SYMMETRIC);
ncclResult_t regResult = ncclCommWindowRegister(comm, buffer.ptr, size, &buffer.window, NCCL_WIN_COLL_SYMMETRIC);
if (regResult != ncclSuccess)
{
ncclMemFree(buffer.ptr);
Expand Down Expand Up @@ -694,6 +570,8 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept
mRegisteredComms.erase(comm);
}

#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)

} // namespace tensorrt_llm::common::nccl_util

#endif // ENABLE_MULTI_DEVICE
104 changes: 42 additions & 62 deletions cpp/tensorrt_llm/common/ncclUtils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,65 +43,11 @@

#if ENABLE_MULTI_DEVICE

#ifdef _WIN32
#include <windows.h>
#else
#include <dlfcn.h>
#endif

TRTLLM_NAMESPACE_BEGIN

namespace common::nccl_util
{

//==============================================================================
// NCCL Helper - Dynamic Library Loading
//==============================================================================

// Helper class for dynamically loading NCCL symbols (ncclMemAlloc, ncclCommWindowRegister)
// This allows the code to work with NCCL libraries that may or may not have these symbols
class NCCLHelper
{
public:
static NCCLHelper& getInstance();

// Dynamic loading function type definition
using ncclCommWindowRegisterFunc = ncclResult_t (*)(ncclComm_t, void*, size_t, ncclWindow_t*, int);
using ncclMemAllocFunc = ncclResult_t (*)(void**, size_t);

// Get function pointer for ncclCommWindowRegister
ncclCommWindowRegisterFunc getNCCLCommWindowRegister();

// Get function pointer for ncclMemAlloc
ncclMemAllocFunc getNCCLMemAlloc();

// Check if NCCL library is successfully loaded
bool isLoaded() const;

NCCLHelper(NCCLHelper const&) = delete;
NCCLHelper& operator=(NCCLHelper const&) = delete;
NCCLHelper(NCCLHelper&&) = delete;
NCCLHelper& operator=(NCCLHelper&&) = delete;

private:
NCCLHelper();
~NCCLHelper();

void loadNCCLLibrary();
void* loadLibraryHandle(char const* libName);
void* getSymbolAddress(void* handle, char const* symbolName);

#ifdef _WIN32
HMODULE mLibraryHandle;
#else
void* mLibraryHandle;
#endif

ncclCommWindowRegisterFunc mNCCLCommWindowRegister;
ncclMemAllocFunc mNCCLMemAlloc;
bool mIsLoaded;
};

//==============================================================================
// NCCL Resource Management
//==============================================================================
Expand Down Expand Up @@ -198,10 +144,32 @@ class NcclCommResource
bool mRegistered;
};

//==============================================================================
// NCCL Version Check
//==============================================================================

// Returns true if the compile-time and runtime NCCL versions support window buffers
// (ncclMemAlloc / ncclCommWindowRegister).
inline bool isNcclWindowSupported()
{
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)
int version = 0;
if (ncclGetVersion(&version) != ncclSuccess)
{
return false;
}
return version >= NCCL_VERSION(2, 28, 0);
#else
return false;
#endif
}

//==============================================================================
// NCCL Window Buffer Allocation
//==============================================================================

#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)

// Represents a buffer with an associated NCCL window
struct NCCLWindowBuffer
{
Expand Down Expand Up @@ -311,17 +279,21 @@ class NCCLWindowAllocator
class ScopedNCCLWindowBuffer
{
public:
ScopedNCCLWindowBuffer(ncclComm_t comm, size_t size)
: mComm(comm)
, mBuffer(NCCLWindowAllocator::getInstance().requestBuffer(comm, size))
ScopedNCCLWindowBuffer(std::shared_ptr<ncclComm_t> comm, size_t size)
: mComm(std::move(comm))
, mBuffer{}
{
Comment thread
coderabbitai[bot] marked this conversation as resolved.
if (mComm && *mComm)
{
mBuffer = NCCLWindowAllocator::getInstance().requestBuffer(*mComm, size);
}
}

~ScopedNCCLWindowBuffer()
{
if (mBuffer.isValid())
{
NCCLWindowAllocator::getInstance().releaseBuffer(mComm, mBuffer.ptr);
NCCLWindowAllocator::getInstance().releaseBuffer(*mComm, mBuffer.ptr);
}
}

Expand Down Expand Up @@ -351,15 +323,15 @@ class ScopedNCCLWindowBuffer
ScopedNCCLWindowBuffer& operator=(ScopedNCCLWindowBuffer&&) = delete;

private:
ncclComm_t mComm;
std::shared_ptr<ncclComm_t> mComm;
NCCLWindowBuffer mBuffer;
};

// Creates a PyTorch tensor backed by an NCCL window buffer.
// The tensor will automatically release the buffer back to the pool when destroyed.
// This is analogous to torch_ext::create_userbuffers_tensor() but for NCCLWindowAllocator.
inline std::pair<torch::Tensor, NCCLWindowBuffer> createNCCLWindowTensor(
ncclComm_t comm, at::IntArrayRef shape, torch::ScalarType dtype)
std::shared_ptr<ncclComm_t> comm, at::IntArrayRef shape, torch::ScalarType dtype)
{
// Calculate buffer size
int64_t buffer_size
Expand All @@ -380,9 +352,15 @@ inline std::pair<torch::Tensor, NCCLWindowBuffer> createNCCLWindowTensor(
auto& allocator = NCCLWindowAllocator::getInstance();
NCCLWindowBuffer buffer;

if (!comm || !*comm)
{
TLLM_LOG_DEBUG("[createNCCLWindowTensor] null comm; returning invalid buffer");
return std::make_pair(torch::Tensor(), NCCLWindowBuffer());
}

try
{
buffer = allocator.requestBuffer(comm, buffer_size);
buffer = allocator.requestBuffer(*comm, buffer_size);
}
catch (std::exception const& e)
{
Expand All @@ -398,14 +376,16 @@ inline std::pair<torch::Tensor, NCCLWindowBuffer> createNCCLWindowTensor(
}

// Create custom deleter that releases the buffer
auto deleter = [comm, ptr = buffer.ptr](void*) { NCCLWindowAllocator::getInstance().releaseBuffer(comm, ptr); };
auto deleter = [comm, ptr = buffer.ptr](void*) { NCCLWindowAllocator::getInstance().releaseBuffer(*comm, ptr); };

// Create tensor from the buffer
auto tensor = torch::from_blob(buffer.ptr, shape, strides_vec, deleter, torch::dtype(dtype).device(torch::kCUDA));

return std::make_pair(tensor, buffer);
}

#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)

} // namespace common::nccl_util

TRTLLM_NAMESPACE_END
Expand Down
Loading