feat(optimizer): NVMe streaming of DistributedOptimizer state by yueming-yuan · Pull Request #63 · radixark/Megatron-LM

yueming-yuan · 2026-07-08T00:18:28Z

Why

Step-time GPU residency of optimizer state (fp32 main + Adam moments, 12 B/param) is the binding constraint for training GLM-5.2 744B on 8 GB300 nodes: ~279 GB/rank against 276.6 GB of HBM, before params/grads/activations. Sleep-window offload (miles #1575 + torch_memory_saver #80) cannot help — the state must be resident when the FusedAdam kernel runs.

How

--optimizer-state-nvme-dir DIR moves the state's home to per-bucket files on node-local NVMe ([main | exp_avg | exp_avg_sq], laid out per param shard in group order). step() processes buckets one at a time: materialize (untyped_storage().resize_) -> load -> per-bucket FusedAdam step -> copy updated mains into the param buffer -> write back -> release. GPU residency is bounded by one bucket.

One TE FusedAdam instance per bucket (groups cloned from the master groups, lr/wd re-synced each step): keeps TE step semantics and per-group step counters exact, no kernel math duplicated.
Lazy seal: mains stay GPU-resident until the first step (so reload_model_params and the first colocate sleep behave exactly as today); moments are lazily created by FusedAdam at step 1 and spill with the mains.
Sync chunked I/O through one pinned staging buffer (--optimizer-state-nvme-chunk-mb, default 256). Startup wipe + atexit reclaim bound disk usage across runs.
v1 guards: non-precision-aware bf16 training only; state_dict/sharded_state_dict raise (checkpointing optimizer state lands in a follow-up).

Testing

Qwen3-30B-A3B RL (miles colocate, 8xH200, TP1/PP1/DP8/EP8, fp32 state): two full rollout->train cycles, zero errors.

rank-0 allocated_GB drops 43.1 -> 31.8 after the first step — exactly bf16 params + fp32 grads; main+moments fully off-GPU in steady state.
sleep/wake collapse to 5.2s/1.3s (cpu-backup baseline: 24s/8.9s) since paused state no longer needs a backup.
grad_norm / train_rollout_logprob_abs_diff in line with the baseline run.
Known v1 perf debt: per-param chunked I/O reaches ~2 GB/s; batching small shards through the staging chunk is the follow-up.

🤖 Generated with Claude Code

Update (2026-07-09): hardened for large-model DP1, validated on GLM-5.2 744B @ 8x GB300

Commit 13a21d94a fixes three issues found bringing this up at scale (TP8/PP4/DP1/EP8, 23.25B params/rank, 279 GB state/rank), none reproducible on the 8x H200 Qwen3-30B A/B (TP1/DP8: small buckets, wide margins):

fp32 main params were GPU-resident until the first step() (lazy eviction). The first rollout + forward/backward runs before any step, so tight configs OOM before the store ever flushes. Now evicted at store init.
Checkpoint load crashed with illegal memory access: load_checkpoint -> reload_model_params -> _copy_model_params_to_main_params writes into the evicted storage. The refresh now routes through the store (materialize -> copy -> flush -> release), which also keeps main semantics correct (post-load weights are the truth).
Bucket-granular streaming inherits DDP bucket sizes. With expert-DP = 1, Megatron leaves an entire expert buffer as a single bucket; late-PP ranks then materialize ~90 GiB main (+ lazy m/v) in one shot and OOM. Specs are now chunked to a fixed 200M-numel budget, decoupled from DDP bucketing.

Validation (GLM-5.2 744B, 8x GB300, full miles RL steps):

step-loop watermark flat at 142 GiB across 120 sub-buckets; release verified (storage size 0)
actor_train 287s (step 0, warm-up) / 216s (step 1, full main+m/v reads) vs ~170s 16-node baseline: NVMe exposure ~46s, mostly absorbed by page cache
train_rollout_logprob_abs_diff 0.0584 / 0.0548, in the 16-node reference band
net: full 744B RL training on half the nodes at ~1.25x step time

Update (2026-07-09): checkpoint save/load (same-topology resume)

Commits ff23ee27a + 34d73a3c2 replace the v1 state_dict/sharded_state_dict raise with a minimal resume path — the flat bucket files on scratch are the state, so checkpointing is a per-rank file copy:

Save: store.save_to() copies each bucket file into <ckpt>/nvme_opt_state/rank%04d_opt{instance}_{uid}/ plus a manifest.json (per-bucket numel, entry numels, Adam step counts). Hooked into save_checkpoint after the regular optimizer path; dist-ckpt keeps handling model weights untouched (state_dict returns a marker, sharded_state_dict returns {}).
Load: store.load_from() asserts the manifest layout matches the freshly built specs (same-topology resume only), copies files back, restores Adam step counts, and pre-creates zero-storage moment tensors so the first step streams everything in. Skipped for --finetune / --no-load-optim; a checkpoint without nvme_opt_state logs a notice and starts fresh (old checkpoints stay loadable).
Bug found by the load-path assert: chained dense/expert optimizers can share distributed_optimizer_instance_id, so both stores landed in the same directory, each wiping the other's files. Runtime survived by POSIX unlink-fd semantics (each store still held fds to its unlinked files) — any path-based access (like save) hit corrupted state. Fixed with a per-process construction counter in the directory name (opt0_0 / opt0_1), deterministic across runs.

Validation (GLM-5.2 744B, 8x GB300, TP8/PP4/DP1/EP8):

save: 64 dirs = 32 ranks x {dense 5 buckets, expert 90 buckets}, ~8 TB, manifest layout correct
resume from that checkpoint: no fresh-start branch, no layout assert, post-resume train step completes with train_rollout_logprob_abs_diff 0.0547 (in band)

Update (2026-07-24): fp32-param models supported, narrow moments, verified on Qwen3.5-35B-A3B

Changes

Models with native-fp32 params now work. The store asserted model_fp32_groups was empty, which rejected any model holding a parameter in fp32 by design — a MoE router's expert_bias (mcore's own _maintain_float32_expert_bias) or a GDN/Mamba A_log. That is not a niche case: it blocked the whole Qwen3.5 family. Those params now step under a small always-resident Adam instead of being streamed; their shard_fp32_groups entries alias the model params' storage, so stepping updates the model directly with no copy-back.
Per-segment storage dtype (--optimizer-state-nvme-moment-dtype, default fp32). The moments tolerate less precision than the master copy: bf16 cuts streaming volume by a third (12 B/param to 8), fp8 by half. fp32 stays bit-identical to keeping the state on GPU. The fp8 options warn — exp_avg_sq needs per-block scaling to survive 8-bit storage, which this does not implement. The manifest records the dtypes and resume verifies them, so bytes written as bf16 cannot be read back as fp32.
Restructured the store. It was one class holding nine mutable fields patched by three builders, with the fp32 itemsize and segment layout hardcoded in five places. Now: _Bucket owns its entries, file, Adam, byte layout and residency and is built in one shot; _Stager is the only code touching host memory or the filesystem, so a GPUDirect Storage backend is a change to that one class (the layout starts every transfer on a 4K boundary, which is what cuFile and O_DIRECT need); plan_buckets is a pure function and testable on its own.
File-handle hygiene: O_CLOEXEC on the bucket files, and posix_fallocate falls back to ftruncate on mounts that do not support it (tmpfs, some NFS).

Verification

Qwen3.5-35B-A3B RL on 8xH200 (colocate, TP1/PP1/DP8/EP8, 3 rollouts x 2 train steps each), on radixark/miles:dev-202607240207 with torch_memory_saver pinned at PR NVIDIA#80's head. This model is the point of the exercise: the previous version of the store could not build on it at all.

Case	Config	Result
control	`--offload-train-target=cpu`	SUCCEEDED, 6/6 steps NORMAL
actor offload to disk	`--offload-train-target=disk`	SUCCEEDED, 6/6 steps NORMAL
both axes	`--offload-train-target=disk --optimizer-state-nvme-dir`	SUCCEEDED, 6/6 steps NORMAL
streaming alone, colocated	`--optimizer-state-nvme-dir`, no actor offload	FAILED — see below
streaming alone, not colocated	4 GPU train + 4 GPU rollout, `--optimizer-state-nvme-dir`, no actor offload	SUCCEEDED, 6/6 steps NORMAL
bf16 moments	both axes + `--optimizer-state-nvme-moment-dtype bf16`	SUCCEEDED, 6/6 steps NORMAL (2 runs)
fp8e4m3 moments	both axes + `--optimizer-state-nvme-moment-dtype fp8e4m3`	SUCCEEDED, 6/6 steps NORMAL (2 runs), accuracy ~1.4x bf16's deviation

What the logs show, beyond exit status:

The lazy-moments path is exercised, not assumed. Streaming volume per step: read 15.4 GB, wrote 46.1 GB on the first step, then read 46.1 GB, wrote 46.1 GB from the second on. 15.4 is exactly 46.1/3 — the first fetch pulls main only, because Adam has not created the moments yet, and all three segments move from then on.
Narrow moments shrink the state as the arithmetic says, and accuracy holds. bf16 takes the state from 46.1 GB to 30.8 GB per rank and fp8e4m3 to 23.1 GB, with both directions of the steady-state step scaling the same way. Accuracy is unaffected for bf16 (rollout-vs-train logprob |diff| 4.9-5.6e-4 across two independent runs, the same band as fp32 and the cpu-offload control) and measurably but not dramatically worse for fp8e4m3 (6.9-7.5e-4, reproducibly ~1.4x bf16, no NaN, 6/6 steps NORMAL).
Step wall-clock is not measurable on this rig, so this PR makes no claim about it. With 8 ranks sharing one md-RAID array, the same configuration's steady-state step ranged 18.8-34.5s (bf16) and 11.5-23.6s (fp8e4m3), and a single step differed 2.4x between ranks. The intervals for different dtypes overlap completely. An earlier revision of this section claimed bf16 cut the step 45%; that was one run against another and does not survive the variance — withdrawn.
A staging bug worth noting for anyone reusing this code. Copying a GPU tensor straight into a differently typed view of pinned host memory does not take the DMA path: it cost ~0.44s per 800 MB against ~0.01s for a same-dtype copy, i.e. narrow dtypes were paying 40x on the host copy and fp8 could end up slower end to end than bf16 despite moving fewer bytes. Casting on the device first and moving same-dtype bytes fixes it. Measured in isolation (no disk), across entry sizes from 400M down to 20K elements, the fix is 2.1-49x for bf16 and 2.8-107x for fp8, i.e. it never regresses even for tiny shards. Round-trip stays exact.
Weights survive the round trip. rollout-vs-train logprob: −0.30930 vs −0.30997 (|diff| 6.7e-4) with both axes on, against 6e-4 for the cpu-offload control.
The fp32-resident branch runs on this model (sub-MB — a per-head A_log — but the branch that used to be a hard assert).
Both ChainedOptimizer members get their own store (dense 1 bucket / 3.5 GB, expert 21 buckets / 46.1 GB), confirming the per-process uid in the directory name.
Streaming stands on its own when not colocated. With 4 GPUs training and 4 serving rollout, each rank carries 41 buckets / 92.2 GB of state (double the 8-GPU colocate case, as expected) and all 6 steps are NORMAL with the actor never offloaded. The same 1/3-then-full read pattern shows up here too (read 30.8 GB, wrote 92.2 GB, and 30.8 = 92.2/3). This is a supported configuration rather than the motivating one. The point of streaming is memory, not colocation: --offload-train frees the training GPUs' HBM for the whole rollout window whether or not the engine shares those GPUs, and streaming keeps the 12 B/param of state out of HBM even 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 need neither; the runs that need both are the ones tight enough to sit at DP=1. GLM-5.2 744B above is that shape: --offload-train-target=disk and streaming together, on half the nodes the run would otherwise need. Disaggregated setups use the same pair, sized so the engines get the HBM the actor gives up.
Under --colocate streaming alone is not sufficient. The store still builds fine (21 buckets / 46.1 GB), then sglang's resume_memory_occupation OOMs in csrc/core.cpp func=resume: without --offload-train the actor's weights and grad buffers stay resident through the rollout and the engine has nowhere to put its KV cache. So the pairing requirement is specific to colocation, not to streaming — documented in miles under docs/advanced/disk-offload.md.

Update (2026-07-24): store moved to miles; this PR is now one file

Commit 0581166, paired with radixark/miles#1793. The implementation lives in miles_plugins/optimizers/nvme_stream.py. What is left here is megatron/training/checkpointing.py, +19 lines; distrib_optimizer.py, optimizer_config.py and arguments.py are back to untouched.

Moved out: the store, its construction, the three --optimizer-state-nvme-* args and their OptimizerConfig fields, and the five if self._nvme_state_store is not None branches. Those five are now instance bindings in miles' _bind():

entry point	binding
`step_with_ready_grads`	`store.step()` plus the param all-gather tail, copied from here and kept in sync by hand
`_copy_model_params_to_main_params`	replaced by a `reload_model_params` binding, so the `_impl` split is gone — `_copy_model_params_to_main_params` has exactly one caller (`MixedPrecisionOptimizer.reload_model_params`) and `ChainedOptimizer` delegates to its children
`state_dict` / `load_state_dict` / `sharded_state_dict`	the same empties as before

The _copy_main_params_to_model_params_for extraction is reverted. It looked like the cheap thing to keep here — a pure refactor, and the range-map walk's invariants belong to whoever changes them. Two things argued the other way:

Its only live caller under streaming was the store. Once step_with_ready_grads is bound to store.step(), _copy_main_params_to_model_params is never reached, so the extracted method existed in this repo purely to serve an out-of-repo caller.
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 data-parallel 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, which is exactly the bug that was open against this PR. The # quantized in the above quantize_param_shard function comment also lost its referent the moment "above" became a different method.

miles adapts the copy instead, pinned to 4716f754#L2469-L2519 in a comment, and rejects fp8 params and the MXFP8 param all-gather in the store constructor. That closes the fp8 gap and takes this PR's optimizer diff to zero.

What stays and why. The two checkpoint hooks are duck-typed through the _nvme_state_store attribute and import nothing from miles. They have to be inside save_checkpoint / load_checkpoint to cover every call path: miles has two save call sites and more will be added, and a missed one shows up as a checkpoint with no optimizer state that silently resumes from zero.

Checkpoint layout policy moved to the store. save_to / load_from take the checkpoint base and derive their own per-rank subdirectory from the same relative_dir the live scratch directory uses. Previously the identity was spelled twice and differently — rank{rank:04d}_opt{instance}_{uid} flat here versus rank{rank}/opt{instance}_{uid} nested in the store — and _nvme_state_checkpoint_dir had to reach through store.dist_opt to build it. That function is gone; the contract is four methods (step, refresh_main_from_model_params, save_to(base), load_from(base)), and nothing here reads store.dist_opt or store.uid. The "this checkpoint has no streamed state" branch moves into load_from, since keeping pre-streaming checkpoints loadable is a policy decision and belongs with the implementation.

Arg surface (miles side, see NVIDIA#1793): one new switch, --stream-optimizer-state-to-disk, replacing --optimizer-state-nvme-dir where "enable" and "where" were the same knob so there could be no default location. --stream-optimizer-state-moments keeps all five dtypes. miles' existing --offload-train-disk-dir / --offload-train-disk-chunk-mb are reused unchanged — they mean the same thing for both mechanisms, which take separate subdirectories of the shared root — so nothing is renamed. Moment dtype is a constructor argument, so the getattr(config, ..., "fp32") silent fallback is gone.

No functional change is intended relative to the validated results above, other than fp8 and MXFP8 now being rejected instead of silently skipped. The bindings and the adapted copy are the same code paths reached a different way, but that still needs an end-to-end run before merge.

Update (2026-07-26): the two remaining hooks are now exercised

Everything this PR still contains is the save_checkpoint / load_checkpoint pair in checkpointing.py, and until now nothing had run them — the earlier validation used --save-interval 20 over 3 steps, so they never fired. Qwen3-30B-A3B RL on 8xH200, streaming with bf16 moments, --save-interval 1 over 2 rollouts, then a resume from the result:

Save hook. Fires for both iterations. Passing the checkpoint base and letting the store derive its own subdirectory produces:

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 — 8 ranks x 2 chained members — with no collisions, which is what dropping _nvme_state_checkpoint_dir was for: the identity used to be spelled rank%04d_opt{instance}_{uid} flat here and rank{rank}/opt{instance}_{uid} nested in the store, and now there is one spelling, in the store. Manifest contents are correct: dtypes {main: float32, exp_avg: bfloat16, exp_avg_sq: bfloat16}, per-group Adam steps [2, 2], dense bucket numel 197,499,904 just under the 200M cap. 228 GB of optimizer state across the 8 ranks, 285 GB for the whole iteration directory.

Timer save_model 57.0 s / 51.8 s, covering the model dist-ckpt plus the store's synchronous shutil.copyfile. That is the cost of this hook sitting before the --async-save block, now measured rather than assumed.

Load hook. 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 used rather than merely copied back, which the streaming volume shows directly. Cold start reads only the mains on the first step (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 the load has filled them, so the first step streams all three segments:

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

Training continues from there: 2 x outcome=NORMAL, train_rollout_logprob_abs_diff 0.0186 against 0.0180-0.0190 for the non-resumed arms, job succeeded, no asserts, no NaN.

The distrib_optimizer.py side of this PR remains empty, and miles-side validation is in radixark/miles#1793.

fp32 main params and Adam moments live in per-bucket NVMe files and are streamed through the GPU bucket-by-bucket during the optimizer step via per-bucket FusedAdam instances. Bounds step-time GPU residency to one bucket. Checkpointing optimizer state is guarded (not supported yet).

- Evict fp32 main params to NVMe at store init: with lazy eviction they stay GPU-resident through the entire first rollout + forward/backward and OOM tight configs before the first step can flush them. - Route _copy_model_params_to_main_params through the store (materialize, copy, flush, release): checkpoint load refreshes main params, which otherwise writes into evicted storage (illegal memory access). - Chunk state specs to a fixed 200M-numel budget: DDP bucketing leaves an entire expert buffer as a single bucket when expert-DP is 1, so bucket-granular streaming materializes ~90 GiB at once on those ranks. Validated on GLM-5.2 744B, 8x GB300 (TP8/PP4/DP1/EP8): two full RL steps, flat 142 GiB watermark through the 120-bucket step loop, actor_train 287s/216s (16-node baseline ~170s), train_rollout_logprob_abs_diff in the reference band.

…topology resume)

… instance_id

Replace the shared-directory + pid-liveness-check reclaim (fragile: pid reuse after a dead process could make an orphaned file look alive forever) with a per-rank/per-cell exclusive subdirectory, the same convention as Megatron-LM's NVMeOptimizerStateStore (radixark/Megatron-LM#63): wipe-then-recreate on startup, plain rmtree on exit, no cross-process coordination needed since directories are never shared. actor_factory.py now computes TMS_DISK_BACKUP_DIR per rank (cell{i}_rank{j}) via a per-actor runtime_env override instead of one shared path for every actor; actor.py's reclaim reads back that same env var so the Python-side cleanup always targets exactly what the C++ hook wrote to.

Route them through a small always-resident Adam instead of hard-rejecting any model that has them; NVMe-stream only the bf16/fp16 buckets, which is where the actual memory pressure is.

Per review: the implementation now lives in miles (miles/backends/megatron_utils/nvme_state_store.py) and is imported lazily at construction, same pattern as the R3 routing-replay hook in moe/router.py. Megatron keeps only the thin config fields and the duck-typed routing hooks in DistributedOptimizer/checkpointing; miles owns the store and can unit-test it.

The move commit deleted megatron/core/optimizer/nvme_state_store.py but left the import behind, so enabling --optimizer-state-nvme-dir raised ImportError.

Streams a DistributedOptimizer's fp32 main params + Adam moments through per-bucket files during optimizer.step(), bounding GPU residency to one bucket, for models whose optimizer state does not fit GPU even transiently. Native-fp32 model params (router expert_bias, GDN A_log) stay GPU-resident via a small always-resident Adam. radixark/Megatron-LM#63 imports this lazily from DistributedOptimizer, so miles owns and tests it.

The step is bound by streaming volume, and the moments tolerate less precision than the master copy: bf16 storage cuts 12 bytes per parameter to 8. Defaults to fp32, which is bit-identical to keeping them on GPU.

The store only manipulates DistributedOptimizer internals -- bucket layout, param groups, Adam state, main-param residency -- so it belongs next to them rather than across a repo boundary.

miles' side of disk offload is the torch_memory_saver wiring: it decides the per-rank directories, the env the actor is launched with, and reclaim. Streaming optimizer state is DistributedOptimizer's own business (radixark/Megatron-LM#63).

Move the store out of megatron/core/optimizer/nvme_state_store.py into miles_plugins/optimizers/nvme_stream.py, and turn the five DistributedOptimizer entry points it overrides into instance bindings here instead of `if self._nvme_state_store is not None` branches in Megatron. Pairs with radixark/Megatron-LM#63, which keeps only the _copy_main_params_to_model_params_for extraction and the two checkpoint hooks. Arg surface: --stream-optimizer-state-to-disk replaces the dir-as-switch, and --disk-offload-dir / --disk-offload-chunk-mb are shared with --offload-train-target=disk rather than duplicated per mechanism. Moment dtype is a constructor argument now, so the silent getattr fallback is gone. Streaming asserts against the paths that read the main params or the master optimizer's state directly (--reset-optimizer-states, --save-local-weight-checksum, --enable-witness) instead of letting them return emptiness, and against colocation without --offload-train, which used to surface as an sglang OOM.

Move the store out of megatron/core/optimizer/nvme_state_store.py into miles_plugins/optimizers/nvme_stream.py, and turn the five DistributedOptimizer entry points it overrides into instance bindings here instead of `if self._nvme_state_store is not None` branches in Megatron. Pairs with radixark/Megatron-LM#63, which keeps only the _copy_main_params_to_model_params_for extraction and the two checkpoint hooks. One new switch: --stream-optimizer-state-to-disk, replacing Megatron's --optimizer-state-nvme-dir, where "enable" and "where" were the same knob so there could be no default location. --offload-train-disk-dir and --offload-train-disk-chunk-mb are reused as-is -- they mean the same thing for both mechanisms, which take separate subdirectories of the shared root -- so their default and validation move to a branch that either mechanism triggers. Streaming's main setting is a run that is not colocated at all, where --offload-train-target stays "cpu" and the default would otherwise never be filled in. --stream-optimizer-state-moments carries the storage dtype, which is a constructor argument now, so the silent getattr fallback on the config field is gone. Streaming asserts against the paths that read the main params or the master optimizer's state directly (--reset-optimizer-states, --save-local-weight-checksum, --enable-witness) instead of letting them return emptiness, and against colocation without --offload-train, which used to surface as an sglang OOM mid-run.

Move the store out of Megatron into miles_plugins/optimizers/nvme_stream.py, and turn the five DistributedOptimizer entry points it overrides into instance bindings here instead of `if self._nvme_state_store is not None` branches there. Pairs with radixark/Megatron-LM#63, whose optimizer diff is now empty: the only thing left on that side is the two checkpoint hooks, which have to sit inside save_checkpoint/load_checkpoint to cover every call path. The per-bucket copy of mains into the param buffer is adapted here from _copy_main_params_to_model_params rather than extracted there, with the source pinned by commit and line range. Extracting it would have handed out a partial operation across a repo boundary: the fp8 branch depends on quantize_param_shard() having run over the whole fp8 param set, which is a DP collective and cannot be split per bucket, so a subset-taking method silently skips fp8 updates unless the caller knows to compensate. Keeping the copy next to its only caller makes that precondition local, and __init__ now rejects fp8 params (and the MXFP8 param all-gather) outright. One new switch: --stream-optimizer-state-to-disk, replacing Megatron's --optimizer-state-nvme-dir, where "enable" and "where" were the same knob so there could be no default location. --offload-train-disk-dir and --offload-train-disk-chunk-mb are reused as-is -- they mean the same thing for both mechanisms, which take separate subdirectories of the shared root -- so their default and validation move to a branch that either mechanism triggers. Streaming's main setting is a run that is not colocated at all, where --offload-train-target stays "cpu" and the default would otherwise never be filled in. --stream-optimizer-state-moments carries the storage dtype, which is a constructor argument now, so the silent getattr fallback on the config field is gone. Streaming also asserts against the paths that read the main params or the master optimizer's state directly (--reset-optimizer-states, --save-local-weight-checksum, --enable-witness) instead of letting them return emptiness, and against colocation without --offload-train, which used to surface as an sglang OOM mid-run.

The store and its wiring move to miles_plugins/optimizers/nvme_stream.py (radixark/miles#1793), which binds the five entry points that used to be `if self._nvme_state_store is not None` branches in distrib_optimizer.py onto each DistributedOptimizer instance. The --optimizer-state-nvme-* args and their OptimizerConfig fields go with it; miles owns the flags now and passes them to the constructor. distrib_optimizer.py is back to untouched. The earlier _copy_main_params_to_model_params_for extraction is reverted too: its only live caller under streaming was the store, and it handed out a partial operation, since the fp8 branch depends on quantize_param_shard() having run over the whole fp8 param set -- a DP collective that cannot be split per bucket. miles adapts the copy instead, with the source pinned by commit and line range, and rejects fp8 params up front. What is left here is the two checkpoint hooks, duck-typed through the _nvme_state_store attribute and importing nothing from miles. They have to be inside save_checkpoint/load_checkpoint to cover every call path: miles has two save call sites and a missed one shows up as a checkpoint with no optimizer state that silently resumes from zero. save_to/load_from now take the checkpoint base and derive their own per-rank subdirectory from the same layout the live scratch directory uses, so _nvme_state_checkpoint_dir goes away and nothing here 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, where the policy belongs.

Move the store out of Megatron into miles_plugins/optimizers/nvme_stream.py, and turn the five DistributedOptimizer entry points it overrides into instance bindings here instead of `if self._nvme_state_store is not None` branches there. Pairs with radixark/Megatron-LM#63, whose optimizer diff is now empty: the only thing left on that side is the two checkpoint hooks, which have to sit inside save_checkpoint/load_checkpoint to cover every call path. The per-bucket copy of mains into the param buffer is adapted here from _copy_main_params_to_model_params rather than extracted there, with the source pinned by commit and line range. Extracting it would have handed out a partial operation across a repo boundary: the fp8 branch depends on quantize_param_shard() having run over the whole fp8 param set, which is a DP collective and cannot be split per bucket, so a subset-taking method silently skips fp8 updates unless the caller knows to compensate. Keeping the copy next to its only caller makes that precondition local, and __init__ now rejects fp8 params (and the MXFP8 param all-gather) outright. One new switch: --stream-optimizer-state-to-disk, replacing Megatron's --optimizer-state-nvme-dir, where "enable" and "where" were the same knob so there could be no default location. --offload-train-disk-dir and --offload-train-disk-chunk-mb are reused as-is -- they mean the same thing for both mechanisms, which take separate subdirectories of the shared root -- so their default and validation move to a branch that either mechanism triggers. Streaming's main setting is a run that is not colocated at all, where --offload-train-target stays "cpu" and the default would otherwise never be filled in. --stream-optimizer-state-moments carries the storage dtype, which is a constructor argument now, so the silent getattr fallback on the config field is gone. Streaming also asserts against the paths that read the main params or the master optimizer's state directly (--reset-optimizer-states, --save-local-weight-checksum, --enable-witness) instead of letting them return emptiness, and against colocation without --offload-train, which used to surface as an sglang OOM mid-run.

Move the store out of Megatron into miles_plugins/optimizers/nvme_stream.py, and turn the five DistributedOptimizer entry points it overrides into instance bindings here instead of `if self._nvme_state_store is not None` branches there. Pairs with radixark/Megatron-LM#63, whose optimizer diff is now empty: the only thing left on that side is the two checkpoint hooks, which have to sit inside save_checkpoint/load_checkpoint to cover every call path. The per-bucket copy of mains into the param buffer is adapted here from _copy_main_params_to_model_params rather than extracted there, with the source pinned by commit and line range. Extracting it would have handed out a partial operation across a repo boundary: the fp8 branch depends on quantize_param_shard() having run over the whole fp8 param set, which is a DP collective and cannot be split per bucket, so a subset-taking method silently skips fp8 updates unless the caller knows to compensate. Keeping the copy next to its only caller makes that precondition local, and __init__ now rejects fp8 params (and the MXFP8 param all-gather) outright. --stream-optimizer-state-to-disk is the switch, replacing Megatron's --optimizer-state-nvme-dir where "enable" and "where" were the same knob so there could be no default location. It requires --offload-train-target=disk: both mechanisms answer the same pressure and are only deployed as a pair, so that is the only combination covered. --offload-train-disk-dir and --offload-train-disk-chunk-mb are therefore reused as-is, with each mechanism taking its own subdirectory of the shared root, and their validation is untouched. --stream-optimizer-state-moment-dtype carries the on-disk dtype, which is a constructor argument now, so the silent getattr fallback on the config field is gone. It is a serialization format, not a compute precision, which is what distinguishes it from --exp-avg-dtype: the step still hands fp32 to FusedAdam, and those flags require the precision-aware optimizer that streaming excludes. Streaming also asserts against the paths that read the main params or the master optimizer's state directly (--reset-optimizer-states, --save-local-weight-checksum, --enable-witness) instead of letting them return emptiness.

yueming-yuan changed the base branch from main to miles-main July 8, 2026 00:20

Zhichenzzz approved these changes Jul 8, 2026

View reviewed changes

yueming-yuan added 3 commits July 8, 2026 23:26

NVMe store checkpointing: save/load bucket files with manifest (same-…

ff23ee2

…topology resume)

Disambiguate store directories: chained optimizer instances can share…

34d73a3

… instance_id

Zhichenzzz added 4 commits July 20, 2026 20:25

Trim redundant NVMe optimizer-state comments

17ca373

Support models with native-fp32 params (router expert_bias, GDN A_log)

c8b5785

Route them through a small always-resident Adam instead of hard-rejecting any model that has them; NVMe-stream only the bf16/fp16 buckets, which is where the actual memory pressure is.

Point the NVMe store import at miles

34cd1bf

The move commit deleted megatron/core/optimizer/nvme_state_store.py but left the import behind, so enabling --optimizer-state-nvme-dir raised ImportError.

Zhichenzzz force-pushed the yueming/nvme-streaming-optimizer branch from 65d190d to 34cd1bf Compare July 24, 2026 19:04

Zhichenzzz added 3 commits July 24, 2026 12:23

Add --optimizer-state-nvme-moment-dtype

1a2581a

The step is bound by streaming volume, and the moments tolerate less precision than the master copy: bf16 storage cuts 12 bytes per parameter to 8. Defaults to fp32, which is bit-identical to keeping them on GPU.

Drop the docstring from _copy_main_params_to_model_params_for

86edf22

Keep the NVMe state store in Megatron

3bf3831

The store only manipulates DistributedOptimizer internals -- bucket layout, param groups, Adam state, main-param residency -- so it belongs next to them rather than across a repo boundary.

Zhichenzzz marked this pull request as draft July 24, 2026 23:48

yueming-yuan mentioned this pull request Jul 25, 2026

feat(optimizer): NVMe optimizer-state streaming as a miles plugin radixark/miles#1793

Open

yueming-yuan force-pushed the yueming/nvme-streaming-optimizer branch from 2e0ca7e to 0581166 Compare July 25, 2026 03:37

yueming-yuan marked this pull request as ready for review July 25, 2026 03:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(optimizer): NVMe streaming of DistributedOptimizer state#63

feat(optimizer): NVMe streaming of DistributedOptimizer state#63
yueming-yuan wants to merge 12 commits into
miles-mainfrom
yueming/nvme-streaming-optimizer

yueming-yuan commented Jul 8, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yueming-yuan commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why

How

Testing

Update (2026-07-09): hardened for large-model DP1, validated on GLM-5.2 744B @ 8x GB300

Update (2026-07-09): checkpoint save/load (same-topology resume)

Update (2026-07-24): fp32-param models supported, narrow moments, verified on Qwen3.5-35B-A3B

Changes

Verification

Update (2026-07-24): store moved to miles; this PR is now one file

Update (2026-07-26): the two remaining hooks are now exercised

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yueming-yuan commented Jul 8, 2026 •

edited

Loading