Stream optimizer main initialization to NVMe

[zianglih]

Summary

@HumansAnd

ci-megatron-pr: #86

Extend --stream-optimizer-state-to-disk to cover distributed-optimizer construction. Miles enables the paired Megatron deferred-main mode, then initializes FP32 mains one existing runtime bucket at a time directly in their final files.

• Paired Megatron change: radixark/Megatron-LM#86, head 5dcbc1020a94ea973316163e8db6719886823e63.
• Miles head: 54f083eb0530e47fce951d2fba243c0d64038cb8.
• Scope: this changes only initial main-parameter construction. Existing optimizer steps, checkpoint save/load, model-parameter reload, scratch paths, store numbering, validation, and failure handling remain unchanged.
• Aggregate diff: five files, +89/-6; no added try, except, or finally blocks.

Design

• Construction: Megatron creates each same-shaped CUDA FP32 main handle and immediately releases its backing storage. Tensor identity, shape, dtype, and device remain available to the distributed optimizer.
• Initialization: after construction, Miles materializes one existing runtime main bucket, copies the still-live BF16 model shards into it, flushes it to its final bucket file, releases CUDA storage, synchronizes and evicts that initialized file range from page cache, then advances.
• No temporary state file: initialized mains are written once to their final offsets. The final files still reserve main plus optimizer moments using the existing streaming implementation.
• Granularity: initialization reuses the existing 200,000,000-element soft bucket limit. A nominal FP32 main bucket is about 0.745 GiB; an indivisible parameter entry can make a bucket larger. This avoids per-parameter I/O while preventing accumulation of the full FP32 main state.
• Residual HBM requirement: the largest individual FP32 shard must still fit briefly during Megatron construction, followed by one main-only runtime bucket during Miles initialization. The change removes cumulative main-state residency; it does not make the largest shard allocation-free.
• Reload behavior: _bind.reload_model_params retains the existing implementation and still invokes Megatron's normal reload callback. This PR intentionally does not bucketize or otherwise change checkpoint/model reload.

Validation

Exact-head public-image B300 tests

• Image: docker.io/radixark/miles:dev-202608210021@sha256:e2e25e4121c595d352933fab81ce05cc8fb2aef8c8478a16abc704b51e330003.
• Hardware/runtime: one C2 node; tests pinned to one NVIDIA B300; PyTorch 2.11.0+cu130; CUDA 13.0.
• Exact revisions: Miles 54f083eb0530e47fce951d2fba243c0d64038cb8; Megatron 5dcbc1020a94ea973316163e8db6719886823e63.
• Import boundary: PYTHONPATH contained only the synced public Miles and Megatron source roots; both imports were asserted to resolve there. The image needed test-only mock==5.2.0 for the Megatron unit-test module.
• Artifact: 20260821T051726Z.tightened-clean.log, SHA-256 511b5a047716f8618d75f9aff89d96a0e00b1929093f56906f2499f6e636b1c9.

Commands:

public_miles_root=/hai-workspace/upstream-main-init-miles
public_megatron_root=/hai-workspace/upstream-main-init-megatron
export PYTHONPATH="${public_megatron_root}:${public_miles_root}" PYTHONNOUSERSITE=1

cd "${public_miles_root}"
CUDA_VISIBLE_DEVICES=0 python -m pytest -q -o addopts='' --disable-warnings \
  tests/fast-gpu/test_nvme_optimizer_main_init.py

cd "${public_megatron_root}"
CUDA_VISIBLE_DEVICES=0 WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 \
  MASTER_ADDR=127.0.0.1 MASTER_PORT=29634 \
  python -m pytest -q -o addopts='' --disable-warnings \
    tests/unit_tests/test_optimizer.py::test_distributed_optimizer_can_defer_main_param_initialization

Raw output:

image=docker.io/radixark/miles:dev-202608210021@sha256:e2e25e4121c595d352933fab81ce05cc8fb2aef8c8478a16abc704b51e330003
miles_revision=54f083eb0530e47fce951d2fba243c0d64038cb8
megatron_revision=5dcbc1020a94ea973316163e8db6719886823e63
torch 2.11.0+cu130 cuda 13.0
miles /hai-workspace/upstream-main-init-miles/miles/__init__.py
megatron_optimizer /hai-workspace/upstream-main-init-megatron/megatron/core/optimizer/__init__.py
.                                                                        [100%]
1 passed, 22 warnings in 0.47s
.                                                                        [100%]
1 passed, 27 warnings in 2.54s

The Miles test verifies exact BF16-to-FP32 bytes in the final file and zero CUDA backing storage after initialization. The Megatron test verifies stable model/optimizer tensor identity, logical shape, FP32 CUDA metadata, and zero backing storage.

Repository checks

At the exact heads above:

$ git diff --check <upstream-base>
(no output)

$ python3 -m py_compile <all changed Python files>
(no output)

$ ruff check <all changed Python files>
All checks passed!

Whitespace-sensitive and whitespace-insensitive aggregate stats are identical. Formatter-only changes were removed instead of running file-wide formatters over upstream baseline drift.

Exact-head full-Miles four-node constructor validation

Run 20260821T060606Z exercised the production Miles initialization path at the exact heads above:

miles/train.py --debug-train-only -> RayTrainGroup -> 32 MegatronTrainRayActor.init calls -> initialize_model_and_optimizer -> paired deferred construction and NVMe initialization.

• Image/hardware: docker.io/radixark/miles:dev-202608210021@sha256:e2e25e4121c595d352933fab81ce05cc8fb2aef8c8478a16abc704b51e330003; four C2 nodes; 32 NVIDIA B300 GPUs.
• Exact revisions: Miles 54f083eb0530e47fce951d2fba243c0d64038cb8; Megatron 5dcbc1020a94ea973316163e8db6719886823e63; harness bbb6b2df9643b7a06ec44d3f0a1488677d4491d6.
• Workload: GLM-5.2, TP2 / PP1 / CP4 / EP32 / DP4 / expert-DP1; DeepEP flex with 20 SMs; sequence length 131,072; BF16 model parameters; FP32 accumulated gradients, mains, and moments; random initialization from an intentionally empty load directory.
• Constructor result: all 32 actors completed initialization in 281.0-327.9 seconds. The native logs contain 64 unique stores and 64 matching main-initialization records: one dense and one expert store for every rank.
• Storage: each rank initialized 2.3 GiB of dense mains and 84.4 GiB of expert mains directly inside final 7.0 GiB and 253.1 GiB stores. The final footprint is approximately 260.1 GiB/rank, 2.032 TiB/node, and 8.130 TiB across the four nodes.
• Outcome: Ray exited zero, with no CUDA OOM, host OOM, ENOSPC, or timeout. The scoped optimizer directory was removed successfully on all four nodes.

Reproduction:

cd /Users/ziangli/playground/projects/upstream-optimizer-main-init-streaming/public_miles_debug_train_glm5_pp1_streaming
./launch_c2.sh provision
./launch_c2.sh sync
./launch_c2.sh ray
./launch_c2.sh prepare
./launch_c2.sh run

Final gate object:

{
  "status": "pass",
  "checks": {
    "ray_job_exit_zero": true,
    "no_timeout": true,
    "no_cuda_oom": true,
    "no_host_oom": true,
    "no_disk_full": true,
    "native_store_records_64": true,
    "native_initialization_records_64": true,
    "unique_store_paths_64": true,
    "unique_initialization_paths_64": true,
    "store_and_initialization_paths_match": true,
    "all_32_ranks_have_dense_and_expert_stores": true,
    "zero_optimizer_streaming_steps": true,
    "zero_training_iterations": true,
    "scoped_optimizer_cleanup_succeeded": true
  }
}

Artifact integrity:

SHA256SUMS   50d1778efe9f7d36b41f9d33102810ed5ace5a9598cd3d7cc0b27a793c361932
result.json  773068eb0d0630f4756caac8e8ed11a95b2c895af26837e801b058c2664431fd
manifest     aab10542b8867817da3773df4cac8098287f78488ca764bdcd148342e9b96400
driver.log   c35c3b736a660f7d70752b32aadcea796f35419c6e44fcaf8d84e30b2659877b

This is an initialization-only production-Miles test: --num-rollout 0 deliberately produced no forward pass, optimizer step, training iteration, or checkpoint operation.

Predecessor two-step GLM-5.2 scale evidence

Run 20260821T001150Z used the same direct-final, bucketwise initialization mechanism on a broader predecessor of this tightened diff: Miles 1697c45ec3973376fdf7c0f0e3fa41aaf8f48682 with Megatron d7f63a091dbe566c2fd8dfd7b7b34cb8821235c7.

• Image/hardware: public Miles image docker.io/radixark/miles:dev-202608200558@sha256:89d2278b083c32d4e794df27aa266cb085a74d0b0bce0b09eab05eddaa972322; four C2 nodes; 32 NVIDIA B300 GPUs.
• Workload: GLM-5.2, TP2 / PP1 / CP4 / EP32 / DP4 / expert-DP1; DeepEP flex with 20 SMs; sequence length 131,072; MBS1; GBS32; BF16 model parameters; FP32 accumulated gradients, mains, and moments; two optimizer steps; random initialization/mock data; no load/save.
• Constructor result: 32/32 initialization records and 64/64 dense/expert stores; maximum PyTorch lifetime allocation 225.12 GiB/GPU; minimum post-constructor CUDA free memory 82.46 GiB/GPU; no constructor OOM.
• Storage: 86.7 GiB of mains inside final 7.0 GiB dense plus 253.1 GiB expert stores per rank, approximately 2.032 TiB/node. There was no temporary full-main file.
• Step result: all 128 (rank, store, step) events completed; four launchers returned zero; no CUDA OOM, host OOM, ENOSPC, timeout, skipped iteration, or NaN.

iteration 1: 1,277,252.1 ms reported mean, 127.8 TFLOP/s/GPU
iteration 2: 1,140,865.4 ms reported mean, 143.0 TFLOP/s/GPU
corrected second interval: 1,004,478.7 ms
corrected throughput: 4,175.60 global tokens/s, 162.42 TFLOP/s/GPU

This scale run is evidence for the mechanism and dense/expert chain, not an exact-current-head run. The current tightened heads are covered by the focused public-image tests above; the removed code concerned checkpoint, path, lifecycle, validation, and failure handling rather than the successful constructor mechanism.

Limitations

• The largest individual FP32 shard and one main-only runtime bucket must fit in HBM.
• This PR inherits the existing streaming feature's supported optimizer, precision, topology, and checkpoint contracts; it adds no new argument guards or compatibility promises.
• Existing checkpoint save/load and model-parameter reload behavior is unchanged. The scale run was fresh/random/no-load and did not exercise a torch_dist save/resume round trip.
• The exact current heads have focused one-B300 validation and the 32-GPU production-Miles constructor validation above. The two-step 32-GPU run used the broader predecessor revisions stated above.
• Buffered application storage timings include GPU/host transfer, optimizer work, page cache, and writeback; they are not physical NVMe bandwidth or durability measurements.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.