Add fully_shard_optimizer for mixed-precision FSDP - #5411

Merged

wujingyue merged 12 commits into

NVIDIA:mainfrom

wujingyue:fsdp/optimizer

Jul 21, 2026

Commits on Jul 9, 2026

Implement experimental fully_shard_optimizer hook adapter
wujingyue
committed
5b8177f
Refine experimental fully_shard_optimizer adapter
wujingyue
committed
a74ec3f
Refine experimental FSDP optimizer behavior tests
wujingyue
committed
f1a5934
Document fully_shard_optimizer adapter alternatives
wujingyue
committed
3dce5f4
Fix fully_shard_optimizer lint formatting
wujingyue
committed
a211993
Document optimizer precision alternatives
wujingyue
committed
98615f2
Move optimizer failure test
wujingyue
committed
e283ece
Rename Adam optimizer failure test
wujingyue
committed
8fc3947
Add FusedAdam MFSDP optimizer test
wujingyue
committed
239fdde
Test MFSDP optimizer grad dtype mismatch
wujingyue
committed
71a0866
Remove custom MFSDP optimizer test
wujingyue
committed
dfbce46

Commits on Jul 20, 2026

Relax MFSDP optimizer dtype assertion
wujingyue
committed
542c84a