Repository navigation
feat(optimizer): cap the host page cache of NVMe-streamed optimizer state - #4143
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
|
@claude review always |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_main_init.py |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_streamed_grads.py |
|
✅ |
|
✅ |
77bbece to
3c1dc76
Compare
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_main_init.py |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_streamed_grads.py |
|
/rerun-test tests/e2e/megatron/test_qwen3_4B_offload_disk_stream.py |
|
✅ |
|
✅ |
|
✅ |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_main_init.py |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_streamed_grads.py |
|
/rerun-test tests/e2e/megatron/test_qwen3_4B_offload_disk_stream.py |
|
✅ |
|
✅ |
|
✅ |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_streamed_grads.py |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_main_init.py |
|
/rerun-test tests/e2e/megatron/test_qwen3_4B_offload_disk_stream.py |
|
✅ |
|
✅ |
|
✅ |
…tate The NVMe store writes each bucket with buffered pwritev, so after a step up to the whole fp32 main and moment state sits in the host page cache. A colocated run reclaims that cache lazily while engines make pinned allocations, and memory guards that count file cache kill workers over it. Syncing on every flush bounds it but puts each bucket's writeback on the step's critical path. --stream-optimizer-state-page-cache-mb caps it per rank. A flush hands its bucket file to one writer thread, which runs fdatasync then POSIX_FADV_DONTNEED in flush order; both release the GIL, so each writeback overlaps the next buckets' fetch and Adam step. The training thread waits only for the oldest file, and only while more than the cap is not yet written back; writer errors re-raise there. The window is shared by a rank's chained stores. Unset keeps the current behavior, including initialization dropping each bucket it writes, now through the same helper. The flag is rejected without streaming, below zero, and with Muon, whose disk state the chunked offloader owns. sync_file_range is not used because it is a no-op on overlayfs.
…oint copies Review of the capped window found gaps. A writeback failure in a step's last buckets only surfaced on a later flush, or never after the final step, because nothing drained the window in production. And save_to/load_from copy every bucket file through the page cache outside the window, so a capped resume left the whole restored state dirty and cached until each bucket was re-flushed. step() and capped initialization now drain the window before returning, so their writeback and its errors stay inside them, and save_to/load_from release each bucket file to the window after copying it. A flush is charged the page-aligned size of the segments it wrote, which is what they occupy in the cache, and flushed() is renamed release() since a checkpoint save only reads the file. The residency tests probe evictability with raw syscalls, skip on tmpfs where pages cannot leave the cache, and a new GPU test drives a capped store through step, save and load.
…he cache With a cap set, save_to and load_from already release the store's own bucket files to the window, but the other end of each copy, the checkpoint's bucket file, still stayed in the page cache: a resume read the whole saved state into it and a save wrote the whole state into it, right before the rollout handoff the cap is meant to protect. Under a cap, each checkpoint-side copy is now written back and dropped as soon as it is copied, through a short-lived fd since the window needs one that outlives the call. The store test checks that the checkpoint's bucket files are uncached after both a save and a load; it fails without this change.
…apshots setup_optimizer_state_streaming reads --stream-optimizer-state-page-cache-mb directly, but make_nvme_setup_args builds a Namespace without it, so test_nvme_setup_backend_ownership failed with AttributeError before reaching its load-policy assertions. The fixture now sets it to None, the uncapped default. The new ClusterConfig field also lands in every worker's serialized --config args, so the miles-run chart and helm-backend launch snapshots gain "stream_optimizer_state_page_cache_mb":null next to the moment dtype.
Only step() and initialize_main_from_model_params drained the window. The other three operations on the store's files left up to the cap pending: refresh_main_from_model_params (the --load-main-params-from-ckpt initialization), load_from (the whole restored state is dirty), and save_to. Their writeback errors then surfaced inside the next step, after it had fetched those buckets, and the capped-initialization promise did not hold for the refresh path. All five now finish through one _finish_writeback(), so none of them leaves a store file cached or a writeback error behind. _Bucket.flush takes the window as a required keyword, so a new flush call cannot skip the cap by omission. Tests: the page-cache helpers move to nvme_setup_fixtures.py, so the GPU files no longer import private helpers from another test module. A CPU test holds the writer thread and checks that releases within the cap return while writeback is still running; it fails if release() waits on writeback. The GPU store test drops its manual drains after save and load and also covers a refresh, so it checks what production leaves behind.
Without a cap, _finish_writeback() finishes no writeback at all, and in save_to it sits right before the manifest write, where it reads like a durability barrier that only the capped path provides. It is now _drain_page_cache_window(), which is what it does in both modes. The GPU store test's docstring also names the four operations it covers instead of claiming every path, since capped initialization is not among them.
A capped save dropped each checkpoint copy synchronously, so its fdatasync finished before the next bucket's copy started and the whole save became serial writeback. The window now takes a checkpoint copy by path and its writer opens, drops, and closes it, so one bucket's writeback overlaps the next copy; resume does the same for the checkpoint files it reads.
…m branch The cap's checks ran before the block that still allows torchtitan streaming. The non-negative check now sits in the Adam branch, and one assert after the Megatron branch rejects the cap everywhere else.
047d06d to
a95f88f
Compare
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_streamed_grads.py |
|
/rerun-test tests/fast-gpu/test_nvme_optimizer_main_init.py |
|
/rerun-test tests/e2e/megatron/test_qwen3_4B_offload_disk_stream.py |
|
✅ |
|
✅ |
|
✅ |
Resolves the serialized-args snapshots against radixark#4143's new optimizer arguments, keeping this branch's p2p_transport and NIXL seed flag.
Re-opens #4122 (merged by mistake, reverted in #4142).
Summary
Add
--stream-optimizer-state-page-cache-mbto cap the streamed Adam state held in host page cache.Motivation
--stream-optimizer-state-to-diskwrites each bucket with bufferedpwritev, so after a step up to the whole fp32 state sits in the host page cache. A colocated run reclaims it lazily while engines make pinned allocations, and guards that count file cache kill workers over it. Syncing every flush bounds it but puts each bucket's writeback on the step's critical path.Usage
Unset keeps today's behavior: no writer thread, the kernel manages the cache, and initialization still drops each bucket it writes. The cap is per rank, needs Megatron Adam streaming, and is rejected with Muon or other backends.
Design Notes
fdatasyncthenPOSIX_FADV_DONTNEEDin flush order, so each writeback overlaps the next buckets' work; the training thread waits only for the oldest file, only while more than the cap is not yet written back.sync_file_rangeis not used, as it is a no-op on overlayfs (measured, kernel 6.8).Verification
pre-commit run --fileson the changed files: passes._PageCacheWindowwith torch/Megatron stubbed out: the new within-cap guard's check passes, and fails against a synchronousrelease().Review Focus
_PageCacheWindow.release: waits only past the cap, only for the oldest files.NVMeOptimizerStateStore._drain_page_cache_window: ends every operation on the store's files.