Support bounded distributed-optimizer main initialization

[zianglih]

What does this PR do ?

@HumansAnd

Add one opt-in constructor primitive for external optimizer-state backends: retain stable distributed-optimizer FP32 main-param tensor handles without retaining the cumulative FP32 main state in HBM.

• Default behavior: unchanged. OptimizerConfig.defer_main_param_initialization defaults to False; the ordinary path still uses shard_model_param.clone().float().
• Deferred behavior: in the ordinary mixed-precision main-param branch, Megatron creates the same-shaped CUDA FP32 tensor, immediately releases its backing storage, and keeps its logical shape, dtype, device, and object identity.
• Storage boundary: Megatron owns only the handle. It does not know about files, NVMe, buckets, checkpoints, or Miles. The caller must initialize every deferred handle before use.
• Peak-memory contract: each FP32 shard exists briefly before its storage is released, so the largest individual shard must still fit. The option removes cumulative main-param residency; it does not make the largest shard allocation-free.
• Aggregate diff: three files, +47/-0; no unrelated formatter changes and no added try, except, or finally blocks.

This extends the integration point used by optimizer-state streaming from radixark/Megatron-LM#63. NVIDIA/Megatron-LM#6544 is related draft work on chunked optimizer-state offloading; this PR deliberately leaves storage policy outside Megatron core.

Paired PR

• Miles: radixark/miles#2653, head 54f083eb0530e47fce951d2fba243c0d64038cb8.
• Miles automatically enables this constructor mode when --stream-optimizer-state-to-disk is on, then initializes one existing runtime bucket at a time directly in its final files.

Validation

base: radixark/Megatron-LM:miles-main @ 2e64d788b67fe8155774200c87606e1e449c494b
head: 5dcbc1020a94ea973316163e8db6719886823e63
paired Miles head: 54f083eb0530e47fce951d2fba243c0d64038cb8

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.
• 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 Megatron test verifies that model and optimizer retain the same main-param tensor objects with nonzero logical size, FP32 CUDA metadata, and zero backing storage. The Miles test verifies exact BF16-to-FP32 bytes after initializing those handles into final storage and confirms the CUDA storage is released again.

Repository checks

$ git diff --check upstream/miles-main
(no output)

$ python3 -m py_compile megatron/core/optimizer/distrib_optimizer.py \
    megatron/core/optimizer/optimizer_config.py tests/unit_tests/test_optimizer.py
(no output)

$ ruff check megatron/core/optimizer/distrib_optimizer.py \
    megatron/core/optimizer/optimizer_config.py tests/unit_tests/test_optimizer.py
All checks passed!

File-wide Black/isort were not applied because the upstream file is not in their current canonical form and doing so introduces unrelated import and line-wrap changes. The submitted diff keeps the baseline formatting untouched.

Exact-head full-Miles four-node constructor validation

Run 20260821T060606Z exercised this exact Megatron head through the production Miles initialization path:

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: Megatron 5dcbc1020a94ea973316163e8db6719886823e63; Miles 54f083eb0530e47fce951d2fba243c0d64038cb8; 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 exercised the same deferred-handle and direct-final bucket mechanism on a broader predecessor: Megatron d7f63a091dbe566c2fd8dfd7b7b34cb8821235c7 with Miles 1697c45ec3973376fdf7c0f0e3fa41aaf8f48682.

• 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 params; 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/steps: approximately 2.032 TiB/node; all 128 (rank, store, step) events completed; 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 mechanism and dense/expert-chain evidence, not an exact-current-head run. The current heads are covered by the focused public-image tests above; the tightened commits removed validation, formatter, checkpoint, path, lifecycle, and failure-handling changes while retaining the successful constructor mechanism.

Remaining limitations

• This boolean is a constructor primitive, not a complete offload implementation. Enabling it without an external backend that initializes the handles is invalid.
• The branch is consumed only by the ordinary mixed-precision main-param path. Existing FP8, precision-aware, and native-FP32 branches are unchanged; callers own compatibility validation.
• The largest individual FP32 shard must fit briefly in HBM.
• Exact-current-head validation covers the focused one-B300 test and the 32-GPU production-Miles constructor run above. The two-step 32-GPU run used the broader predecessor revisions stated above and did not exercise checkpoint save/resume.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.