NVIDIA / NVIDIA/Megatron-LM

Enable Muon optimizer training with MFSDP v2

Open
#5,533 7 comments 0 reactions 1 assignee Claimed by @janEbert View on GitHub
MFSDPv2 PoR
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.