[https://nvbugs/6525010][fix] Fix TestLlama3_3_70BInstruct::test_nvfp… - #17040
[https://nvbugs/6525010][fix] Fix TestLlama3_3_70BInstruct::test_nvfp…#17040liji-nv wants to merge 5 commits into
Conversation
…4_tp4 UB reuse The B200 full accuracy sequence constructs several PyTorch engines in one process. UserBuffersManager updated its configured allocation size for each engine but retained CUDA VMM regions registered by earlier engines because UserBufferAllocator::deallocate was a no-op. This allowed the Llama 3.1 8B TP4 case to leave a 64 MiB userbuffer in the singleton pool and the following Llama 3.3 70B TP4 case to reuse it for a 128 MiB request. The undersized mapping produced an invalid CUDA access that surfaced synchronously at the explicit graph-output copy when CUDA_LAUNCH_BLOCKING was enabled. Add collective userbuffer unregistration that synchronizes ranks, unbinds multicast mappings, unmaps and frees virtual address ranges, releases local and imported allocation handles, and clears communicator metadata. Track logical releases independently so callers may release in any order while physical registrations are retired in the required LIFO order. Expose UserBuffersManager shutdown through the Python binding and invoke it after model tensors and CUDA graphs are released during engine cleanup. Reject manager reinitialization while an earlier engine pool remains, preventing silent cross-engine reuse. Extend the multi-GPU allocation tests to shut down each manager instance and repeatedly reinitialize it more than MAX_REGIONS times before changing allocation size. This verifies that shutdown unregisters physical regions instead of only dropping pool bookkeeping. Remove the B200 waive for TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True]. The exact four-case sequence passed twice in one sbatch allocation with CUDA_LAUNCH_BLOCKING=1 (Slurm job 1554986), and the new userbuffer lifecycle regression passed in both iterations. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62665 [ run ] triggered by Bot. Commit: |
|
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 (5)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughUser-buffer allocation and collective registration now support explicit release. A manager shutdown API is exposed to Python and integrated into model-engine and executor shutdown. Tests cover cleanup, repeated reinitialization, region reuse, and shutdown failures. ChangesUser-buffer lifecycle
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant PyTorchModelEngine
participant UserBuffersManager
participant UserBufferAllocator
PyExecutor->>PyTorchModelEngine: request user-buffer shutdown
PyTorchModelEngine->>UserBuffersManager: shut down initialized manager
UserBuffersManager->>UserBufferAllocator: deallocate tracked buffers
UserBufferAllocator-->>UserBuffersManager: release allocator entries
UserBuffersManager-->>PyTorchModelEngine: return teardown status
PyTorchModelEngine-->>PyExecutor: report shutdown result
🚥 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.
Actionable comments posted: 1
🧹 Nitpick comments (2)
cpp/tensorrt_llm/kernels/userbuffers/userbuffers-host.cpp (1)
436-483: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTeardown correctly mirrors registration in reverse; strict LIFO invariant enforced via
TLLM_CHECK.The unmap/free/release sequence, the
mc_offsetrewind (with an assertion tying it back to the recordedmc_ptr[handle]), the hostfree()ofuchandles/peer_ptr, and thecudaMemsetclearing the peer-pointer table all correctly reverse the corresponding steps inregister_user_buffer_collective. This function is collective (an initialcudaDeviceSynchronize/ub_barrier, and a trailingub_barrier), so all ranks must call it in lockstep — see the related comment onUserBuffersManager::shutdown()for a cross-rank ordering concern this introduces.Consider adding a short comment above this function documenting that callers must release buffers strictly in reverse-allocation order (enforced by
TLLM_CHECK(handle == comm->free_region - 1)), since that invariant is easy to miss for future callers of this low-level API.🤖 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/kernels/userbuffers/userbuffers-host.cpp` around lines 436 - 483, Add a brief comment immediately above unregister_user_buffer_collective documenting that buffers must be released in strict reverse-allocation order. Reference the existing handle == comm->free_region - 1 TLLM_CHECK as the enforcement of this LIFO invariant, without changing teardown behavior.tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
2625-2698: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCleanup ordering and idempotent-retry design look correct.
Releasing cuda graphs / model refs / legacy
ub_buffersbeforerelease_gc()+ conditionalshutdown_userbuffers_manager()is the right order, and leaving_userbuffers_manager_initialized/_cleanup_doneunset on failure correctly enables a safe retry on the nextcleanup()call.Minor: the docstring's "Raises" section only calls out the RuntimeError chained from failed legacy userbuffer deallocations;
ub.shutdown_userbuffers_manager()(line 2696) can independently raise (e.g. if a manager-tracked buffer is still marked in-use) and isn't mentioned. Worth a one-line addition for future maintainers debugging a raised exception here.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 2625 - 2698, Update the cleanup() docstring’s Raises section to also document that ub.shutdown_userbuffers_manager() may raise independently, including when manager-tracked buffers remain in use; keep the existing legacy userbuffer deallocation RuntimeError description unchanged.
🤖 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 `@cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp`:
- Around line 84-98: Ensure PyTorchModelEngine.cleanup() invokes
ub.shutdown_userbuffers_manager() on every rank even when an earlier legacy
ub_buffers cleanup operation raises. Preserve the collective teardown ordering
by capturing or deferring the exception, completing manager shutdown across all
ranks, then propagating the cleanup error afterward.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/userbuffers/userbuffers-host.cpp`:
- Around line 436-483: Add a brief comment immediately above
unregister_user_buffer_collective documenting that buffers must be released in
strict reverse-allocation order. Reference the existing handle ==
comm->free_region - 1 TLLM_CHECK as the enforcement of this LIFO invariant,
without changing teardown behavior.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2625-2698: Update the cleanup() docstring’s Raises section to also
document that ub.shutdown_userbuffers_manager() may raise independently,
including when manager-tracked buffers remain in use; keep the existing legacy
userbuffer deallocation RuntimeError description unchanged.
🪄 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: e147c6f8-859e-4922-9d23-7b0f5ad908be
📒 Files selected for processing (10)
cpp/tensorrt_llm/kernels/userbuffers/ub_allocator.cppcpp/tensorrt_llm/kernels/userbuffers/ub_allocator.hcpp/tensorrt_llm/kernels/userbuffers/userbuffers-host.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffers.hcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cppcpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.hcpp/tensorrt_llm/nanobind/userbuffers/bindings.cpptensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/waives.txttests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
… shutdown PyTorchModelEngine.cleanup raised immediately when a legacy directly-owned userbuffer failed to deallocate. On a multi-rank engine, that rank skipped UserBuffersManager shutdown while peers could enter its collective unregister sequence, risking mismatched teardown or a hang. Collect legacy deallocation failures, release tensor-owned resources, and invoke shutdown_userbuffers_manager on every initialized rank before propagating the original cleanup error. Preserve failed legacy buffers and the incomplete cleanup state so deterministic callers can retry safely. Add a focused cleanup regression that injects a legacy deallocation failure and verifies manager shutdown still runs before the error is returned, while the failed buffer and retry state are retained. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
153-181: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise the retry completion path.
Coverage summary: added
PyTorchModelEngineTestCase.test_cleanup_shuts_down_manager_after_legacy_userbuffer_error; no tests were modified or removed. The test covers deferred errors, shutdown, flag reset, and failed-buffer retention, but the coverage verdict is insufficient because it never callscleanup()again to verify the failed buffer is retried and_cleanup_donebecomesTrue.As per path instructions, test-code changes must include a coverage summary and verdict.
Suggested extension
- with patch( + with patch( "tensorrt_llm._torch.pyexecutor.model_engine.ub.ub_deallocate", - side_effect=legacy_error), patch( + side_effect=[legacy_error, None]) as deallocate, patch( "tensorrt_llm._torch.pyexecutor.model_engine." "ub.shutdown_userbuffers_manager") as shutdown_manager, \ patch("tensorrt_llm._torch.pyexecutor.model_engine.release_gc"): with self.assertRaisesRegex( RuntimeError, "Failed to deallocate one or more userbuffers"): engine.cleanup() + engine.cleanup() + self.assertEqual(deallocate.call_count, 2) shutdown_manager.assert_called_once_with() + self.assertIsNone(engine.ub_buffers) self.assertFalse(engine._userbuffers_manager_initialized) - self.assertEqual(len(engine.ub_buffers), 1) - self.assertFalse(engine._cleanup_done) + self.assertTrue(engine._cleanup_done)🤖 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/_torch/executor/test_pytorch_model_engine.py` around lines 153 - 181, Extend test_cleanup_shuts_down_manager_after_legacy_userbuffer_error to call engine.cleanup() a second time after the expected first-call exception. Verify the retained userbuffer is retried successfully, _cleanup_done becomes True, and the manager shutdown behavior remains correct; update the test’s coverage summary and verdict to reflect the completed retry-path coverage.Source: Path instructions
🤖 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 `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 153-181: Extend
test_cleanup_shuts_down_manager_after_legacy_userbuffer_error to call
engine.cleanup() a second time after the expected first-call exception. Verify
the retained userbuffer is retried successfully, _cleanup_done becomes True, and
the manager shutdown behavior remains correct; update the test’s coverage
summary and verdict to reflect the completed retry-path coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e482eb37-b5e9-4507-85ab-f84f54a943b3
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
|
PR_Github #62729 [ run ] triggered by Bot. Commit: |
|
PR_Github #62665 [ run ] completed with state |
|
The unregister mechanics look careful — the LIFO drain via Two concerns, both about callers, and both stemming from 1.
Worth either confirming that path is collective-safe, or restricting the actual unregister to the deterministic 2. The new LIFO precondition contradicts the retry contract
Non-blocking: |
|
PR_Github #62729 [ run ] completed with state
|
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The core fix is correct and directly addresses the no-op deallocate root cause of nvbugs/6525010, but the new teardown runs cross-rank collective barriers whose failure modes are not covered by tests, so it merits confirmation before merge.
Concerns
- [MAJOR]
cpp/tensorrt_llm/kernels/userbuffers/userbuffers-host.cpp:452(anduserbuffersManager.cppshutdown()) - collective teardown can hang or crash across ranks- What is wrong:
unregister_user_buffer_collectivedoescudaDeviceSynchronize+ub_barrier(comm->comm_intra)at entry and a trailingub_barrier.UserBuffersManager::shutdown()first asserts!buffer.second("still in use") for every buffer viaTLLM_CHECK_WITH_INFObefore any teardown. - How it fails: these are collective ops requiring every rank in lockstep and identical LIFO order. If one rank still has a buffer in-use (or releases in a different order) while a peer does not, the peers pass the check and block on
ub_barrierwaiting for the rank that threw before the barrier — a hang. AnyTLLM_CHECKfailure insideunregisteron one rank aborts it mid-collective while peers wait forever. - Suggested fix: guarantee
cleanup()/shutdown()is invoked symmetrically on all ranks with identical buffer state, and validate the in-use invariant before any rank enters a barrier (ideally collectively), so no rank can abort after a peer has entered the barrier.
- What is wrong:
Minor notes (non-blocking)
cpp/tensorrt_llm/kernels/userbuffers/userbuffers-host.cpp:439-TLLM_CHECK(handle > 0)forbids unregistering handle 0; if a real userbuffer ever lands on handle 0 the LIFO pop aborts. Document that region 0 is reserved, or usehandle >= 0.tests/unittest/_torch/executor/test_pytorch_model_engine.py:153- test never callscleanup()a second time, so the retried-buffer completion path (ub_bufferscleared,_cleanup_doneTrue) is untested.tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py:205- reinit test returnsTruewhenub_supported()is False; usepytest.skipso a non-run is not reported as a pass.
QA view
- Test coverage: partial - the reinit multi-GPU test exercises shutdown/reinit past
MAX_REGIONS, and the unit test covers deferred-error ordering, but the retry-completion path and the shutdown-while-in-use abort path are uncovered. - SM coverage: the userbuffer VMM/multicast teardown is architecture-independent, so the 2-GPU unit test is a fair functional check; the arch-specific nvfp4/B200 failure is validated only by the un-waived
TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True]running on B200 in CI, which cannot be confirmed from the diff. - Test code: reinit test silently passes on unsupported HW; single-tensor-per-iteration does not exercise multi-buffer LIFO within one manager instance.
- Test time: significant - removing the waive re-enables a full B200 70B nvfp4 accuracy case; the reinit test adds smaller multi-GPU time.
- Needs
/qa-verify: yes - a B200 waive is removed for an arch-specific fix; QA should confirm the un-waived accuracy test passes reliably on B200 in CI and that the shutdown/cleanup collective path does not hang.
Does this actually fix 6525010?
Yes. The bug is stale reuse of an undersized 64MiB VMM region because deallocate was a no-op; the fix makes deallocate call unregister_user_buffer_collective (LIFO), adds shutdown() to retire the pool, calls it during engine cleanup, and rejects re-init while a prior pool exists. The >→>= MAX_REGIONS change additionally closes an off-by-one. I could not independently rerun the B200 sequence, so this rests on the un-waived integration test.
Possible new issues
- Cross-rank deadlock/crash on asymmetric shutdown as described above.
mc_offsetrewind assumes every torn-down region hadUB_MEM_MC_CREATED; mixed MC/non-MC regions could desync the multicast offset.- The
>=capacity change reduces registrable regions by one versus prior behaviour.
What I could not verify
- Whether region 0 is reserved at init (relevant to the
handle > 0assertion). - Whether
cleanup()/shutdown()is always called symmetrically across ranks with identical buffer state (the collective safety hinges on this). - Runtime behaviour of the un-waived B200 accuracy test in CI.
Automated review by NVCortex Lite, run by @fredricz-20070104.
…eardown Userbuffer deallocation now performs CUDA synchronization and MPI barriers, so invoking it from rank-asynchronous destructor cleanup can hang or run after MPI finalization. It is also unsafe to retry after a buffer has entered the unregister drain because CUDA VMM teardown is not transactional. Move userbuffer release into deterministic terminal executor shutdown. Release CUDA graphs first, synchronize the device, visit target and draft engines in a fixed order on every rank, and defer errors until all engine collectives have been attempted. Keep destructor cleanup non-collective and make failed unregister fail-stop instead of retrying partially released buffers. Preserve userbuffer managers during the temporary executor shutdown used for KV-cache capacity estimation. The final executor reuses those engines, and prematurely shutting down the manager reset its capacity to zero, causing later TP4 allocations to fail. Encoder shutdown now follows the same explicit graph-before-userbuffer ordering. Add focused coverage for deferred legacy cleanup errors, non-collective destructor cleanup, all-engine collective ordering, and nonterminal manager preservation. Improve the manager capacity assertion to report requested and available bytes. Validated on one 4xB200 allocation with CUDA_LAUNCH_BLOCKING=1. Two consecutive iterations passed the focused unit tests, multi-GPU manager reinitialization test, and all four integration cases, including TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True]. Slurm job 1563971 completed with exit code 0. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
161-175: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAssert the cleanup order.
This test proves that
shutdown_userbuffers_manager()runs beforeshutdown_userbuffers()returns. It does not prove thatub_deallocate()runs before manager shutdown.Record both side effects and assert the deallocation call occurs before manager shutdown. This protects the deferred-error contract.
🤖 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/_torch/executor/test_pytorch_model_engine.py` around lines 161 - 175, Update the test around engine.shutdown_userbuffers() to record calls to ub_deallocate and shutdown_userbuffers_manager in a shared order list, then assert the deallocation call occurs before manager shutdown. Keep the existing exception and state assertions, and ensure both mocked side effects append distinct markers when invoked.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/encoder_executor.py`:
- Around line 59-68: Make EncoderExecutor.shutdown idempotent by adding a
completion guard, such as _cleanup_done, before accessing self.model_engine.
Return immediately on subsequent calls, and mark cleanup complete only after the
existing release, synchronization, and shutdown_userbuffers flow finishes while
preserving the finally-based model_engine deletion.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2670-2679: Update the shutdown error handling around the
userbuffer deallocation and manager shutdown checks so that when both ub_errors
and manager_error are present, the raised RuntimeError includes diagnostic
details from both failures. Preserve the existing single-failure messages and
exception chaining behavior for cases where only one failure occurs, while
keeping _userbuffers_shutdown_failed set for any failure.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 1533-1537: Update the userbuffer shutdown error handling around
userbuffer_errors to log every collected failure, including errors beyond
userbuffer_errors[0], before raising the RuntimeError. Preserve the existing
aggregated exception and first-error chaining while ensuring each engine
shutdown failure is surfaced through the established logging mechanism.
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 111-124: Annotate both new test functions with monkeypatch:
pytest.MonkeyPatch and -> None, and annotate the nested fail_target_userbuffers
function with -> None. Keep the existing test behavior and coverage unchanged.
---
Outside diff comments:
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 161-175: Update the test around engine.shutdown_userbuffers() to
record calls to ub_deallocate and shutdown_userbuffers_manager in a shared order
list, then assert the deallocation call occurs before manager shutdown. Keep the
existing exception and state assertions, and ensure both mocked side effects
append distinct markers when invoked.
🪄 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: 02e081bd-ddb3-4ed5-bfbd-fe8feb192c26
📒 Files selected for processing (7)
cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpptensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/encoder_executor.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
🚧 Files skipped from review as they are similar to previous changes (1)
- cpp/tensorrt_llm/kernels/userbuffers/userbuffersManager.cpp
1 has been fixed, the cleanup now runs explicitly in shutdown. 2 is not a real concern. We do not actually expect any retry. Any retry means the whole Buffer is not in good state. |
We verified the current teardown path and believe rank state remains symmetric by construction:
Therefore, the manager’s in-use check should either fail on every rank before unregister begins or pass on every rank, after which all ranks drain the same LIFO sequence. The described hang remains theoretically possible after an existing rank-local fatal condition, such as a CUDA error, process loss, asymmetric internal-API use, or another bug that has already violated the collective contract. A preflight collective would not handle a rank that never enters teardown and would add another collective without addressing general MPI fault tolerance. Given that we found no supported execution path that can produce divergent manager state, we do not consider this an actionable blocker for the current change. Collective state validation could still be considered separately as defensive hardening. |
…down Make EncoderExecutor shutdown idempotent so repeated terminal cleanup does not access an already-deleted model engine. The completion guard is set only after graph release, device synchronization, and the userbuffer shutdown attempt have run, while retaining unconditional engine deletion. Preserve complete diagnostics when teardown fails. Combined legacy-buffer and manager failures now report both details while chaining the first deallocation error, and PyExecutor logs every target or draft engine shutdown failure before raising its aggregate error. Extend focused coverage for shutdown ordering, nonterminal manager preservation, repeated encoder shutdown, multi-engine error logging, and combined manager/deallocation failures. Add the requested pytest type annotations without changing existing test behavior. Validation: all pre-commit hooks passed. On B200, Slurm job 1567861 completed with exit code 0 and all seven focused teardown tests passed. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
|
Re-reviewed at
One residual thing, non-blocking: Not approving: shared C++ runtime with no other approval yet, and 4 threads still open. |
|
/bot run --disable-fail-fast |
|
PR_Github #62935 [ run ] triggered by Bot. Commit: |
|
Re-reviewed
Each is correct only if its predicate is rank-uniform. Where it isn't, the failure mode is a hang, not an error: the skipping rank exits while its peers sit in the barrier. (1) is the sharpest — a per-rank deallocation failure is exactly the case where ranks diverge, and that's the path that raises without the barrier. I don't think any of these is reachable today ( |
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
|
PR_Github #62935 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63036 [ run ] triggered by Bot. Commit: |
|
PR_Github #63036 [ run ] completed with state
|
4 participants
…4_tp4 UB reuse
The B200 full accuracy sequence constructs several PyTorch engines in one process. UserBuffersManager updated its configured allocation size for each engine but retained CUDA VMM regions registered by earlier engines because UserBufferAllocator::deallocate was a no-op.
This allowed the Llama 3.1 8B TP4 case to leave a 64 MiB userbuffer in the singleton pool and the following Llama 3.3 70B TP4 case to reuse it for a 128 MiB request. The undersized mapping produced an invalid CUDA access that surfaced synchronously at the explicit graph-output copy when CUDA_LAUNCH_BLOCKING was enabled.
Add collective userbuffer unregistration that synchronizes ranks, unbinds multicast mappings, unmaps and frees virtual address ranges, releases local and imported allocation handles, and clears communicator metadata. Track logical releases independently so callers may release in any order while physical registrations are retired in the required LIFO order.
Expose UserBuffersManager shutdown through the Python binding and invoke it after model tensors and CUDA graphs are released during engine cleanup. Reject manager reinitialization while an earlier engine pool remains, preventing silent cross-engine reuse.
Extend the multi-GPU allocation tests to shut down each manager instance and repeatedly reinitialize it more than MAX_REGIONS times before changing allocation size. This verifies that shutdown unregisters physical regions instead of only dropping pool bookkeeping.
Remove the B200 waive for TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True]. The exact four-case sequence passed twice in one sbatch allocation with CUDA_LAUNCH_BLOCKING=1 (Slurm job 1554986), and the new userbuffer lifecycle regression passed in both iterations.
Dev Engineer Review
UserBufferAllocator.shutdown_userbuffers_manager()through nanobind.TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True].QA Engineer Review
test_userbuffers_manager_reinitializeand_run_userbuffers_manager_reinitialize.test_pytorch_model_engine.py.test_py_executor.py.tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.tests/integration/test_lists/waives.txt.Verdict: needs follow-up (CBTS coverage mapping is unavailable for the new regression tests).
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.