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

[yueming-yuan]

Stacked on #1575 (base will retarget to main when that merges). Pairs with radixark/Megatron-LM#63, which is updated alongside this.

Why

--offload-train-target=disk (#1575) spills the paused actor at phase boundaries. That cannot help when the optimizer state does not fit the GPU while the step runs — by the time the Adam kernel launches, everything is resident again. Streaming solves that instead: the fp32 main params and Adam moments live in per-bucket files, and each step brings in one bucket, updates it, and writes it back.

The mechanism was written and validated in Megatron-LM#63 (GLM-5.2 744B on 8xGB300, Qwen3.5-35B-A3B on 8xH200). This PR moves the implementation into miles and shrinks what Megatron has to carry.

What moves

megatron/core/optimizer/nvme_state_store.py -> miles_plugins/optimizers/nvme_stream.py, plus the wiring. The five if self._nvme_state_store is not None branches inside DistributedOptimizer become instance bindings in _bind():

entry point	binding
step_with_ready_grads	store.step() + the param all-gather tail
reload_model_params	routed through store.refresh_main_from_model_params
state_dict / load_state_dict / sharded_state_dict	empties; save_to/load_from carry the real bytes

Megatron keeps one thing: the two checkpoint hooks in checkpointing.py, which have to be inside save_checkpoint/load_checkpoint because miles has two save call sites and a missed one shows up as a checkpoint with no optimizer state that silently resumes from zero. Its optimizer diff is empty — #63 is down to one file, +19 lines.

The per-bucket copy is adapted here, not extracted there

step() has to get each bucket's updated mains into the param buffer before flush() releases their storage, so it needs a subset version of _copy_main_params_to_model_params. The first version of this work extracted one in Megatron and called it. That was wrong on two counts:

1. Its only live caller under streaming was the store. Once step_with_ready_grads is bound to store.step(), Megatron's own _copy_main_params_to_model_params is never reached, so the extracted method existed there purely to serve an out-of-repo caller.
2. It handed out a partial operation. The fp8 branch relies on quantize_param_shard() having already run over the whole fp8 param set, and that is a DP collective which cannot be hoisted into a per-subset method — plan_buckets cuts on each rank's own shard sizes, so bucket counts differ across ranks and the calls would not line up. A caller taking only the subset method therefore skips fp8 weight updates with no error. The # quantized in the above quantize_param_shard function comment also lost its referent the moment "above" became a different method.

So _copy_main_to_model_params lives here, pinned by commit and line range in a comment (4716f754#L2469-L2519), following the same convention as miles/utils/seqlen_balancing.py and the miles_plugins/models/* kernels. The precondition it cannot satisfy is enforced next to it: the constructor rejects fp8 params and the MXFP8 param all-gather. Drift on a Megatron bump is the cost — the pinned link makes it auditable, and the assert world_range.size == shard_main_param.nelement() is the tripwire.

miles' own coupling is two lazy imports; grep -rn nvme_stream miles/ lists all of it.

Arg surface: one new switch, nothing renamed

--offload-train --offload-train-target disk      # actor, from #1575
--stream-optimizer-state-to-disk                 # optimizer state, new
--offload-train-disk-dir DIR                     # from #1575, now used by both
--offload-train-disk-chunk-mb N                  # from #1575, now used by both
--stream-optimizer-state-moments {fp32,bf16,...} # advanced, default fp32

Streaming requires --offload-train-target=disk, asserted at startup. Both mechanisms answer the same pressure — not enough HBM for the model — and are only ever deployed as a pair, so that is the only combination covered. Streaming on its own does run (#63's matrix has a passing row for it) and it is the one case where offload has nothing to do, because nothing else wants the training GPUs during rollout; but nothing validates it, and failing at startup beats someone discovering that at 3am. Since the assert makes streaming imply disk offload, --offload-train-disk-dir and --offload-train-disk-chunk-mb are reused unchanged with each mechanism taking its own subdirectory of the shared root, and #1575's validation block is untouched — the diff against it is two help strings.

The one thing that does not carry over is Megatron's --optimizer-state-nvme-dir, where "enable" and "where" were the same knob, so there could be no default location. It also read too close to the three flags miles already inherits (--offload-optimizer-states, --optimizer-cpu-offload, --rl-offload-optimizer-during-inference), two of which are mutually exclusive with this. stream encodes the actual distinction: offload happens at phase boundaries, streaming happens inside optimizer.step().

Moments default to fp32, so enabling streaming does not change results — it is bit-identical to GPU-resident state and only trades step time for memory. bf16 stays an explicit perf opt-in. The fp8 options are kept (#63 measured fp8e4m3 at ~1.4x bf16's deviation, 6/6 steps NORMAL) but documented as not recommended. The dtype is a constructor argument now, so getattr(config, "optimizer_state_nvme_moment_dtype", "fp32") and its silent fallback are gone.

Guards

Three miles-side readers go through the main params or the master optimizer's per-param state, both of which the store owns — mains have 0-sized storage between steps and the master Adam never holds state. They were unguarded and would have returned emptiness rather than failing:

• --reset-optimizer-states (walks chained_optimizer.optimizer.state.values()) — the reset would silently do nothing
• --save-local-weight-checksum (reads param.main_param)
• --enable-witness (writes main_param.data[local_idx])

The disk-offload requirement above also subsumes the --colocate without --offload-train case, which #63 hit as an sglang resume_memory_occupation OOM mid-run. Plus the usual exclusions (--optimizer-cpu-offload, --offload-optimizer-states, precision-aware optimizer, multi-LoRA, non-Adam, non-distributed optimizer).

_bind also asserts reuse_grad_buf_for_mxfp8_param_ag is off: the bound step_with_ready_grads drops Megatron's branch to _copy_main_params_to_param_buffer.

Checkpointing

save_to/load_from now take the checkpoint base and derive their own per-rank subdirectory from the same relative_dir the live scratch directory uses, so the layout is spelled once instead of once per repo (Megatron had rank0000_opt0_0 flat while the store had rank0/opt0_0 nested). Megatron no longer reads store.dist_opt or store.uid; the duck-typed contract is four methods. The "this checkpoint has no streamed state" case moves into load_from, which is a policy decision — it keeps pre-streaming checkpoints loadable — and now lives in one place.

Known limits, documented in docs/advanced/disk-offload.md

• Same-topology resume only. The on-disk layout follows this rank's DP shard with no global parameter identity, so changing TP/PP/DP/EP fails the layout assert rather than resharding. This is a regression against the torch_dist path for runs that enable streaming; fixing it means emitting per-param ShardedTensors, which sharded_state_dict's build-the-whole-plan-up-front shape does not currently allow.
• Checkpoint save is synchronous. The hook sits before the --async-save block in Megatron, so the state copy blocks. Byte volume is unchanged from today (dist-ckpt writes the same 12 B/param, and bf16 moments write less) — what is lost is the async parallel writer. Separate PR.

Guard added

--offload-train-target=disk sets TMS_INIT_ENABLE_CPU_BACKUP=0 alongside TMS_INIT_ENABLE_DISK_BACKUP=1. A torch_memory_saver built before fzyzcjy/torch_memory_saver#80 ignores the disk variable but still honours the cpu one, so pause() frees the actor with no backup anywhere and resume() returns uninitialized memory. Nothing fails at that point — the weights are garbage and it surfaces a rollout later as nan log-probs and then found NaN in local grad norm for bucket #0, which reads like a model bug. actor_factory now asserts the LD_PRELOADed binary knows the env var (checked against the binary, not the Python API, since the binary is what reads it). radixark/miles:dev currently ships 0.0.9.post1 and trips it.

Testing

Qwen3-30B-A3B RL on 8xH200 (h200-sci-k8s devbox), colocate, TP4/EP8/DP8, bf16, --num-rollout 3, dapo-math-17k, eval off. Both arms identical except for streaming. Because streaming is mutually exclusive with --optimizer-cpu-offload / --use-precision-aware-optimizer, which the launcher's ("H100", 1) preset adds, those three flags were dropped from both arms — so the baseline holds the whole optimizer state on GPU, which is the intended contrast.

3/3 train steps outcome=NORMAL valid_step=true on all 8 ranks in both arms. Weights round-trip correctly: train_rollout_logprob_abs_diff 0.01803 / 0.01864 (baseline) vs 0.01804 / 0.01875 / 0.01896 (streaming).

Peak is torch.cuda.max_memory_allocated() bracketing the actor_train timer, max over the 8 ranks; steps 1-2 only, since step 0 is lower until Adam creates the moments.

	baseline (disk offload only)	streaming, bf16 moments	delta
peak allocated	81.14 / 80.84 GB	38.86 / 38.64 GB	-42.3 GB (-52%)
peak reserved	95.08 / 95.16 GB	54.00 / 52.64 GB	-41.8 GB
actor_train	39.9 / 42.9 s	55.6 / 59.4 s	+16.1 s (+39%)
sleep (offload)	45.6 / 53.1 s	8.9 / 8.3 s	-5.3x
wake_up	27.4 s	3.7 s	-7.4x

Both numbers check out against what the store reports, so they are not coincidence:

• Memory. The store holds 28.5 GB/rank on disk (dense 1.5 + expert 27.0) at 8 B/param with bf16 moments, i.e. 3.56 G params/rank. GPU-resident at 12 B/param that is 42.8 GB; measured reduction 42.3 GB, within 1.2%.
• Time. Self-reported I/O per step is 0.9 s (dense) + ~14.3 s (expert) = ~15.2 s/rank against a +16.1 s train-time delta; the rest is the per-bucket Adam launches and the copy.

The docs' "trades step time for memory" is too pessimistic for the colocated pairing. With the state already on disk the paused actor has ~43 GB/rank less to spill, so sleep/wake collapse and the cycle gets faster:

per rollout (train + sleep + wake)   baseline: 41.4 + 49.4 + 27.4 = 118.2 s
                                    streaming: 57.5 +  8.6 +  3.7 =  69.8 s   (-41%)

The lazy-moments path is exercised rather than assumed: the expert store reads 13.5 GB on the first step (= 27.0 x 4/8, mains only) and 27.0 GB from the second on. Both ChainedOptimizer members get their own store (opt0_0 dense / opt0_1 expert), confirming the per-process uid disambiguation.

Not covered: fp32 moments (should be bit-identical, worth running as the correctness control); checkpoint save/load, since --save-interval 20 never fires in 3 steps, so the synchronous copy noted above is still unmeasured. One streaming attempt lost its rollout engine at step 3 to sglang's own couldn't get a response from detokenizer for last 20 s liveness check; it did not recur on rerun and host RAM was never under pressure (103 of 2009 GB), so it is not attributed to streaming, but not ruled out either.

Second arm: fp32 moments, 4xB300, fp8 rollout

Closes the fp32-moments control the arm above lists as not covered, on different hardware, topology and rollout precision: 4xB300, TP4/PP1/CP1/EP4, colocate, --rollout-fp8 with bf16 training, --offload-train-target disk, --num-rollout 3, dapo-math-17k, 8192 response length, eval off. (This arm also dropped --use-fault-tolerance from the launcher, which turned out to be unnecessary: _FT_DEFAULT_COMPONENTS is ["rollout"], so "train" not in ft_components and the witness / weight-checksum auto-enables that streaming rejects never fire. The H200 arm above kept FT on and confirms it: ft_components ['rollout'], enable_witness False, save_local_weight_checksum False, no assert. Streaming is only incompatible with an explicit --ft-components train, where refusing is the intended behaviour. Both arms here were symmetric, so the comparison is unaffected.)

Three 3-step runs, perf 0/1/2 in each, no NaN, no traceback:

actor offload	streaming	train_rollout_logprob_abs_diff (steps 0/1/2)
cpu	off	0.0329 / 0.0336 / 0.0350
disk	off	0.0325 / 0.0340 / 0.0349
disk	on, fp32 moments	0.0328 / 0.0336 / 0.0351

Agreement step by step across all three arms, so the fp32 path behaves as documented.

Metric here is the train actor's allocated_GB from miles' own print_memory at the end of the train phase, plus a 2 s nvidia-smi sampler windowed to actor_train start..end; steps 1-2 only.

	baseline (disk offload only)	streaming, fp32 moments	delta
train-actor allocated_GB	132.1 GB	46.7 GB	-85.4 GB (-65%)
device peak, train phases	173.8-177.9 GiB	85.4-88.5 GiB	~-89 GiB
actor_train	99.0 / 106.9 s (mean 103.0)	139.3 / 148.9 s (mean 144.1)	+41.1 s (+40%)

Both deltas check out against what the store reports:

• Memory. The store holds 85.4 GB/rank on disk (dense 4.4 + expert 81.0) at 12 B/param with fp32 moments; the measured reduction is 85.4 GB, exact.
• Time. Self-reported I/O is 2.4 s (dense) + 37.2 s (expert) = 39.6 s/rank against a +41.1 s train-time delta, so 96% of the added time is the synchronous streaming I/O.
• Lazy moments, fp32 flavour. The expert store reads 27.0 GB on the first step (= 81.0/3, mains only, three equal-sized fp32 segments) and 81.0 GB from the second on. With the bf16 arm's 4/8 ratio that pins the per-segment dtype accounting from both directions.

Two things this arm needed that the H200 arm did not. torch_memory_saver at #80 — radixark/miles:dev ships 0.0.9.post1 and tripped exactly the silent-corruption failure described under Guards, surfacing as found NaN in local grad norm in step 0's backward until the matching build was installed. And a one-line sglang fix (sgl-project/sglang#32385) for an unrelated uninitialised-scale bug that makes fp8 rollout unusable for any model with moe_intermediate / 128 % 4 != 0, Qwen3-30B-A3B included.

Measurement note: the device-level peak taken over the whole run is ~197 GiB in both arms and shows nothing, because --sglang-mem-fraction-static 0.7 dominates it. The numbers above window the sampler to the train phases, where the colocated engine is offloaded.

Checkpoint save/load is covered separately, below.

Checkpoint save and resume

Run on the 8xH200 box, same arms as above, --save-interval 1 over 2 rollouts to make the hook fire, then a resume from the produced checkpoint. This is the first exercise of save_to / load_from and therefore of the two checkpointing.py hooks in radixark/Megatron-LM#63.

Save. Both iterations got their state, and the layout is what relative_dir intends:

iter_0000001/nvme_opt_state/rank00000/opt0_0/{bucket00000.bin, manifest.json}
                                     /opt0_1/{bucket00000.bin .. bucket00017.bin, manifest.json}
                            rank00001/...   (8 ranks)

16 manifests, i.e. 8 ranks x 2 chained members, no collisions — the point of moving the subdirectory naming out of Megatron and into the store, where the live scratch directory already spelled it. Per rank the store is 1 bucket / 1.5 GB (dense) plus 18 buckets / 27.0 GB (expert) at 8 B/param, so 228 GB of optimizer state across the 8 ranks; the whole iter_0000001 directory is 285 GB. Manifest contents check out: dtypes {main: float32, exp_avg: bfloat16, exp_avg_sq: bfloat16}, per-group Adam steps [2, 2], and numel 197,499,904 for the dense bucket, just under the 200M cap.

Cost of the synchronous copy, previously unmeasured: Timer save_model 57.0 s and 51.8 s. That covers the model dist-ckpt plus the 228 GB shutil.copyfile of the store, so it is an upper bound on the streaming part rather than an isolated figure.

Resume. Both members load, and the manifest dtype and layout asserts pass:

NVMe optimizer state loaded: 18 buckets <- .../iter_0000001/nvme_opt_state/...
NVMe optimizer state loaded:  1 buckets <- .../iter_0000001/nvme_opt_state/...

The restored state is actually used, not just copied back. On a cold start the first step reads only the mains — 4/8 of the total with bf16 moments — because Adam has not created the moments yet. After a resume allocate_moments() has pre-created them and load_from has filled them, so the very first step streams all three segments:

	first step after	expert store read	wrote
cold start	—	13.5 GB (= 27.0 x 4/8, mains only)	27.0 GB
resume	iter_0000001	27.0 GB (full)	27.0 GB

Training continues normally from there: 2 x outcome=NORMAL, train_rollout_logprob_abs_diff 0.0186, in the same band as the arms above, job exits succeeded with no asserts and no NaN.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.