TF XLA greedy generation
Author: Rocketknight1
Pull request: yes

This is a first attempt at refactoring greedy generation in TF to use XLA. It has a couple of main problems still:

- Very hardcoded for GPT-2 right now
- No stopping based on EOS tokens yet (EDIT: Fixed!)
- Some really ugly array update hackery caused by the lack of tools like `dynamic_update_slice` (EDIT: Fixed!)
- `prepare_inputs_for_generation` assumes greedy generation, this is wrong (EDIT: Fixed)
- Outputs are slightly wrong compared to a manual loop - I'm still investigating (EDIT: Fixed!)

However, the outputs are coherent and the speedup compared to the old method is approximately 25X.

**UPDATE (from Patrick)**

PR should be good to go. See description below.

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_15786). All of your documentation changes will be reflected on that endpoint.

Outputs now correctly match a manual loop, up to minor float32 errors. Code has been simplified with the addition of direct calls to XLA, and we're now at >30X speedup versus the original code.

@Rocketknight1 

I've been testing the code by running the following code example - put it in a `run.py` file:

```python
from transformers import TFGPT2LMHeadModel, GPT2Tokenizer
import tensorflow as tf
import sys

use_xla = bool(int(sys.argv[1]))

physical_devices = tf.config.list_physical_devices('GPU')
for device in physical_devices:
    tf.config.experimental.set_memory_growth(device, True)

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = TFGPT2LMHeadModel.from_pretrained("gpt2")

input_ids = tokenizer("hello there can you continue", return_tensors="tf").input_ids

outputs = model.generate(input_ids, use_xla=use_xla)

print("Output", tokenizer.batch_decode(outputs))
```

XLA-case `python run.py 1`
Non-XLA-case `python run.py 0`

Can you try it out? Note that it's a pretty hacky solution that I've quickly done - we can surely make it cleaner. But overall I really don't see the problem with make xla_generate compatible for the non-xla case. 

Could you take a look, try it out and see whether my idea is too limited?

Pr is ready to go. Ran all of the following slow tests:

```python
RUN_SLOW=1 pytest -sv \
tests/bart/test_modeling_tf_bart.py \
tests/t5/test_modeling_tf_t5.py \
tests/gpt2/test_modeling_tf_gpt2.py \
tests/vision_encoder_decoder/test_modeling_tf_vision_encoder_decoder.py \
tests/encoder_decoder/test_modeling_tf_encoder_decoder.py \
tests/speech_to_text/test_modeling_tf_speech_to_text.py
```

as well as

```bash
RUN_SLOW=1 pytest -sv tests/rag/test_modeling_tf_rag.py
```
(this one is tested on brutasse as it requires a heavy index).

Everything passes now.

The PR is now ready for merge - we now have a XLA-compileable `greedy_search` method which fully replaced the old `greedy_search` generate method. The XLA-compilable `greedy_search` behaves exactly the same way as the previous one for eager execution, but also allows for XLA compilation. 
This was made possible at the expense of some code overhead (see function decorator flags), and a couple of `if use_xla` statements.

Since XLA-compileable `generate` has model-dependent cache tensors, we need to add a **single** new function to every model we want to make XLA-compilable, called `_update_model_kwargs_for_xla_generation`. One is added for GPT2 now so that GPT2 now works with XLA. The method is a bit hacky at the moment and could definitely be cleaned up in a future PR - I've copy-pasted some code from the previous state of the PR. XLA gpt2 generate gives the same result as non-XLA generate though.

IMO, the design is more or less fixed now as we have 0 breaking changes, no duplicated (non-XLA and XLA) generate methods, and relatively clean code.

Some important things that should be discussed here though:

1 - At the moment the way the cache and attention_mask is padded to make it XLA-compatible is a bit hacky (see code in `_update_model_kwargs_for_xla_generation`), but it allows for 0 breaking changes. It might be preferable to change the way `past_key_values` is implemented inside the model to potentially make generate even faster and less "hacky". This would mean however that `past_key_values` would change its tensor shape which is a breaking change though since we do return `past_key_values` when one passes `use_cache=True`. So might not be worth it

2 - XLA can be enabled with a flag `use_xla=True/False` which defaults to `False` at the moment. Once we make all models XLA compileable we could default the flag to `True`. TF has a somewhat weird way of enabling jit-XLA compilation with the function decorator which is why we have so many functions inside functions now in this PR. I do think it's the most intuitive way though is to allow the user to pass a `use_xla=True/False` flag. But I'm no TF expert - what do you think @sgugger @Rocketknight1 @gante ?

3. XLA generate does **not** cut before `max_length` is reached. It'll just output 0's in the end for now (think we should change this in a follow-up PR). The reason is that cutting the final outputs for some reason changes the axes and it's also "cleaner"/more "barebone" IMO to just output the whole `max_length` tensor. We should just make sure that the tensor is always padded with 0's. Would be nice to do this in a follow-up PR though.

In terms of next steps cc @Rocketknight1 @gante:

- a) GPT2 is mostly used with `do_sample=True`. So we should refactor this method analogs to this one. Should not be difficult now that it's done for greedy search.
- b) We should add more models (should be trivial to make GPT-J, GPT-Neo, ...) XLA-compatible. Also shouldn't be too hard for T5 now. 
- c) We now should make the logits processors XLA-compileable - I left a TODO in the PR. It won't be easily possible to make **all** logits processors XLA-compileable, but we should aim to do it for the important ones. Again one should just look at how it's done for Flax and more or less copy-paste the design

Might be a bit difficult to review the git diff here. Would be nice if @Rocketknight1 and @gante, you could go through the new file and carefully check if everything makes sense - happy to answer any question you might have.

> 2 - XLA can be enabled with a flag use_xla=True/False which defaults to False at the moment. Once we make all models XLA compileable we could default the flag to True. TF has a somewhat weird way of enabling jit-XLA compilation with the function decorator which is why we have so many functions inside functions now in this PR. I do think it's the most intuitive way though is to allow the user to pass a use_xla=True/False flag. But I'm no TF expert - what do you think @sgugger @Rocketknight1 @gante ?

There is an alternative solution -- using the `tf.config` functions to trigger whether we want [XLA](https://www.tensorflow.org/api_docs/python/tf/config/optimizer/set_jit) or [eager execution](https://www.tensorflow.org/api_docs/python/tf/config/run_functions_eagerly) on functions decorated with `tf.function`. This approach has pros and cons, with respect to the current one.

Pros:
1. No need to define inner functions with `@xla_compile`. Everything that might get XLA compilation gets a `@tf.function`, and the exec mode is set by calling the appropriate config at the beginning of `generate()`
2. Repeated `generate()` calls do not redefine the functions, so they should benefit from previous XLA warmup -- not sure if this is a significant advantage or not 🤔  (could be, for instance, on hosted APIs that expect repeated calls on the same function)

Cons:
1. Setting any of these `tf.config` methods requires an empty TF session. This means we can't call a generate without and with XLA without resetting the session. Also, if we were to call it at the beginning of `generate()`, no compiled TF model could exist (again, because it needs an empty session).

Personally, I'm slightly leaning towards this approach (with the `tf.config` calls) if we could make it work, but no strong feelings :) 

> > 2 - XLA can be enabled with a flag use_xla=True/False which defaults to False at the moment. Once we make all models XLA compileable we could default the flag to True. TF has a somewhat weird way of enabling jit-XLA compilation with the function decorator which is why we have so many functions inside functions now in this PR. I do think it's the most intuitive way though is to allow the user to pass a use_xla=True/False flag. But I'm no TF expert - what do you think @sgugger @Rocketknight1 @gante ?
> 
> There is an alternative solution -- using the `tf.config` functions to trigger whether we want [XLA](https://www.tensorflow.org/api_docs/python/tf/config/optimizer/set_jit) or [eager execution](https://www.tensorflow.org/api_docs/python/tf/config/run_functions_eagerly) on functions decorated with `tf.function`. This approach has pros and cons, with respect to the current one.
> 
> Pros:
> 
> 1. No need to define inner functions with `@xla_compile`. Everything that might get XLA compilation gets a `@tf.function`, and the exec mode is set by calling the appropriate config at the beginning of `generate()`
> 2. Repeated `generate()` calls do not redefine the functions, so they should benefit from previous XLA warmup -- not sure if this is a significant advantage or not thinking  (could be, for instance, on hosted APIs that expect repeated calls on the same function)
> 
> Cons:
> 
> 1. Setting any of these `tf.config` methods requires an empty TF session. This means we can't call a generate without and with XLA without resetting the session. Also, if we were to call it at the beginning of `generate()`, no compiled TF model could exist (again, because it needs an empty session).
> 
> Personally, I'm slightly leaning towards this approach (with the `tf.config` calls) if we could make it work, but no strong feelings :)

Happy to switch to using `tf.config(...)` if you guys prefer! I'm not familiar with what is intuitive in TF so trusting you guys here. But I'm also not too happy with all those subfunctions and complicated function wrappers, so this approach makes a lot of sense. 
Just to understand better, do you think it make more sense to set `tf.config.optimizer.set_jit(enabled=use_xla)` right after `generate(...)` or should we move it out completely of `generate()` and let the user set it before calling generate? In this case we wouldn't need a `use_xla` flag here at all no? Not sure what's better here

> Just to understand better, do you think it make more sense to set tf.config.optimizer.set_jit(enabled=use_xla) right after generate(...) or should we move it out completely of generate() and let the user set it before calling generate? In this case we wouldn't need a use_xla flag here at all no? Not sure what's better here

There would be no need for the `use_xla` flag in `tf.function`. I'd say to keep it in `generate()`, as it should be the outermost function that may need XLA inside. @Rocketknight1 what do you think (current approach vs `tf.config`?)

@Rocketknight1 and I were talking on Slack, and @Rocketknight1 came up with a third option, which is better in our opinion :D 

The current approach has the disadvantages described above -- it has nested functions to implement the decorator on and, because of that, it may lose the ability to cache XLA compiled functions (= performance loss on repeated calls)

The approach I suggested relies on clearing sessions under some circumstances, which is not user-friendly.

The third option:
 1. Let's not use `tf.function()` as a decorator, but as a function. We could do something like `xla_greedy_generate = tf.function(greedy_generate, jit_compile=True)` -- have a function called by `__init__()` that initializes all needed XLA functions;
 2. `tf.function()` doesn't accept `**kwargs`, so we couldn't rely on a `use_xla` flag. Instead, we could use `tf.executing_eagerly()` to determine the mode of operation, inside `greedy_generate`.
 
This third option would be the best of both worlds -- we would avoid nested functions, ensure we can cache XLA compilation, and avoid relying on functions that need clear sessions. @patrickvonplaten what do you think? I'd be happy to give it a go if you're busy with other things :)

> @Rocketknight1 and I were talking on Slack, and @Rocketknight1 came up with a third option, which is better in our opinion :D
> 
> The current approach has the disadvantages described above -- it has nested functions to implement the decorator on and, because of that, it may lose the ability to cache XLA compiled functions (= performance loss on repeated calls)
> 
> The approach I suggested relies on clearing sessions under some circumstances, which is not user-friendly.
> 
> The third option:
> 
> 1. Let's not use `tf.function()` as a decorator, but as a function. We could do something like `xla_greedy_generate = tf.function(greedy_generate, jit_compile=True)` -- have a function called by `__init__()` that initializes all needed XLA functions;
> 2. `tf.function()` doesn't accept `**kwargs`, so we couldn't rely on a `use_xla` flag. Instead, we could use `tf.executing_eagerly()` to determine the mode of operation, inside `greedy_generate`.
> 
> This third option would be the best of both worlds -- we would avoid nested functions, ensure we can cache XLA compilation, and avoid relying on functions that need clear sessions. @patrickvonplaten what do you think? I'd be happy to give it a go if you're busy with other things :)

Ah yes, I see! Using `tf.function(greedy_search)` is def better than the function decorator logic. Not 100% sure whether overwriting greedy search at `model.__init__(...)` is the way to go though. Wouldn't it in that case just be better to just mention in the docs that the generate method can be compiled as follows:

```
jitted_generate = tf.function(model.generate, jit_compile=True)

output = jitted_generate(input_ids)
```

? We could make use of `tf.executing_eagerly()` inside the code to figure out if we need XLA or not no?