[None][test] Probe: full-warp hidden_size in allreduce CI smoke test#16333
[None][test] Probe: full-warp hidden_size in allreduce CI smoke test#16333MrGeva wants to merge 9 commits into
Conversation
Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
The fused all-reduce residual/RMSNorm kernel derives its block size from the number of 128-bit accesses. For the AutoDeploy test model, hidden_size=64 with a 16-bit dtype produces an eight-thread block. blockReduceSumV2 assumes a complete warp, so its shuffle named inactive lanes and triggered undefined CUDA behavior. On B200 this caused the TWOSHOT CUDA-graph warmup to hang. The production fix uses an active-mask warp reduction for sub-warp blocks while retaining the existing reduction for full-warp and multi-warp launches. The native regression test covers the exact 8192 x 64 TWOSHOT fused-residual-RMSNorm geometry, and the complete B200 custom-ops directory now passes. Remove the temporary fault-handler, logger tee, environment logging, and forced TWOSHOT debug level used to diagnose the CI hang. Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
… single-warp RMSNorm reduction blockReduceSum() unconditionally added a __syncthreads() plus shared-memory cross-warp step, even for single-warp blocks (blockDim.x <= 32), where the preceding shuffle-based warpReduceSum() already reconverges every participating lane. For the ONESHOT strategy, whose Lamport all-reduce relies on each thread independently busy-waiting on its own peer buffer, that unnecessary barrier turns per-token progress into a block-wide synchronization point and inflates tail latency under contention, observed as a ~300s CI hang for hidden_size=64/ONESHOT (the geometry used by the AutoDeploy multigpu smoke test's tiny Llama config) that then poisoned the following TWOSHOT case. Skip the shared-memory/__syncthreads() path entirely when the block fits in a single warp, matching the original validated fix's behavior for that case, while keeping the broader multi-warp partial-block coverage. Also add a ONESHOT variant of the existing hidden_size=64 native regression case, mirroring the TWOSHOT one, since the gap in that coverage is what let this regression through. Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Diagnostic-only change: overrides the tiny test model's hidden_size from 64 (which launches an 8-thread, partial-warp block in the fused all-reduce RMSNorm kernel) to 256 (exactly one full warp), so the smoke test no longer exercises the partial-warp reduction path at all. Two prior CI runs of the base branch (agent/fix-ad-allreduce-ci) failed with an identical CUBLAS_STATUS_EXECUTION_FAILED error in an unrelated cuBLAS GEMM (torch_linear_simple), on the same node (nsc-svg-slurm-1-gpu-120), at a different allreduce strategy each time. This branch isolates whether that failure is specific to the partial-warp kernel path or is unrelated node/infra flakiness that also reproduces on the full-warp path. Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1" |
|
PR_Github #58992 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThe change adds warp- and block-level reduction helpers, expands AllReduce CUDA coverage, moves end-to-end strategy testing into a multi-GPU smoke test, and removes obsolete custom-ops benchmark waives. ChangesAllReduce validation
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Pytest
participant FlashInfer
participant DatasetPrep
participant CliRunner
participant ThroughputCLI
Pytest->>FlashInfer: prewarm JIT kernels
Pytest->>DatasetPrep: generate synthetic dataset
Pytest->>CliRunner: invoke benchmark with allreduce_strategy
CliRunner->>ThroughputCLI: run multi-GPU throughput test
ThroughputCLI-->>CliRunner: return exit status
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu (1)
413-437: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse the kernel norm for FP8
Outvariants
ForkARResidualRMSNormOutFP8Quant, quantizem_norm_outinstead ofref_output;compare(..., atol=0)is exact, so the independent host reference can drift at FP8 rounding boundaries. Keepref_outputonly for the non-Outpatterns whereHasNormOut<Pattern>is false.🤖 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/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu` around lines 413 - 437, Update the FP8 branch in the quantization test to use m_norm_out as the quantization input for kARResidualRMSNormOutFP8Quant patterns, while retaining ref_output for non-Out patterns where HasNormOut<Pattern> is false. Preserve the existing exact comparison behavior and output validation.
🧹 Nitpick comments (3)
tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py (3)
34-38: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCustom
TimeoutErrorshadows the builtin, which can hide unrelated timeouts.
except TimeoutErrorat Line 306 only catches this module's local class. Ifrunner.invoke(...)ever raises a genuine builtinTimeoutError(e.g. from an unrelated socket/IO timeout deep in the LLM API stack), it will propagate uncaught rather than being handled by this block — and worse, it will print as "TimeoutError" in both cases, making CI failures hard to disambiguate. Ruff'sA001(builtin shadowing) isn't enforced in this repo's select list, so this won't be caught by CI lint, but it's still worth a distinct name for debuggability.♻️ Proposed rename
-class TimeoutError(Exception): +class BenchmarkTimeoutError(Exception): """Exception raised when a test times out.""" - - passAnd update the
raise/exceptsites accordingly.Based on learnings, "Ruff's
[tool.ruff.lint].selectenables onlyD, E, F, I, PLE, W" and does not include theA(flake8-builtins) category, so this will not be flagged by CI lint.Also applies to: 306-306
🤖 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 `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py` around lines 34 - 38, Rename the module-local TimeoutError class to a distinct name such as TestTimeoutError, then update all corresponding raise and except references in the test, including the handling block near runner.invoke. Preserve the existing timeout behavior while allowing builtin TimeoutError exceptions to remain distinguishable.Source: Learnings
88-122: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueLock file handle isn't released via
finally; an exception between acquire and close leaks the flock for the process lifetime.
lock_fis opened andflock'd at Lines 94-95, but thetorch.no_grad()block andlock_f.close()at Line 122 aren't wrapped intry/finally. If anything outside the two narrowly-scoped innertry/except Exceptionblocks raises (e.g.torch.cuda.empty_cache(), or an uncaught error type), the lock is never explicitly released for the remainder of the process, which could stall a concurrent xdist worker waiting on the sameflashinfer_jit_prewarm.lock. CPython's refcounting will eventually close the fd on GC, but that's incidental, not guaranteed.♻️ Proposed fix
lock_path = pathlib.Path(tempfile.gettempdir()) / "flashinfer_jit_prewarm.lock" lock_f = open(lock_path, "w") fcntl.flock(lock_f.fileno(), fcntl.LOCK_EX) - # Create dummy tensors to trigger kernel JIT compilation - with torch.no_grad(): - device = torch.device("cuda:0") - ... - torch.cuda.empty_cache() - if lock_f is not None: - lock_f.close() + try: + # Create dummy tensors to trigger kernel JIT compilation + with torch.no_grad(): + device = torch.device("cuda:0") + ... + torch.cuda.empty_cache() + finally: + if lock_f is not None: + lock_f.close()🤖 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 `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py` around lines 88 - 122, Wrap the lock-protected prewarm section in a try/finally so the file handle acquired in the lock_f setup is always closed, including when torch.no_grad(), torch.cuda.empty_cache(), or other uncaught operations raise. Keep the existing page and sampling prewarm behavior unchanged, and perform the lock_f.close() cleanup in the finally path only when a handle was successfully opened.
133-133: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
model_namehardcoded twice.
"meta-llama/Meta-Llama-3.1-8B-Instruct"is duplicated between theshared_datasetfixture andtest_allreduce_strategies. Extracting a module-level constant avoids drift if the model is ever changed.Also applies to: 219-219
🤖 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 `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py` at line 133, Replace the duplicated model string in the shared_dataset fixture and test_allreduce_strategies with a single module-level constant, and reference that constant in both locations so future model changes stay synchronized.
🤖 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.
Inline comments:
In `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`:
- Around line 220-233: Coverage needs follow-up outside this diagnostic change:
after confirming whether failures are infrastructure-related, restore
partial-warp coverage in the test configuration around get_small_model_config,
either by reverting hidden_size to 64 or parametrizing the smoke test to
exercise both hidden_size=64 and 256. Preserve the full-warp diagnostic case
while ensuring the original partial-warp regression path is covered before
merging.
- Around line 211-215: The docstring for the test’s allreduce_strategy argument
must include SYMM_MEM alongside the existing listed strategies. Update that Args
description without changing the parametrization or test behavior.
- Around line 40-63: Replace the signal.alarm-based timeout context manager
around the strategy test with subprocess isolation for each CliRunner.invoke()
call. Run every strategy invocation in its own process, enforce the timeout from
the parent, terminate or kill the child when it expires, and propagate the
child’s result or failure without allowing a hung NCCL/CUDA call to affect
subsequent parametrized strategies.
---
Outside diff comments:
In `@cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu`:
- Around line 413-437: Update the FP8 branch in the quantization test to use
m_norm_out as the quantization input for kARResidualRMSNormOutFP8Quant patterns,
while retaining ref_output for non-Out patterns where HasNormOut<Pattern> is
false. Preserve the existing exact comparison behavior and output validation.
---
Nitpick comments:
In `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`:
- Around line 34-38: Rename the module-local TimeoutError class to a distinct
name such as TestTimeoutError, then update all corresponding raise and except
references in the test, including the handling block near runner.invoke.
Preserve the existing timeout behavior while allowing builtin TimeoutError
exceptions to remain distinguishable.
- Around line 88-122: Wrap the lock-protected prewarm section in a try/finally
so the file handle acquired in the lock_f setup is always closed, including when
torch.no_grad(), torch.cuda.empty_cache(), or other uncaught operations raise.
Keep the existing page and sampling prewarm behavior unchanged, and perform the
lock_f.close() cleanup in the finally path only when a handle was successfully
opened.
- Line 133: Replace the duplicated model string in the shared_dataset fixture
and test_allreduce_strategies with a single module-level constant, and reference
that constant in both locations so future model changes stay synchronized.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5cb3e0ca-e070-4e84-8de6-be0710f94b5f
📒 Files selected for processing (6)
cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cucpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cucpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cutests/integration/test_lists/waives.txttests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.pytests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py
💤 Files with no reviewable changes (2)
- tests/integration/test_lists/waives.txt
- tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
| @contextmanager | ||
| def timeout(seconds): | ||
| """Context manager that raises TimeoutError if code block exceeds time limit. | ||
|
|
||
| Args: | ||
| seconds: Maximum time in seconds to allow the code block to run | ||
|
|
||
| Raises: | ||
| TimeoutError: If the code block execution exceeds the time limit | ||
| """ | ||
|
|
||
| def timeout_handler(signum, frame): | ||
| raise TimeoutError(f"Test execution exceeded {seconds} seconds timeout") | ||
|
|
||
| # Set the signal handler and alarm | ||
| old_handler = signal.signal(signal.SIGALRM, timeout_handler) | ||
| signal.alarm(seconds) | ||
| try: | ||
| yield | ||
| finally: | ||
| # Restore the old signal handler and cancel the alarm | ||
| signal.alarm(0) | ||
| signal.signal(signal.SIGALRM, old_handler) | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target test and nearby context.
python3 - <<'PY'
from pathlib import Path
path = Path("tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py")
text = path.read_text().splitlines()
for start in (1, 25, 35, 40, 60, 280, 300):
end = min(len(text), start + 35)
print(f"\n--- {path}:{start}-{end} ---")
for i in range(start, end + 1):
if i <= len(text):
print(f"{i:4d}: {text[i-1]}")
PY
# Find relevant symbols/usages in the test file.
rg -n "timeout|SIGALRM|CliRunner|invoke|threadleak|multiprocessing|subprocess|pytest\.mark" tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py
# Locate adjacent test helpers or fixtures that may affect teardown/isolation.
rg -n "threadleak|CliRunner|allreduce|strategy|timeout" tests/unittest/auto_deploy/multigpu/smoke -g '*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 13637
Run each strategy in a separate process
signal.alarm() around CliRunner.invoke() still depends on the main Python process regaining control, so a hang inside NCCL/CUDA can outlive the timeout and leave the parametrized test in a bad state for the next strategy. Move each invocation to a subprocess and kill it on timeout instead of relying on SIGALRM.
🤖 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 `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`
around lines 40 - 63, Replace the signal.alarm-based timeout context manager
around the strategy test with subprocess isolation for each CliRunner.invoke()
call. Run every strategy invocation in its own process, enforce the timeout from
the parent, terminate or kill the child when it expires, and propagate the
child’s result or failure without allowing a hung NCCL/CUDA call to affect
subsequent parametrized strategies.
| Args: | ||
| llm_root: Root directory fixture | ||
| shared_dataset: Shared dataset fixture (prepared once for all test runs) | ||
| allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL) | ||
| """ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Docstring omits SYMM_MEM.
The Args section lists strategies "(AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL)" but the parametrize list (Line 191) also includes SYMM_MEM.
📝 Proposed fix
- allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL)
+ allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL, SYMM_MEM)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| Args: | |
| llm_root: Root directory fixture | |
| shared_dataset: Shared dataset fixture (prepared once for all test runs) | |
| allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL) | |
| """ | |
| Args: | |
| llm_root: Root directory fixture | |
| shared_dataset: Shared dataset fixture (prepared once for all test runs) | |
| allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL, SYMM_MEM) | |
| """ |
🤖 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 `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`
around lines 211 - 215, The docstring for the test’s allreduce_strategy argument
must include SYMM_MEM alongside the existing listed strategies. Update that Args
description without changing the parametrization or test behavior.
| # Override hidden_size to a multiple of one warp's worth of 128-bit accesses (32 threads * | ||
| # 8 fp16 elements/access = 256) so the fused all-reduce/RMSNorm kernel never launches a | ||
| # partial-warp block. This isolates whether CI failures are specific to the partial-warp | ||
| # code path or are unrelated infra flakiness that also affects the full-warp path. | ||
| config = get_small_model_config( | ||
| model_name, | ||
| model_kwargs={ | ||
| "num_hidden_layers": 1, | ||
| "hidden_size": 256, | ||
| "intermediate_size": 256, | ||
| "num_attention_heads": 2, | ||
| "num_key_value_heads": 1, | ||
| }, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
QA note: this diagnostic change removes coverage of the partial-warp kernel path.
As per the PR objectives, bumping hidden_size from 64 to 256 intentionally moves the kernel launch from a partial warp (8 threads) to a full warp, which is the point of this diagnostic probe and matches the stated intent that this PR is not meant to merge. As a QA follow-up: if this diagnostic confirms the CI failures are infra-related rather than partial-warp-specific, this smoke test should be restored to (or parametrized to include) hidden_size=64 before merging, otherwise the original partial-warp regression coverage introduced for the earlier fix (PR #16111) is silently dropped. Flagging as a required follow-up outside this PR rather than a blocking defect here.
Based on path instructions, "tests/**: Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM" and to state "whether coverage is sufficient, insufficient, or needs follow-up outside the PR."
🤖 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 `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`
around lines 220 - 233, Coverage needs follow-up outside this diagnostic change:
after confirming whether failures are infrastructure-related, restore
partial-warp coverage in the test configuration around get_small_model_config,
either by reverting hidden_size to 64 or parametrizing the smoke test to
exercise both hidden_size=64 and 256. Preserve the full-warp diagnostic case
while ensuring the original partial-warp regression path is covered before
merging.
Source: Path instructions
|
PR_Github #58992 [ run ] completed with state |
2 participants
Description
Diagnostic-only PR, not intended to merge. Branched from
agent/fix-ad-allreduce-ciatd5163b1c(PR #16111).Changes the AutoDeploy allreduce-strategy smoke test's tiny model
hidden_sizefrom 64 to 256, so the fused all-reduce/RMSNorm kernel'sblockReduceSumalways launches exactly one full warp (32 threads) instead of the partial-warp (8-thread) block that PR #16111's kernel fix specifically targets. Everything else (kernel code, test structure, strategies tested) is identical to #16111's current head.Motivation
Two consecutive
DGX_B200-4_GPUs-AutoDeploy-1CI runs of #16111 (builds 47476, 47504) failed with an identicalCUBLAS_STATUS_EXECUTION_FAILEDerror in an unrelated cuBLAS GEMM (torch_linear_simple/trtllm_cublas_mm), on the same physical node (nsc-svg-slurm-1-gpu-120), at a different allreduce strategy each time (AUTO, thenNCCL) with all other strategies passing around it.This PR isolates whether that failure is specific to the partial-warp kernel path touched by #16111, or is unrelated node/infra flakiness that also reproduces on the untouched full-warp path.
Test Coverage
Same
test_allreduce_strategiessmoke test as #16111, all 6 strategies, just withhidden_size=256.PR Checklist
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests