Layer-sharded Muon: one Newton-Schulz home per weight over (GTP × TP)
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
### Problem
Both existing GTP+Muon execution modes leave large costs on the table at scale:
**duplicated** mode all-gathers every weight and runs the full-matrix Newton-Schulz
redundantly on every rank of the (GTP × TP) domain; **distributed** mode sheds the
redundancy but pays O(5N) Gram all-reduces per step and its Gram-space polynomial
does not shard further.
### Proposal
A third mode, opt-in via `--use-layer-sharding-muon`: since a transformer has far
more 2D weights than the domain has ranks, shard the *work* at matrix granularity
instead of the math. Each weight gets one NS home (LPT bin-packing by NS cost);
all_to_all exchanges assemble the full matrix there, the identical full-matrix NS
runs with zero communication and zero redundancy, and reverse exchanges scatter the
result back. **Bit-identical to duplicated mode in the default configuration**
(verified `torch.equal` at the optimizer level, and 20/20 iterations bit-identical
lm-loss in a deterministic real-model A/B).
Measured (32 GPUs GB300, TP2 × gtp_remat8, EP8 × egtp2, production matrix shapes):
optimizer step 632 → 157 ms vs duplicated (356 ms for distributed); NCCL calls per
step 76/276 → 4.
This builds directly on the LayerWise distributed optimizer and complements #6379
(NS-cost shard balancing) and #6381 (`--muon-use-syrk`). Implementation is complete
and tested on an internal fork (nv-mistralai-megatron#266); PR to follow. Design doc
attached/linked.
Contributor guide
Research direction
The issue names no repository files or tests. Start by reading the existing GTP+Muon duplicated and distributed execution modes, then review the attached or linked design document and the internal implementation reference. Done means the opt-in layer-sharding mode achieves the stated communication and optimizer behavior while preserving bit-identical results in the reported checks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100