Repository navigation

Add preliminary Muon+M-FSDP support - #4486

Draft

janEbert wants to merge 32 commits into

NVIDIA:mainfrom

janEbert:muon-m-fsdp

Draft

Add preliminary Muon+M-FSDP support#4486
janEbert wants to merge 32 commits into
NVIDIA:mainfrom
janEbert:muon-m-fsdp

Commits on May 21, 2026
- Make Muon work with M-FSDP
  janEbert
  committed
  2c2c68b
- Make Muon+M-FSDP work up to ZeRO-3
  janEbert
  committed
  13117e2
- Add Muon+M-FSDP tests
  janEbert
  committed
  d5cb02b
- Support further features
  janEbert
  committed
  c0aeafc
- Apply Cory suggestions
  janEbert
  committed
  04d9bd4
- LayerwiseDistOpt not necessary for Megatron-FSDP.
  
  cspades
  authored and
  janEbert
  committed
  a096591
- Remove extra AG calls during Muon step.
  
  cspades
  authored and
  janEbert
  committed
  4f87d37
- Check for FSDP instances via type
  janEbert
  committed
  fab856e
- Add Muon+M-FSDP numerics test
  janEbert
  committed
  a48caf6
- Do not fill with redundant zeros
  janEbert
  committed
  540452a
- Cache boundary detection
  janEbert
  committed
  45e4cae
- Add batched all-gathers
  janEbert
  committed
  c488e77
- Fix boundary detection
  janEbert
  committed
  e9087bb
- Fix throughput regression
  janEbert
  committed
  c9cbe63
- Test whether empty shards skip Newton-Schulz
  janEbert
  committed
  034b4a7
- Improve docstring clarity
  janEbert
  committed
  e87008f
- Try to avoid DTensor copies
  janEbert
  committed
  b82c4d5
- Add error checks and tests for H(F)SDP handling
  janEbert
  committed
  e5d4348
- Add boundary detection caching for H(F)SDP
  janEbert
  committed
  f13ec75
- Fix H(F)SDP boundary detection
  janEbert
  committed
  6366f78
- Fix M-FSDP DTensor grad statistics
  janEbert
  committed
  2626020
- Fix hybrid M-FSDP delayed outer grad reduction
  janEbert
  committed
  829a30b
- Do not redundantly install weights with >1 optims
  janEbert
  committed
  75acf1a
- Add Muon M-FSDP gather fast paths
  janEbert
  committed
  868bb14
- Test Muon M-FSDP gather fast paths
  janEbert
  committed
  9475c40
- Expose Muon M-FSDP gather options via config
  janEbert
  committed
  b7a1420
- Add Muon+M-FSDP comm+comp overlap
  janEbert
  committed
  4b110fd
- Forward weight-install method to M-FSDP adapter
  janEbert
  committed
  eec7a28
- Improve Muon M-FSDP overlap scheduling
  janEbert
  committed
  399f252
- Fix Muon HFSDP boundary index agreement
  janEbert
  committed
  3c5a382
- Optimize Muon M-FSDP batched gather reconstruction
  janEbert
  committed
  6803ac5
- Add flat gather support
  janEbert
  committed
  4edce9c

Commits

Commits on May 21, 2026