[https://nvbugs/6501404][fix] Request the output window only when `windowBuffer0.isValid() ||… - #16911
[https://nvbugs/6501404][fix] Request the output window only when `windowBuffer0.isValid() ||…#16911trtllm-agent wants to merge 3 commits into
Conversation
Walkthrough
ChangesSymmetric all-reduce output handling
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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
…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>
Signed-off-by: handongl <handongl@nvidia.com>
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."
| unittest/_torch/misc/test_autotuner.py::test_cutedsl_nvfp4_heuristic_matches_full_sweep SKIP (https://nvbugs/6490028) | ||
| unittest/_torch/misc/test_share_tensor.py::TestShareTensor::test_share_tensor_different_dtypes SKIP (https://nvbugs/6418021) | ||
| unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912) | ||
| 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.
|
/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 kill |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run |
|
PR_Github #65090 [ kill ] triggered by Bot. Commit: |
|
PR_Github #65090 [ kill ] completed with state |
|
PR_Github #65091 [ run ] triggered by Bot. Commit: |
8 participants
Summary
windowBuffer0.isValid() || bufferSizeBytes >= minRegistrationThreshold, reusing the existing invalid-buffer fallback to torch::empty_like; gating on the threshold rather than the raw capability flags keeps the fast path enabled where it works and honors TLLM_NCCL_MIN_REGISTRATION.Test plan
Links
Dev Engineer Review
runNCCLAllReduceSymmetricto request the output window only whenwindowBuffer0.isValid() || bufferSizeBytes >= minRegistrationThreshold.torch::empty_likefallback when symmetric-memory allocation is unavailable.outputTensor.data_ptr()directly toncclAllReduce.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.