Stream optimizer main initialization to NVMe

[zianglih]

Summary

@HumansAnd

ci-megatron-pr: #86

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

• Paired Megatron change: radixark/Megatron-LM#86 at c0dad034d0eb0e57fde0267d38589f42eb7213b1.
• Upstream integration: upstream now gives Muon its own file-backed chunked-state path. This PR leaves that path intact and enables deferred main initialization only when the selected optimizer is not Muon.
• Scope: only initial Adam/DistributedOptimizer main-parameter construction changes. Existing optimizer steps, checkpoint save/load, Muon state handling, model-parameter reload, store numbering, and fallback handling remain unchanged.
• Diff: one commit, five files, +89/-6; no compatibility fallback or broad exception handling was added.

Design

• Megatron creates each same-shaped CUDA FP32 main handle and immediately releases its backing storage while retaining tensor identity and metadata.
• Miles materializes one existing main bucket, copies the still-live BF16 model shards into it, flushes directly to the final bucket file, releases CUDA storage, synchronizes and evicts that initialized file range, then advances.
• Initialization reuses the existing 200,000,000-element soft bucket limit and final store layout. There is no temporary full-main file or parallel fallback path.
• The largest individual FP32 shard must still fit briefly during Megatron construction, followed by one main-only runtime bucket during Miles initialization.

Validation

Miles base:    ca25d697c0743fde52d27a5ebd88a3883b7fd838
Miles head:    ab760c086fde48c7c40b10030970008f6ee64e65
Megatron base: 235952df607b3820716e5e67728a5ab470ca33ae
Megatron head: c0dad034d0eb0e57fde0267d38589f42eb7213b1

• Image: docker.io/radixark/miles:dev-202608301223; OCI index sha256:22c82828fca9215e22c8493bfe53fba12ca2b63feb0e6343c052fc48f16e7d47; linux/amd64 manifest sha256:cb5beb813caa5aaaf47b144bc4fabe7b75cae1c1cf65bf7cac5b6840c5b110e3.
• Environment: one C2 node, 8 x NVIDIA B300 SXM6 AC (275040 MiB each), driver 590.48.01, PyTorch 2.13.0+cu130, CUDA 13.0, image CUDA_VERSION=13.0.3. Focused GPU tests used GPU 0; each E2E used four GPUs.
• Source boundary: the image contained both exact bases. C2 tested Miles code at 0401d4e3f2baf20855bd4ce821e20d4569664fbe; the submitted head differs only by qualifying one documentation sentence to the Adam/DistributedOptimizer path, and every changed Python blob is identical. Megatron was tested at the exact head above. The full Miles command prepends /root/Megatron-LM, so the two changed Megatron runtime files were overlaid on its exact image base; their hashes matched the synced Megatron head.

Focused tests

cd /hai-workspace/upstream-main-init-miles-rebase-0830
PYTHONPATH=$PWD:/root/Megatron-LM \
python -m pytest -q -o addopts= --disable-warnings \
  tests/fast/optimizers/test_nvme_stream.py

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$PWD:/root/Megatron-LM \
python -m pytest -q -o addopts= --disable-warnings \
  tests/fast-gpu/test_nvme_optimizer_main_init.py

cd /hai-workspace/upstream-main-init-megatron-rebase-0830
CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$PWD \
torchrun --master-addr=127.0.0.1 --master-port=29687 \
  --nnodes=1 --nproc-per-node=1 -m pytest -q -o addopts= --disable-warnings \
  tests/unit_tests/test_optimizer.py::test_distributed_optimizer_can_defer_main_param_initialization

........                                                                 [100%]
8 passed, 20 warnings in 0.16s
.                                                                        [100%]
1 passed, 20 warnings in 0.42s
.                                                                        [100%]
1 passed, 26 warnings in 2.95s

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

The effective dispatch contract was also checked directly:

optimizer=adam is_muon=False defer_main_param_initialization=True
optimizer=dist_muon is_muon=True defer_main_param_initialization=False

Four-GPU Qwen3-4B E2E

After one explicit-loopback checkpoint conversion, the existing E2E execute() paths ran with TP2/PP1/CP1, four B300s, two rollouts/optimizer steps, colocated SGLang, disk train offload, and --stream-optimizer-state-to-disk:

cd /hai-workspace/upstream-main-init-miles-rebase-0830
PYTHONPATH=$PWD:/root/Megatron-LM python -c \
  'from tests.e2e.megatron.test_qwen3_4B_offload_disk_stream import execute; execute()'
PYTHONPATH=$PWD:/root/Megatron-LM python -c \
  'from tests.e2e.megatron.test_qwen3_4B_muon_offload_disk import execute; execute()'

Adam/DistributedOptimizer, Ray job raysubmit_G7VgaVvS5pNDL9Xn:

rank 0: initialized 3.7 GB; steps read/write 3.7/11.2 GB in 6.6s, 11.2/11.2 GB in 5.5s
rank 1: initialized 3.7 GB; steps read/write 3.7/11.2 GB in 6.6s, 11.2/11.2 GB in 5.5s
rank 2: initialized 3.7 GB; steps read/write 3.7/11.2 GB in 6.6s, 11.2/11.2 GB in 4.7s
rank 3: initialized 3.7 GB; steps read/write 3.7/11.2 GB in 5.0s, 11.2/11.2 GB in 4.8s
disk offload armed for 4 ranks under /root/train_offload_disk_stream
optimizer state streaming ran on 4 ranks
Job 'raysubmit_G7VgaVvS5pNDL9Xn' succeeded

Muon upstream control, Ray job raysubmit_1DNkhNExcJ1nLigM:

each rank: setup=1, file-backed Muon steps=4, disk-offload armed=1
each rank: two state partitions (3.38 GB and 0.73 GB) for each of two optimizer steps
each rank: Adam NVMe main-init records=0, Adam NVMe streaming records=0
disk offload armed for 4 ranks under /root/train_offload_muon_disk
Muon optimizer state was file-backed on 4 ranks
Job 'raysubmit_1DNkhNExcJ1nLigM' succeeded

Both jobs completed both training steps. No CUDA OOM, host OOM, ENOSPC, invalid step, or Ray failure was observed.

Repository checks

$ git diff --check ca25d697c0743fde52d27a5ebd88a3883b7fd838
(no output)
$ python3 -m py_compile <all changed Python files>
(no output)
$ ruff check <all changed Python files>
All checks passed!

Limitations

• The image checkpoint-conversion helper initially selected an unresolvable pod hostname. It was stopped before checkpoint creation; the same conversion completed with explicit loopback torchrun, then both existing E2E execution paths ran unchanged.
• SGLang logged recovered post-warm-up /freeze_gc connection races before reporting ready. Subsequent requests and both Ray jobs succeeded.
• The run covered fresh load, initialization, two optimizer steps, disk offload/wake-up, and weight updates for Adam and Muon. It did not exercise checkpoint save/resume.
• This PR inherits the existing streaming feature's optimizer, precision, topology, and checkpoint contracts; it adds no new fallback or compatibility promise.

[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.