Repository navigation
[https://nvbugs/6501404][fix] Request the output window only when bufferSizeBytes >= minRegistrationThreshold - #16911
trtllm-agent wants to merge 6 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. ℹ️ Recent review info
Walkthrough
ChangesSymmetric all-reduce output handling
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1
✨ Finishing Touches: 💡 1
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/thop/allreduceOp.cpp (1)
564-569: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the window allocation result
const.
windowOutputandwindowBuffer1are never reassigned.Proposed fix
- auto [windowOutput, windowBuffer1] = createNCCLWindowTensor(rawComm, input.sizes(), input.scalar_type()); + auto const [windowOutput, windowBuffer1] + = createNCCLWindowTensor(rawComm, input.sizes(), input.scalar_type());As per coding guidelines, “declare unmodified variables
const.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/thop/allreduceOp.cpp` around lines 564 - 569, Declare the structured-binding variables windowOutput and windowBuffer1 as const in the createNCCLWindowTensor result within the surrounding allreduce operation, preserving the existing validity check and outputTensor assignment.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cpp/tensorrt_llm/thop/allreduceOp.cpp`:
- Around line 564-569: Declare the structured-binding variables windowOutput and
windowBuffer1 as const in the createNCCLWindowTensor result within the
surrounding allreduce operation, preserving the existing validity check and
outputTensor assignment.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 99f0fbf3-1e74-44af-92b9-b0dbc4a1acd8
📒 Files selected for processing (1)
cpp/tensorrt_llm/thop/allreduceOp.cpp
f74413d to
2593d10
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Making the output allocation follow the same gate as the input is the right consistency fix regardless of the bug — with a 64-byte tensor and a ~290 KB threshold at 2 ranks, the old code skipped registration for the input and then ran a full collective allocateAndRegisterBuffer for the output, which is both inconsistent and wasteful.
What I don't follow is the causal story. The failure in nvbugs/6501404 was on 2×H100, where NVLink is present, so minRegistrationThreshold is never SIZE_MAX on that path and the mechanism the new comment describes never fires there. allocateAndRegisterBuffer is also written specifically so every rank reaches the min-allreduce even when ncclMemAlloc fails asymmetrically. So this change plausibly removes a collective from the hot path, but it isn't shown to remove the one that hung — and the linked bug records that the failure could not be reproduced.
So I'd land the gating change on its own merits and keep the waiver until there's evidence the hang is actually gone. A concrete way to get that evidence without holding up this PR: open a separate draft PR that (a) removes the waiver, (b) adds instrumentation around the symmetric allreduce path (log per-rank windowBuffer0.isValid(), bufferSizeBytes, minRegistrationThreshold, and entry/exit of allocateAndRegisterBuffer), and (c) runs only the failing stage with the test list trimmed to test_row_linear_norm_fusion (and neighbors if needed) so repeated runs are cheap on capacity and turnaround. Re-run that until the hang reproduces; the instrumentation then tells you which rank diverged and why. Once you have a pre-fix hang and a post-fix pass on the same stage, dropping the waiver here is easy to justify.
| // ncclAllReduce plus cudaStreamSynchronize inside allocateAndRegisterBuffer cannot | ||
| // complete, so allocating here unconditionally hangs every rank. | ||
| torch::Tensor outputTensor; | ||
| if (windowBuffer0.isValid() || bufferSizeBytes >= minRegistrationThreshold) |
There was a problem hiding this comment.
This if guards a collective. createNCCLWindowTensor → requestBuffer → allocateAndRegisterBuffer does an ncclAllReduce on the sync flag plus ncclCommWindowRegister, so every rank has to make the same decision here or the ones that enter will wait forever for the ones that didn't.
Of the two operands, only one is safe in that role. bufferSizeBytes >= minRegistrationThreshold is computed from the same inputs on every rank, so it's uniform. windowBuffer0.isValid() is not: it comes from allocator.searchBuffer(comm, input.data_ptr()) or from a pool best-fit inside requestBuffer, both of which depend on rank-local allocator state. If the input ends up registered on rank 0 but not on rank 1 while the size is below the threshold, rank 0 calls the collective and rank 1 skips it — a hang that the old unconditional call could not produce.
If the intent is "the input is window-backed, so the output should be too", either gate on the rank-uniform condition alone, or add a comment explaining why windowBuffer0.isValid() is guaranteed to agree across ranks.
| void* outputPtr = windowBuffer1.isValid() ? windowBuffer1.ptr : outputTensor.data_ptr(); | ||
| if (!windowBuffer1.isValid()) | ||
| // Use a window-backed output buffer under the same threshold gate as the input above. | ||
| // minRegistrationThreshold is SIZE_MAX without NVLink/MNNVL, where the collective |
There was a problem hiding this comment.
The comment explains the fix with a mechanism that can't apply to the reported failure. It says minRegistrationThreshold is SIZE_MAX without NVLink/MNNVL and that the collective inside allocateAndRegisterBuffer therefore can't complete — but the bug reproduced on 2×H100, which has NVLink, so that branch was never taken. allocateAndRegisterBuffer is also written so that all ranks reach the min-allreduce even when ncclMemAlloc fails on only some of them.
Suggest describing what the change actually does instead of asserting an unproven hang cause, e.g.: "Allocate an output window only when the input path also registered a window. This keeps window use all-or-nothing and lets small messages skip the registration collective entirely."
| @@ -351,7 +351,6 @@ unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_ | |||
| unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" SKIP (https://nvbugs/6464169) | |||
There was a problem hiding this comment.
The bug this waiver points at was never reproduced, and the fix above doesn't clearly explain the observed hang on 2×H100. Un-waiving here risks the pre-merge job going intermittently red again with no new information.
I'd keep the waiver until there's a run that hangs without this patch and passes with it. To get there cheaply: put the un-waive plus instrumentation of the symmetric allreduce path (per-rank windowBuffer0.isValid(), bufferSizeBytes, minRegistrationThreshold, entry/exit of allocateAndRegisterBuffer) in a throwaway draft PR, trim the test list for the failing stage down to this test, and run that stage repeatedly until it hangs. That keeps the capacity cost and turnaround per attempt low, and the logs will show which rank diverged. Then drop the waiver here with the before/after runs linked.
There was a problem hiding this comment.
@brnguyen2 is it okay to merge the PR first. If the hang issue still exists, we can re-open the bug.
There was a problem hiding this comment.
Fair — the later commit marking the test post_merge changes my objection. With it out of the blocking pre-merge glob, un-waiving no longer risks turning pre-merge intermittently red, which was the whole basis for keeping the waiver. Merge it.
Two asks so this doesn't just go quiet: keep the bug open until a post-merge run on the 2-GPU H100 stage has actually exercised the un-waived test, and treat that stage as the signal rather than the pre-merge result on this PR. If it hangs again there, the reproduction data from that run is what we lacked the first time.
|
/bot run |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot kill |
|
/bot run |
|
PR_Github #64957 [ run ] triggered by Bot. Commit: |
|
PR_Github #64957 [ run ] completed with state
|
|
/bot run |
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
Semantic conflict reviewThe verdict of record is the Latest recorded state: Review completed: inconclusive for head Best-effort AI judgment for the recorded revisions. PASS, FAIL and INCONCLUSIVE may be incomplete or incorrect. PR authors and reviewers should independently verify the evidence and relevant behavior. This semantic review and its status/workflow are advisory, not required merge checks under current repository rules; other merge requirements still apply. Advisory status does not make a confirmed defect safe to ignore.
Processed request and reply comments are minimized to reduce timeline noise; they remain expandable for audit. |
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
a7161f4 to
4749fee
Compare
4749fee to
18b7cab
Compare
|
Bug engineer ownership @kris1025, as the engineer assigned to the linked bug, you are responsible for driving this bot PR to resolution. Please coordinate review, address feedback and required code changes, get CI passing, and verify the original failure before driving it to merge. If the PR is obsolete or superseded, close it with a reason. Keep this PR updated with blockers and the next step. |
dfca738 to
42aa8ad
Compare
…ration threshold runNCCLAllReduceSymmetric allocated its window-backed output buffer with an unconditional createNCCLWindowTensor, bypassing the minRegistrationThreshold gate that the input path a few lines above already applies. That threshold is set to SIZE_MAX when neither NVLink nor MNNVL is supported, so on such topologies the collective ncclAllReduce plus cudaStreamSynchronize inside allocateAndRegisterBuffer never completes and every rank hangs, which the CI stage reports as "Test terminated unexpectedly". Apply the same gate to the output allocation: request a window buffer only when the input already obtained one or the message is at least as large as the threshold. Gating on the threshold rather than on mIsNVLINKSupported / mIsMNNVLSupported keeps the symmetric-memory fast path enabled wherever it works and still honors TLLM_NCCL_MIN_REGISTRATION. The pre-existing invalid-buffer fallback to torch::empty_like covers the skipped case, so no new error path is introduced. Signed-off-by: handongl <handongl@nvidia.com>
test_row_linear_norm_fusion needs 2 GPUs and has been flaky under nvbugs/6501404. Mark it @pytest.mark.post_merge so it is dropped from the blocking pre_merge 2-GPU glob (unittest/_torch/multi_gpu -m "not post_merge") on l0_dgx_h100, and add a matching 2-GPU post_merge condition block that collects it via -m "post_merge" so coverage is retained in the non-blocking post_merge stage. Signed-off-by: linquanh <linquanh@nvidia.com>
Signed-off-by: linquanh <linquanh@nvidia.com>
Signed-off-by: linquanh <linquanh@nvidia.com>
Signed-off-by: handongl <handongl@nvidia.com> [nvbugs/6501404][chore] Sort ray_orchestrator waive entry The file-contents-sorter pre-commit hook requires waives.txt to be byte-sorted. The newly added ray_orchestrator waive was placed after the sampler entry; reorder it before to satisfy the hook. Signed-off-by: linquanh <linquanh@nvidia.com>
…RMS_NORM release Rebasing onto main brings in the RMS_NORM release of the internal output window in runNCCLAllReduceSymmetric (NVIDIA#20232), which reads windowBuffer1 after the fused fallback. The threshold-gated output allocation declared windowBuffer1 inside the gate, so that read was out of scope. Declare windowBuffer1 at function scope. It stays invalid unless the gate allocates a valid window, which is exactly when outputTensor is window-backed and must be released, so the release keeps its meaning. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
42aa8ad to
baeb776
Compare
Summary
This PR updates the NCCL symmetric allreduce path so the output tensor no longer unconditionally requests a window-backed allocation.
Previously,
runNCCLAllReduceSymmetric()applied the registration threshold when deciding whether to copy the input into a new window, but always calledcreateNCCLWindowTensor()for the output. Small messages could skip that input allocation while still entering the output window allocation/registration path. Window allocation is a collective, so taking that path on some ranks and not others is the failure mode this change avoids.The new behavior requests a window-backed output tensor only when:
bufferSizeBytes >= minRegistrationThresholdIf no valid output window is obtained, the existing
torch::empty_like(...)fallback is used. The threshold still honorsTLLM_NCCL_MIN_REGISTRATION.This is not the same predicate as the input path. The input can already be window-backed for a sub-threshold message:
searchBuffer()returns a validwindowBuffer0when the caller reused a previously registered tensor, and that pointer is used assendbufregardless of size. The output gate does not includewindowBuffer0.isValid(). In that casencclAllReducesees a registered send buffer and a plain recv buffer. All ranks still take the same branches when SPMD inputs match, so this is a send/recv mix on the symmetric kernel (possible degrade / fallback), not an unpairedcreateNCCLWindowTensorhang.Test coverage
test_row_linear_norm_fusionaspost_merge.unittest/_torch/multi_gpu -m "post_merge" TIMEOUT (90)unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2]Test plan
test_row_linear_norm_fusion[2-hidden:16-seqlen:2].createNCCLWindowTensorfor the output.Links
Dev Engineer Review
runNCCLAllReduceSymmetricto request the output window only whenbufferSizeBytes >= minRegistrationThreshold. This is intentionally size-only and does not mirror the input'swindowBuffer0.isValid()reuse path.torch::empty_likefallback when symmetric-memory allocation is unavailable.outputTensor.data_ptr()directly toncclAllReduce.TLLM_NCCL_MIN_REGISTRATIONbehavior.unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2].QA Engineer Review
tests/integration/test_lists/waives.txt.