On Fri, 8 Sept 2023 at 18:44, Ben Thompson ***@***.***> wrote: It would be nice to have a flag to enable flash attention in models where that would make sense. This is helpful for performance and memory usage in larger models. In my case working with Pythia 12B, I get ~50% better performance and ~4x larger batch sizes when using flash attention. I also find numerical stability in float16 to be better using flash attention, probably because the model was trained using flash attention. The downside of using flash attention in TransformerLens is that the we would not have access to intermediate quantities in the attention calculation like the attention matrix itself. This is why I would suggest having a default-off flag so that users can choose whether they need those intermediate values to be available. In addition, when only a small subset of attention intermediates are needed, it's much faster to just cache the input to the attention layer and then recompute those intermediates when needed. Thanks! — Reply to this email directly, view it on GitHub <#378>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ASRPNKOCL65GY7N52PH2PQDXZNKRXANCNFSM6AAAAAA4QWO7HE> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>