Flax model much slower than Pytorch

[yorickvanzweeden]

Hi,

I am running the same QuestionAnswering model using Pytorch and Flax, but I am seeing quite a bad performance for Flax.

Model	Time
Pytorch (GPU)	30 ms
Flax (GPU)	4 sec
Flax (TPU)	15 sec

I have warmed the jit function, and I think I have put the model and data on the device in both cases. What am I doing wrong here?

Notebook: https://colab.research.google.com/drive/18ouY4S9rYtlEC1dht_w8WU4fFKnP9yhP#scrollTo=FLfV3ymlwRJC

Environment info

• transformers version: 4.16.2
• Platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.12
• PyTorch version (GPU?): 1.10.0+cu111 (True)
• Tensorflow version (GPU?): 2.7.0 (True)
• Flax version (CPU?/GPU?/TPU?): 0.4.0 (gpu)
• Jax version: 0.2.25
• JaxLib version: 0.1.71
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Who can help

@patrickvonplaten

Information

Model I am using (Bert, XLNet ...): BERT (FlaxBertForQuestionAnswering, BertForQuestionAnswering)

The problem arises when using:

• the official example scripts: (give details below)
• [ X] my own modified scripts: (give details below)

The tasks I am working on is:

• an official GLUE/SQUaD task: (give the name)
• [ X] my own task or dataset: (give details below)

To reproduce

Steps to reproduce the behavior:

1. Run this notebook on Google Colab (GPU runtime)

Expected behavior

About the same performance for Pytorch as for Flax.

[patrickvonplaten]

cc @patil-suraj think we need to improve our docs here.

@yorickvanzweeden the very first time you run run_model_jit(jax_inputs_batched), the model is compiled which takes time. However the second time you run it, it should be pretty fast. Could you maybe run the code say 20 times for both PyTorch and Flax and then check if the Flax is still as slow?

[yorickvanzweeden]

@patrickvonplaten Thank you for your swift response!

I noticed caching on the inputs, so I have two inputs. When I do something like this, I still get about 4 seconds versus 40 ms.

# PyTorch
%timeit -r 20 -n 1 model(**inputs_gpu) # <-- 20x repeat
> 1 loop, best of 20: 20.5 ms per loop
%timeit -r 1 -n 1 model(**inputs_gpu2)
> 1 loop, best of 1: 39.1 ms per loop

# Flax
%timeit -r 20 -n 1 run_model_jit(jax_inputs_batched) # <-- 20x repeat
> 1 loop, best of 20: 44.7 µs per loop
%timeit -r 1 -n 1 run_model_jit(jax_inputs_batched2)
> 1 loop, best of 1: 3.83 s per loop 

Also, please note that I am using FlaxBertForQuestionAnsweringModule, and not FlaxBertForQuestionAnswering directly. I would like to use the module directly.

But also when using FlaxBertForQuestionAnswering.from_pretrained, I get the second input in seconds (5.75 sec).

I updated the notebook for the 20x repeat, and using the model rather than the module.

[patrickvonplaten]

Hey @yorickvanzweeden,

Cool, so it seems that once the flax inputs are cached, the forward pass is very fast no?

Now the problem is that for the second input jax_inputs_batched2 the code is recompiled again. The most important thing when dealing with jax.jit(...) is to remember the following if you jax.jit(...) to create a "jitted" function, e.g. run_model_jit(...) and then run the function, e.g. run_model_jit(jax_inputs_batched2), the following happens under the hood:

1. JAX checks if it can find a compiled binary that corresponds to the function you jitted run_model_jit(...) and the input shape you forwarded.
2. If there is no match, then JAX compiles the function + input shape and we repeat step 1. If there is a match then the compiled binary is executed.

In your case for run_model_jit(jax_inputs_batched2) JAX does find a compiled binary, but the compiled binary corresponds to run_model_jit + jax_inputs_batched.shape and while you again use run_model_jit, jax_inputs_batched2.shape != jax_inputs_batched.shape which means that the function is compiled again. To remedy this you should make sure that all your inputs have the exact same shape, e.g. by padding the input. This way you won't trigger a new recompilation.

[yorickvanzweeden]

Hi @patrickvonplaten,

Thank you for the explanation. The different shapes indeed resulted in recompilation. It is now about 15% faster than PyTorch :)

Thanks for the help!