[Feature request] incorporation of flash-attention 2 about dnabert_2 HOT 3 OPEN

I've encountered a similar issue. I believe the recently released flash-attention version 2 significantly improves the execution time for DNA-BERT2. By the way, I greatly appreciate your work; it has been very helpful to me.

from dnabert_2.

Hey,

Thank you very much for your interest in our work and for this great suggestion!

However, I can't do it right away since I have been quite busy recently. But you are very welcome to submit PRs if you find a good way to do it. I will work on it after a few weeks.

from dnabert_2.

I think this might be harder than one would expect, since the current triton implementation (only one that allows for AliBi) is pretty well known to be broken, unless you have a specific dev version of triton. There's a lot of discussion on the flash repo pertaining to this. If someone really wants to give it a go, this might be promising to use instead.
But currently, implementation is a bit more involved than just switching out a few lines on DNABERT_2.

from dnabert_2.