[Dev] Support LayerWiseDistributedOptimizer with torch_dist checkpoint format - #1928

Merged

skyw merged 23 commits into

BoxiangW:boxiangw/dist_ckpt_with_dist_muon

Nov 4, 2025

Commits on Nov 3, 2025

Init

BoxiangW
authored and
FDecaYed
committed
87e173d
fixed-DP

BoxiangW
authored and
FDecaYed
committed
b852d50
Dist Adam works with dist_ckpt with fixed DP size

BoxiangW
authored and
FDecaYed
committed
d8f526b
Lint

BoxiangW
authored and
FDecaYed
committed
d3d423d
Address comments

BoxiangW
authored and
FDecaYed
committed
b734507
Add docstring

BoxiangW
authored and
FDecaYed
committed
52e3960
Fix bug

BoxiangW
authored and
FDecaYed
committed
34e9d0b
Add changing DP support

BoxiangW
authored and
FDecaYed
committed
a9e4c0c
Fix PR

BoxiangW
authored and
FDecaYed
committed
46c5bb7
Fix lint

BoxiangW
authored and
FDecaYed
committed
719a8e4
Fix lint

BoxiangW
authored and
FDecaYed
committed
fd76b42
Fix import

BoxiangW
authored and
FDecaYed
committed
ffb9a70
Fix test

BoxiangW
authored and
FDecaYed
committed
111ce35
Fix lint and copyright

BoxiangW
authored and
FDecaYed
committed
3655f8e
lint

BoxiangW
authored and
FDecaYed
committed
3df09ed
Update muon.py
FDecaYed
authored
b987681
Update optimizer.py
FDecaYed
authored
6c1aed9
Update muon.py
FDecaYed
authored
1018a43
Merge branch 'dev' into boxiangw/dist_ckpt_with_dist_muon
BoxiangW
authored
b6c156d

Commits on Nov 4, 2025

Fix dependency issue
BoxiangW
committed
59d3a8e
Revert print
BoxiangW
committed
12a6250
Add ckpt guard
BoxiangW
committed
eabae18
Fix bug
BoxiangW
committed
31e116f