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: unchanged. OptimizerConfig.defer_main_param_initialization defaults to False; the ordinary path still uses shard_model_param.clone().float().
• Deferred path: Megatron creates the same-shaped CUDA FP32 tensor, immediately releases its backing storage, and retains its logical shape, dtype, device, and object identity. The caller must initialize every deferred handle before use.
• Boundary: Megatron owns only the handle. Files, NVMe buckets, writeback, and lifecycle remain outside core.
• Peak memory: the largest individual FP32 shard must still fit briefly; this removes cumulative main-param residency, not the largest single allocation.
• Scope: only the ordinary mixed-precision distributed-optimizer branch changes. Existing CPU/chunked offload, quantized, FP8, precision-aware, and native-FP32 paths are untouched. No fallback or exception-handling path was added.
• Upstream integration: Miles now has a separate file-backed Muon state path. The paired Miles change enables this primitive only for non-Muon optimizer streaming; Muon does not enter the deferred-main path.
• Diff: one commit, three files, +47/-0.

Paired Miles integration: radixark/miles#2653 at ab760c086fde48c7c40b10030970008f6ee64e65.

Validation

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

• 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 integration run used four GPUs.
• Source boundary: the image contained both exact bases. Megatron was tested at the exact submitted head. C2 tested Miles code at 0401d4e3f2baf20855bd4ce821e20d4569664fbe; the submitted Miles head differs only by qualifying one documentation sentence, and all changed Python blobs are identical. The full Miles command prepends /root/Megatron-LM, so its two changed Megatron runtime files were overlaid on the exact image base; their hashes matched the synced Megatron head.

Focused tests:

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

cd /hai-workspace/upstream-main-init-miles-rebase-0830
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

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

The Megatron test verifies stable model/optimizer tensor identity, nonzero logical shape, FP32 CUDA metadata, and zero backing storage. The Miles test verifies exact BF16-to-FP32 bytes after direct-final initialization and confirms the CUDA storage is released again.

The existing Miles Qwen3-4B disk-stream E2E then ran with TP2/PP1/CP1, four B300s, two rollouts/optimizer steps, colocated SGLang, disk train offload, and --stream-optimizer-state-to-disk:

Ray job raysubmit_G7VgaVvS5pNDL9Xn: SUCCEEDED
direct main-param initialization: 4/4 rank worker logs
disk-offload reclaim armed:       4/4 rank worker logs
NVMe streaming steps:             8 records across 4/4 ranks
each rank: initialized 3.7 GB directly to its final file
step 1: each rank read 3.7 GB and wrote 11.2 GB
step 2: each rank read 11.2 GB and wrote 11.2 GB
disk offload armed for 4 ranks under /root/train_offload_disk_stream
optimizer state streaming ran on 4 ranks

Upstream's separate Muon E2E then passed as a control: Ray job raysubmit_1DNkhNExcJ1nLigM succeeded, all 4 ranks armed disk offload and logged four file-backed Muon state records, while all 4 recorded zero Adam NVMe main-init and streaming records. The direct dispatch probe agreed:

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

Both jobs completed both steps normally. There was no CUDA OOM, host OOM, ENOSPC, invalid step, or Ray failure.

Repository checks at the submitted head:

$ git diff --check 235952df607b3820716e5e67728a5ab470ca33ae
(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!

Current-run boundaries: checkpoint conversion required an explicit loopback rendezvous because the C2 pod hostname was not resolvable; it then completed and both existing E2E execution paths ran unchanged. SGLang logged recovered post-warm-up /freeze_gc connection races before reporting ready. The run covered fresh load, initialization, two optimizer steps, offload/wake-up, and weight updates for Adam and Muon, but not 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 this PR.