[DO NOT MERGE] Draft mfsdp_refactor to dev branch#5665

Draft

shjwudp wants to merge 36 commits into

shjwudp:mfsdp_refactor_dev_upstream

Commits on May 27, 2026

feat(mfsdp-v2): mixed-precision infrastructure and MXFP8 support (by Tongliu)

Autumn1998
authored and
shjwudp
committed
feat(mfsdp-v2): checkpoint infrastructure, NVFP4 support, and FSDP v2 core (by Jack Chang)
shjwudp
committed

Commits on May 29, 2026

Commits on May 31, 2026

feat(mfsdp-v2): ZeRO-1/2/3 sharding strategies, MixedPrecisionPolicy rename

Autumn1998
authored and
shjwudp
committed

Commits on Jun 1, 2026

feat(mfsdp-v2): sync module states after checkpoint load, harden uneven dtensor split, improve checkpoint inspector
shjwudp
committed

Commits on Jun 4, 2026

Commits on Jun 5, 2026

Commits on Jun 8, 2026

Commits on Jun 11, 2026

Commits on Jun 12, 2026

Commits on Jun 23, 2026

Integrate M-FSDP v2 with 1F1B pipeline overlap; add lazy grad buffers, CPU offload, and design docs
shjwudp
authored

Commits on Jun 30, 2026

CUDA graph capture via vendored te-graph-runtime with `capture_time_hooks` for MFSDP v2 (#21)
shjwudp
authored

Commits on Jul 2, 2026