Draft [2/N] Megatron FSDP: Introduce Megatron-FSDP2 with per-module fully_shard() API - #4849

Draft

shjwudp wants to merge 68 commits into

NVIDIA:mainfrom

shjwudp:mfsdp_refactor_main_stage2

Commits on Apr 21, 2026

Add param group

Autumn1998
authored and
shjwudp
committed
9d450b4
1. Add fsdp2/m-fsdp toy example for API align
shjwudp
committed
c76b84e
init fully_shard v2 FSDPModule and fsdp hooks
shjwudp
committed
2fcea79
use post_backward do post_accumulate_grad job
shjwudp
committed
463ae77
fix import name
shjwudp
committed
480c53b
fix set of bugs
shjwudp
committed
79fb44e
fix FSDPModule param name mapping
shjwudp
committed
cafb667
Add get_state_dict support
shjwudp
committed
e4ebb50
init fully_shard api mcore_fsdp_adapter support
shjwudp
committed
6b815e7
fix bug in mcore_fsdp_adapter + fully_shard_v2
shjwudp
committed
54eedd7
add mcore fsdp + fully_shard_v2 ut
shjwudp
committed
3ba54b9
Add _materialize_meta_module function
shjwudp
committed
2571a40

Commits on Apr 22, 2026

fix chunk_size_factor & add mock param_and_grad_buffer
shjwudp
committed
8e85301

Commits on Apr 23, 2026

Squash merge mfsdp_refactor_main_debug into mfsdp_refactor_main

shjwudp
and
Copilot
committed
66ede6e

Commits on Apr 24, 2026

Commits on Apr 26, 2026

feat: add unshard prefetch and reduce-scatter overlap for fully_shard_rewrite
shjwudp
committed
06682e2

Commits on Apr 28, 2026

fix: fix overlap bugs in fully_shard_rewrite and add design doc
shjwudp
committed
950d304

Commits on May 6, 2026

fix(mfsdp): convergence fix, memory optimization, and polish for overlap path
shjwudp
committed
4844fe8

Commits on May 10, 2026

feat(mfsdp): llama3-8b convergence fixes, debug tooling, and docs
shjwudp
committed
58aadfb

Commits on May 12, 2026

Add FSDP v2 MXFP8 mixed precision support
Autumn1998
committed
8a2515b

Commits on May 13, 2026

feat(mfsdp): activation recompute support, bucket allocator, and convergence fixes
shjwudp
committed
9d4ac5c
rename fully_shard_rewrite
shjwudp
committed
7916ae0

Commits on May 14, 2026

Commits on May 15, 2026

Commits on May 18, 2026

Commits on May 19, 2026

Commits on May 20, 2026

fix(mfsdp): handle fp8 fine-grained recompute hooks
Autumn1998
committed
8ee871c
docs(mfsdp): clarify fine-grained hook owner
Autumn1998
committed
39e23c7
Merge pull request #8 from Autumn1998/tongliu/fsdp-v2-mixed-precision
shjwudp
authored
a5d1114
feat: Megatron FSDP v2 checkpoint save/load and online format conversion
shjwudp
committed
66cac17
fix(mfsdp): use fp32 main grads and inplace NCCL reduce-scatter
shjwudp
committed
3e22618
fix dist_grad dtype check
shjwudp
committed
ab550d3

Commits on May 23, 2026

refactor(mfsdp): overhaul checkpoint.py for load_torch_dist_into_fsdp_v2
shjwudp
committed
2e0ff19

Commits on May 25, 2026

Commits on May 26, 2026

Commits on May 27, 2026

Commits on May 28, 2026

Refactor FSDP v2 unshard and reduce_grad
Autumn1998
committed
00507fd
Refine FSDP v2 replicated buffer refresh
Autumn1998
committed
123d7f9

Commits on May 29, 2026