Fix TEGroupedMLP pre-backward unshard in fine-grained FSDP hooks for … - #4989

Open

rapatel wants to merge 7 commits into

rapatel:ripatel/fsdp_1f1b_pre_bwd_fix_dev

Commits on Jun 30, 2026

Fix TEGroupedMLP pre-backward unshard in fine-grained FSDP hooks for 1F1B EP overlap.
rapatel
committed
c1fdf20
Symmetric fine-grained FSDP unshard for TEGroupedMLP via fine_grained_recurse_module_types
rapatel
committed
3294394
Rename submodule unshard helper arg to pass_direction Literal[forward, backward].
rapatel
committed
b1278dd
Lazy-import TEGroupedMLP to fix fsdp_dtensor checkpoint load.
rapatel
committed
e1ecaa3
Fix fine-grained FSDP unshard logic and add SharedExpertMLP to recurse allowlist.
rapatel
committed
141ae98
Applying formatting.
rapatel
committed
9bd4b68
Fix line-too-long linting error.
rapatel
committed
dc9caf9