feat(optimizer): NVMe optimizer-state streaming as a miles plugin

docs/advanced/disk-offload.md

@@ -38,12 +38,75 @@ matching `TMS_DISK_BACKUP_*` environment, and reclaiming the files at startup an
 Because pause and resume happen at phase boundaries, everything is resident again by the
 time the optimizer step runs. If the binding constraint is instead that the optimizer
 state does not fit the GPU *during* the step, actor offload cannot help; that case is what
-Megatron's `--optimizer-state-nvme-dir` addresses (see
-[radixark/Megatron-LM#63](https://github.com/radixark/Megatron-LM/pull/63)), and the two
-compose.
+`--stream-optimizer-state-to-disk` addresses, and the two compose.
+
+## Streaming the optimizer state
+
+```bash
+--offload-train --offload-train-target disk \
+--stream-optimizer-state-to-disk \
+--offload-train-disk-dir /scratch/miles_offload
+```
+
+The two together are what this is for. Actor offload gets the paused actor out of HBM for
+the whole rollout window; streaming keeps the largest part of it — 12 bytes per parameter
+of fp32 main params and Adam moments — from being in HBM in the first place, including
+while the step runs, which offload cannot do. Data parallelism divides the optimizer state,
+so a run with GPUs to spare shards it small enough to stay resident and needs neither; the
+runs that need both are the ones tight enough to sit at DP=1. GLM-5.2 744B on 8 GB300 nodes
+is the shape of it: 279 GB of state per rank against 276.6 GB of HBM, on half the nodes the
+run would otherwise need.
+
+The fp32 main params and Adam moments live in per-bucket files instead of on the GPU.
+Each step brings in one bucket, updates it, and writes it back, so peak residency is one
+bucket rather than the whole state. Buckets are capped at 200M elements independently of
+DDP's bucket sizes, which reach tens of GB at DP=1. Native-fp32 model params (a router's
+`expert_bias`, a GDN/Mamba `A_log`) stay GPU-resident under a small separate Adam: they
+are tiny, and unlike the bf16 path their optimizer shards alias the model params directly.
+
+`fp32` storage is bit-identical to keeping the state on GPU, so turning this on does not
+change results — it trades step time for memory. The step is I/O bound and the moments
+tolerate less precision than the master copy, so they can be stored narrower:
+
+```bash
+--stream-optimizer-state-moment-dtype bf16
+```
+
+That cuts streaming volume by a third (12 bytes per param to 8). A checkpoint records the
+dtypes it was written with and a resume verifies them, so bytes written as bf16 can never
+be read back as fp32. The fp8 options work but are not recommended: `exp_avg_sq` needs
+per-block scaling to survive 8-bit storage, which this does not implement.
+
+Two limits to know about. Resume is same-topology only — the on-disk layout follows this
+rank's DP shard, so changing TP/PP/DP/EP fails the layout assert rather than resharding.
+And the optimizer state is copied to the checkpoint directory synchronously, outside
+`--async-save`, so expect checkpoint saves to take noticeably longer.
+
+The two also help each other. With the optimizer state already on disk there is that much
+less to move when the actor is paused: on Qwen3-30B-A3B, sleep/wake went from 24s/8.9s to
+5.2s/1.3s once the paused actor no longer carried the state.
 
 ## Choosing
 
-- Host RAM holds the paused actor: keep the default `--offload-train-target=cpu`. It is
-  faster and needs no scratch space.
-- Host RAM does not hold it: switch to `--offload-train-target=disk`.
+`--offload-train` is the base mechanism, and it is not tied to colocation: the actor is
+resident only during `train()` and sleeps for the whole rollout window either way. What
+changes is who takes the freed HBM. Colocated, it is the engine on the same GPUs, so
+offload is mandatory and defaults on. Disaggregated, it is whatever else you fit on the
+training GPUs, which is the point when you are sizing engines against a fixed cluster.
+
+- Paused actor fits in host RAM: keep `--offload-train-target=cpu` (default). Fastest, no
+  scratch space.
+- It does not fit: `--offload-train-target=disk`.
+- On top of that, the optimizer state does not fit the GPU *during the step*, which offload
+  cannot help with: add `--stream-optimizer-state-to-disk`, and consider
+  `--stream-optimizer-state-moment-dtype bf16` to claw back some of the I/O cost.
+
+Streaming requires `--offload-train-target=disk`; miles asserts it. The two answer the same
+pressure and are only deployed as a pair, so that is the only combination covered.
+Streaming on its own does run — it is the one case where offload has nothing to do, because
+nothing else wants the training GPUs during rollout — but nothing validates it.
+
+Both mechanisms share `--offload-train-disk-dir` and `--offload-train-disk-chunk-mb`. Point the
+directory at real node-local NVMe: a tmpfs mount, which `/tmp` is on many systems, keeps
+the data in RAM and defeats both. The chunk is a pinned host staging buffer and each
+mechanism allocates its own, so enabling both costs 2x that per rank.

miles/backends/megatron_utils/model.py

@@ -181,6 +181,12 @@ def setup_model_and_optimizer(
             model_chunks=model,
             use_gloo_process_groups=args.enable_gloo_process_groups,
         )
+
+    if args.stream_optimizer_state_to_disk:
+        from miles_plugins.optimizers.nvme_stream import setup_optimizer_state_streaming
+
+        setup_optimizer_state_streaming(args, optimizer)
+
     opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer)
     return model, optimizer, opt_param_scheduler
 

miles/ray/train/actor_factory.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 
 import ray
 from ray.util.placement_group import PlacementGroup
@@ -48,6 +49,10 @@ def allocate_gpus_for_actor(
         env_vars["LD_PRELOAD"] = dynlib_path
         env_vars["TMS_INIT_ENABLE"] = "1"
         if args.offload_train_target == "disk":
+            assert b"TMS_INIT_ENABLE_DISK_BACKUP" in Path(dynlib_path).read_bytes(), (
+                f"{dynlib_path} has no disk backend; reinstall torch_memory_saver at the commit "
+                f"docker/Dockerfile pins."
+            )
             env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "0"
             env_vars["TMS_INIT_ENABLE_DISK_BACKUP"] = "1"
             env_vars["TMS_DISK_BACKUP_CHUNK_MB"] = str(args.offload_train_disk_chunk_mb)

miles/utils/arguments.py

@@ -138,13 +138,48 @@ def add_cluster_arguments(parser):
                     "NVMe (--offload-train-disk-dir) for the case where even CPU RAM cannot hold it."
                 ),
             )
+            parser.add_argument(
+                "--stream-optimizer-state-to-disk",
+                action="store_true",
+                help=(
+                    "Stream the fp32 main params and Adam moments through per-bucket files on "
+                    "node-local NVMe during optimizer.step(), bounding GPU residency to one bucket. "
+                    "For when the optimizer state does not fit the GPU *while the step runs*: "
+                    "--offload-train-target=disk cannot help there, because pause/resume happen at "
+                    "phase boundaries and everything is resident again by the time Adam launches. "
+                    "Bit-identical to keeping the state on GPU, at the cost of disk traffic every "
+                    "step. Distinct from --offload-optimizer-states and --optimizer-cpu-offload, "
+                    "and mutually exclusive with both."
+                ),
+            )
+            parser.add_argument(
+                "--stream-optimizer-state-moment-dtype",
+                type=str,
+                default="fp32",
+                choices=["fp32", "bf16", "fp16", "fp8e4m3", "fp8e5m2"],
+                help=(
+                    "On-disk dtype for the streamed Adam moments; the fp32 master copy is always "
+                    "fp32. This is a serialization format, not a compute precision: the step still "
+                    "hands fp32 tensors to FusedAdam, and the cast happens on the way to and from "
+                    "disk. That makes it distinct from --exp-avg-dtype / --exp-avg-sq-dtype, which "
+                    "change what the optimizer holds and require --use-precision-aware-optimizer, "
+                    "so they cannot be combined with streaming at all. "
+                    "The step is I/O bound and the moments tolerate less precision than the master "
+                    "copy, so bf16 cuts streaming volume by a third (12 bytes per param to 8). "
+                    "fp32 is bit-identical to keeping the moments on GPU. The fp8 options need "
+                    "per-block scaling for exp_avg_sq to be sound, which this does not implement, "
+                    "and are not recommended."
+                ),
+            )
             parser.add_argument(
                 "--offload-train-disk-dir",
                 type=str,
                 default=None,
                 help=(
-                    "Node-local directory for train disk-offload files when "
-                    "--offload-train-target=disk. Should be fast local NVMe (e.g. /scratch). "
+                    "Node-local directory for the disk-offload files, used by both "
+                    "--offload-train-target=disk and --stream-optimizer-state-to-disk (each under "
+                    "its own subdirectory). Should be fast local NVMe (e.g. /scratch); a tmpfs "
+                    "mount, which /tmp is on many systems, keeps the data in RAM and defeats both. "
                     "Files are per-process and overwritten in place every step (bounded size); "
                     "defaults to $SCRATCH/miles_train_offload_<uid>."
                 ),
@@ -154,8 +189,10 @@ def add_cluster_arguments(parser):
                 type=int,
                 default=256,
                 help=(
-                    "Chunk size (MiB) for streaming the training actor GPU<->disk in disk-offload "
-                    "mode. Bounds the pinned host staging buffer regardless of the total offloaded size."
+                    "Chunk size (MiB) for the GPU<->disk transfers, i.e. the pinned host staging "
+                    "buffer, which bounds host memory regardless of how much is moved. Used by both "
+                    "--offload-train-target=disk and --stream-optimizer-state-to-disk, and each "
+                    "allocates its own, so enabling both costs 2x this per rank."
                 ),
             )
 
@@ -2853,6 +2890,38 @@ def miles_validate_args(args):
             f"chunk={args.offload_train_disk_chunk_mb}MB"
         )
 
+    if args.stream_optimizer_state_to_disk:
+        assert args.use_distributed_optimizer, "--stream-optimizer-state-to-disk requires the distributed optimizer"
+        assert (
+            args.optimizer == "adam"
+        ), f"--stream-optimizer-state-to-disk requires --optimizer adam, got {args.optimizer}"
+        assert not args.multi_lora, "--stream-optimizer-state-to-disk does not support multi-LoRA"
+        assert not args.optimizer_cpu_offload, "--stream-optimizer-state-to-disk excludes --optimizer-cpu-offload"
+        assert (
+            not args.offload_optimizer_states
+        ), "--stream-optimizer-state-to-disk excludes --offload-optimizer-states"
+        assert (
+            not args.use_precision_aware_optimizer
+        ), "--stream-optimizer-state-to-disk requires mcore to hold the fp32 main params"
+        assert not args.reset_optimizer_states, (
+            "--reset-optimizer-states walks the master optimizer's state, which the NVMe store "
+            "leaves empty, so the reset would silently do nothing"
+        )
+        assert not args.save_local_weight_checksum, (
+            "--save-local-weight-checksum reads param.main_param, whose storage the NVMe store "
+            "resizes to 0 between steps"
+        )
+        assert (
+            not args.enable_witness
+        ), "--enable-witness reads the master optimizer's per-param state, which the NVMe store owns"
+        assert args.offload_train_target == "disk", (
+            "--stream-optimizer-state-to-disk is only supported alongside " "--offload-train-target=disk; pass both."
+        )
+        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}"
+        )
+
     if args.async_max_concurrent_samples is not None:
         assert args.async_max_concurrent_samples >= args.n_samples_per_prompt, (
             f"--async-max-concurrent-samples ({args.async_max_concurrent_samples}) must be at least "