NVIDIA / NVIDIA/Megatron-LM

[Prototype] MatrixFSDP-style TensorAtomic placement for MFSDP v2 Muon

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.