[Prototype] MatrixFSDP-style TensorAtomic placement for MFSDP v2 Muon
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Goal
Prototype the placement strategy described in [MatrixFSDP](https://arxiv.org/abs/2607.05895) in MFSDP v2 using a DBuffer `TensorAtomic` placement.
This is a subissue of #5533 and explores candidate approach 3: keep each Muon matrix tensor-atomic on one data-parallel rank so Newton–Schulz can run locally without optimizer-step communication.
## Background
MatrixFSDP changes the ZeRO-3 shard placement rather than distributing the matrix optimizer itself:
- Each 2D matrix is owned in full by one data-parallel rank.
- Other ranks hold empty resident shards for that matrix.
- Non-matrix parameters are packed into tail-owner roles and continue to use Adam-style optimization.
- Forward and backward materialize and reshard parameters.
- The complete matrix gradient lands on its owner, allowing Muon to run locally without an optimizer-step matrix collective.
Prototype this idea in MFSDP v2 and evaluate its correctness, performance, memory behavior, and fit with the DBuffer placement model. The implementation approach is intentionally left open so the prototype can explore the most suitable planning, communication, and integration strategy.
## References
- MatrixFSDP: https://arxiv.org/abs/2607.05895
- Parent issue: #5533
- Initial DBuffer implementation: #4835
Contributor guide
Research direction
Start with parent issue #5533, the MatrixFSDP paper, and the initial DBuffer implementation in #4835. Prototype TensorAtomic placement for MFSDP v2, then evaluate correctness, performance, memory behavior, and fit with the DBuffer placement model; the payload does not name specific files or tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100