(T5) tf.function wrapped model.generate() does not produce the same result as non-wrapped model.generate()
Author: JEF1056
Pull request: no
Labels: name: WIP

name: bug

### System Info

```shell
- `transformers` version: 4.19.0
- Platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.13
- Huggingface_hub version: 0.6.0
- PyTorch version (GPU?): 1.11.0+cu113 (False)
- Tensorflow version (GPU?): 2.8.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No
```


### Who can help?

@patrik @gante 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

1. Run the code below:
```
from transformers import TFT5ForConditionalGeneration, AutoTokenizer
import tensorflow as tf

tokenizer = AutoTokenizer.from_pretrained("t5-small")
model = TFT5ForConditionalGeneration.from_pretrained("t5-small")

input_ids = tokenizer("translate English to German: I need to convert this into a savedmodel.", padding='max_length', return_tensors="tf").input_ids

# This is normal model.generate()
outputs_0 = model.generate(input_ids)
print(outputs_0)
print(tokenizer.batch_decode(outputs_0))

# This is wrapped with tf.function()
wrapped = tf.function(model.generate)
outputs_1 = wrapped(input_ids)
print(outputs_1)
print(tokenizer.batch_decode(outputs_1))

assert outputs_0 == outputs_1, "Results are not equal."
```
2. Observe results. This is what i get on my machine:
```
tf.Tensor(
[[    0  1674  2171    67     7    16   236 20819    15     7  8731   561
  18980    29     5     1]], shape=(1, 16), dtype=int32)
['<pad> Ich muss dies in ein gespeichertes Modell umwandeln.
']
tf.Tensor(
[[    0  1674  2171    67     7    16   236 20819    15     7     7     7
      7     7     7     7     7     7     7     7]], shape=(1, 20), dtype=int32)
['<pad> Ich muss dies in ein gespeichertesssssssssss']
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
[<ipython-input-74-27b3c9e72d47>](https://localhost:8080/#) in <module>()
     18 print(tokenizer.batch_decode(outputs_1))
     19 
---> 20 assert outputs_0 == outputs_1, "Results are not equal."

AssertionError: Results are not equal.
```
This issue gets worse the larger the max_length option in model.generate() gets.

### Expected behavior

```shell
The results are expected to be equal.
```


Comments:
@JEF1056 thank you for pointing it out 👍 I can reproduce the issue with transformers==4.19 and 4.18 (4.17 and older versions do not support the `tf.function` wrapper), with and without input padding.

I will look into it and let you know of any findings :) In the recent past, we found numerical instabilities in some compiled functions on CPU, so it may be related.

Interesting. I've tried with a GPU and used the master branch of the repository as well and results are identical. Unfortunately, I don't have a lot of experience with XLA so i'm not sure if I can be of much help.

Thanks for picking up the issue though!

@JEF1056 upon further digging in debugging mode, I couldn't find any unexpected behavior. I did encounter the expected behavior, which explains the mismatch:
- `tf.function` compiles and optimizes the graph, which rearranges FP32 operations. Rearranging FP32 operations leads to very minor numerical differences (see [here](https://stackoverflow.com/questions/48957828/floating-point-arithmetic-why-would-order-of-addition-matter) an explanation). We can see this in the encoder forward pass, which has differences in the order of 1e-6;
- Generation is in essence a sequence of forward passes, where past hidden outputs are fed as inputs. What starts as a tiny difference quickly builds up into a larger difference, which at some point results in different tokens.

The reverse can also be observed: if we pick an input with a stronger signal or a more powerful model, we see smaller differences at a token level. Consider the following example, and try it out with `t5-small`, `t5-base`, and `t5-large`:

```
from transformers import TFT5ForConditionalGeneration, AutoTokenizer
import tensorflow as tf

tokenizer = AutoTokenizer.from_pretrained("t5-large")
model = TFT5ForConditionalGeneration.from_pretrained("t5-large")

input_ids = tokenizer("translate English to German: This is a very long sentence that is easy to translate because it has common words.", padding='max_length', return_tensors="tf").input_ids

# This is normal model.generate()
outputs_0 = model.generate(input_ids)
print(outputs_0)
print(tokenizer.batch_decode(outputs_0))

# This is wrapped with tf.function()
wrapped = tf.function(model.generate)
outputs_1 = wrapped(input_ids)
print(outputs_1)
print(tokenizer.batch_decode(outputs_1))
```

The outputs are not exactly the same, but they are sensible. In any case, `tf.function` + generation is something we are working at the moment, stay tuned for further updates (which may alleviate this issue) :D

Also -- @Rocketknight1 @patrickvonplaten this issue [`tf.function` resulting in different FP32 ops -> different generate results] is something I'm seeing constantly and, curiously, the `tf.function` generation outputs qualitatively worse text. I wonder if there is any way to mitigate this issue and/or if we should try to contact the TF team 🤔 

@gante I see, and that does make sense. I presume that the issue occurs in T5 and not GPT2 becasue T5 is encoder-decoder, while GPT2 is decoder only?
```
from transformers import TFGPT2LMHeadModel, AutoTokenizer
import tensorflow as tf
import numpy as np

model_name = 'gpt2'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = TFGPT2LMHeadModel.from_pretrained(model_name)

input_ids = tokenizer("This is a sentence that ", return_tensors="tf").input_ids

# This is normal model.generate()
outputs_0 = model.generate(input_ids, max_length=100)
print(outputs_0)
print(tokenizer.batch_decode(outputs_0))

# This is wrapped with tf.function()
wrapped = tf.function(lambda x: model.generate(x, max_length=100))
outputs_1 = wrapped(input_ids)
print(outputs_1)
print(tokenizer.batch_decode(outputs_1))

assert np.array_equal(outputs_0, outputs_1), "Results are not equal."
```
The above code works fine.

For context, I've been trying to convert t5 to tensorflowjs and though I have got it working before by creating my own generate function in tf.js, having huggingface's generate function directly as part of the savedmodel would really improve its speed.
In the current t5 tf.function generate() wrapper, XlaDynamicUpdateSlice and TensorListConcatV2 ops are used, which pobably won't ever be supported by tf.js, do you know if there is any way to implement tf.function without these?

@gante from my experience small numerical differences in the magnitude of 1e-6 should not lead to different tokens being generated (in flax they don't). Also the generate outputs of https://github.com/huggingface/transformers/issues/17222#issue-1234576122 look quite bad. Since we don't seem to have this issue in GPT2, could it be that something encoder-decoder specific is not done correctly in `tf.function` ? E.g. the `encoder_attention_mask` or the cache? @JEF1056 could you try running the code also with `use_cache=False` to see if we still get a difference ?

To me it looks like a bug still from our side

@patrickvonplaten

With use_cache:
```
from transformers import TFT5ForConditionalGeneration, AutoTokenizer
import tensorflow as tf

tokenizer = AutoTokenizer.from_pretrained("t5-small")
model = TFT5ForConditionalGeneration.from_pretrained("t5-small")

input_ids = tokenizer("translate English to German: This is a very short sentence.", padding='max_length', return_tensors="tf").input_ids

# This is normal model.generate()
outputs_0 = model.generate(input_ids)
print(outputs_0)
print(tokenizer.batch_decode(outputs_0))

# This is wrapped with tf.function()
wrapped = tf.function(model.generate)
outputs_1 = wrapped(input_ids)
print(outputs_1)
print(tokenizer.batch_decode(outputs_1))

assert outputs_0 == outputs_1, "Results are not equal."
```
```
tf.Tensor([[    0   644   229   236  1319  7755    49 20144     5     1]], shape=(1, 10), dtype=int32)
['<pad> Das ist ein sehr kurzer Satz.
']
tf.Tensor(
[[   0  644  229  236 1319 7755    5    1    0    0    0    0    0    0
     0    0    0    0    0    0]], shape=(1, 20), dtype=int32)
['<pad> Das ist ein sehr kurz.
<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>']
```

Without use_cache:
```
from transformers import TFT5ForConditionalGeneration, AutoTokenizer
import tensorflow as tf

tokenizer = AutoTokenizer.from_pretrained("t5-small")
model = TFT5ForConditionalGeneration.from_pretrained("t5-small")

input_ids = tokenizer("translate English to German: This is a very short sentence.", padding='max_length', return_tensors="tf").input_ids

# This is normal model.generate()
outputs_0 = model.generate(input_ids, use_cache=False)
print(outputs_0)
print(tokenizer.batch_decode(outputs_0))

# This is wrapped with tf.function()
wrapped = tf.function(lambda x: model.generate(x, use_cache=False))
outputs_1 = wrapped(input_ids)
print(outputs_1)
print(tokenizer.batch_decode(outputs_1))

assert outputs_0 == outputs_1, "Results are not equal."
```
```
tf.Tensor([[    0   644   229   236  1319  7755    49 20144     5     1]], shape=(1, 10), dtype=int32)
['<pad> Das ist ein sehr kurzer Satz.
']
tf.Tensor(
[[   0  644  229  236 1319 7755    5    1    0    0    0    0    0    0
     0    0    0    0    0    0]], shape=(1, 20), dtype=int32)
['<pad> Das ist ein sehr kurz.
<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>']
```

Results appear to be the same.
I would also like to bring up that after wrapping generate() in tf.function, it no longer appears to respect stop tokens, which should be something that can be fixed on huggingface's side.

I think @gante 's point that small numerical differences magnified by generating an autoregressive fashion is correct, as results diverge wildly the longer the generated sentence is.

@JEF1056, 

Note that the model does generate the EOS token. It's expected to get padding in the end since shapes need to be static when using XLA. You can remove the `<pad>` tokens by doing:

```py
print(tokenizer.batch_decode(outputs_1, skip_special_tokens=True))
```

Still interested in a more in-detail analysis where exactly the differences start to creep in and at what point they become very significant. XLA generation does not look good enough to me to not be a bug

I see, I didn't realize that XLA shapes need to be static (seems waste compute for longer sequences though). That might pose some difficulties in getting generate() to work with GPT2 (decoder models in general) since inputs can't be padded to length in that case.



For some additional context @JEF1056: we have already found (and mitigated) an XLA/non-XLA mismatch (see https://github.com/tensorflow/tensorflow/issues/55682), so we can't rule out conversion problems. Sadly, they are hard to detect, as it relies on numerical debugging with XLA.

I'm going to continue developing XLA + generate, keeping this issue in the backlog -- there is a chance future changes fix the issue we are seeing, or that I stumble across the root cause naturally as I enable the use of `tf.function` on other models.

Thank you for raising the issue, and let us know if your see related problems 🙏  

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/main/CONTRIBUTING.md) are likely to be ignored.

This issue is fixed, XLA T5 should be working properly :)