Pipeline the layer-sharded Muon all_to_all to hide exchange latency (and cut peak memory)
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
Layer-sharded Muon (LSH) currently performs **one** all_to_all covering every matrix a
DP shard owns, then runs Newton-Schulz on all matrices homed to this rank, then **one**
reverse all_to_all. This serializes communication and compute, and holds every
intermediate buffer live simultaneously.
Splitting the exchange into `C` chunks and overlapping chunk *k*'s Newton-Schulz with
chunk *k+1*'s all_to_all would recover the ~17% of LSH step time currently spent stalled
in NCCL, and gives a memory knob that would otherwise require falling back to
`duplicated`.
This affects the expert (EGTP) axis, where a DP shard owns 23-24 matrices. Dense (GTP)
profiles mostly own a single matrix, so there is nothing to chunk there.
## Current behaviour
`layer_sharded_all_to_all_fwd` / `_bwd` (`megatron/core/optimizer/layer_sharded_a2a.py`)
take the full `momentum_list` and build one flat send buffer over all of it:
```python
send_parts = []
for g_prime in range(gtp_size):
if params_for_rank[g_prime]:
chunk = torch.cat([m.contiguous().flatten() for _, m in params_for_rank[g_prime]])
send_parts.append(chunk)
send_buf = torch.cat(send_parts) if send_parts else ...
torch.distributed.all_to_all_single(recv_buf, send_buf, ...)
```
## Linked PR
Layer sharded Muon (LSH): #6683
Contributor guide
Research direction
Read megatron/core/optimizer/layer_sharded_a2a.py, starting with layer_sharded_all_to_all_fwd and _bwd, and review the linked Layer sharded Muon work in #6683. Compare the expert EGTP path with the dense GTP case, then verify that chunked communication overlaps Newton-Schulz while preserving forward and backward behavior and reducing peak memory.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 48/100