Investigate and implement Flash Attention #5

Activity

1. 0x000011b

   added this to AI/ML Model Backlogon Feb 19, 2023
2. 0x000011b

   converted this from a draft issue on Feb 19, 2023
3. 0x000011b

   added

   enhancementNew feature or requestNew feature or request

   on Feb 19, 2023
4. 0x000011b

   moved this from 📋 Backlog to 🏗 In progress in AI/ML Model Backlogon Mar 3, 2023
5. 0x000011b

   assigned

   TearGosling

   and

   0x000011b

   on Mar 3, 2023
6. 0x000011b commented on Mar 6, 2023

   

   0x000011b

   on Mar 6, 2023

   Author

   More actions

   Me and @TearGosling played around with optimizing NeoX by using components from xFormers after I profiled the training code. Results:

   • Usage of the MLP component there causes performance drop and lots of warnings, so we're ignoring this one
   • Usage of flash attention results in:
     • 👍 A decent speedup (~17% IIRC)
     • 👎 A significant increase in training and evaluation loss
     • 😐 No noticeable VRAM savings
   • Usage of the rotary embedding implementation results in:
     • 👍 A decent speedup (when used together with flash attention, gets us over a 20% throughput increase IIRC)
     • 😐 A very minor increase in training loss

   Next steps:

   • Figure out whether the higher loss values with flash attention actually translate to a worse model, or whether they can be accounted for by some other factor we're unaware of.
   • No noticeable VRAM savings with flash attention doesn't make a lot of sense. Investigate whether this is due to e.g. memory fragmentation (since we're not properly pre-allocating memory for dynamic tensor sizes).
7. 0x000011b commented on May 2, 2023

   

   0x000011b

   on May 2, 2023

   Author

   More actions

   Flash attention via xFormers landed on #8. "Lack" of VRAM savings was actually just fragmentation messing us up.
8. 0x000011b

   closed this as completedon May 2, 2023
9. github-project-automation

   moved this from 🏗 In progress to ✅ Done in AI/ML Model Backlogon May 2, 2023