Skip to main content
Colocated RL keeps a training actor and a rollout engine on the same GPUs, so the actor must get out of the way while the engine generates. miles offloads the actor during that window, and by default the backup lives in pinned host memory. For large models that copy does not fit, and disk offload is the alternative.

​
Usage

Instead of a pinned host copy, the paused actor is streamed to per-rank files through a fixed-size pinned staging buffer, so host memory stays bounded by --offload-train-disk-chunk-mb regardless of how much is offloaded. Each rank writes to its own directory under --offload-train-disk-dir (defaults to $SCRATCH/miles_train_offload_<uid>), the files are overwritten in place every step, and they are removed when the actor exits. Point the directory at real node-local NVMe. A tmpfs mount (including /tmp on many systems) keeps the backup in RAM and defeats the purpose.

​
How it works

This runs on torch_memory_saver, which hooks the allocator, so it does not care what the memory holds — weights, gradient buffers and optimizer state all move as one block when the actor is paused, and come back on resume. miles’ part is choosing the per-rank directory, launching each actor with the matching TMS_DISK_BACKUP_* environment, and reclaiming the files at startup and exit. 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 --stream-optimizer-state-to-disk addresses, and the two compose.

​
Streaming the optimizer state

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. For the Adam/DistributedOptimizer path, streaming also bounds initialization: Megatron releases each fp32 main shard’s storage after creating its tensor handle, then Miles fills one existing 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. 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:
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. Three 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. A checkpoint written before streaming was enabled cannot be resumed with it: the streamed state is the only optimizer state read, so miles refuses rather than silently restarting Adam from zero — pass --no-load-optim to accept a fresh optimizer state. 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

--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 on its own is the disaggregated case: nothing else wants the training GPUs during rollout, so there is no actor to park and --offload-train-target is never read. Pass --stream-optimizer-state-to-disk alone there. Under --offload-train, though, the two go together — a run that cannot hold the optimizer state on the GPU for the step will not hold a pinned host copy of the whole actor either — and miles asserts that pairing. 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.