[TRTLLM-11237][fix] [fix] Synchronize NCCL memory allocation error handling by nv-lschneider · Pull Request #12125 · NVIDIA/TensorRT-LLM

nv-lschneider · 2026-03-11T22:47:40Z

Summary by CodeRabbit

Bug Fixes
- Improved memory allocation error handling in multi-GPU distributed operations to prevent potential deadlocks by ensuring synchronized failure checks across all GPUs.
- Enhanced stability for collective communication operations when allocation failures occur on any GPU.

Description

This adds a global sync before the window registration, so in case a rank OOMs registration fails gracefully.
This adds extra overhead in the registration path, but registration was slow before and we are re-using buffers to get this out of the critical path.

Test Coverage

I am actively looking for the test that was breaking. Help!

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

coderabbitai · 2026-03-11T22:52:42Z

📝 Walkthrough

Walkthrough

This PR modifies NCCL memory allocation error handling to implement a synchronized, collective check across ranks instead of per-rank failure handling. When ncclMemAlloc fails on any rank, a coordinated barrier using ncclAllReduce detects this and ensures all ranks either succeed together or fail gracefully. Fallback logic is updated to prefer NCCL_SYMMETRIC with coordinated failure handling.

Changes

Cohort / File(s)	Summary
NCCL Memory Allocation Synchronization `cpp/tensorrt_llm/common/ncclUtils.cpp`	Implements collective rank synchronization for ncclMemAlloc failure detection. Allocates device-side rank-sync flag, uses ncclAllReduce to detect allocation failures on any rank, and performs coordinated cleanup and buffer return if any rank failed. Replaces immediate per-rank exception throwing with synchronized abort path.
AllReduce Operation Cleanup `cpp/tensorrt_llm/thop/allreduceOp.cpp`	Removes TODO comment noting pending use of NCCL_SYMMETRIC after memory allocation issue resolution. No functional changes.
Torch Fallback Strategy Update `tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`	Updates AllReduceRunner fallback behavior when tactic is -1 to select NCCL_SYMMETRIC instead of plain NCCL. Integrates with NCCLWindowAllocator's cross-rank barrier for handling asymmetric allocation failures.

Sequence Diagram

sequenceDiagram
    participant Rank0 as Rank 0
    participant Rank1 as Rank 1
    participant RankN as Rank N
    participant NCCL as NCCL Collective
    
    Rank0->>Rank0: ncclMemAlloc<br/>(may fail)
    Rank1->>Rank1: ncclMemAlloc<br/>(may fail)
    RankN->>RankN: ncclMemAlloc<br/>(may fail)
    
    Rank0->>Rank0: cudaMalloc rankSyncFlag
    Rank1->>Rank1: cudaMalloc rankSyncFlag
    RankN->>RankN: cudaMalloc rankSyncFlag
    
    Rank0->>Rank0: Set flag = (alloc failed ? 1 : 0)
    Rank1->>Rank1: Set flag = (alloc failed ? 1 : 0)
    RankN->>RankN: Set flag = (alloc failed ? 1 : 0)
    
    Rank0->>NCCL: ncclAllReduce(MIN)
    Rank1->>NCCL: ncclAllReduce(MIN)
    RankN->>NCCL: ncclAllReduce(MIN)
    
    NCCL-->>Rank0: reduced_flag
    NCCL-->>Rank1: reduced_flag
    NCCL-->>RankN: reduced_flag
    
    alt Any rank failed
        Rank0->>Rank0: Free allocation<br/>Return empty buffer
        Rank1->>Rank1: Free allocation<br/>Return empty buffer
        RankN->>RankN: Free allocation<br/>Return empty buffer
    else All succeeded
        Rank0->>Rank0: Proceed to<br/>ncclCommWindowRegister
        Rank1->>Rank1: Proceed to<br/>ncclCommWindowRegister
        RankN->>RankN: Proceed to<br/>ncclCommWindowRegister
    end

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title [TRTLLM-11237][fix] Synchronize NCCL memory allocation error handling accurately describes the main change: synchronizing error handling for NCCL memory allocation across ranks.
Description check	✅ Passed	The description explains the core change (global sync before window registration for graceful OOM handling) and acknowledges overhead trade-offs, but lacks detailed test coverage information and has incomplete checklist items.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

cpp/tensorrt_llm/common/ncclUtils.cpp (1)

411-411: Consider declaring localAllocOk as const.

This variable is not modified after initialization. As per coding guidelines, it should be declared as const int.

-    int localAllocOk = (allocResult == ncclSuccess) ? 1 : 0;
+    int const localAllocOk = (allocResult == ncclSuccess) ? 1 : 0;

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@cpp/tensorrt_llm/common/ncclUtils.cpp` at line 411, The variable localAllocOk
in ncclUtils.cpp is never modified after initialization; change its declaration
to const int localAllocOk = (allocResult == ncclSuccess) ? 1 : 0; so it is
immutable—locate the usage in the nccl allocation handling block where
localAllocOk is declared and update its type to const int (symbol:
localAllocOk).

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@cpp/tensorrt_llm/common/ncclUtils.cpp`:
- Around line 438-443: Add explicit CUDA error checks for the two cudaMemcpy
calls surrounding the NCCL all-reduce: the host->device copy that writes
localAllocOk into rankSyncFlag and the device->host copy that reads allAllocOk
back. After each cudaMemcpy (the calls that use rankSyncFlag, localAllocOk, and
allAllocOk), check the returned cudaError_t and handle failures (e.g., log via
processLogger/your logger or stderr, set a safe failure value for allAllocOk,
free rankSyncFlag, and abort/return an error code) before proceeding to
ncclAllReduce or using the result; ensure the cleanup (cudaFree(rankSyncFlag))
still runs on error paths.

---

Nitpick comments:
In `@cpp/tensorrt_llm/common/ncclUtils.cpp`:
- Line 411: The variable localAllocOk in ncclUtils.cpp is never modified after
initialization; change its declaration to const int localAllocOk = (allocResult
== ncclSuccess) ? 1 : 0; so it is immutable—locate the usage in the nccl
allocation handling block where localAllocOk is declared and update its type to
const int (symbol: localAllocOk).

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ac070788-1367-4d41-8e65-1ee4a4e39f6e

📥 Commits

Reviewing files that changed from the base of the PR and between ae2bd16 and a66cebb.

📒 Files selected for processing (3)

cpp/tensorrt_llm/common/ncclUtils.cpp
cpp/tensorrt_llm/thop/allreduceOp.cpp
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

💤 Files with no reviewable changes (1)

cpp/tensorrt_llm/thop/allreduceOp.cpp

nv-lschneider · 2026-03-11T23:28:06Z

/bot run --add-multi-gpu --disable-fail-fast

tensorrt-cicd · 2026-03-11T23:36:32Z

PR_Github #38641 [ run ] triggered by Bot. Commit: e3c5b1b Link to invocation

tensorrt-cicd · 2026-03-12T07:23:02Z

PR_Github #38641 [ run ] completed with state SUCCESS. Commit: e3c5b1b
/LLM/main/L0_MergeRequest_PR pipeline #29970 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

nv-lschneider · 2026-03-12T22:04:29Z

/bot run --add-multi-gpu --disable-fail-fast

tensorrt-cicd · 2026-03-12T22:11:43Z

PR_Github #38785 [ run ] triggered by Bot. Commit: ebe11a0 Link to invocation

tensorrt-cicd · 2026-03-13T05:19:35Z

PR_Github #38785 [ run ] completed with state SUCCESS. Commit: ebe11a0
/LLM/main/L0_MergeRequest_PR pipeline #30099 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

nv-lschneider · 2026-03-13T22:28:42Z

/bot run --add-multi-gpu --disable-fail-fast

tensorrt-cicd · 2026-03-13T22:34:50Z

PR_Github #38914 [ run ] triggered by Bot. Commit: ebe11a0 Link to invocation

tensorrt-cicd · 2026-03-13T23:58:05Z

PR_Github #38914 [ run ] completed with state SUCCESS. Commit: ebe11a0
/LLM/main/L0_MergeRequest_PR pipeline #30220 completed with status: 'SUCCESS'

CI Report

Link to invocation

nv-lschneider · 2026-03-16T20:20:15Z

/bot run --add-multi-gpu --disable-fail-fast

tensorrt-cicd · 2026-03-16T20:26:52Z

PR_Github #39127 [ run ] triggered by Bot. Commit: a653d9a Link to invocation

tensorrt-cicd · 2026-03-17T03:01:55Z

PR_Github #39127 [ run ] completed with state FAILURE. Commit: a653d9a
/LLM/main/L0_MergeRequest_PR pipeline #30386 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

nv-lschneider · 2026-03-17T21:43:02Z

/bot run --add-multi-gpu --disable-fail-fast

tensorrt-cicd · 2026-03-17T21:51:12Z

PR_Github #39319 [ run ] triggered by Bot. Commit: f4fa84a Link to invocation

tensorrt-cicd · 2026-03-18T04:28:26Z

PR_Github #39319 [ run ] completed with state SUCCESS. Commit: f4fa84a
/LLM/main/L0_MergeRequest_PR pipeline #30565 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

…cateAndRegisterBuffer Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

…apper configs Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

…for ncclMemAlloc/cudaMalloc cleanup Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

…sterBuffer Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

… failures Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

…d of throwing Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

…edefining Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

…imiK2VL WAR Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

nv-lschneider · 2026-04-01T15:28:38Z

/bot run --add-multi-gpu --disable-fail-fast

tensorrt-cicd · 2026-04-01T15:35:04Z

PR_Github #41219 [ run ] triggered by Bot. Commit: 2ff284d Link to invocation

tensorrt-cicd · 2026-04-01T23:35:29Z

PR_Github #41219 [ run ] completed with state SUCCESS. Commit: 2ff284d
/LLM/main/L0_MergeRequest_PR pipeline #32179 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

nv-lschneider · 2026-04-02T13:16:33Z

/bot run --add-multi-gpu --disable-fail-fast

tensorrt-cicd · 2026-04-02T13:22:35Z

PR_Github #41436 [ run ] triggered by Bot. Commit: 2ff284d Link to invocation

tensorrt-cicd · 2026-04-02T14:37:10Z

PR_Github #41436 [ run ] completed with state SUCCESS. Commit: 2ff284d
/LLM/main/L0_MergeRequest_PR pipeline #32367 completed with status: 'SUCCESS'

CI Report

Link to invocation

lucaslie

lgtm, thank you for improving the implementation

…ndling (NVIDIA#12125) Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

…e sticky ncclMemAlloc errors NCCLWindowAllocator::allocateAndRegisterBuffer used a per-rank cudaMalloc'd flag + ncclAllReduce(min) to make a per-rank ncclMemAlloc failure non-hanging. The intent was that all ranks reach the collective allreduce regardless of ncclMemAlloc's outcome, so the failing rank can signal "abort" instead of silently skipping the collective. However the cudaMalloc for the rank-sync flag ran AFTER ncclMemAlloc. When ncclMemAlloc returns ncclUnhandledCudaError (the symptom in the MGMN Llama-3.3-70B TP=2 run), the CUDA runtime can be left with a sticky last-error. The next runtime CUDA call (cudaMalloc for the rank-sync flag) then inherits that error, TLLM_CUDA_CHECK throws, and the failing rank unwinds out of allocateAndRegisterBuffer without ever reaching the collective ncclAllReduce(min). Healthy ranks proceed into the collective and block, producing the hang the bug reports. Fix: - Pre-allocate rankSyncFlag before ncclMemAlloc, so the flag is always ready on every rank irrespective of what ncclMemAlloc does. - After a failed ncclMemAlloc, drain any sticky CUDA last-error via cudaGetLastError() so the subsequent cudaMemcpy/ncclAllReduce on the CUDA stream see a clean device state. Both changes are confined to allocateAndRegisterBuffer. On the success path they are no-ops; on the failure path they keep the failing rank on the same collective control flow as healthy ranks, restoring the graceful fallback to plain NCCL allreduce that PR NVIDIA#12125 originally introduced. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>

nv-lschneider requested a review from a team as a code owner March 11, 2026 22:47

nv-lschneider requested a review from hyukn March 11, 2026 22:47

github-actions Bot assigned nv-lschneider Mar 11, 2026

coderabbitai Bot changed the title ~~[TRTLLM-11237][fix] @coderabbitai title~~ [TRTLLM-11237][fix] [fix] Synchronize NCCL memory allocation error handling Mar 11, 2026

coderabbitai Bot reviewed Mar 11, 2026

View reviewed changes

Comment thread cpp/tensorrt_llm/common/ncclUtils.cpp Outdated

schetlur-nv requested review from schetlur-nv March 11, 2026 22:53

nv-lschneider force-pushed the lschneider/nccl-symmetric-oom-resistant branch from e3c5b1b to ebe11a0 Compare March 12, 2026 22:02

nv-lschneider requested a review from Tabrizian March 14, 2026 00:13

Tabrizian reviewed Mar 16, 2026

View reviewed changes

Comment thread cpp/tensorrt_llm/common/ncclUtils.cpp Outdated

Comment thread tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

nv-lschneider force-pushed the lschneider/nccl-symmetric-oom-resistant branch from a653d9a to f4fa84a Compare March 17, 2026 21:42

nv-lschneider requested review from a team as code owners March 17, 2026 21:42

nv-lschneider requested a review from HuiGao-NV March 17, 2026 21:42

nv-lschneider added 13 commits April 1, 2026 10:28

fix naming, defensive default, and check ncclAllReduce return in allo…

71da53a

…cateAndRegisterBuffer Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

add error checks on cudaMemcpy calls for rank-sync flag

042ddc8

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

check every NCCL and cuda call

0e08f36

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

preallocate NCCL window buffers at 2x max model size for headroom

cc6311f

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

expose text_config attrs in KimiK25Config for model_config compat

c604a8e

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

gracefully skip allreduce pre-alloc when hidden_size absent on VLM wr…

19785df

…apper configs Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

make allreduce pre-alloc config querying graceful with warning

fb1c927

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

add TLLM_CUDA_CHECK_WARN/TLLM_NCCL_CHECK_WARN macros and RAII guards …

abf52c5

…for ncclMemAlloc/cudaMalloc cleanup Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

use TLLM_NCCL_CHECK/TLLM_CUDA_CHECK macros throughout allocateAndRegi…

047fad7

…sterBuffer Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

harden collective error handling: warn on H2D, throw on allreduce/D2H…

a458fd5

… failures Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

warn and return empty buffer on ncclCommWindowRegister failure instea…

a2d27fb

…d of throwing Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

fix TLLM_NCCL_CHECK redefinition: use multiDeviceUtils.h instead of r…

098b0d6

…edefining Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

[None][fix] use get_text_config() for allreduce hidden_size; remove K…

2ff284d

…imiK2VL WAR Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

nv-lschneider force-pushed the lschneider/nccl-symmetric-oom-resistant branch from cd42163 to 2ff284d Compare April 1, 2026 15:28

lucaslie approved these changes Apr 2, 2026

View reviewed changes

Tabrizian merged commit bba6ca4 into NVIDIA:main Apr 2, 2026
5 checks passed

karen-sy pushed a commit to karen-sy/TensorRT-LLM that referenced this pull request Apr 7, 2026

[TRTLLM-11237][fix] [fix] Synchronize NCCL memory allocation error ha…

00db7a1

…ndling (NVIDIA#12125) Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

Uh oh!

Conversation

nv-lschneider commented Mar 11, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nv-lschneider commented Mar 11, 2026

Uh oh!

tensorrt-cicd commented Mar 11, 2026

Uh oh!

tensorrt-cicd commented Mar 12, 2026

Uh oh!

nv-lschneider commented Mar 12, 2026

Uh oh!

tensorrt-cicd commented Mar 12, 2026

Uh oh!

tensorrt-cicd commented Mar 13, 2026

Uh oh!

nv-lschneider commented Mar 13, 2026

Uh oh!

tensorrt-cicd commented Mar 13, 2026

Uh oh!

tensorrt-cicd commented Mar 13, 2026

Uh oh!

Uh oh!

Uh oh!

nv-lschneider commented Mar 16, 2026

Uh oh!

tensorrt-cicd commented Mar 16, 2026

Uh oh!

tensorrt-cicd commented Mar 17, 2026

Uh oh!

nv-lschneider commented Mar 17, 2026

Uh oh!

tensorrt-cicd commented Mar 17, 2026

Uh oh!

tensorrt-cicd commented Mar 18, 2026

Uh oh!

nv-lschneider commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

nv-lschneider commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

lucaslie left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

nv-lschneider commented Mar 11, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Mar 11, 2026 •

edited

Loading