Fix Muon optimizer under ZeRO CPU offload and bound gather buffers - #8464
Open
jinyouzhi wants to merge 13 commits into
Open
Fix Muon optimizer under ZeRO CPU offload and bound gather buffers#8464jinyouzhi wants to merge 13 commits into
jinyouzhi wants to merge 13 commits into
Conversation
Reconstruct full Muon gradients and momentum before Newton-Schulz updates while preserving CPUAdam auxiliary updates across ZeRO stages. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Limit cached GPU scratch buffers with LRU eviction and explicit cleanup, and cover buffer lifecycle behavior in tests. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…e eviction - Guard _apply_distributed_muon_update in ZeRO-3 when offload_optimizer is enabled to prevent device mismatch during backward and duplicate Newton-Schulz updates. - Unwrap single-buffer outputs in _muon_all_gather_partitions under multi-GPU ZeRO-1/2 to avoid returning nested lists. - Initialize _muon_allgather_buffers as an OrderedDict in ZeRO-1/2 so popitem(last=False) works during LRU cache eviction. - Clear cached all-gather buffers in DeepSpeedZeroOptimizer.destroy(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
jinyouzhi
requested review from
loadams,
tjruwase and
tohtana
as code owners
September 8, 2026 16:33
# Conflicts: # deepspeed/runtime/zero/stage_1_and_2.py
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b139b8f82c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| @@ -236,6 +236,9 @@ | |||
| self.dtype = self.optimizer.param_groups[0]['params'][0].dtype | |||
There was a problem hiding this comment.
This non-merge commit has no Signed-off-by trailer, so it does not satisfy the repository's mandatory DCO requirement and will be rejected by the corresponding CI check; recreate the commit with git commit --signoff.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
Comment on lines
+2468
to
+2469
| if self._swappable_optimizer_subgroup(sub_group_id): | ||
| self._optimizer_states_and_gradient_swap_in(sub_group_id) |
There was a problem hiding this comment.
When ZeRO-3 optimizer state is NVMe-swappable and save_muon_momentum_buffer_in_memory=true, this unconditional swap-in (paired with the unconditional swap-out below) ignores the option's documented promise to keep Muon momentum in CPU/GPU memory. It also performs this NVMe round trip before _prepare_sub_group() performs the normal optimizer-state swap again, adding two avoidable transfers per step for the large configurations this option targets; use the resident muon_momentum_buffer_partitioned_groups_flat path when the flag is enabled, as the previous Muon update path did.
Useful? React with 👍 / 👎.
…memory is set In ZeRO-3 CPU offload path, bypass NVMe swap-in/swap-out and use the resident muon_momentum_buffer_partitioned_groups_flat when save_muon_momentum_buffer_in_memory is enabled, avoiding unnecessary NVMe round-trips. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…caling Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…residency - Fix gradient clipping norm accounting under loss scaling by storing scaled update norm in norm_for_param_grads for ZeRO-1/2/3. - Exclude resident ZeRO-3 Muon momentum buffer from OptimizerStateSwapInfo to prevent eviction by NVMe swapper. - Ensure swappable optimizer subgroups properly swap in and write back updated gradients and states in ZeRO-3 CPU offload. - Expand TestMuonOffloadLossScaling to ZeRO-1/2/3 with clipping equivalence validation across loss scales. - Add TestMuonZero3NVMeMomentumResidency for multi-step persistence of resident momentum under NVMe offload. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…d swapper - Retain unswapped gradient fragment ownership in OptimizerStateSwapInfo across Muon writeback and step until swap_out_optimizer_state. - Guard swapped gradient writing in writeback_optimizer_state_and_gradients when swapped_gradients is empty. - Implement writeback_optimizer_state_and_gradients and release_swap_buffers in PipelinedOptimizerSwapper. - Ensure synchronous swap-in without async prefetch during _apply_muon_updates_cpu_offload. - Expand TestMuonZero3NVMeMomentumResidency to cover both non-pipelined and pipelined NVMe swapping. - Add test_zero3_nvme_aggregate_unswapped_fragments for swappable subgroups composed of sub-MiB unswapped gradient fragments. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
- Add test_zero3_nvme_mixed_fragments_numerical_equivalence to TestMuonZero3NVMeMomentumResidency. - Construct a single subgroup containing both >= 1 MiB (swapped) and < 1 MiB (unswapped) parameters. - Verify multi-step numerical equivalence against a non-NVMe CPU offload reference across both partitioned and pipelined swappers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
… reference - Fix Muon parameter selection in ZeRO-3 by checking getattr(p, 'use_muon', False) and p.ds_shape instead of p.ndim. - Assert all 3 mixed-fragment parameters (large, proj, small) are identified and tracked. - Replace shared CPU-offload reference with an independent pure-PyTorch full-gradient reference maintaining momentum across steps. - Fix parameter selection in TestMuonOffloadLossScaling for ZeRO-3. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
- Add TestMuonZero3NVMeMultiRankMixedFragments with world_size=2 covering cross-rank all-gather, partition slicing, and DP gradient averaging. - Size parameters so per-rank partition includes both >= 1 MiB (swapped) and < 1 MiB (unswapped) fragments in the same subgroup. - Construct independent high-precision reference holding FP32 master weights/momentum and averaging per-rank FP16 gradients. - Compare actual update tensor (init - final) directly against reference update using ref_update.norm() as denominator. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
ZeRO-3 allocated the Muon momentum buffer with communication_data_type, so under an fp16 config the accumulator was rounded to fp16 on every step even though the update itself was computed in fp32. ZeRO-1/2 derive the buffer from the fp32 master partition via zeros_like(flat_param), and the NVMe optimizer swapper already stores state in master_weights_and_grads_dtype, so fp16 momentum was both lossy and inconsistent with the rest of the stack. Allocate the momentum buffer in the master dtype, gather it (and the gradients it is blended with) in that dtype, and promote the gradient before Newton-Schulz in the non-offload path. The oracles in the ZeRO-3 NVMe tests created fp16 momentum too, so they reproduced and accepted the same drift; they now keep fp32 momentum. The 2-rank test could not observe ZeRO-3 padding because every Muon matrix had an element count divisible by the world size. Add an odd-numel matrix, assert that it really produces a padded final partition, and cover the padding sensitive reconstruction and slicing paths. Both NVMe tests also allowed 25% update-relative error, which was wide enough to hide a tail reconstruction or second-step momentum bug. The error was not inherent: a sum() loss over a linear stack yields rank-1 gradients, and Newton-Schulz amplifies fp16 noise in the near-null directions by roughly the fifth power of its slope at zero. Switching to a squared loss over a wide batch (scaled out of the fp16 subnormal range) and averaging DP gradients in the fp16 communication dtype, as DeepSpeed does, drops the observed error from 0.12-0.24 to under 0.02, so the bound is now 0.05. Verified on 8x Intel Battlemage (XPU): 161 passed in tests/unit/ops/muon, and 121 passed / 61 skipped in the ZeRO NVMe checkpointing and tensor fragment suites. Also refine the surrounding Muon work: drop an unused local in the ZeRO-1/2 offload path, collapse the triplicated momentum lookup and writeback branches in _apply_distributed_muon_update, and hoist the duplicated gradient writeback in the two optimizer swappers into a shared OptimizerSwapper helper. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
The in-memory Muon momentum buffer (save_muon_momentum_buffer_in_memory) is held twice: in optimizer.state and in the muon_momentum_buffer_partitioned_groups_flat cache that the Muon update path actually reads. Neither checkpoint load path kept the two in sync: * _rigid_load_state_dict() calls Optimizer.load_state_dict(), which rebinds optimizer.state to fresh tensors. The cache kept serving the pre-load (zero) momentum and wrote it straight back over the restored state on the next step. * NVMe optimizer offload skips the ZeRO optimizer state dict entirely and rebuilds state by copying swap files. A resident buffer is excluded from swapping by design, so it never reaches those files and the restored value was dropped outright. Copy the checkpointed momentum into the resident buffer and re-bind optimizer.state to it in both paths, so a resumed run continues from the momentum it saved. Add an interrupted-vs-uninterrupted checkpoint test over the non-offload path plus the partitioned and pipelined NVMe swappers. It asserts the restored momentum, the shared tensor identity, and that the post-resume weight update matches the uninterrupted run. Without the fix all three variants fail, and the end-to-end update error is 0.91 instead of 0. Validated on 8x Intel Battlemage XPU: tests/unit/ops/muon/test_muon.py 164 passed; test_nvme_checkpointing.py + test_zero_optimizer.py 59 passed (10 pre-existing missing-fixture collection errors). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
In PR #8278, we extended the auxiliary optimizer in
MuonWithAuxAdamto support optimized backends such asCPUAdam. However, Muon with CPU offload under ZeRO stages had critical correctness issues and performance/memory pitfalls:param.gradis on accelerator devices. Calling_apply_distributed_muon_updateduring backward caused cross-device runtime errors and double-updating.TypeError: zeros_like()crashes on the first step.Changes
CPUAdam) intact._get_norm_groups()and guard_apply_distributed_muon_updateduring backward whenoffload_optimizeris enabled._muon_allgather_buffersLRU cache (capped at 256MB) usingOrderedDictto prevent VRAM spikes and memory fragmentation.destroy().Tests
pytest tests/unit/ops/muon/test_muon.py -k 'not TestMuonConfigs'TestMuonOptimizerOffload(ZeRO stages 1, 2, and 3 with CPU offload)TestMuonAllGatherBufferLifecycleTestMuonZero12NumericalCorrectnessyapf,flake8,codespell,check-torchdist,check-license).