feat(optimizer): NVMe streaming of DistributedOptimizer state

[yueming-yuan]

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):

1. 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.
2. 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).
3. 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)