Add experimental Megatron-FSDP fully_shard implementation - #5387

Merged

wujingyue merged 24 commits into

NVIDIA:mainfrom

wujingyue:fsdp/minimal

Jun 24, 2026

Commits on Jun 17, 2026

WIP: add experimental minimal FSDP path
wujingyue
committed
844c625
Refine minimal experimental FSDP path
wujingyue
committed
832af46
Refine experimental FSDP buffer APIs
wujingyue
committed
afbe9fa
Refine experimental FSDP gradient contract
wujingyue
committed
63987db
Require matching FSDP main grad dtype
wujingyue
committed
14e807b
Reuse FSDP model weights for matching main weights
wujingyue
committed
680f17b
Sync FSDP model weights before unshard
wujingyue
committed
5f06688
Preserve autograd grad dtype in minimal FSDP
wujingyue
committed
09909a2
Compare minimal FSDP loss curve with baseline
wujingyue
committed
c4d13a8
Adapt minimal FSDP to split DBuffer API
wujingyue
committed
f1c2e57
Split minimal FSDP runtime modules
wujingyue
committed
195b4ca
Split minimal FSDP module mixin
wujingyue
committed
13055d8
Remove experimental FSDP meta parameter support
wujingyue
committed
8759750
Clarify FSDP version counter preservation
wujingyue
committed
23c47a8
Rename experimental FSDP runtime types
wujingyue
committed
08e84c6
Document main_grad allocation lifetime
wujingyue
committed
ea39f57
Document post-backward reshard storage choice
wujingyue
committed
e80d570

Commits on Jun 21, 2026

Remove experimental FSDP distributed pytest markers
wujingyue
committed
99b1261
Remove DBuffer distributed pytest markers
wujingyue
committed
c12e418
Remove experimental FSDP CUDA skip guards
wujingyue
committed
f959479
Document experimental fully_shard mixin attachment
wujingyue
committed
adb4e12
Document fully_shard mixin behavior
wujingyue
committed
06286bc

Commits on Jun 22, 2026

Cover FSDP loss parity with microbatches
wujingyue
committed
0dfb83b

Commits on Jun 24, 2026

Merge branch 'main' into fsdp/minimal
wujingyue
authored
e3216c4