Merge main into amitw/miles-nixl-upstream-pr

Resolves the serialized-args snapshots against radixark#4143's new optimizer arguments, keeping this branch's p2p_transport and NIXL seed flag.

Author: yueming-yuan

docs/advanced/disk-offload.md

@@ -74,7 +74,25 @@ each fp32 main shard's storage after creating its tensor handle, then Miles fill
 runtime bucket at a time directly in the final files. There is no temporary state file. Peak
 initialization HBM is the largest individual shard during construction and one main-only bucket
 afterward (normally at most about 200M fp32 elements, except for an oversized entry). Each
-initialized range is synchronized and evicted from page cache before continuing.
+initialized bucket is evicted from the page cache once its writeback completes.
+
+The step's writes are buffered, so by default the host page cache keeps whatever of the state
+it has room for. That is the fastest setting when host RAM can hold the whole state: the next
+step reads it from memory, and the kernel finishes the writeback during rollout. A colocated
+run pays for it elsewhere, since that cache is reclaimed lazily and races every pinned
+allocation at the rollout/train handoff. Cap it per rank:
+
+```bash
+--stream-optimizer-state-page-cache-mb 8192
+```
+
+A writer thread then writes back and drops each flushed bucket in order, so a bucket's
+writeback overlaps the buckets after it instead of stalling its own flush. A flush waits only
+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
+put both ends of each bucket's copy, the store's file and the checkpoint's, through the same
+cap, so one bucket's writeback overlaps the next copy, and likewise finish it before returning.
 
 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/rollout/filter_hub/sample_filters.py

@@ -36,20 +36,23 @@ def mask_truncated_zero_reward(args: Namespace, data: list[list[Sample]]) -> Non
     """
     global _warned_dict_reward
     total = masked = 0
-    for sample in iter_samples(data):
-        total += 1
-        if not args.reward_key and isinstance(sample.reward, dict):
-            # No scalar to compare without --reward-key; leave the sample in the loss.
-            if not _warned_dict_reward:
-                logger.warning(
-                    "mask_truncated_zero_reward: rewards are dicts and --reward-key is not set, "
-                    "so those samples are never masked"
-                )
-                _warned_dict_reward = True
-            continue
-        reward = sample.get_reward_value(args) if sample.reward is not None else None
-        if is_cut_off(sample) and (reward is None or reward <= 0):
-            sample.remove_sample = True
-            masked += 1
+    # ``data`` is a list of groups. With session server v2, a group entry is itself a list of
+    # samples (one per session branch), so open each group with ``iter_samples``.
+    for group in data:
+        for sample in iter_samples(group):
+            total += 1
+            if not args.reward_key and isinstance(sample.reward, dict):
+                # No scalar to compare without --reward-key; leave the sample in the loss.
+                if not _warned_dict_reward:
+                    logger.warning(
+                        "mask_truncated_zero_reward: rewards are dicts and --reward-key is not set, "
+                        "so those samples are never masked"
+                    )
+                    _warned_dict_reward = True
+                continue
+            reward = sample.get_reward_value(args) if sample.reward is not None else None
+            if is_cut_off(sample) and (reward is None or reward <= 0):
+                sample.remove_sample = True
+                masked += 1
     if masked:
         logger.info(f"mask_truncated_zero_reward: left {masked}/{total} cut-off zero-reward samples out of the loss")

miles/utils/args/configs/cluster.py

@@ -221,6 +221,21 @@ class ClusterConfig(BaseConfig):
             ),
         ),
     ] = "fp32"
+    stream_optimizer_state_page_cache_mb: A[
+        int | None,
+        Arg(
+            help=(
+                "Per-rank cap (MiB) on the streamed Adam state that may sit in the host page cache "
+                "waiting to be written back. With a cap, a writer thread writes back and drops each "
+                "flushed bucket off the step's critical path, a flush waits only while more than the "
+                "cap is still waiting, and the step finishes its writeback before returning, so none "
+                "of the store's files stay cached past a step and the next step reads them from disk. "
+                "Unset leaves the cache to the kernel, which is fastest when host RAM holds the whole "
+                "state, but a colocated run then races that cache at every pinned allocation. "
+                "Initialization never leaves its writes cached either way."
+            ),
+        ),
+    ] = None
     offload_train_disk_dir: A[
         str | None,
         Arg(

miles/utils/arguments.py

@@ -1405,7 +1405,14 @@ def miles_validate_args(args):
                 "offloader; pass --chunked-optimizer-state-offload and a non-zero "
                 "--optimizer-state-offload-fraction"
             )
+            assert args.stream_optimizer_state_page_cache_mb is None, (
+                "--stream-optimizer-state-page-cache-mb caps the Adam bucket store; Muon's disk state is "
+                "file-backed mappings the chunked offloader owns"
+            )
         else:
+            assert (
+                args.stream_optimizer_state_page_cache_mb is None or args.stream_optimizer_state_page_cache_mb >= 0
+            ), "--stream-optimizer-state-page-cache-mb must be >= 0"
             assert (
                 args.use_distributed_optimizer
             ), "--stream-optimizer-state-to-disk requires the distributed optimizer"
@@ -1448,9 +1455,16 @@ def miles_validate_args(args):
         assert (
             not args.enable_witness
         ), "--enable-witness reads the master optimizer's per-param state, which the NVMe store owns"
+        page_cache_mb = args.stream_optimizer_state_page_cache_mb
         logger.info(
             f"Streaming optimizer state to disk, dir={args.offload_train_disk_dir}, "
-            f"chunk={args.offload_train_disk_chunk_mb}MB, moments={args.stream_optimizer_state_moment_dtype}"
+            f"chunk={args.offload_train_disk_chunk_mb}MB, moments={args.stream_optimizer_state_moment_dtype}, "
+            f"page cache {'left to the kernel' if page_cache_mb is None else f'capped at {page_cache_mb}MB'}"
+        )
+    else:
+        assert args.stream_optimizer_state_page_cache_mb is None, (
+            "--stream-optimizer-state-page-cache-mb caps the Megatron Adam bucket store; it needs "
+            "--stream-optimizer-state-to-disk and is not supported on other backends"
         )
 
     if args.async_max_concurrent_samples is not None:

miles_plugins/optimizers/nvme_stream.py

@@ -38,6 +38,8 @@
 import os
 import shutil
 import time
+from collections import deque
+from concurrent.futures import Future, ThreadPoolExecutor
 from types import MethodType
 from typing import TYPE_CHECKING, NamedTuple
 
@@ -101,6 +103,55 @@ def _rw_full(op, fd: int, offset: int, buf) -> None:
         done += n
 
 
+def _drop_from_page_cache(fd: int) -> None:
+    """Write the file back and drop it from the host page cache; DONTNEED drops only clean pages."""
+    os.fdatasync(fd)
+    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 the page cache of the store's files and their checkpoint copies; no flush waits on its own writeback.
+
+    A writer thread drops each released file in release order; both syscalls release the GIL, so a file's
+    writeback overlaps the buckets after it. A release waits only while more than the cap is still waiting to be
+    dropped, and only for the oldest files. The bucket being read for the next step comes on top of the cap.
+    """
+
+    def __init__(self, cap_bytes: int):
+        self._cap_bytes = cap_bytes
+        self._writer = ThreadPoolExecutor(max_workers=1, thread_name_prefix="nvme-page-cache")
+        self._waiting: deque[tuple[int, Future]] = deque()
+        self._waiting_bytes = 0
+
+    def release(self, file: int | str, nbytes: int) -> None:
+        """Queue a file just written or read, ``nbytes`` of it now in the page cache, to be dropped.
+
+        ``file`` is a store file's fd, or the path of a checkpoint copy, which the writer opens and closes itself.
+        """
+        drop = _drop_path_from_page_cache if isinstance(file, str) else _drop_from_page_cache
+        self._waiting.append((nbytes, self._writer.submit(drop, file)))
+        self._waiting_bytes += nbytes
+        while self._waiting and (self._waiting_bytes > self._cap_bytes or self._waiting[0][1].done()):
+            self._retire_oldest()
+
+    def drain(self) -> None:
+        while self._waiting:
+            self._retire_oldest()
+
+    def _retire_oldest(self) -> None:
+        nbytes, dropped = self._waiting.popleft()
+        self._waiting_bytes -= nbytes
+        dropped.result()  # a failed writeback surfaces here, on the training thread
+
+
 def plan_buckets(entries_by_ddp_bucket: dict, limit: int = BUCKET_NUMEL_LIMIT) -> list[list[_Entry]]:
     planned, current, numel = [], [], 0
     for _, entries in sorted(entries_by_ddp_bucket.items(), key=lambda kv: kv[0]):
@@ -169,12 +220,15 @@ def __init__(self, path: str, entries: list[_Entry], adam, stager: _Stager, dtyp
         self.numel = sum(e.main_param.numel() for e in entries)
 
         self.offsets: dict[str, list[int]] = {}
+        self.segment_nbytes: dict[str, int] = {}  # page-aligned, which is what a segment occupies in the cache
         at = 0
         for segment in SEGMENTS:
             self.offsets[segment] = []
+            start = at
             for entry in entries:
                 self.offsets[segment].append(at)
                 at += _align(entry.main_param.numel() * dtypes[segment].itemsize)
+            self.segment_nbytes[segment] = at - start
         self.nbytes = at
         self.fd = _allocate_file(path, at)
         self.moments_ready = False
@@ -198,8 +252,10 @@ def _move(self, segments, *, to_disk: bool) -> int:
     def fetch(self) -> int:
         return self._move(SEGMENTS if self.moments_ready else SEGMENTS[:1], to_disk=False)
 
-    def flush(self, segments=SEGMENTS) -> int:
+    def flush(self, segments=SEGMENTS, *, page_cache: _PageCacheWindow | None) -> int:
         moved = self._move(segments, to_disk=True)
+        if page_cache is not None:
+            page_cache.release(self.fd, nbytes=sum(self.segment_nbytes[segment] for segment in segments))
         self.moments_ready = self.moments_ready or tuple(segments) == SEGMENTS
         return moved
 
@@ -227,9 +283,11 @@ def __init__(
         chunk_mb: int,
         moment_dtype: str = "fp32",
         allow_fresh_state: bool = False,
+        page_cache: _PageCacheWindow | None = None,
     ):
         self.dist_opt = distrib_optimizer
         self._allow_fresh_state = allow_fresh_state
+        self._page_cache = page_cache
         self.uid = NVMeOptimizerStateStore._next_uid
         NVMeOptimizerStateStore._next_uid += 1
         config = distrib_optimizer.config
@@ -444,10 +502,11 @@ def step(self) -> bool:
             bucket.adam.step()
             self._detach_grads(bucket.entries)
             self._copy_main_to_model_params(bucket.entries)
-            written += bucket.flush()
+            written += bucket.flush(page_cache=self._page_cache)
         if self._fp32_adam is not None:
             self._sync_lr_wd(self._fp32_adam, self._fp32_group_indices)
             self._fp32_adam.step()
+        self._drain_page_cache_window()
         logger.info(
             f"NVMe streaming step: {len(self.buckets)} buckets, read {read / 1024**3:.1f} GB, "
             f"wrote {written / 1024**3:.1f} GB in {time.monotonic() - started:.1f}s"
@@ -465,9 +524,12 @@ def initialize_main_from_model_params(self) -> int:
                 assert param_range.size == entry.main_param.nelement()
                 source_shard = entry.model_param.view(-1)[param_range.start : param_range.end]
                 entry.main_param.copy_(source_shard)
-            written += bucket.flush(segments=("main",))
-            os.fdatasync(bucket.fd)
-            os.posix_fadvise(bucket.fd, 0, bucket.offsets["exp_avg"][0], os.POSIX_FADV_DONTNEED)
+            written += bucket.flush(segments=("main",), page_cache=self._page_cache)
+            if self._page_cache is None:
+                # Initialization writes the whole main state in one pass, just before the checkpoint load and
+                # the engines' pinned allocations, so even with no cap it leaves none of that cached.
+                _drop_from_page_cache(bucket.fd)
+        self._drain_page_cache_window()
         return written
 
     @torch.no_grad()
@@ -476,7 +538,14 @@ def refresh_main_from_model_params(self, copy_fn) -> None:
             bucket.materialize_main()
         copy_fn()
         for bucket in self.buckets:
-            bucket.flush(segments=("main",))
+            bucket.flush(segments=("main",), page_cache=self._page_cache)
+        self._drain_page_cache_window()
+
+    def _drain_page_cache_window(self) -> None:
+        # Every operation on the store's files drains the window before it returns, so no file stays cached past it
+        # and a failed writeback fails that operation rather than a later one, or none after the last.
+        if self._page_cache is not None:
+            self._page_cache.drain()
 
     @torch.no_grad()
     def save_to(self, base_dir: str) -> None:
@@ -495,7 +564,14 @@ 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; both go through the window, so one bucket's
+                # writeback overlaps the next bucket's copy.
+                self._page_cache.release(bucket.fd, nbytes=bucket.nbytes)
+                self._page_cache.release(copy, nbytes=bucket.nbytes)
+        self._drain_page_cache_window()
         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:
@@ -536,11 +612,16 @@ 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)
+                self._page_cache.release(source, nbytes=bucket.nbytes)
             for group, step in zip(bucket.adam.param_groups, meta["steps"], strict=True):
                 if step:
                     group["step"] = step
             bucket.allocate_moments()
+        self._drain_page_cache_window()
         fp32_state = os.path.join(dirpath, "fp32_resident_optimizer.pt")
         if self._fp32_adam is not None and os.path.isfile(fp32_state):
             self._fp32_adam.load_state_dict(torch.load(fp32_state))
@@ -559,6 +640,10 @@ def setup_optimizer_state_streaming(args, optimizer) -> None:
 
     dir_root = optimizer_state_dir_root(args)
     purge_rank_dir(dir_root)
+    # One window per rank: the chained dense and expert stores share this process's page cache.
+    page_cache = None
+    if args.stream_optimizer_state_page_cache_mb is not None:
+        page_cache = _PageCacheWindow(cap_bytes=args.stream_optimizer_state_page_cache_mb * 1024 * 1024)
     stores = []
     for dist_opt in optimizer.chained_optimizers:
         assert isinstance(
@@ -576,6 +661,7 @@ def setup_optimizer_state_streaming(args, optimizer) -> None:
             args.offload_train_disk_chunk_mb,
             args.stream_optimizer_state_moment_dtype,
             allow_fresh_state=args.backend.no_load_optim,
+            page_cache=page_cache,
         )
         written = store.initialize_main_from_model_params()
         logger.info(