Add experimental Megatron-FSDP fully_shard implementation - #4976

Closed

wujingyue wants to merge 13 commits into

NVIDIA:pull-request/4835from

wujingyue:fsdp/minimal

Commits on Jun 16, 2026

WIP: add experimental minimal FSDP path
wujingyue
committed
8b43c26
Refine minimal experimental FSDP path
wujingyue
committed
0646346
Refine experimental FSDP buffer APIs
wujingyue
committed
f5521fc
Refine experimental FSDP gradient contract
wujingyue
committed
dd38bc4
Require matching FSDP main grad dtype
wujingyue
committed
948caff
Reuse FSDP model weights for matching main weights
wujingyue
committed
5217e08
Sync FSDP model weights before unshard
wujingyue
committed
b2ac8c0
Preserve autograd grad dtype in minimal FSDP
wujingyue
committed
6ad5fbc
Compare minimal FSDP loss curve with baseline
wujingyue
committed
de302ea
Adapt minimal FSDP to split DBuffer API
wujingyue
committed
878b9a3
Split minimal FSDP runtime modules
wujingyue
committed
345e96e
Split minimal FSDP module mixin
wujingyue
committed
2521db9
Remove experimental FSDP meta parameter support
wujingyue
committed
19a814e