Fix Muon optimizer under ZeRO CPU offload and bound gather buffers - #8464
Conversation
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 @@ def __init__( | |||
| 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 👍 / 👎.
| 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 👍 / 👎.
|
hi @jinyouzhi whenever this is ready for review just turn it from draft and I'll start review, thanks! |
daa6261 to
e377eaf
Compare
Thank you for your remind! Just rebase with latest master and ready to review. |
…ch (deepspeedai#8600) Fixes deepspeedai#8443. ## The problem ZeRO-3 applies Muon inside the gradient reduce (`_apply_distributed_muon_update`, called from `__avg_scatter_contiguous_grads`), and that runs every micro-batch. With `gradient_accumulation_steps: n`, the momentum advances `n` times per optimizer step, and Newton-Schulz orthogonalizes each micro-batch's partial gradient instead of the accumulated one. ZeRO-1/2 apply Muon at the accumulation boundary and are correct. On 2 GPUs, fp32, with the same 8 samples per step either way (one micro-batch of 8 at `gas=1`, four of 2 at `gas=4`), three steps, relative difference in the weights: | | master | this PR | | --- | ---: | ---: | | ZeRO-2, `gas=1` vs `gas=4` | 3.6e-4 | 3.6e-4 | | ZeRO-3, `gas=1` vs `gas=4` | **1.3e-1** | 3.6e-4 | | Newton-Schulz calls, ZeRO-3, 2 matrices, 2 steps, `gas=4` | 16 | 4 | ZeRO-3 now lands on exactly ZeRO-2's figure. ## The change This is option 1 from the discussion in deepspeedai#8443. It is scoped to ZeRO-3 without optimizer offload. - The reduce path no longer runs Muon when optimizer offload is off. The partitions accumulate the raw averaged gradient, as they do for every other optimizer. - `step()` calls `_apply_muon_to_accumulated_grads()` after the overflow check and before the gradient norm. For each Muon sub-group, it: - all-gathers each parameter's accumulated gradient partitions, in chunks bounded by `reduce_bucket_size` as the reduce buckets were; - runs the existing round-robin Muon update once; - writes each rank's slice back into its partition. - The per-sub-group body of `_apply_distributed_muon_update` is moved into `_muon_update_sub_group`. It takes the full-shape gradients explicitly, because at step time the parameters are partitioned and `param.grad` can't hold them. The reduce path calls it with `param.grad` as before. - The gradient norm is still taken over the Muon update, as before. Clipping semantics are unchanged (deepspeedai#8439 / deepspeedai#7776 are separate). - Because the update now runs after the overflow check, a step the loss scaler discards no longer touches the momentum. That is the ZeRO-3 counterpart of deepspeedai#8435. - Collectives: each Muon parameter's gradient and momentum are gathered once per step instead of once per micro-batch. At `gas=n` that is `n` times fewer. The optimizer-offload path is unchanged; deepspeedai#8464 is working on it. jinyouzhi added a pointer to this shape in deepspeedai#8464, and the overlap is limited to `_apply_distributed_muon_update`. ## Testing On 2×H20: - New file `tests/unit/v1/ops/muon/test_muon_zero3_grad_accum.py`. Both tests fail on master and pass here. - `test_newton_schulz_runs_once_per_matrix_per_step`: Newton-Schulz calls summed over ranks come to 2 × steps, not 2 × steps × gas. - `test_gradient_accumulation_matches_one_large_micro_batch[2, 3]`: `gas=1` and `gas=4` agree to within half-precision Newton-Schulz noise at both stages. - `tests/unit/v1/ops/muon/` plus `tests/unit/runtime/zero/test_per_head_muon.py`: 260 passed. --------- Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
alanhuangyoo
left a comment
There was a problem hiding this comment.
Nice, this fixes it. I ran ZeRO-1/2/3 with CPU offload against ZeRO-0 on 2×H20 (same batch on every rank, no clipping) and it matches bit for bit, while master is about 5e-2 off at ZeRO-1/2.
One thing to sort out: #8609 does the universal checkpoint part too, and _flat_optimizer_states here skips any state that isn't partition-sized. Without offload, ZeRO-1/2 keep Muon's momentum whole per param, so it gets dropped quietly on conversion (#8609 handles that). Would you be ok leaving the UC changes to #8609 so this one stays about offload? Your offload momentum is partition-sized, so it converts fine there as is.
Thank you very much for your thorough review. I’m confident that your tests and suggestions will greatly accelerate fixing this issue. I completely agree with your proposed approach to handling the momentum conversion. This was originally intended as a supplementary improvement while validating the offload fix, so dividing the work this way should allow us both to stay focused on our respective areas. |
13cf46d to
3690e87
Compare
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>
…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>
…ll site Add a docstring note pinpointing exactly where a future boundary-gated, step-time rewrite of _apply_distributed_muon_update should land, mirroring the CPU-offload path's approach. Purely documentation, no behavior change; intended to keep any future deepspeedai#8443 fix small and reduce merge-conflict surface regardless of which change lands first. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
The CPU-offload Muon update paths in stage_1_and_2.py and stage3.py called muon_update() without forwarding param.muon_num_heads, so parameters configured for per-head Muon were silently orthogonalized as one full matrix under optimizer offload instead of per-head, changing optimization behavior relative to non-offloaded execution. Pass num_heads=getattr(param, "muon_num_heads", None) in both call sites, and add a regression test comparing trained parameters (not losses, which are insensitive to this difference once Adam's per-element normalization is applied on top of the Muon update) between per-head and whole-matrix orthogonalization under CPU offload across ZeRO stages 1/2/3. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Skip optimizer-state swap-in when the Muon momentum buffer is configured to remain resident in memory. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Skip the reduce-path Muon update when optimizer offload is enabled. CPU-offloaded gradients are handled by the step-time offload path, and must not enter the non-offload all-gather path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
|
Great, thanks! I'll make sure #8609 still handles the offload layout once this lands. |
| if tensor is not None: | ||
| # Callers can pin an optimizer state in memory (e.g. the Muon momentum buffer under | ||
| # save_muon_momentum_buffer_in_memory) by tagging the tensor, which excludes it from swapping. | ||
| if not getattr(tensor, "swappable", True) or getattr(tensor, "is_resident", False): |
There was a problem hiding this comment.
what is the scenario when swappable is used and what is the scenario for is_resident?
There was a problem hiding this comment.
Good question. is_resident is redundant, I’ll remove it and keep swappable=False.
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
| swap_out_tensors(aio_handle, swap_buffers, swap_paths) | ||
| assert len(swap_buffers) == aio_handle.wait() | ||
| if swap_info.unswapped_gradients: | ||
| swap_info.write_unswapped_gradients(src_buffer=parameter.grad) |
There was a problem hiding this comment.
do we need to call swap_info.release_unswapped_gradients() after this line?
There was a problem hiding this comment.
No, not at this intermediate writeback point. These small gradient fragments have no NVMe backing file: write_unswapped_gradients() updates their CPU copies with the Muon-transformed gradients, and _prepare_sub_group() subsequently swaps the subgroup in again for the actual optimizer step. Clearing them here would lose those updates; when all fragments are unswapped it can also leave the next swap-in without a gradient buffer.
They are already released by the normal post-step swap_out_optimizer_state() in both PartitionedOptimizerSwapper and PipelinedOptimizerSwapper.
I have added a lifecycle comment and strengthened the all-unswapped-fragments regression locally: two training steps are compared with an independent full-gradient Muon reference for both swappers. Both cases pass on an RTX 5090 D using real NVMe I/O.
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Load swappable gradients independently of Muon momentum residency. Cover resident and swapped momentum with the NVMe numerical regression and document the distinction.\n\nValidated on RTX 5090 D 32 GB GPUs with real NVMe I/O, using partitioned and pipelined swappers. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Pack only owned ZeRO-1/2 slices and chunk collectives to a 64 MiB per-rank scratch budget. Process full gradients and momentum in bounded batches while preserving whole-matrix Newton-Schulz updates. Validate multi-step equivalence and actual collective buffer sizes on 1, 2, and 4 RTX 5090 D 32 GB GPUs; retain the independent full-gradient numerical regressions. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Explain why intermediate gradient writeback must retain CPU fragments until the optimizer consumes them. Compare two all-unswapped training steps with an independent full-gradient reference for both NVMe swappers. Validated both swapper implementations with real NVMe I/O on an RTX 5090 D 32 GB GPU. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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).