FlaxBertModel examples (and fast attention)
Author: StefanoSalvatori
Pull request: no

Are there any examples that show how to use the FlaxBertModel?

Would it be possible to replace the current SelfAttention module with the one proposed here https://github.com/google-research/google-research/tree/master/performer/fast_self_attention?

Comments:
It should work - we'll run some experiments soon :-) https://github.com/huggingface/transformers/pull/8358 @TevenLeScao

Ok, great!

For the moment i'm running some experiments myself with FlaxBertModel but i'm getting unexpected behaviors: it seems that the flax implementation is slower than the torch one; I tried to run this simple code in google colab with GPU environment 

```python
!pip install --upgrade pip
!pip install --upgrade jax jaxlib==0.1.57+cuda101 -f https://storage.googleapis.com/jax-releases/jax_releases.html
!pip install flax
!pip install transformers

import flax
import jax
from transformers import BertModel
from transformers import FlaxBertModel
from transformers import AutoTokenizer
import time

baseModel ='bert-base-uncased'
torchBert = BertModel.from_pretrained(baseModel)
flaxBert = FlaxBertModel.from_pretrained(baseModel)

tokenizer = AutoTokenizer.from_pretrained(baseModel)
torchInput = tokenizer.encode("Random sentence", truncation=True, padding=True, return_tensors='pt')
flaxInput = tokenizer.encode("Random sentence", truncation=True, padding=True, return_tensors='jax')

start_time = time.time()
for _ in range(10):
  torchBert(torchInput)
print("torch - ", time.time()-start_time)

start_time = time.time()
for _ in range(10):
  flaxBert(flaxInput)
print("flax - ", time.time()-start_time)
```

and i'm getting the following output

```
torch -  0.6615538597106934
flax -  5.129613161087036
```

What am i missing?

You should probably use `jax.jit` to speed it up

Indeed, as @patrickvonplaten says, using `jax.jit` on `flaxBert` will speed things up considerably. This will first shape-compile the function using XLA, and every time you call the function again (provided the shapes are the same), it will run the compiled version directly. I've demonstrated in this Colab: https://colab.research.google.com/drive/1davNsnV34KDZOyJ9i8zZfvxAVjBJC4dp?usp=sharing

Make sure you set the accelerator to GPU/TPU! (Runtime -> Change runtime type)

Here's a summary:

```
>>> %timeit torchBert(torchInput)
1 loop, best of 5: 75.3 ms per loop
>>> %timeit flaxBert(flaxInput)
1 loop, best of 5: 1.41 s per loop
>>> %timeit jitted_flax(flaxInput)
100 loops, best of 5: 11.2 ms per loop
```

Note that this excluded the compilation time for the first time we called `jitted_flax`. Including this will increase the overall execution time, but since it has to be done only once this is negligible as you execute this function more often.

To learn more about JAX's jit, this quickstart is quite useful: https://jax.readthedocs.io/en/latest/notebooks/quickstart.html

Thank you for you comment @marcvanzee. Indeed after @patrickvonplaten's reply I checked Flax and Jax documentation more carefully confirming that JIT compilation could solve performance issues. It still surprises me though that pytorch is quite fast even without JIT compilation while the same is not true for Flax.  Frankly i didn't even know that JIT existed in pytorch so i'd be curious too to se how it compares to Flax.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.