[None][feat] Add an opt-in raw-weight cache to the HF weight loader - #16054
Conversation
📝 WalkthroughWalkthroughAdds an opt-in LRU cache for raw HuggingFace checkpoint weight tensors, keyed by file path/size/mtime, controlled via environment variables. Refactors safetensors/bin loading through a new ChangesHF Weight Loader Caching
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant HfWeightLoader
participant WithWeightCache as _with_weight_cache
participant WeightCache as _WEIGHT_CACHE
participant PrefetchLoad as _prefetch_and_load
HfWeightLoader->>WithWeightCache: request weights (cache key)
WithWeightCache->>WeightCache: lookup key
alt cache hit
WeightCache-->>WithWeightCache: cached weights
WithWeightCache->>WithWeightCache: optional MPI barrier (barrier_on_hit)
WithWeightCache-->>HfWeightLoader: ConsumableWeightsDict
else cache miss
WithWeightCache->>WeightCache: evict LRU entries if needed
WithWeightCache->>PrefetchLoad: load_fn()
PrefetchLoad->>PrefetchLoad: prefetch heuristic + MPI barrier
PrefetchLoad->>PrefetchLoad: _load_weights_in_parallel
PrefetchLoad-->>WithWeightCache: loaded weights
WithWeightCache->>WeightCache: store weights
WithWeightCache-->>HfWeightLoader: ConsumableWeightsDict
end
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}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (1)
144-146: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the
load_fnparameter.
load_fnis unannotated; it's invoked asload_fn()and returns aConsumableWeightsDict. Add a preciseCallabletype so mypy/pyright can check callers.♻️ Proposed annotation
- def _with_weight_cache(self, weight_files: List[str], - use_consolidated: bool, barrier_on_hit: bool, - load_fn) -> ConsumableWeightsDict: + def _with_weight_cache( + self, weight_files: List[str], use_consolidated: bool, + barrier_on_hit: bool, + load_fn: Callable[[], ConsumableWeightsDict] + ) -> ConsumableWeightsDict:As per coding guidelines: "Always annotate functions" and "Prefer specifying argument types in
Callables."🤖 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 144 - 146, The _with_weight_cache method currently leaves load_fn untyped even though it is called as a zero-argument loader returning a ConsumableWeightsDict. Update the signature of _with_weight_cache in weight_loader.py to annotate load_fn with an appropriate Callable type that reflects no inputs and a ConsumableWeightsDict return, so type checkers can validate callers and the existing return annotation stays consistent.Source: Coding guidelines
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py (2)
110-137: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNo assertion that the cache-hit barrier is exercised.
Per the PR summary, "The safetensors hit path uses the same local MPI barrier behavior as the miss path" — i.e.,
_with_weight_cachecallslocal_mpi_barrier()on a hit whenbarrier_on_hit=True. This test drives exactly that scenario (safetensors hit) but never mocks/assertslocal_mpi_barrierwas invoked, so a regression that skips the barrier on hit (a distributed deadlock risk across ranks) wouldn't be caught here.Consider patching
tensorrt_llm._torch.models.checkpoints.hf.weight_loader.local_mpi_barrierand asserting it's called once on the second (second = loader.load_weights(...)) call.🤖 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 110 - 137, Add coverage for the safetensors cache-hit barrier path in test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper by mocking local_mpi_barrier from HfWeightLoader’s weight_loader module and asserting it is called once on the second load_weights invocation. Keep the existing assertions around _load_weights_in_parallel and cached raw weight reuse, and verify the hit path exercises the same barrier behavior as the miss path when barrier_on_hit is enabled.
6-203: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFingerprint-invalidation test appears to be missing.
The PR objectives state unit tests cover "hit, miss, eviction-before-load, and fingerprint invalidation behavior," but this diff only adds tests for reuse-on-hit, eviction-before-load-on-miss, and disabled-by-default. There's no test that mutates a cached file's
mtime/size and asserts the resulting cache key changes / triggers a miss (per PR summary: "Cache entries are invalidated using file fingerprints based on absolute path, size, andmtime_ns"). This is a coverage gap on a core correctness property of the cache (stale-data avoidance), so I'd rate current coverage as insufficient relative to the stated test matrix.Suggest adding
test_weight_cache_invalidated_on_fingerprint_changetotests/unittest/_torch/models/checkpoints/hf/test_weight_loader.pythat touches/rewrites the safetensors file (changing size or mtime) between twoload_weights()calls and asserts_load_weights_in_parallelis invoked twice.As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM ... suggest concrete list file names and 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/_torch/models/checkpoints/hf/test_weight_loader.py` around lines 6 - 203, Coverage is insufficient because there is no test for cache invalidation when a safetensors file fingerprint changes. Add a new test in HfWeightLoader coverage, e.g. test_weight_cache_invalidated_on_fingerprint_change, that mutates the checkpoint file between two load_weights() calls by changing size or mtime and then verifies _load_weights_in_parallel runs again instead of reusing the cached entry. Use the existing HfWeightLoader and _WEIGHT_CACHE setup to assert the fingerprint-based cache key changes and forces a miss.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 `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 144-146: The _with_weight_cache method currently leaves load_fn
untyped even though it is called as a zero-argument loader returning a
ConsumableWeightsDict. Update the signature of _with_weight_cache in
weight_loader.py to annotate load_fn with an appropriate Callable type that
reflects no inputs and a ConsumableWeightsDict return, so type checkers can
validate callers and the existing return annotation stays consistent.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py`:
- Around line 110-137: Add coverage for the safetensors cache-hit barrier path
in test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper by mocking
local_mpi_barrier from HfWeightLoader’s weight_loader module and asserting it is
called once on the second load_weights invocation. Keep the existing assertions
around _load_weights_in_parallel and cached raw weight reuse, and verify the hit
path exercises the same barrier behavior as the miss path when barrier_on_hit is
enabled.
- Around line 6-203: Coverage is insufficient because there is no test for cache
invalidation when a safetensors file fingerprint changes. Add a new test in
HfWeightLoader coverage, e.g.
test_weight_cache_invalidated_on_fingerprint_change, that mutates the checkpoint
file between two load_weights() calls by changing size or mtime and then
verifies _load_weights_in_parallel runs again instead of reusing the cached
entry. Use the existing HfWeightLoader and _WEIGHT_CACHE setup to assert the
fingerprint-based cache key changes and forces a miss.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2710e5a7-954e-4daa-83eb-68afe7807c40
📒 Files selected for processing (2)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py
|
/bot run --disable-fail-fast |
|
PR_Github #57985 [ run ] triggered by Bot. Commit: |
|
PR_Github #57985 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58015 [ run ] triggered by Bot. Commit: |
|
PR_Github #58015 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #58418 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #58543 [ run ] triggered by Bot. Commit: |
|
PR_Github #58418 [ run ] completed with state |
|
PR_Github #58543 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58576 [ run ] triggered by Bot. Commit: |
|
PR_Github #58576 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58602 [ run ] triggered by Bot. Commit: |
|
PR_Github #58602 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58620 [ run ] triggered by Bot. Commit: |
|
PR_Github #58620 [ run ] completed with state
|
Workers that load the same checkpoint repeatedly (e.g. across LLM instances sharing one MPI worker pool) can keep the raw checkpoint tensors in CPU RAM instead of re-reading and re-deserializing them: - Off by default; enabled with TRTLLM_HF_WEIGHT_CACHE=1. - Keyed by file fingerprints (abs path, size, mtime_ns) so any checkpoint change invalidates the entry. - LRU with TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES (default 1); eviction happens BEFORE the new load so CPU never holds the old cached and new loading weights at once (a ~2x peak otherwise). - Cache hits return a fresh ConsumableWeightsDict wrapper sharing the read-only tensors; the safetensors hit path joins the same local MPI barrier the miss path is about to enter (rank-local caches can diverge after eviction). - One _with_weight_cache() wrapper serves both the safetensors and bin/pth paths. Covered by unit tests (hit/miss/eviction/fingerprint invalidation). Split out of NVIDIA#15777 per review (independent optimization with its own memory and invalidation considerations). Signed-off-by: qgai <qgai@nvidia.com>
_with_weight_cache already evicts before load_fn() (the load-bearing eviction for the memory peak); re-running it inside _cache_loaded_weights was derivable-as-done work. _clear_weight_cache becomes a staticmethod like its sibling helpers. Signed-off-by: qgai <qgai@nvidia.com>
…ective divergence Two review findings on the HF raw-weight cache (thanks @QiJune): 1. Shared-tensor mutation: the cache shares raw tensors across loads, and a consumer mutating one in place (e.g. Nemotron-H's preprocess applies exp_()/neg_() to a view of the raw tensor; .to(torch.float32) is a no-op for fp32 checkpoints) would poison every later cache hit. The cache now records a sampled fingerprint per stored tensor and re-verifies it on every hit: a mutated entry is detected, logged with the offending keys, dropped and reloaded from disk. Models with in-place preprocessing (currently Nemotron-H, per a mapper audit) thus effectively opt out of caching - every hit self-heals into a fresh load - until their preprocessing is made non-mutating. 2. Collective divergence: with divergent rank-local cache states the first collective was Barrier on a hit rank but Allreduce (inside _get_local_available_host_memory) on a miss rank - a deadlock. The hit path now mirrors the miss path's exact sequence (Allreduce then Barrier); the bin/pth path performs no collectives on either side. The barrier_on_hit parameter is renamed mirror_load_collectives to match. Unit tests: in-place mutation between loads -> detected + clean reload; recorded collective sequences for a forced hit and a forced miss are identical ([allreduce, barrier]). Signed-off-by: qgai <qgai@nvidia.com>
|
/bot run --disable-fail-fast |
27769d5 to
37ecb41
Compare
|
PR_Github #58677 [ run ] triggered by Bot. Commit: |
|
PR_Github #58677 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58741 [ run ] triggered by Bot. Commit: |
|
PR_Github #58741 [ run ] completed with state |
4 participants
Summary
Independent optimization split out of #15777 per review. Workers that load the same checkpoint repeatedly (e.g. across LLM instances sharing one MPI worker pool) can keep the raw checkpoint tensors in CPU RAM instead of re-reading and re-deserializing them.
TRTLLM_HF_WEIGHT_CACHE=1. The default path pays nothing (no fingerprinting).TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES(default 1). Eviction happens BEFORE the new load so CPU never holds the old cached and the new loading weights at once (a ~2x transient peak otherwise).ConsumableWeightsDictwrapper sharing the read-only tensors; the safetensors hit path joins the same local MPI barrier the miss path is about to enter (rank-local caches can diverge after eviction). One_with_weight_cache()wrapper serves both the safetensors and bin/pth paths.Measured
Qwen3-235B-A22B nvfp4 on 4xB200: warm reload 106.6s -> 76.2s with the cache (page-cache-warm baseline); DeepSeek-V3-Lite accuracy suites show the same per-case reload saving when workers are shared.
Companion PRs
Test plan
Summary by CodeRabbit
New Features
Bug Fixes
Tests