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

Merged

rapatel merged 7 commits into

NVIDIA:mainfrom

rapatel:ripatel/fsdp_1f1b_pre_bwd_fix_main

Jun 30, 2026

Commits on Jun 30, 2026

Fix TEGroupedMLP pre-backward unshard in fine-grained FSDP hooks for 1F1B EP overlap.
rapatel
committed
64f8231
Symmetric fine-grained FSDP unshard for TEGroupedMLP via fine_grained_recurse_module_types
rapatel
committed
2819a17
Rename submodule unshard helper arg to pass_direction Literal[forward, backward].
rapatel
committed
0ad5777
Lazy-import TEGroupedMLP to fix fsdp_dtensor checkpoint load.
rapatel
committed
50672eb
Fix fine-grained FSDP unshard logic and add SharedExpertMLP to recurse allowlist.
rapatel
committed
c0422cb
Applying formatting.
rapatel
committed
00ddbc6
Fix line-too-long linting error.
rapatel
committed
cc78b37