[yulangz]

Description

Adapt CPU-staged Muon for the separate WeftMoE engine integration without copying that integration into this branch.

• Allow expert updates to invoke an FP8 writer on the optimizer compute stream instead of BF16 writeback, before recording completion and recycling the staging slot.
• Accept the intentionally released BF16 expert storage only when the registered writer/storage identity matches; avoid redundant single-rank parameter gathering.
• Expose get_full_parameter_master(param) for complete locally owned CPU FP32 matrices. Residency callers drain the optimizer first; unowned/partial parameters return None.
• Bind an installed WeftMoE adapter after staged Muon construction. Without an adapter, this hook is a no-op.
• Extend Muon tests for released-storage writeback and complete-master ownership; document how to combine the integrations.

The Muon update algorithm and native kernels are unchanged.

Related Issue

Targets feature/muon-optimizer-offload from #1743. Complements #1891 and replaces the Muon-specific portion of #1890. This PR adds no AWEX/SGLang adapter implementation, example screenshots, dependencies, or lockfile changes. Both integration PRs are needed for end-to-end Muon + WeftMoE.

Type of Change

• New feature

Validation

• Muon-only tree: 50 passed across test_gpu_staged_muon.py, test_gpu_staged_muon_checkpoint.py, and test_megatron_optimizer_config.py using the Muon-capable Megatron-Core 0.19.0 environment.
• pre-commit run --all-files passed; lockfiles match the Muon target branch.
• The two PR branches merge automatically with both AdamW and Muon binding hooks retained.
• Combined-tree 8-GPU adapter smoke passed for Muon BF16 and FP8 (TP2/EP8, two layers, full recomputation, two optimizer steps, offload/resume and workspace release). BF16 exercised direct restoration from complete CPU master matrices; FP8 exercised released expert storage and inner-master checkpoint roundtrip.

The combined-tree checkpoint check covers inner optimizer master/momentum state and BF16 materialization, not full Megatron distributed-checkpoint save/resume. No new full-model or multi-node RL run was launched.

Checklist

• Relevant tests pass
• Pre-commit hooks pass
• Optional adapter hooks preserve standalone Muon behavior
• No dependency or lockfile changes
• Created by a coding agent using the create-pr workflow

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.