FlaxCLIPModel is 10x faster than CLIPModel during inference, and 100x slower getting gradients?
Author: josephrocca
Pull request: no

## Environment info
- `transformers` version: 4.10.0
- Platform: Linux-5.11.0-7620-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.9.0+cu102 (True)
- Tensorflow version (GPU?): 2.5.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (gpu)
- Jax version: 0.2.17
- JaxLib version: 0.1.68
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
@patil-suraj

## Information

Model I am using: FlaxCLIPModel and CLIPModel

The problem arises when using the official example scripts (though the gradient code is not from the examples).

## To reproduce

I may be doing something terribly, terribly wrong here, and I initially thought it must have something to do with me misunderstanding `jit` compilation/caching, or something like that, but having put together some very minimal examples, and tested on my own machine, and on a colab notebook, I'm still seeing a huge disparity in performance between Pytorch and JAX in both inference and in getting gradients.

Here's a notebook which minimally replicates the performance disparity:
 * https://colab.research.google.com/gist/josephrocca/3b4aca3ebf66e8b46a035eb0f3711150 (make sure to select GPU runtime)
 * https://gist.github.com/josephrocca/3b4aca3ebf66e8b46a035eb0f3711150

And, for quick reference, here are the relevant code snippets from that notebook for **inference** (note that in the notebook I run `jit(jax_model)` before `%timeit` to pre-cache the compilation):
```python
def pt_forward():
    inputs = processor(text=["red"], images=[image], return_tensors="pt", padding=True)
    outputs = pt_model(**inputs)
    return outputs.logits_per_image

%timeit pt_forward()
# 81.2 ms ± 559 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

```python
def jax_forward():
    inputs = processor(text=["red"], images=[image], return_tensors="jax", padding=True)
    outputs = jit(jax_model)(**inputs)
    return outputs.logits_per_image
    
%timeit jax_forward().block_until_ready()
# 7.83 ms ± 1.11 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

And for getting **gradients** of the inputs with respect to the output (note that in the notebook I run `jit(grad(jax_loss))` before `%timeit` to pre-cache the compilation):
```python
for p in pt_model.parameters():
    p.requires_grad = False

def pt_loss(inputs):
    outputs = pt_model(**inputs)
    logits_per_image = outputs.logits_per_image
    return logits_per_image

def pt_test():
    text = "green"
    image = Image.new('RGB', (224, 224), color=(0, 255, 0))
    inputs = processor(text=[text], images=[image], return_tensors="pt", padding=True)
    inputs.pixel_values.requires_grad = True
    loss = pt_loss(inputs)
    loss.backward()
    return inputs.pixel_values.grad

def jax_loss(pixel_values, input_ids, attention_mask):
    inputs = {"pixel_values":pixel_values, "input_ids":input_ids, "attention_mask":attention_mask}
    outputs = jax_model(**inputs)
    logits_per_image = outputs.logits_per_image
    return logits_per_image[0][0]

def jax_test():
    text = "green"
    image = Image.new('RGB', (224, 224), color=(0, 255, 0))
    inputs = processor(text=[text], images=[image], return_tensors="jax", padding=True)
    grads = jit(grad(jax_loss))(inputs.pixel_values, inputs.input_ids, inputs.attention_mask)
    return grads
```
```python
%timeit pt_test()
# 129 ms ± 1.04 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
```python
%timeit jax_test().block_until_ready()
# 11.2 s ± 34.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```
## Expected behavior

I'd have expected CLIPModel and FlaxCLIPModel to be *roughly* comparable in performance.


Comments:
Hi there, 

It seems you have frozen the PyTorch model's parameters so no gradients are calculated for them, which could explain the huge speed difference. 

Also, I would avoid jitting inside a function and then call it multiple times, it may cause re-compilation (I'm not sure).
We could first `jit` the function and then call that inside the main function, like

```python
forward =  jit(jax_model) # jit here

def jax_forward():
    inputs = processor(text=["red"], images=[image], return_tensors="jax", padding=True)
    outputs = forward(**inputs) # call the jitted function
    return outputs.logits_per_image

@patil-suraj

> It seems you have frozen the PyTorch model's parameters so no gradients are calculated for them, which could explain the huge speed difference.

Do you mean this part?

```python
for p in pt_model.parameters():
     p.requires_grad = False
```

If so, I tried commenting that out and Pytorch is still 1-2 OOM faster. The reason I added that code was because my understanding is that `jax.grad` only gets gradients of the first parameter of the function, [i.e. [default `argnums` value is `0`](https://jax.readthedocs.io/en/latest/jax.html#jax.grad)] so since JAX is only getting gradients of the `pixel_values` input, I wanted to make it so Pytorch is also only getting those too, for fairness.

> Also, I would avoid jitting inside a function and then call it multiple times, it may cause re-compilation (I'm not sure).

I'm quite sure that's not the case, since `jax_forward()` is much, much slower the first time I call it, and then after that it's way *faster* than `pt_forward()`. I tested this again just now in the Colab to confirm: ~800ms first run of `jax_forward()` vs ~30ms on successive runs, while `pt_forward()` always takes ~260ms.

[The docs for jax.jit](https://jax.readthedocs.io/en/latest/jax.html#jax.jit) explicitely mention that there is a compilation cache, and a lot of examples from intro tutorials use the a pattern like `jit(jax_model)(**inputs)` instead of `jit_jax_model(**inputs)` - one just needs to be aware of the caching rules - especially (from docs): "Static arguments are included as part of a compilation cache key" (which makes perfect sense).

Also worth mentioning that I thought that for JAX to be so fast with inference, perhaps it's caching outputs by default, but I couldn't find any documentation suggesting that, and I think I've eliminated that possibility with a test like this:

```python
i = 0
def jax_forward():
    global i
    i += 1
    inputs = processor(text=[f"red {i}"], images=[image], return_tensors="jax", padding=True)
    outputs = jit(jax_model)(**inputs)
    return outputs.logits_per_image
```

Let me know if there are any other tests you'd like me to run. The Colab notebook I linked in my original post is quite minimal and self-contained so if you just switch runtime to GPU and then `Runtime -> Run all` (after making some code changes to test a hypothesis), then it should fairly painlessly give you an answer.

@patil-suraj It occurred to me that you might have meant to mention `jit_grad_jax_loss = jit(grad(jax_loss))`, rather than the `forward` functions, since it's the grad part that's 100x slower in Flax, and so I tried pre-compiling `jit(grad(jax_loss))`, and you were right! Now Flax is 10x faster in **both** inference and getting the gradients! 🤔 I'll have to do some more reading of the JAX/Flax docs to work out what's going on here. I guess something is causing the `jit` compilation cache to be bypassed. Maybe related to [this JAX issue](https://github.com/google/jax/issues/2095). Either way, **thank you for your help!** :pray: 

```python
jit_grad_jax_loss = jit(grad(jax_loss))

def jax_test_pre_jit():
    text = "green"
    image = Image.new('RGB', (224, 224), color=(0, 255, 0))
    inputs = processor(text=[text], images=[image], return_tensors="jax", padding=True)
    grads = jit_grad_jax_loss(inputs.pixel_values, inputs.input_ids, inputs.attention_mask)
    return grads
```

I also just tested `FlaxBertModel` against `BertModel` and Flax is again 10x faster during inference:

* https://gist.github.com/josephrocca/8f3797f1f7605eccddd7b3d438052734
* https://colab.research.google.com/gist/josephrocca/8f3797f1f7605eccddd7b3d438052734

Could it be that Flax is just a lot faster than Pytorch with these models? I'd have thought the speed difference would be smaller.

Happy for you to close this if the observed inference speed difference is not suspicious to you :+1:

**Edit**: [I've asked a question](https://github.com/google/jax/discussions/7871) about `jit(f)` being cached, but `jit(grad(f))` not being cached (using a much simpler code example) over on the JAX discussion forum. **Edit 2**: The answer turns out to be pretty simple: `grad(f)` [isn't cached](https://github.com/google/jax/issues/2095) (but there's discussion in that issue of giving `grad` a cache), and thus returns a new function each time, and so `jit` must recompile it every time (since it caches compilations based on the function reference rather than the function source string).

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.