[dev] Add chunked optimizer-state and master-weight offload

[hxbai]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

The existing optimizer-state offload path restores the complete optimizer state to GPU for every optimizer step, so it does not reduce peak optimizer-update memory.

This PR replaces that implementation with chunked GPU optimizer execution backed by pinned CPU canonical storage. Selected optimizer tensor states and master weights remain on CPU between updates. Tensor states are restored, updated, and offloaded in bounded chunks, while selected master weights use one full restore window.

External optimizers remain unaware of offloading; no Transformer Engine or Emerging Optimizers changes are required.

Acknowledgements and related work

This work benefited from our collaboration with @yanring (Zijie Yan) at Periodic Labs, including design discussions and his implementation with chunked optimizer-state and master-weight offloading.

See #6544, for a complementary implementation that chunks both optimizer states and master weights.

Supported configurations

• Adam with DistributedOptimizer
• BF16 Muon with compact-layout LayerWiseDistributedOptimizer
• BF16 and FP8 parameter gather, including MXFP8 grad-buffer reuse
• Precision-aware FusedAdam
• Partial optimizer-state offload
• Distributed checkpoint save/load
• Checkpoint compatibility between offloaded and non-offloaded configurations

User interface

The new options are:

--chunked-optimizer-state-offload
--optimizer-state-offload-chunk-size-mb <size>
--optimizer-state-offload-fraction <fraction>

⚠️ 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 @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• 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

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

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hxbai]

/claude strict-review

[claude]

Strict Review Summary

Findings: CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 1

Scope

This PR replaces the previous OptimizerStateOffloader (full-state reload every step, peak-neutral) with a chunked GPU optimizer-execution path backed by pinned CPU canonical storage (ChunkedOptimizerStateOffloader). It touches the core optimizer stack (optimizer.py, distrib_optimizer.py, layer_wise_optimizer.py), config/argument plumbing, checkpointing, and the training loop, with ~1900 lines of new/updated unit tests.

Assessment

The implementation is careful and well-documented. I traced the areas most prone to silent-corruption bugs and found them sound:

• CUDA stream lifecycle: D2H/H2D ordering uses wait_stream/record_stream correctly; H2D is ordered after both the compute stream (staging-buffer allocation) and the D2H stream (CPU-buffer production) before reading either side. The two-slot staging pool bounds host run-ahead as documented.
• Master-weight rebinding: canonical CPU/CUDA binding transitions are guarded by assert_master_weights_resident at every external master reader, and self-heal via ensure_master_for_param_sync.
• Checkpoint save/load: load_state_dict_without_device_cast correctly avoids torch's automatic device cast that would reconstruct full state on CUDA; precision-aware FusedAdam scaled-state restore is handled through TE's public setter in bounded chunks; sync/async and torch_dist restrictions are enforced at both arg-validation and save_checkpoint time.
• Chained/LayerWise integration: one shared D2H/H2D stream pair across nested managers; bounded cross-child prefetch pipeline; no dangling references to the removed offload API. All new identifiers have real use paths; typing imports are all consumed.
• Process-group guidance: no new direct parallel_state.get_*_group() reads introduced.

Notable points (non-blocking)

1. [IMPORTANT Compatibility] The deprecated --offload-optimizer-states alias now carries new hard checkpoint requirements (torch_dist, synchronous optimizer-state save). Existing scripts using torch format or --async-save will fail hard on resume. Suggest surfacing the reason in the runtime FutureWarning.
2. [SUGGESTION] _group_values_equal intentionally accepts CUDA tensor-valued group fields on shape/dtype/device only (to avoid a per-chunk host sync). Correct for supported optimizers; a targeted test for the subset-dependent-CUDA-field case would lock in the contract.

Risk level: Low-Moderate. Large surface area, but confined to an opt-in feature path (no behavior change when the flags are off), with strong test coverage and defensive validation.

[hxbai]

/ok to test a215785

[hxbai]

/ok to test 2c6276a

[hxbai]

/ok to test 8bfbea0

[Wohox]

LGTM, thanks!

[hxbai]

/ok to test 2de7461