fix(optimizer): drop the checkpoint's copy of the capped state from the 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.

Author: guapisolo

docs/advanced/disk-offload.md

@@ -91,8 +91,7 @@ writeback overlaps the buckets after it instead of stalling its own flush. A flu
 while more than the cap is still waiting to be written back; the bucket being read comes on top
 of it. The step finishes its writeback before it returns, so none of the store's files stay
 cached past a step and each step reads its state back from disk. Checkpoint save and resume
-copy the store's files through the same cap; the checkpoint's own copy is cached like any other
-checkpoint file.
+copy the store's files through the same cap and drop the checkpoint's copy from the cache too.
 
 With `fp32` storage the state round-trips bit for bit and only the grad norm's summation
 order differs from keeping it on GPU, so turning this on changes results by rounding at

miles_plugins/optimizers/nvme_stream.py

@@ -154,6 +154,14 @@ def _drop_from_page_cache(fd: int) -> None:
     os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
 
 
+def _drop_path_from_page_cache(path: str) -> None:
+    fd = os.open(path, os.O_RDONLY | os.O_CLOEXEC)
+    try:
+        _drop_from_page_cache(fd)
+    finally:
+        os.close(fd)
+
+
 class _PageCacheWindow:
     """Caps how much of the store's files the host page cache holds, without a flush waiting on its own writeback.
 
@@ -588,9 +596,13 @@ def save_to(self, base_dir: str) -> None:
             ],
         }
         for bucket in self.buckets:
-            shutil.copyfile(bucket.path, os.path.join(dirpath, os.path.basename(bucket.path)))
+            copy = os.path.join(dirpath, os.path.basename(bucket.path))
+            shutil.copyfile(bucket.path, copy)
             if self._page_cache is not None:
+                # Both ends of the copy are a full bucket in the cache: the store's file through the window, the
+                # checkpoint's at once, since its fd does not outlive this loop.
                 self._page_cache.release(bucket.fd, nbytes=bucket.nbytes)
+                _drop_path_from_page_cache(copy)
         if self._fp32_adam is not None:
             torch.save(self._fp32_adam.state_dict(), os.path.join(dirpath, "fp32_resident_optimizer.pt"))
         with open(os.path.join(dirpath, "manifest.json"), "w") as f:
@@ -631,9 +643,11 @@ def load_from(self, base_dir: str) -> bool:
         for bucket, meta in zip(self.buckets, manifest["buckets"], strict=True):
             assert meta["numel"] == bucket.numel
             assert meta["entry_numels"] == [e.main_param.numel() for e in bucket.entries]
-            shutil.copyfile(os.path.join(dirpath, meta["file"]), bucket.path)
+            source = os.path.join(dirpath, meta["file"])
+            shutil.copyfile(source, bucket.path)
             if self._page_cache is not None:
                 self._page_cache.release(bucket.fd, nbytes=bucket.nbytes)
+                _drop_path_from_page_cache(source)
             for group, step in zip(bucket.adam.param_groups, meta["steps"], strict=True):
                 if step:
                     group["step"] = step

tests/fast-gpu/test_nvme_optimizer_streamed_grads.py

@@ -190,6 +190,16 @@ def test_capped_store_leaves_no_page_cache_after_a_step_save_or_load(tmp_path):
     def resident() -> list[int]:
         return [_resident_pages(bucket.fd) for bucket in store.buckets]
 
+    def checkpoint_resident() -> int:
+        total = 0
+        for path in (tmp_path / "ckpt").rglob("bucket*.bin"):
+            fd = os.open(path, os.O_RDWR)
+            try:
+                total += _resident_pages(fd)
+            finally:
+                os.close(fd)
+        return total
+
     try:
         (norm,) = store.grads_for_norm()
         store.record_grad_norm(norm if multi_tensor_scale_tensor_impl is not None else float(norm))
@@ -200,12 +210,14 @@ def resident() -> list[int]:
         store.save_to(str(tmp_path / "ckpt"))
         store._page_cache.drain()
         assert resident() == [0, 0]
+        assert checkpoint_resident() == 0, "the checkpoint's copy is dropped too"
 
         for bucket in store.buckets:
             os.pwrite(bucket.fd, bytes(bucket.nbytes), 0)
         assert store.load_from(str(tmp_path / "ckpt"))
         store._page_cache.drain()
         assert resident() == [0, 0]
+        assert checkpoint_resident() == 0, "reading the checkpoint leaves it uncached"
         for bucket, before in zip(store.buckets, stepped, strict=True):
             torch.testing.assert_close(_read_main(bucket, 0, before.numel()), before, atol=0, rtol=0)
     finally: