Epic: Integration of FSDP2 #325

Activity

1. le1nux

   added

   enhancementNew feature or requestNew feature or request

   on Mar 25, 2025
2. le1nux

   added this to the v0.4.0 milestone on Mar 25, 2025
3. le1nux

   assigned

   le1nux

   and

   flxst

   on Mar 25, 2025
4. le1nux

   added the

   Feature

   issue type on Mar 25, 2025
5. le1nux

   added sub-issues on Mar 27, 2025

   • Add feature to calculate the number of parameters based on FSDP2 device mesh #326
   • get_theoretical_flops_per_token is not model agnostic, even though the function declaration indicates this #327
   • Addressing discussion points in #319 #328
   • Add testing for FSDP2 #329
   • Fix failing multi-gpu tests #330
   • Failing end2end test test_initialization.py must be fixed #331
   • MultiGPU / End2end test test_get_total_number_of_trainable_parameters in test_utils.py fails / succeeds non-deterministically #332
   • LR Scheduler instantiation fails when Optimizer loaded via DCP #315
   • Apply Torch-Compile #133
   • Add torch compile uni test #333
   • Add debug setup to torch.copile component #334
6. le1nux

   added sub-issues on Apr 8, 2025

   • Run manual check for FSDP1 and FSDP2 equivalence #339
   • Add test for experiment_id syncing #340
7. le1nux

   assigned

   mali-git

   and

   fromm-m

   on Apr 8, 2025
8. le1nux

   added sub-issues on Apr 18, 2025

   • Weight tying breaks with meta device initialisation #357
   • [ABL] Checkpoint conversion fails with the FSDP2 integration #358
9. le1nux

   changed the title [-]Integration of FSDP2[/-] [+]Epic: ntegration of FSDP2[/+] on Apr 27, 2025
10. le1nux

    added a sub-issue on Apr 27, 2025

    • FSDP2: Full Activation Checkpointing does not work with torch.compile #361
11. le1nux

    changed the title [-]Epic: ntegration of FSDP2[/-] [+]Epic: Integration of FSDP2[/+] on Apr 27, 2025
12. le1nux

    added a sub-issue on May 1, 2025

    • MFU Calculation has potential precision issues #364
13. le1nux

    closed this as completedon Oct 8, 2025
14. le1nux

    removed sub-issues on Oct 10, 2025

    • Weight tying breaks with meta device initialisation #357
    • [ABL] Checkpoint conversion fails with the FSDP2 integration #358
    • FSDP2: Full Activation Checkpointing does not work with torch.compile #361