[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)
✨ 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."
| 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.
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 |
5483275 to
49c5735
Compare
|
/bot run |
|
PR_Github #65526 [ run ] triggered by Bot. Commit: |
|
PR_Github #65526 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65549 [ run ] triggered by Bot. Commit: |
|
PR_Github #65549 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Two things to settle before this lands.
The comment above the new guard says the threshold is SIZE_MAX without NVLink/MNNVL and that allocating unconditionally therefore hangs. That isn't the configuration this bug was filed from: on a 2-GPU H100 node NVLink is supported, so minRegistrationThreshold is the linear-model value (roughly 146k elements at 2 ranks), and what actually changes for a 32-element allreduce is that the output no longer requests a window at all. Please rewrite the comment to describe that path — as written it points the next reader at an invariant that doesn't apply here.
Second, the PR description covers only the C++ change. The branch also marks all of test_row_linear_norm_fusion post_merge and adds a 2-GPU post_merge block to l0_dgx_h100.yml. Demoting the whole test rather than just the flaky parametrization is the right call — the other parameters plausibly hit the same hang and we've just not caught a pre-merge repro — but it should be stated in the description and commit message so the coverage move is deliberate on the record. It also means the first real signal on this fix comes from a post-merge run, so that stage is worth watching explicitly after merge.
| // 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.
windowBuffer0.isValid() is rank-local state (an allocator pool hit on this rank's input pointer), while requestBuffer inside createNCCLWindowTensor can fall through to allocateAndRegisterBuffer, which is collective. If one rank's input happens to be pool-registered and another's isn't while bufferSizeBytes < minRegistrationThreshold, the ranks disagree about entering that allocation — the same hang class this PR is closing.
The windowBuffer0.isValid() || clause buys nothing that the size test doesn't already cover: whenever the input got a window through this function, the size was already >= minRegistrationThreshold. Dropping it leaves a purely rank-uniform predicate:
if (bufferSizeBytes >= minRegistrationThreshold)| torch::Tensor outputTensor = windowBuffer1.isValid() ? normOut : torch::empty_like(inputTensor); | ||
| 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. |
There was a problem hiding this comment.
This comment describes the !mIsNVLINKSupported && !mIsMNNVLSupported → SIZE_MAX case, which does not hold on the 2-GPU H100 stage this bug came from (NVLink is supported there, so the threshold is the fitted linear-model value — about 146k elements at 2 ranks). For the failing case the effective change is that a small message no longer registers an output window, matching what the input path a few lines above already does. Please state that instead.
kris1025
changed the title
Signed-off-by: linquanh <linquanh@nvidia.com>
|
/bot run |
|
PR_Github #65771 [ run ] triggered by Bot. Commit: |
|
PR_Github #65771 [ run ] completed with state
|
| - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off-cpp_mamba_cache] | ||
| - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-python_mamba_cache] | ||
| - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-cpp_mamba_cache] | ||
| - condition: |
There was a problem hiding this comment.
We don't have a test stage matching those configs. You need to add a new test stage like DGX_H100-2_GPUs-PyTorch-Others-Post-Merge-1 here - https://github.com/NVIDIA/TensorRT-LLM/blob/main/jenkins/L0_Test.groovy#L5443.
There was a problem hiding this comment.
The new stage is added now.
Signed-off-by: linquanh <linquanh@nvidia.com>
|
/bot run |
|
PR_Github #66119 [ run ] triggered by Bot. Commit: |
|
PR_Github #66119 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66467 [ run ] triggered by Bot. Commit: |
12 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 to the input buffer, but always calledcreateNCCLWindowTensor()for the output buffer. That meant small messages could skip input window registration while still entering the output window allocation/registration path.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.This keeps the symmetric-memory fast path available for eligible buffers while allowing small messages to avoid the extra window allocation/registration path. The threshold still honors
TLLM_NCCL_MIN_REGISTRATION.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]no longer hangs.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.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.