Enable Muon optimizer training with MFSDP v2
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Goal
Enable training models with the Muon optimizer when using MFSDP v2.
## Background
Muon applies Newton-Schulz (NS) updates to momentum tensors. For correctness and efficiency, those momentum tensors need to remain tensor-atomic at the point where NS is applied.
With MFSDP v2, parameters and gradients may be dim-0 uneven-sharded across multiple ranks. This means a single tensor's weights and gradients can span ranks, while Muon wants the corresponding momentum tensor to be available as a tensor-atomic unit for the NS step.
## Candidate approaches
1. Keep weights and gradients sharded across ranks, and all-gather momentum for boundary parameters before applying NS.
2. Keep weights and gradients sharded across ranks, but use all-to-all-v for boundary-parameter momentum before applying NS to avoid duplicated computation.
3. Keep weights and gradients tensor-atomic as well. This makes the optimizer communication-free, but shifts communication into forward/backward, which would use all-to-all-v rather than all-gather/reduce-scatter.
## Current status
Jan has prototyped approaches 1 and 2.
Approach 2 appears to have stronger performance data so far, so it should likely be implemented first after a simple baseline is in place.
## Implementation plan
* https://github.com/NVIDIA/Megatron-LM/pull/6597
* Implement peer-to-peer communication
* Add FsdpOrthogonalizedOptimizer
* MCore Integration, including ChainedOptimizer work for adam+muon
## Implementation notes
DBuffer may be useful for the simpler/unbatched communication path, but it may not be the right abstraction for every Muon communication strategy.
Before an optimizer step, there is one `main_grad` DBuffer per `ParameterGroup`. For communication that stays within that structure, it should be straightforward to express redistribution directly, for example with something like:
```python
DBuffer.redistribute(axis, TensorAtomic())
```
However, for a batched all-to-all-v implementation that combines all `main_grad` DBuffers into a few large buffers, DBuffer may become an extra indirection without buying much. In that case, the implementation can use `DBuffer.get_tensor()` and build the required flattened/batched communication layout directly.
This is related to whether the all-to-all-v path should be batched, or whether an unbatched/pipelined approach performs well enough. If the unbatched or pipelined approach works as well or better, it may be preferable to implement the redistribution directly in DBuffer.
## Related issues and PRs
- #3174
- #4486
- #5078
- #5307
- #4691
- #5179
- #5394
- #2163
- #3176
- #4091
Contributor guide
Assessment
This issue has not been assessed yet.