Memory Efficient Attention integration about big_vision

[akolesnikoff]

I have not ran any deep investigations, but currently I expect that it is not worth it:

1. Speed-wise XLA is already supposed to do the optimizations that flash-attention does. However, if you try it and observe that XLA is doing worse job that manual implementation of FLASH attention, we may have a look into it. It can also be different for TPUs vs GPUs.
2. Memory-wise, a single attention matrix for 4k tokens weighs 64 MB. And by wrapping the whole transformer block inside flax.linen.remat one can guarantee that no more than 1 (or 2) attention matrices are materialized at any moment in time. So memory-wise it is only helpful when scaling beyond 4k tokens, which is not a typical vision scenario.

[maximzubkov]

Thank you for the detailed explanation! I'll give it a try in my experiments