NVIDIA / NVIDIA/Megatron-LM

Pipeline the layer-sharded Muon all_to_all to hide exchange latency (and cut peak memory)

Open
#6,719 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.