Add FSDP support to TiledMLP by preventing premature resharding during the tiled backward recompute loop.#1128

Draft

alektebel wants to merge 11 commits into

linkedin:mainfrom

alektebel:fsdp-tiledmlp-support

Commits on Mar 5, 2026

fix on tiled_mlp
alektebel
committed

Commits on Mar 6, 2026

Commits on Mar 15, 2026

test(tiled_mlp): add FSDP compatibility test for LigerTiledSwiGLUMLP
alektebel
committed

Commits on Mar 19, 2026

Fix TiledMLP FSDP compatibility and memory optimizations\n\n- Fix shard offset calculation and FSDP gradient comparison\n- Optimize memory with lazy allocation and in-place accumulation\n- Relax fl…
alektebel
committed

Commits on Mar 20, 2026

Fix TiledMLP FSDP compatibility and memory optimizations\n\n- Fix shard offset calculation and FSDP gradient comparison\n- Optimize memory with lazy allocation and in-place accumulation\n- Relax fl…
alektebel
committed

Commits on Mar 23, 2026

Apply suggestion from @Tcc0403

alektebel
and
Tcc0403
authored

Commits on Mar 28, 2026

docs(test_tiled_mlp): add Axolotl integration notes and improve docstrings

alektebel
and
claude
committed

Commits on Mar 31, 2026