[None][test] Reduce multi-GPU pre-merge time via shared MPI sessions and HF weight cache - #15872
Closed
sunnyqgg wants to merge 21 commits into
Closed
[None][test] Reduce multi-GPU pre-merge time via shared MPI sessions and HF weight cache#15872sunnyqgg wants to merge 21 commits into
sunnyqgg wants to merge 21 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>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57161 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot kill |
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>
…ed MPI sessions Extend the shared-MpiPoolSession grouping (method 1: session reuse; method 2: session reuse + HF weight cache) to the highest-ROI pre-merge stages: - DeepSeekV3Lite (H100): 14 fp8_block_scales_4gpus cases + 2 cuda_graph_padding cases grouped onto the shared 4-GPU session. test_guided_decoding_4gpus is excluded: it uses the bf16 checkpoint and sets TRTLLM_XGUIDANCE_LENIENT via mocker, which cannot reach workers of a pre-spawned pool (env is snapshotted at spawn). - GPTOSS (H100 + GB200): one union w4_4gpus premerge matrix; each platform's test-db yml selects its own case ids. Only non-waived cases are in the matrix -- moving a waived case to a new (grouped) test id would silently un-waive it. test_w4a16 is excluded (OVERRIDE_QUANT_ALGO env cannot reach pre-spawned workers). - Llama3.3-70B (GB200): fp8_tp4 x2, nvfp4_tp4 x2 and fp4_tp2pp2 (fusion=False only: the TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED patch cannot reach pre-spawned workers, and False matches the default). - NemotronV3 Super/Ultra (GB200): block_reuse / parallelism / mtp_ar cases grouped; online_eplb cases deferred pending an EPLB shared-state-on-reused-workers audit. - DeepSeekV32 (B200, new shared 8-GPU session fixture): nvfp4 baselines and piecewise_cuda_graph[baseline] share one FP4-v2 weight-cache entry. DeepSeekR1 groups were dropped: all member cases are currently waived. Also fix a latent bug: test-list matching is exact-nodeid (test_list_parser.modify_by_test_list), so the earlier unbracketed premerge_grouped yml entries never matched any collected test. All grouped entries (including the existing DeepSeekV3Lite nvfp4/bf16 ones) now list fully-parametrized ids; yml order keeps same-checkpoint cases adjacent so the weight-cache LRU (max_entries=1) stays hot. Verified: pytest --collect-only shows a 1:1 match between all 58 grouped yml entries and collected node ids; the 342 original parametrized cases still collect unchanged. Signed-off-by: qgai <qgai@nvidia.com>
sunnyqgg
force-pushed
the
grouped_test_premerge
branch
from
July 2, 2026 07:50
f94c700 to
c96e862
Compare
Closed
3 tasks
sunnyqgg
marked this pull request as draft
July 2, 2026 07:51
Collaborator
|
PR_Github #57166 [ kill ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57161 [ run ] completed with state |
Contributor
📝 WalkthroughWalkthroughThis PR adds an LRU cache for HuggingFace checkpoint weights, introduces MPI session ownership tracking across executor/proxy/LLM components to avoid shutting down externally-owned sessions, and adds grouped-test infrastructure plus test suite refactors enabling shared MPI sessions/LLM instances across many multi-GPU accuracy and unit tests. ChangesWeight caching and MPI session ownership
Grouped multi-GPU test infrastructure and coverage
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant HfWeightLoader
participant WeightCache
participant Disk
HfWeightLoader->>WeightCache: compute cache key (file stat + use_consolidated)
alt cache hit
WeightCache-->>HfWeightLoader: return cached weights
HfWeightLoader->>HfWeightLoader: local_mpi_barrier()
else cache miss
WeightCache->>WeightCache: evict LRU entry
HfWeightLoader->>Disk: load weights in parallel
Disk-->>HfWeightLoader: raw weights
HfWeightLoader->>WeightCache: store loaded weights
end
sequenceDiagram
participant TestCase
participant SharedMpiSession
participant MakeSharedLlm
participant Worker
TestCase->>SharedMpiSession: request shared MpiPoolSession
SharedMpiSession-->>TestCase: session
TestCase->>MakeSharedLlm: build LLM factory(mpi_session)
MakeSharedLlm->>Worker: construct LLM with _mpi_session kwarg
TestCase->>Worker: run inference
TestCase->>Worker: reset_worker_torch_compile_state / clear_worker_weight_cache
Suggested reviewers: 🚥 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}
Contributor
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)
tensorrt_llm/executor/proxy.py (1)
442-446: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
shutdown_abort()on worker-init failure bypasses_owns_mpi_session, can kill a shared/external session.Line 531-532 correctly guards the graceful shutdown path with
_owns_mpi_session, but the abort path at line 444 (triggered when a worker fails during initialization) still unconditionally callsself.mpi_session.shutdown_abort(...). For an externally-owned (e.g. grouped-test sharedMpiPoolSession) session, this forcibly destroys it, cascading the failure to all other consumers of that shared session — the exact outcome this PR's ownership tracking is meant to prevent.🛡️ Proposed fix
if ready_signal != GenerationExecutorProxy.READY_SIGNAL: logger.error(f"Executor worker initialization error: {error_trace}") - self.mpi_session.shutdown_abort(reason=ready_signal) + if self._owns_mpi_session: + self.mpi_session.shutdown_abort(reason=ready_signal) raise RuntimeError( "Executor worker returned error") from ready_signalAlso applies to: 531-532
🤖 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/executor/proxy.py` around lines 442 - 446, The worker-init failure path in GenerationExecutorProxy._get_worker_ready unconditionally calls self.mpi_session.shutdown_abort, which can destroy an externally owned shared MpiPoolSession. Update the abort handling to follow the same _owns_mpi_session ownership guard used in the graceful shutdown path, and only abort the MPI session when this proxy owns it. Keep the logger.error and RuntimeError behavior in place, but avoid calling shutdown_abort on shared/external sessions.
🧹 Nitpick comments (1)
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py (1)
103-208: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage is insufficient for the bin/pth cache path.
The new tests only create
.safetensorsfiles, butHfWeightLoader.load_weights()has a separate.bin/.pthcache branch. Add at least one cache-hit test intests/unittest/_torch/models/checkpoints/hf/test_weight_loader.pyusingmodel.binormodel.pthso both changed loader paths are covered. As per path instructions, coverage for this test file is currently insufficient for the changed branches.🤖 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/models/checkpoints/hf/test_weight_loader.py` around lines 103 - 208, Add a cache-hit test in test_weight_loader.py that exercises the .bin/.pth branch of HfWeightLoader.load_weights, since the current coverage only uses .safetensors and misses the separate cache path. Reuse the existing HfWeightLoader, ConsumableWeightsDict, and mock.patch.object patterns, but create a checkpoint with model.bin or model.pth and assert the second load hits the cache and does not call _load_weights_in_parallel again. Keep the existing cache cleanup via HfWeightLoader._clear_weight_cache.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.
Inline comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 55-71: The HF weight cache disabling logic is incomplete because
_weight_cache_max_entries() can return 0 for max_entries=0 or invalid values,
but load_weights() still performs cache reads through _get_cached_weights() when
_is_weight_cache_enabled() is true. Update the cache gating in load_weights()
(and any shared cache path around _get_cached_weights()/cache population) so a
max entries value of 0 or a parse failure disables both reads and writes,
preventing stale cached weights from being reused or retained in memory.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 77-87: The remaining late imports from _grouped_test_utils in
test_llm_api_pytorch still violate E402, so update the import block to either
add # noqa: E402 to each of the remaining imports or move the entire grouped
import block above patch_mpi_pool_session_for_env(). Use the existing symbols
clear_worker_weight_cache, hf_weight_cache_env, make_shared_llm,
reset_shared_session_torch_compile_state, and shared_mpi_session to locate the
affected imports.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py`:
- Line 104: The cache-enabled test is inheriting an external
TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES value, which can make the test fail for
unrelated environment reasons. In the cache setup around the monkeypatch in
test_weight_loader, explicitly set TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES to a known
positive value or remove it before the weight loader is exercised, so the
behavior is controlled within the test.
---
Outside diff comments:
In `@tensorrt_llm/executor/proxy.py`:
- Around line 442-446: The worker-init failure path in
GenerationExecutorProxy._get_worker_ready unconditionally calls
self.mpi_session.shutdown_abort, which can destroy an externally owned shared
MpiPoolSession. Update the abort handling to follow the same _owns_mpi_session
ownership guard used in the graceful shutdown path, and only abort the MPI
session when this proxy owns it. Keep the logger.error and RuntimeError behavior
in place, but avoid calling shutdown_abort on shared/external sessions.
---
Nitpick comments:
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py`:
- Around line 103-208: Add a cache-hit test in test_weight_loader.py that
exercises the .bin/.pth branch of HfWeightLoader.load_weights, since the current
coverage only uses .safetensors and misses the separate cache path. Reuse the
existing HfWeightLoader, ConsumableWeightsDict, and mock.patch.object patterns,
but create a checkpoint with model.bin or model.pth and assert the second load
hits the cache and does not call _load_weights_in_parallel again. Keep the
existing cache cleanup via HfWeightLoader._clear_weight_cache.
🪄 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: b5ebd729-f52d-4441-acf7-a2d78f3e27e2
📒 Files selected for processing (15)
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_b200.ymltests/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
+55
to
+71
| def _is_weight_cache_enabled() -> bool: | ||
| return os.environ.get(_WEIGHT_CACHE_ENV, | ||
| "0").lower() in ("1", "true", "yes", "on") | ||
|
|
||
| @staticmethod | ||
| def _weight_cache_max_entries() -> int: | ||
| try: | ||
| return max( | ||
| 0, | ||
| int( | ||
| os.environ.get(_WEIGHT_CACHE_MAX_ENTRIES_ENV, | ||
| _DEFAULT_WEIGHT_CACHE_MAX_ENTRIES))) | ||
| except ValueError: | ||
| logger.warning( | ||
| f"Invalid {_WEIGHT_CACHE_MAX_ENTRIES_ENV} value; disabling HF weight cache." | ||
| ) | ||
| return 0 |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Make max_entries=0/invalid disable reads too.
_weight_cache_max_entries() returns 0 for invalid values and logs that the cache is disabled, but load_weights() still calls _get_cached_weights() when TRTLLM_HF_WEIGHT_CACHE is truthy. If the cache was previously populated, stale entries can still be returned and kept in memory despite TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES=0 or an invalid value.
Proposed fix
`@staticmethod`
def _is_weight_cache_enabled() -> bool:
- return os.environ.get(_WEIGHT_CACHE_ENV,
- "0").lower() in ("1", "true", "yes", "on")
+ enabled = os.environ.get(_WEIGHT_CACHE_ENV,
+ "0").lower() in ("1", "true", "yes", "on")
+ if not enabled:
+ HfWeightLoader._clear_weight_cache()
+ return False
+ if HfWeightLoader._weight_cache_max_entries() <= 0:
+ HfWeightLoader._clear_weight_cache()
+ return False
+ return TrueAlso applies to: 125-131
🤖 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/models/checkpoints/hf/weight_loader.py` around lines 55 -
71, The HF weight cache disabling logic is incomplete because
_weight_cache_max_entries() can return 0 for max_entries=0 or invalid values,
but load_weights() still performs cache reads through _get_cached_weights() when
_is_weight_cache_enabled() is true. Update the cache gating in load_weights()
(and any shared cache path around _get_cached_weights()/cache population) so a
max entries value of 0 or a parse failure disables both reads and writes,
preventing stale cached weights from being reused or retained in memory.
Comment on lines
+77
to
+87
| from tensorrt_llm.llmapi._grouped_test_utils import \ | ||
| clear_worker_weight_cache as _clear_worker_weight_cache # noqa: E402 | ||
| from tensorrt_llm.llmapi._grouped_test_utils import \ | ||
| hf_weight_cache_env as _hf_weight_cache_env | ||
| from tensorrt_llm.llmapi._grouped_test_utils import \ | ||
| make_shared_llm as _make_shared_llm | ||
| from tensorrt_llm.llmapi._grouped_test_utils import \ | ||
| reset_shared_session_torch_compile_state as \ | ||
| _reset_shared_session_torch_compile_state | ||
| from tensorrt_llm.llmapi._grouped_test_utils import \ | ||
| shared_mpi_session as _shared_mpi_session |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect what precedes these imports to determine whether E402 is actually triggered here.
sed -n '1,90p' tests/integration/defs/accuracy/test_llm_api_pytorch.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3906
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether this file or test subtree is exempt from E402 elsewhere in config.
rg -n --hidden --glob '!**/.git/**' 'E402|ignore\s*=|per-file-ignores|extend-per-file-ignores' \
pyproject.toml ruff*.toml .pre-commit-config.yaml tests || true
# Inspect the exact import block and the preceding code shape with line numbers.
sed -n '1,110p' tests/integration/defs/accuracy/test_llm_api_pytorch.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 14244
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Run Ruff on the file to see exactly which lines, if any, trigger E402.
ruff check --select E402 tests/integration/defs/accuracy/test_llm_api_pytorch.py || true
# Show the import block with line numbers for mapping diagnostics to statements.
nl -ba tests/integration/defs/accuracy/test_llm_api_pytorch.py | sed -n '70,90p'Repository: NVIDIA/TensorRT-LLM
Length of output: 2994
Add # noqa: E402 to the remaining _grouped_test_utils imports
tests/integration/defs/accuracy/test_llm_api_pytorch.py:79-86 — the last four late imports still trigger E402; either suppress each one or move the whole block above patch_mpi_pool_session_for_env().
🤖 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/integration/defs/accuracy/test_llm_api_pytorch.py` around lines 77 -
87, The remaining late imports from _grouped_test_utils in test_llm_api_pytorch
still violate E402, so update the import block to either add # noqa: E402 to
each of the remaining imports or move the entire grouped import block above
patch_mpi_pool_session_for_env(). Use the existing symbols
clear_worker_weight_cache, hf_weight_cache_env, make_shared_llm,
reset_shared_session_torch_compile_state, and shared_mpi_session to locate the
affected imports.
|
|
||
|
|
||
| def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, monkeypatch): | ||
| monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Isolate TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES in cache-enabled tests.
This test enables the cache but inherits any outer TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES; a developer/CI env value of 0 or invalid text can make the cache test fail for the wrong reason. Set it to a known positive value or delete it before loading.
Proposed fix
def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, monkeypatch):
monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1")
+ monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", raising=False)
HfWeightLoader._clear_weight_cache()📝 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
| monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") | |
| monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") | |
| monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", raising=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 `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py` at line
104, The cache-enabled test is inheriting an external
TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES value, which can make the test fail for
unrelated environment reasons. In the cache setup around the monkeypatch in
test_weight_loader, explicitly set TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES to a known
positive value or remove it before the weight loader is exercised, so the
behavior is controlled within the test.
Collaborator
|
PR_Github #57166 [ kill ] completed with state |
2ez4bz
approved these changes
Jul 2, 2026
|
|
||
| @pytest.fixture(scope="module") | ||
| def shared_mpi_session_8gpu(hf_weight_cache): | ||
| # See shared_mpi_session_4gpu for why hf_weight_cache must be set up before |
Collaborator
There was a problem hiding this comment.
Nit: why not move the comment from shared_mpi_session_4gpu into the _shared_mpi_session helper? (Can be done in a future PR though).
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
3 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
Extends the shared-MPI-session grouping infrastructure from #15873 to the highest-ROI (P0/P1) pre-merge stages. No library changes in this PR — test-only.
New grouped pre-merge variants (
*_premerge_grouped, module-scoped shared 4-GPU/8-GPU sessions + HF weight cache):fp8_block_scales_4gpuscases + 2cuda_graph_paddingcasesw4_4gpus: one union matrix; DGX_H100 selects 6 case ids, GB200 selects 4fp8_tp4x2,nvfp4_tp4x2,fp4_tp2pp2[fusion_off_compile_on]block_reusex3; NemotronV3 Ultra (GB200):block_reuse[ADP4_MTP],parallelism[ADP2_PP2],mtp_arnvfp4baselines x2 +piecewise_cuda_graph[baseline]Grouping eligibility rules applied (documented in code comments):
OVERRIDE_QUANT_ALGO,TRTLLM_XGUIDANCE_LENIENT, gemm-allreduce fusion=True) stay ungrouped — the pool snapshots env at spawn.sampler_async_workercases are ordered last in their group.Changes
tests/integration/defs/accuracy/test_llm_api_pytorch.py—_run_*_caserunner refactors + grouped variants + shared 8-GPU fixturetests/integration/test_lists/test-db/{l0_dgx_h100,l0_gb200_multi_gpus,l0_dgx_b200}.yml— swap grouped case ids for the replaced entries (fully-parametrized ids; same-checkpoint cases adjacent for weight-cache LRU hits)Test plan
pytest --collect-only: all 58 grouped yml entries match collected node ids 1:1; the 342 original parametrized cases still collect unchanged