[https://nvbugs/6501404][fix] Request the output window only when bufferSizeBytes >= minRegistrationThreshold - #16911
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⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Walkthrough
ChangesSymmetric all-reduce output handling
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description clearly explains the issue, solution, behavior for sub-threshold messages, fallback behavior, test coverage, and test plan. The repository checklist is not reproduced, but the core required information is complete. ✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
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 |
|
/bot run |
|
PR_Github #67805 [ run ] triggered by Bot. Commit: |
|
PR_Github #67805 [ run ] completed with state
|
|
/bot run |
|
/bot kill |
c4470d6 to
21c489f
Compare
|
/bot run |
|
PR_Github #68563 [ run ] triggered by Bot. Commit: |
|
PR_Github #68563 [ run ] completed with state
|
crazydemo
left a comment
crazydemo
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The C++ output-window gating is a reasonable consistency fix, but the PR's own validation is broken and it can't merge as-is: it is blocked/CHANGES_REQUESTED, the regression test for the bug it claims to fix stays waived, and it silently adds an unrelated waiver.
Issues
- [MAJOR]
tests/integration/test_lists/waives.txt:371- 6501404 waiver NOT removed despite PR claim; the fix is never exercised - [MAJOR]
tests/integration/test_lists/waives.txt:370- new undocumented waiver for test_row_linear[2-balanced] (6507113) - [MINOR]
cpp/tensorrt_llm/thop/allreduceOp.cpp:570- output gate drops thewindowBuffer0.isValid()branch the input path uses - [MINOR]
cpp/tensorrt_llm/thop/allreduceOp.cpp:588- collective now usesoutputTensor.data_ptr()instead ofwindowBuffer1.ptr
QA view
- Test coverage: missing - the validating test
test_row_linear_norm_fusion[2-hidden:16-seqlen:2]is still SKIPped in waives.txt, so the post_merge marker and new l0_dgx_h100 entry never run it. The changed allreduce path has no active test. - SM coverage: touches the sm90 H100 NVLink symmetric-memory path; the intended test targets sm90 2-GPU but is blocked by the retained waiver, so no arch is actually validated.
- Test code: contradictory (test both post_merge and waived); an unrelated waiver added instead of a fix; whole test demoted to post_merge rather than one parametrisation.
- Test time: small - new 2-GPU post_merge stage with a 90-min timeout, post-merge only.
- Needs
/qa-verify: yes - the fix's regression test does not run, a new waive is introduced, and test-list infra changed.
Does this actually fix 6501404?
Unclear. The change makes the output allocation follow the same threshold gate as the input, which is a sound consistency improvement. But per reviewer brnguyen2 the 2xH100 config has NVLink, so minRegistrationThreshold is never SIZE_MAX on that path and the hang mechanism the comment implies never fires; the original failure was never reproduced. And because the 6501404 test remains waived, there is no run proving the hang is gone. So this lands the gating change on its own merits but does not verify the bug is fixed.
Possible new issues
- If the input can be window-backed (
windowBuffer0.isValid()) for a sub-threshold message, the input uses a registered symmetric pointer while the output falls back to a plain tensor — mixed registered/unregistered buffers on the symmetric path. outputTensor.data_ptr()may not equal the previously-used registeredwindowBuffer1.ptr; if so NCCL writes to an unregistered address.
What I could not verify
- The full input-side condition and whether
windowBuffer0.isValid()can be true for sub-threshold messages (input block not shown in the diff). - Whether
createNCCLWindowTensorguaranteeswindowOutput.data_ptr() == windowBuffer1.ptr. - Runtime behaviour of the symmetric allreduce with a mixed registered/plain buffer pair.
Automated review by NVCortex Lite, run by @crazydemo.
| // registration threshold used by the input path. Smaller messages use a | ||
| // regular output tensor and skip the window allocation path. | ||
| torch::Tensor outputTensor; | ||
| if (bufferSizeBytes >= minRegistrationThreshold) |
There was a problem hiding this comment.
[MINOR] Output gate drops the windowBuffer0.isValid() branch used by input
The output window is now requested only when bufferSizeBytes >= minRegistrationThreshold. The release notes describe the intended condition as windowBuffer0.isValid() || bufferSizeBytes >= minRegistrationThreshold, matching the input path. As written, if the input is window-backed (windowBuffer0 valid) but the message is below the threshold, the input uses a registered symmetric pointer while the output falls back to a plain torch::empty_like pointer. Mixing a registered send buffer with an unregistered recv buffer in runNCCLAllReduceSymmetric is exactly the asymmetry class this code is sensitive to and could degrade or misbehave on the symmetric kernel. Confirm whether the input path can be window-backed for sub-threshold messages; if so, mirror the full condition here.
There was a problem hiding this comment.
The description is updated now.
crazydemo
left a comment
crazydemo
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: A bug fix whose regression test remains waived (so it does not reproduce or validate the original 6501404 hang), plus a newly added unrelated waiver and changes to test-list infrastructure. A human QA should confirm whether the 6501404 test should run and actually passes with the fix, and whether disabling test_row_linear[2-balanced] is intended.
Automated review by NVCortex Lite, run by @crazydemo.
80cda7c to
d9a01cc
Compare
|
/bot run |
|
PR_Github #69417 [ run ] triggered by Bot. Commit: |
|
PR_Github #69417 [ run ] completed with state
|
…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>
fa05136 to
15d744f
Compare
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>
15d744f to
a7161f4
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #70322 [ run ] triggered by Bot. Commit: |
|
PR_Github #70322 [ run ] completed with state
|
14 participants
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.