[https://nvbugs/6272397][fix] Prevent host OOM during checkpoint prefetch - #17430
[https://nvbugs/6272397][fix] Prevent host OOM during checkpoint prefetch#17430moraxu wants to merge 1 commit into
Conversation
…etch Checkpoint prefetch warmed the OS page cache by reading each safetensors shard with a whole-file f.read(), pinning the entire file in anonymous memory per in-flight file. With up to 16 prefetch threads per local rank (128 concurrent multi-GB reads on an 8-GPU node), slow shared storage lets those transient buffers accumulate into hundreds of GB on top of the page cache itself, exhausting host memory: the kernel OOM killer then takes down the whole job mid-prefetch, e.g. during the ~641 GB DeepSeek-R1 FP8 prefetch on H20-3e perf CI runners. Read in fixed-size 64 MB chunks into one bounded buffer per in-flight file instead. This warms the page cache identically while keeping the per-thread footprint constant. Also emit a rate-limited progress heartbeat so a slow-but-healthy prefetch produces observable output instead of tens of minutes of log silence (which output-stall watchdogs punish with SIGKILL, masking the real failure mode). Signed-off-by: Michal Guzek <mguzek@nvidia.com>
|
/bot run |
WalkthroughCheckpoint prefetching now reads files in bounded 64 MiB chunks. Worker threads report byte progress through a synchronized callback, which emits periodic heartbeat logs. Tests cover chunk boundaries, missing files, and progress logging. ChangesPrefetch progress
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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: 2
🤖 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 356-358: Keep the read-loop callback in weight_loader.py lines
356-358 limited to byte accounting via report_progress. Update the worker-future
handling at lines 382-404 to emit rate-limited heartbeat logs while futures
remain pending, including when a blocking read produces no completed bytes. Add
a blocking-read test in
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py lines 327-347
that verifies a heartbeat occurs before read completion and subsequent
heartbeats are rate-limited.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py`:
- Around line 327-347: Extend test_prefetch_files_emits_progress_heartbeat with
a controlled blocking read that takes longer than _PREFETCH_LOG_INTERVAL_SEC,
using a synchronization mechanism to pause and resume the read. Assert that
“Prefetch progress” is logged before the blocked read completes, and verify the
number or timing of logs remains rate-limited rather than emitting continuously.
Keep the existing chunk-based heartbeat assertions intact.
🪄 Autofix
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: 889d2b4a-28ca-457c-8f77-6849b01521ff
📒 Files selected for processing (2)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py
| while num_read := f.readinto(buffer): | ||
| if report_progress is not None: | ||
| report_progress(num_read) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Decouple heartbeat scheduling from completed reads.
A heartbeat check occurs only after readinto() returns. A slow 64 MiB read can block beyond 60 seconds, so the process can remain silent during active prefetch.
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L356-L358: keep this callback for byte accounting only.tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L382-L404: schedule rate-limited logs while worker futures are pending, including when no read completes.tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py#L327-L347: add a blocking-read test that verifies a heartbeat before read completion and verifies rate limiting.
As per path instructions, test coverage must validate changed test behavior.
📍 Affects 2 files
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L356-L358(this comment)tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py#L382-L404tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py#L327-L347
🤖 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 356
- 358, Keep the read-loop callback in weight_loader.py lines 356-358 limited to
byte accounting via report_progress. Update the worker-future handling at lines
382-404 to emit rate-limited heartbeat logs while futures remain pending,
including when a blocking read produces no completed bytes. Add a blocking-read
test in tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py lines
327-347 that verifies a heartbeat occurs before read completion and subsequent
heartbeats are rate-limited.
Source: Path instructions
| def test_prefetch_files_emits_progress_heartbeat(tmp_path, monkeypatch): | ||
| # The heartbeat is what keeps a slow prefetch observable (and alive under | ||
| # output-stall watchdogs); with the log interval forced to zero it must | ||
| # fire for every chunk. | ||
| from tensorrt_llm._torch.models.checkpoints.hf import weight_loader as wl | ||
|
|
||
| monkeypatch.setattr(wl, "_PREFETCH_CHUNK_SIZE_BYTES", 1024) | ||
| monkeypatch.setattr(wl, "_PREFETCH_LOG_INTERVAL_SEC", 0.0) | ||
| files = [] | ||
| for i in range(3): | ||
| file = tmp_path / f"model-0000{i}-of-00003.safetensors" | ||
| file.write_bytes(os.urandom(4 * 1024)) | ||
| files.append(str(file)) | ||
|
|
||
| with mock.patch.object(wl.logger, "info") as info: | ||
| HfWeightLoader().prefetch_files(files) | ||
|
|
||
| progress_logs = [call for call in info.call_args_list if "Prefetch progress" in str(call)] | ||
| # Every chunk logs when the interval is zero: 3 files x 4 KB at a 1 KB | ||
| # chunk size means at least 12 heartbeats (short reads only add more). | ||
| assert len(progress_logs) >= 12 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Test a read that exceeds the heartbeat interval.
Setting _PREFETCH_LOG_INTERVAL_SEC to zero only tests completed-chunk callbacks. It cannot detect that a blocked readinto() call prevents heartbeat evaluation.
Add a controlled blocking-read test. Assert that progress logging continues before the read completes and remains rate-limited.
As per path instructions, test coverage must validate changed test behavior.
🤖 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 327 - 347, Extend test_prefetch_files_emits_progress_heartbeat with a
controlled blocking read that takes longer than _PREFETCH_LOG_INTERVAL_SEC,
using a synchronization mechanism to pause and resume the read. Assert that
“Prefetch progress” is logged before the blocked read completes, and verify the
number or timing of logs remains rate-limited rather than emitting continuously.
Keep the existing chunk-based heartbeat assertions intact.
Source: Path instructions
|
PR_Github #64699 [ run ] triggered by Bot. Commit: |
|
PR_Github #64699 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The diagnosis is right and the fix is minimal and behavior-preserving — main comment is on the mechanism (inline), and mostly asks you to record why readinto was chosen over mmap + MADV_POPULATE_READ in the PR description.
Two smaller points:
- The heartbeat text ends with
(local rank).loggeralready prefixes[RANK n], so the suffix reads as a stray token rather than telling you the totals are this rank's share of the checkpoint. SuggestPrefetch progress: X / Y GB (this rank's share)or dropping it. - The known limitation (a fully hung mount still logs nothing, because the heartbeat is progress-gated) is worth a one-line comment in
prefetch_filesnext toreport_progress, not just in the PR description — the next person to debug a silent stall will read the code, not the PR.
NVBug tag is fine; no docs/changelog owed for an internal log line.
| # those buffers accumulate into hundreds of GB across the local | ||
| # ranks, which can OOM the host. Chunked reads warm the OS page | ||
| # cache identically with a constant per-thread footprint. | ||
| buffer = memoryview(bytearray(_PREFETCH_CHUNK_SIZE_BYTES)) |
There was a problem hiding this comment.
Chunking fixes the OOM, but it still copies every byte out of the page cache into an anonymous buffer only to discard it — ~640GB of pointless memcpy per node, on top of the read itself.
The alternative for a pure "get these pages into the page cache" warm-up is mmap() the file read-only, madvise(MADV_POPULATE_READ) over it, then munmap(): the kernel populates the page cache directly, with no user-space copy and no anonymous buffer at all. The repo already has the pieces — tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py:564 does the ctypes libc.madvise call for MADV_POPULATE_WRITE including the EINVAL/ENOSYS fallback for kernels older than 5.14 (MADV_POPULATE_READ has the same 5.14 requirement), and tensorrt_llm/_torch/mmap_utils.py wraps madvise_range. The chunked readinto loop written here would be the fallback path, so it'd be additive rather than a rewrite, and progress reporting works the same way if you madvise in chunks.
That said, mmap on network filesystems is a known sore spot, and checkpoints here are typically on Lustre — if mmap/madvise was already ruled out for that reason (or for any other), that's a fine answer; please state it in the PR description so the choice is on record and nobody re-litigates it later. If it wasn't considered, it's worth a look given the memcpy volume.
Either way the description should say which it is.
| files.append(str(file)) | ||
|
|
||
| with mock.patch.object(wl.logger, "info") as info: | ||
| HfWeightLoader().prefetch_files(files) |
There was a problem hiding this comment.
prefetch_files shards its input with file_names[local_mpi_rank()::local_mpi_size()], so this test's >= 12 heartbeat count silently depends on the process not being launched under MPI — under local_mpi_size() > 1 this rank only reads 1 of the 3 files and the assertion fails.
Monkeypatch wl.local_mpi_rank -> 0 and wl.local_mpi_size -> 1 so the test asserts on a fixed file set regardless of how it's launched.
3 participants
Dev Engineer Review
QA Engineer Review
tests/integration/test_lists/coverage was reported.Description
What the problem actually was
HfWeightLoader._prefetch_one_filewarmed the OS page cache by callingf.read()— materializing each ~4 GB safetensors shard as a whole Python bytes object, purely to throw it away. With 16 prefetch threads per rank × 8 ranks, that's 128 concurrent multi-GB buffers. On fast storage they live for milliseconds. On slow storage (the H20-3e CI runners: zero of 128 files finished in ~30 minutes), all 128 buffers fill up simultaneously — up to ~0.5 TB of unreclaimable anonymous memory on top of the page cache the prefetch is deliberately creating. The prefetch gate (641GB < 0.9 × available) budgets nothing for this, and reads host-wide/proc/meminfobesides. The kernel OOM killer then takes out a rank mid-prefetch; MPI aborts; the launcher hangs silently; the test harness's 1800 s stall detector SIGKILLs it — and the log's last line isPrefetching model-00113-of-000163.safetensors to memory..., which everyone read as "mysterious OOM during a 641 GB prefetch that should have fit."The fix
Two changes in
weight_loader.py, deliberately behavior-neutral otherwise:Prefetch progress: X.XXGB / Y.YYGB (local rank), at most once per 60 s) so a slow-but-healthy prefetch emits output instead of the fatal silence — any future storage slowness produces a diagnosable log instead of a stall-kill mislabeled as OOM.Known limitation: the heartbeat is progress-gated. A hard-hung mount (no read completing at all) still produces silence and will still be stall-killed — arguably the correct outcome; a timer-thread heartbeat was deliberately left out of scope.
Why this doesn't reproduce on typical dev nodes
To reproduce it, the bug needs two environmental ingredients that a docker memory cap alone can't supply. The killer is the transient read buffers, and buffers only accumulate when storage is slow — a typical dev node prefetches at ~2.7 GB/s, completing files in seconds, so at most a few GB of buffers ever coexist; the affected CI runners read ~10× slower, letting all 128 buffers grow toward full size together. And a memory cap mostly caps page cache, which the kernel happily reclaims — clean page cache cannot OOM a cgroup — so capped repro runs sail through while pointing at the "641 GB > limit" red herring. The bug was never "the checkpoint doesn't fit"; it was "slow storage turns the prefetch's own scratch buffers into half a terabyte." This also makes the failure version-independent: the prefetch code is identical from 1.3.0rc17 through rc23, matching the CI history.
Test Coverage
Unit tests (new)
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py:test_prefetch_one_file_reads_full_file_in_bounded_chunks— full file consumed strictly chunk by chunk, including a trailing partial chunk.test_prefetch_one_file_missing_file_is_noop— missing files silently skipped, as before.test_prefetch_files_emits_progress_heartbeat— heartbeat fires per chunk when the interval is forced to zero.Testing scenario (before/after efficacy demo)
Same 8×H200 node, same container with a 400 GiB memory cap (
--memory=400g --memory-swap=400g), same test:perf/test_perf.py::test_perf[deepseek_r1_0528_fp8-bench-pytorch-float8-input_output_len:1000,1000-reqs:20000-ep:8-tp:8-gpus:8](one of the failing CI cases; ~641 GB checkpoint).dmesg:CONSTRAINT_MEMCG ... Killed process (python3) anon-rss:52992132kB); MPI aborts; the harness stall detector SIGKILLs the hung launcher ~30 minutes later — reproducing the exact CI signature (died with <Signals.SIGKILL: 9>during prefetch) end to end, including the misleading log shape.PASSED, cold-cache 641 GB prefetch included, with heartbeat lines pacing the prefetch and container memory staying two orders of magnitude below the cap.Validation on the affected CI hardware
The failing cases live in the QA weekly lane (
tests/integration/test_lists/qa/llm_perf_core.yml, H20-3e pool), which is not reachable from PR CI. A one-off run of the patched build on an H20-3e runner has been requested in NVBug 6272397; the heartbeat additionally guarantees that any residual failure there produces an attributable log instead of a silent SIGKILL.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.