[None][test] Reuse MPI sessions and cache HF weights across grouped multi-GPU LLM tests - #15873
Closed
sunnyqgg wants to merge 20 commits into
Closed
[None][test] Reuse MPI sessions and cache HF weights across grouped multi-GPU LLM tests#15873sunnyqgg wants to merge 20 commits into
sunnyqgg wants to merge 20 commits into
Conversation
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
…cesses When multiple LLM instances share the same worker processes (e.g. a reused MpiPoolSession), the module-global _pp_comm was reassigned on each new pp>1 LLM, dropping the previous NcclCommunicatorOp and triggering a collective ncclCommDestroy at an unsynchronized point during the next model build. On reused workers this can deadlock in commDestroySync. Since the PP comm is a world communicator depending only on (world_size, rank), reuse the existing one when world_size matches and only refresh the routing mapping. Single-LLM runs are unaffected. Signed-off-by: qgai <qgai@nvidia.com>
- Fix weight-cache env ordering: shared_mpi_session_4gpu now depends on hf_weight_cache so TRTLLM_HF_WEIGHT_CACHE is exported before the MPI pool spawns (MpiPoolSession snapshots env at spawn time); add an assertion to pin the contract. - Add a _make_shared_llm factory (shared_llm_* fixtures) so tests inject the shared MPI session transparently, and migrate the RPC/NVFP4 direct-construction tests onto it. - Mark the grouped NVFP4 test threadleak(enabled=False): the reused session keeps mpi4py's _manager_spawn thread alive by design. - Only compute the HF weight-cache key when the cache is enabled (avoid a per-shard os.stat on the default path). Signed-off-by: qgai <qgai@nvidia.com>
On a reused MpiPoolSession, Dynamo's per-code-object recompile counter is process-global in the worker and never reset across LLMs. Each torch_compile case recompiles the same model.forward under new guards; the count accumulates and eventually trips recompile_limit (16), which is a HARD failure under enable_fullgraph=True (FailOnRecompileLimitHit) and aborts the whole MPI job (observed as a silent crash at the ~10th case). Submit torch._dynamo.reset() to each worker after every case so compile state starts clean per LLM. Test-only; production (non-reused) paths are untouched. Verified: full 13-case grouped run now passes 13/13 on B300. Signed-off-by: qgai <qgai@nvidia.com>
…ckage module The grouped multi-GPU session-reuse helpers (shared_mpi_session, make_shared_llm, mpi_session_kwargs, reset_worker_torch_compile_state, restore_env_var) were duplicated across the unittest and integration test roots, and had already begun to diverge. Consolidate them into tensorrt_llm/llmapi/_grouped_test_utils.py so both roots (which have separate sys.paths and cannot share a module under tests/) import one source. Heavy imports are lazy to avoid circular imports at package init. Test-infra only. Verified: package import + pytest collection resolve (13 tests collected). Signed-off-by: qgai <qgai@nvidia.com>
Move the HF-weight-cache env enable/restore logic into a reusable contextmanager (hf_weight_cache_env) in _grouped_test_utils.py, keeping only a thin pytest fixture wrapper in the test. The fixture itself stays test-side so the shipped package never imports pytest. Docstring documents the load-bearing ordering (must be entered before the MpiPoolSession spawns). Verified: contextmanager sets/restores env correctly; 13 tests collected. Signed-off-by: qgai <qgai@nvidia.com>
Replace the repeated 'take shared_mpi_session_Ngpu + spread **_mpi_session_kwargs(...)' boilerplate in every harness-based test with module-scoped mpi_kwargs_2gpu / mpi_kwargs_4gpu fixtures that yield the ready-to-spread kwargs dict. Tests now just spread **mpi_kwargs_Ngpu and never reference the raw session or the helper; _mpi_session_kwargs is used only inside the two fixtures. RPC tests keep the shared_llm_2gpu factory. No behavior change. Signed-off-by: qgai <qgai@nvidia.com>
- Collapse the 13 hand-spelled _NVFP4_4GPU_PREMERGE_CASES dicts (which repeated 4 constant keys x13 + one lpc flip) into a compact tuple matrix of the varying axes plus a builder that fills the constants. Same ids and field values, ~130 lines -> ~30. - make_shared_llm's closure now reuses mpi_session_kwargs instead of re-implementing the 'inject _mpi_session if not None' conditional. Signed-off-by: qgai <qgai@nvidia.com>
…ardown The per-worker HF raw-weight cache previously relied on the shared session's worker processes exiting to be cleared. Make invalidation explicit: on shared_llm_4gpu teardown (which runs before the session shuts down its workers), submit clear_worker_weight_cache() to each worker. This closes the 'setup cache / teardown invalidate' loop instead of leaning on process death. Signed-off-by: qgai <qgai@nvidia.com>
…ss-scoped sessions Group test_llm_multi_gpu_pytorch.py into TestLlmMultiGpu2gpu (@gpu2) and TestLlmMultiGpu4gpu (@gpu4), sharing a small _MultiGpuLlmTests base whose class-scoped mpi_session/mpi_kwargs/shared_llm fixtures build one MpiPoolSession of self.n_gpus workers. The session is created lazily and torn down at class end, so the 2-GPU and 4-GPU pools no longer coexist (as they did with the previous module-scoped fixtures). Update the two explicit nodeid references for test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 (l0_dgx_h100.yml, waives.txt) to include the class; -m gpu2/gpu4 and -k filters are unaffected. Signed-off-by: qgai <qgai@nvidia.com>
Upstream's independently-curated pre-merge smoke set (superseded by the grouped test during the rebase) included two CUTEDSL configs. Add them to _NVFP4_4GPU_PREMERGE_MATRIX so the grouped test preserves that coverage; _run_nvfp4_4gpus_case already skips CUTEDSL off SM100/103. Signed-off-by: qgai <qgai@nvidia.com>
The case dicts embedded id=row[0] only to strip it back out at call time (ids=[case['id']...] + a case_kwargs comprehension). Keep the id in the matrix (row[0]) for the parametrize ids, make each case a pure _run_nvfp4_4gpus_case kwargs dict, and splat it directly (**case). No behavior change. Signed-off-by: qgai <qgai@nvidia.com>
Rename HfWeightLoader.clear_weight_cache -> _clear_weight_cache: its only callers are tests (via _grouped_test_utils.clear_worker_weight_cache and the weight-loader unit tests), and it matches the sibling private cache helpers (_get_cached_weights, _cache_loaded_weights). Signals it is test-infra, not a supported production API. Signed-off-by: qgai <qgai@nvidia.com>
…ession Extract _run_bfloat16_4gpus_case from test_bfloat16_4gpus and add test_bfloat16_4gpus_premerge_grouped, which reuses the SAME module-scoped shared_llm_4gpu MPI session as the nvfp4 group (both are 4-worker DeepSeek-V3-Lite, so world_size matches). This saves one more session spawn in the gb200 pre-merge stage. The shared runner resets per-worker torch.compile/Dynamo state between cases (same FailOnRecompileLimitHit guard as the nvfp4 group). The HF weight cache is keyed by checkpoint file fingerprint, so the bf16 model gets its own entry; with max_entries=1 the LRU evicts the other model on the first cross-model miss, bounding CPU cache to one model's raw tensors. Replace the three individual bf16 pre-merge entries in l0_gb200_multi_gpus.yml with the grouped test; the post-merge individual test_bfloat16_4gpus[...] entries are unaffected (method still exists). Signed-off-by: qgai <qgai@nvidia.com>
… peak On a cache miss, evict LRU entries down to the cap BEFORE loading the new weights, so the old (cached) and new (loading) raw tensors are never held in CPU at the same time. Previously the eviction happened only at store time (after the load), so switching models on a reused MPI-session worker transiently held both models' weights (~2x CPU peak). CPU now stays bounded by max_entries models even during a cross-model load. Add a unit test asserting the cache is already empty when the new load runs (max_entries=1), proving eviction precedes allocation. Signed-off-by: qgai <qgai@nvidia.com>
/simplify cleanup on the bf16-grouping + weight-cache-eviction changes: - Hoist the identical per-case torch.compile-reset finally block (duplicated verbatim in _run_nvfp4_4gpus_case and _run_bfloat16_4gpus_case) into reset_shared_session_torch_compile_state(make_llm) in _grouped_test_utils, so both call it in one line and a future grouped test can't drift. - Trim the over-long _evict_to_make_room docstring (10 lines -> 2). No behavior change. Signed-off-by: qgai <qgai@nvidia.com>
Test-list matching is exact-nodeid (test_list_parser.modify_by_test_list), so an unbracketed yml entry for a parametrized test never matches any collected test and is silently ignored as an invalid filter. Expand the grouped DeepSeek-V3-Lite nvfp4/bf16 entries to per-case bracketed ids; same-checkpoint cases stay adjacent so the HF weight-cache LRU (max_entries=1) stays hot across cases. Signed-off-by: qgai <qgai@nvidia.com>
Closed
3 tasks
Contributor
📝 WalkthroughWalkthroughThis PR adds NCCL PP communicator reuse, an in-process LRU cache for HuggingFace checkpoint weights, MPI session ownership tracking across executor/proxy/LLM classes to prevent shutting down externally-owned sessions, and shared MPI/LLM test utilities enabling grouped multi-GPU test execution with updated test lists. ChangesMPI ownership, weight cache, and grouped test infrastructure
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant HfWeightLoader
participant Cache
participant Disk
HfWeightLoader->>Cache: compute cache key, lookup
alt cache hit
Cache-->>HfWeightLoader: cached raw weights
HfWeightLoader->>HfWeightLoader: local_mpi_barrier()
HfWeightLoader-->>HfWeightLoader: return fresh ConsumableWeightsDict
else cache miss
HfWeightLoader->>Cache: evict_to_make_room()
HfWeightLoader->>Disk: load_weights_in_parallel()
Disk-->>HfWeightLoader: raw weights
HfWeightLoader->>Cache: store weights (LRU)
end
sequenceDiagram
participant BaseLLM
participant Proxy as ExecutorProxy
participant MpiSession
BaseLLM->>BaseLLM: capture mpi_session, set _owns_mpi_session
BaseLLM->>Proxy: create executor with mpi_session
Proxy->>Proxy: set _owns_mpi_session
BaseLLM->>BaseLLM: shutdown()
alt owns session
BaseLLM->>MpiSession: shutdown()
else externally provided
BaseLLM-->>BaseLLM: skip shutdown
end
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/llmapi/test_llm.py (1)
1824-1925: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
llm.shutdown()isn't exception-safe — risks leaking executor state on a shared MPI session.
llm.shutdown()(line 1925) only runs if every assert in thefor output in llm.generate(...)loop and the streamingasyncio.run(main())path succeeds. Unlikellm_get_stats_test_harness/llm_get_stats_async_test_harness(which usewith LLM_CLASS(...) as llm:) and_test_llm_capture_request_error(which wraps intry/finally), this harness has no equivalent guard. Sincetest_llm_multi_gpu_pytorch.py::TestLlmMultiGpu2gpu::test_llm_return_logprobs_streaming_tp2now calls this harness with a shared, class-scoped MPI session (**mpi_kwargs), a failed assertion here would skipllm.shutdown(), leaving that LLM's executor un-shut-down on workers that are reused by every other test in the same class — risking hangs/failures in subsequent tests, not just a resource leak in isolation.🐛 Proposed fix
llm = LLM_CLASS( llama_model_path, kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4, **kv_cache_args_extra), tensor_parallel_size=tp_size, gather_generation_logits=True, **llm_args_extra, **extra_llm_kwargs, ) + try: + _run_llm_return_logprobs_body(llm, ...) + finally: + llm.shutdown()Simplest minimal fix: wrap the existing body (from
prompts = [...]throughasyncio.run(main())) intry: ... finally: llm.shutdown()instead of appending the bare call at the end.🤖 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/llmapi/test_llm.py` around lines 1824 - 1925, llm_return_logprobs_test_harness is not exception-safe because llm.shutdown() only runs on the happy path, so a failing assert or streaming error can leave the shared LLM executor alive in the MPI session. Wrap the body that starts at prompts and includes llm.generate and the asyncio.run(main()) streaming path in a try/finally, and move llm.shutdown() into the finally block so it always executes even when the harness fails.
🧹 Nitpick comments (2)
tensorrt_llm/llmapi/_grouped_test_utils.py (2)
1-133: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo direct unit test coverage for this utility module.
The ordering contract in
hf_weight_cache_env/shared_mpi_session_4gpu(env must be exported before spawn) and the dict-building logic inmpi_session_kwargs/restore_env_varare load-bearing but only exercised indirectly through multi-GPU integration tests. A small CPU-only unit test (e.g.tests/unittest/llmapi/test_grouped_test_utils.py) assertingrestore_env_varcorrectly pops/restores env vars andmpi_session_kwargsreturns{}forNonewould catch regressions cheaply without requiring GPUs.🤖 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/llmapi/_grouped_test_utils.py` around lines 1 - 133, Add lightweight CPU-only unit coverage for the helper functions in _grouped_test_utils: create a test module that exercises restore_env_var with both removal and restoration cases, and mpi_session_kwargs with None and a non-None sentinel session. Keep the tests focused on these symbols so they validate the env restore behavior and empty-dict fallback without needing MPI or GPUs.
27-27: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing return/parameter type annotations on several public functions.
hf_weight_cache_env,shared_mpi_session,mpi_session_kwargs,reset_shared_session_torch_compile_state, andmake_shared_llmare missing return type annotations and/or parameter annotations (e.g.mpi_session,make_llmare untyped;hf_weight_cache_env/shared_mpi_sessionhave no return type). Since lazy imports are used to avoid circular imports, anif TYPE_CHECKING:block could supply the types (e.g.MpiPoolSession,Iterator[None]) without introducing an eager import.♻️ Example fix
+from typing import TYPE_CHECKING, Iterator, Optional + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession, MpiSession + `@contextmanager` -def hf_weight_cache_env(max_entries: str = "1"): +def hf_weight_cache_env(max_entries: str = "1") -> Iterator[None]: ... -def shared_mpi_session(n_workers: int): +def shared_mpi_session(n_workers: int) -> Iterator["Optional[MpiPoolSession]"]: ... -def mpi_session_kwargs(mpi_session) -> dict: +def mpi_session_kwargs(mpi_session: "Optional[MpiSession]") -> dict: ... -def make_shared_llm(mpi_session): +def make_shared_llm(mpi_session: "Optional[MpiSession]") -> Callable[..., "LLM"]: ...As per coding guidelines, "Always annotate functions with return types (use
Noneif no return); annotate class members and variables when necessary" and "do not use# type: ignoreannotations" / avoidtyping.Any.Also applies to: 49-49, 68-71, 91-91, 117-117
🤖 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/llmapi/_grouped_test_utils.py` at line 27, Add explicit type annotations to the public helpers in `_grouped_test_utils.py`: give `hf_weight_cache_env` and `shared_mpi_session` return types, annotate `mpi_session`/`make_llm` parameters and the return types in `mpi_session_kwargs`, `reset_shared_session_torch_compile_state`, and `make_shared_llm`. Use an `if TYPE_CHECKING:` block for imported types like `MpiPoolSession` and `Iterator[None]` so you can keep lazy imports without circular dependencies.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.
Inline comments:
In `@tensorrt_llm/_torch/distributed/communicator.py`:
- Around line 1249-1260: The PPCommNCCL reuse check in communicator.py only
compares world_size, but NcclCommunicatorOp is created with both world_size and
rank, so reuse can keep an nccl_comm bound to the wrong rank. Update the reuse
guard in the PPCommNCCL branch to require both the existing
_pp_comm.mapping.world_size and _pp_comm.mapping.rank to match the incoming
mapping before reassigning _pp_comm.mapping.
---
Outside diff comments:
In `@tests/unittest/llmapi/test_llm.py`:
- Around line 1824-1925: llm_return_logprobs_test_harness is not exception-safe
because llm.shutdown() only runs on the happy path, so a failing assert or
streaming error can leave the shared LLM executor alive in the MPI session. Wrap
the body that starts at prompts and includes llm.generate and the
asyncio.run(main()) streaming path in a try/finally, and move llm.shutdown()
into the finally block so it always executes even when the harness fails.
---
Nitpick comments:
In `@tensorrt_llm/llmapi/_grouped_test_utils.py`:
- Around line 1-133: Add lightweight CPU-only unit coverage for the helper
functions in _grouped_test_utils: create a test module that exercises
restore_env_var with both removal and restoration cases, and mpi_session_kwargs
with None and a non-None sentinel session. Keep the tests focused on these
symbols so they validate the env restore behavior and empty-dict fallback
without needing MPI or GPUs.
- Line 27: Add explicit type annotations to the public helpers in
`_grouped_test_utils.py`: give `hf_weight_cache_env` and `shared_mpi_session`
return types, annotate `mpi_session`/`make_llm` parameters and the return types
in `mpi_session_kwargs`, `reset_shared_session_torch_compile_state`, and
`make_shared_llm`. Use an `if TYPE_CHECKING:` block for imported types like
`MpiPoolSession` and `Iterator[None]` so you can keep lazy imports without
circular dependencies.
🪄 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: c5cbf39f-195c-45aa-ac2b-60dfe3341aab
📒 Files selected for processing (14)
tensorrt_llm/_torch/distributed/communicator.pytensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytensorrt_llm/executor/executor.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/rpc_proxy.pytensorrt_llm/llmapi/_grouped_test_utils.pytensorrt_llm/llmapi/llm.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/test-db/l0_gb200_multi_gpus.ymltests/integration/test_lists/waives.txttests/unittest/_torch/models/checkpoints/hf/test_weight_loader.pytests/unittest/llmapi/test_llm.pytests/unittest/llmapi/test_llm_multi_gpu_pytorch.py
Comment on lines
+1249
to
+1260
| elif isinstance(_pp_comm, PPCommNCCL) and \ | ||
| _pp_comm.mapping.world_size == mapping.world_size: | ||
| # Reuse the existing world NCCL communicator across LLM instances that | ||
| # share the same worker processes (e.g. a reused MpiPoolSession). The | ||
| # underlying comm depends only on (world_size, rank) -- it is a world | ||
| # communicator, independent of the pp/tp/ep layout -- so only the | ||
| # routing mapping needs refreshing. Recreating it would drop the old | ||
| # comm and trigger a collective ncclCommDestroy at an unsynchronized | ||
| # point during the next model build, which can deadlock on reused | ||
| # workers. Single-LLM (production) runs are unaffected: _pp_comm starts | ||
| # as None, so the first call still constructs a fresh PPCommNCCL. | ||
| _pp_comm.mapping = mapping |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Include rank in the communicator reuse guard.
NcclCommunicatorOp is constructed with both world_size and rank (Lines 1180-1183), and the new comment also says the communicator depends on (world_size, rank). Reusing when only world_size matches can leave _pp_comm.nccl_comm bound to an old rank while _pp_comm.mapping is refreshed to a new one.
Proposed fix
- elif isinstance(_pp_comm, PPCommNCCL) and \
- _pp_comm.mapping.world_size == mapping.world_size:
+ elif (
+ isinstance(_pp_comm, PPCommNCCL)
+ and _pp_comm.mapping.world_size == mapping.world_size
+ and _pp_comm.mapping.rank == mapping.rank):📝 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.
Suggested change
| elif isinstance(_pp_comm, PPCommNCCL) and \ | |
| _pp_comm.mapping.world_size == mapping.world_size: | |
| # Reuse the existing world NCCL communicator across LLM instances that | |
| # share the same worker processes (e.g. a reused MpiPoolSession). The | |
| # underlying comm depends only on (world_size, rank) -- it is a world | |
| # communicator, independent of the pp/tp/ep layout -- so only the | |
| # routing mapping needs refreshing. Recreating it would drop the old | |
| # comm and trigger a collective ncclCommDestroy at an unsynchronized | |
| # point during the next model build, which can deadlock on reused | |
| # workers. Single-LLM (production) runs are unaffected: _pp_comm starts | |
| # as None, so the first call still constructs a fresh PPCommNCCL. | |
| _pp_comm.mapping = mapping | |
| elif ( | |
| isinstance(_pp_comm, PPCommNCCL) | |
| and _pp_comm.mapping.world_size == mapping.world_size | |
| and _pp_comm.mapping.rank == mapping.rank): | |
| # Reuse the existing world NCCL communicator across LLM instances that | |
| # share the same worker processes (e.g. a reused MpiPoolSession). The | |
| # underlying comm depends only on (world_size, rank) -- it is a world | |
| # communicator, independent of the pp/tp/ep layout -- so only the | |
| # routing mapping needs refreshing. Recreating it would drop the old | |
| # comm and trigger a collective ncclCommDestroy at an unsynchronized | |
| # point during the next model build, which can deadlock on reused | |
| # workers. Single-LLM (production) runs are unaffected: _pp_comm starts | |
| # as None, so the first call still constructs a fresh PPCommNCCL. | |
| _pp_comm.mapping = mapping |
🤖 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/distributed/communicator.py` around lines 1249 - 1260,
The PPCommNCCL reuse check in communicator.py only compares world_size, but
NcclCommunicatorOp is created with both world_size and rank, so reuse can keep
an nccl_comm bound to the wrong rank. Update the reuse guard in the PPCommNCCL
branch to require both the existing _pp_comm.mapping.world_size and
_pp_comm.mapping.rank to match the incoming mapping before reassigning
_pp_comm.mapping.
jieli-matrix
approved these changes
Jul 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Infrastructure to cut multi-GPU test startup overhead: groups of
LLM-API tests can reuse oneMpiPoolSessionacross cases (skipping per-case MPI spawn), and cases sharing a checkpoint can reuse an opt-in per-worker HF raw-weight cache (skipping repeated disk loads). All mechanisms are off by default and have no impact on production single-LLM runs.Library changes:
LLM/executor proxies honor an injected shared MPI session via_owns_mpi_sessionownership semantics: a session the LLM does not own is never shut down by it. The live session is kept offLlmArgs(args are pickled to workers;MpiSessionis not pickleable).init_pp_commreuses the existing world NCCL communicator across LLM instances on reused workers when world_size matches; recreating it would trigger a collectivencclCommDestroyat an unsynchronized point and can deadlock. Single-LLM runs are unaffected (_pp_commstarts as None).HfWeightLoadergains an LRU raw-weight cache gated byTRTLLM_HF_WEIGHT_CACHE(+TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES), keyed by checkpoint file fingerprints (path/size/mtime). Cache miss evicts BEFORE loading to avoid holding old+new weights at once (~2x CPU peak); cache hits still take the local MPI barrier so ranks with divergent hit/miss states cannot deadlock.Test infrastructure:
tensorrt_llm/llmapi/_grouped_test_utils.py(new, in-package because the unittest and integration test roots cannot share a module): shared-session generator, LLM factory that transparently injects the session, per-case workertorch._dynamo.reset()(recompile counts accumulate on reused workers and triprecompile_limit=16, a hard failure underfullgraph=True), and explicit worker weight-cache invalidation on teardown.TestLlmMultiGpu2gpu/TestLlmMultiGpu4gpuwith class-scoped sessions.Changes
tensorrt_llm/llmapi/llm.py,executor/{executor,proxy,rpc_proxy}.py— shared-session ownership semanticstensorrt_llm/_torch/distributed/communicator.py— PP NCCL communicator reusetensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py— opt-in HF raw-weight cachetensorrt_llm/llmapi/_grouped_test_utils.py— shared grouped-test helpers (new)tests/integration/defs/accuracy/test_llm_api_pytorch.py— shared 4-GPU fixtures + DeepSeek-V3-Lite grouped pre-merge teststests/unittest/llmapi/test_llm_multi_gpu_pytorch.py— gpu2/gpu4 class reorg with class-scoped sessionstests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py— weight-cache unit tests (new)tests/integration/test_lists/test-db/{l0_gb200_multi_gpus,l0_dgx_h100}.yml,waives.txt— grouped entries + path updatesTest plan
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py)pytest --collect-only: every grouped test-db entry matches a collected node idNotes
Follow-up PR #15872 (stacked on this branch) extends the grouping to the P0/P1 pre-merge sets (DeepSeek-V3-Lite fp8 on H100, GPT-OSS-120B, Llama3.3-70B, NemotronV3, DeepSeek-V3.2) and must merge after this one.
Summary by CodeRabbit
New Features
Bug Fixes