TF: XLA generation not working properly in some models
Author: gante
Pull request: no
Labels: name: TensorFlow

name: Good Second Issue

This issue is used to track TensorFlow XLA generation issues, arising from #17857. There are three categories of issues, sorted in descending order by severity:

### Key model issues

These are heavily-used models, whose quality should be prioritized.

- [x] T5 -- The quality of the results decreases with `max_length`. See [here](https://github.com/huggingface/transformers/pull/17857/files#r906702367).
- [x] GPT-J -- fails simple generate tests with numerical issues

### Models failing basic tests

These models are failing `test_xla_generate_fast` -- a short greedy generation.

- [ ] LED
- [ ] Speech2Text
- [ ] XLNet
- [ ] XGLM

### Models failing complex tests

These are models failing `test_xla_generate_slow` -- a long beam search generation.

- [x] Bart
- [x] Blenderbot
- [x] Marian
- [x] mbart
- [x] OPT
- [x] Pegasus

Comments:
@gante do you require any help with this issue? Happy to contribute

Hi @anmolsjoshi 👋 

If you are comfortable with debugging XLA, absolutely :) My recommendation would be to pick a model from "Models failing complex tests" (the others might require significant architecture changes), and to start debugging. The number 1 suspect is always the position embeddings, which may not be handling the case where `past` is padded. Let me know if you are up to it, and which model would you like to take!  

Hi @gante, I did have a bit of a poke around. I think the complex tests all fail for the same reason: those models have a setting `max_position_embeddings` that is set to 20 by default during testing and which is too short for the “slow” tests. Here’s a simple fix for those: https://github.com/dsuess/transformers/commit/4a3e27164ae941fcd649b8565d7d92a4552d689f. I’ll give the other ones a shot now

Hello @gante, may I ask if there is anything that I can contribute?


Hi JuheonChu 👋 Actually yes! I have a few unchecked models at the top, but I wouldn't recommend spending time there unless you plan to use those architectures -- they are infrequently used.

However, two popular models are currently failing their XLA tests with beam search:
- Marian
- OPT

You can see the failing test if you install from `main` (`pip install --upgrade git+https://github.com/huggingface/transformers.git`) and run it e.g. for OPT `NVIDIA_TF32_OVERRIDE=0 RUN_SLOW=1 py.test -vv tests/models/opt/test_modeling_tf_opt.py::TFOPTModelTest::test_xla_generate_slow`

I haven't dived in yet, so I don't know what's the cause for the failure. You'll have to hop into debug mode and see what is breaking :)

Can @katiele47 and I try working on them? 

@JuheonChu of course!

> @JuheonChu of course!
@gante  Are we figuring out the cause of the testing failures based on the clues as follows?

![Error 1](https://user-images.githubusercontent.com/35699839/219711119-9459c1d9-22c4-4673-9b49-2e4815515a96.png)
![Error 2](https://user-images.githubusercontent.com/35699839/219711143-135b5069-62ca-4622-823a-6df3fe572318.png)
![Error 3](https://user-images.githubusercontent.com/35699839/219711156-00580d72-2cc2-4acc-a5a0-7bccf7392097.png)


@JuheonChu yes. My suggestion would be to attempt to find where the numerical differences start from (between the XLA and the non-XLA version), using a debugger. Please note that you can't print variables with `jit_compile=True`, so you should set it to `False`. From there, the root cause is typically apparent.

Be warned, these sort of tasks sometimes are very time-consuming to complete :)

> @JuheonChu yes. My suggestion would be to attempt to find where the numerical differences start from (between the XLA and the non-XLA version), using a debugger. Please note that you can't print variables with `jit_compile=True`, so you should set it to `False`. From there, the root cause is typically apparent.
> 
> Be warned, these sort of tasks sometimes are very time-consuming to complete :)

Thank you very much for your valuable guidance! We will try and keep you updated!

Hi @gante, I've attempted to reproduce the failed XLA test on the OPT model using your suggested commands. The cause of error I had was somehow different from @JuheonChu's. Would you be able to verify if the following is the expected failing test output? If not, I assume it could be due to my local repo. Thanks!
<img width="1015" alt="Screen Shot 2023-02-21 at 11 20 44 PM" src="https://user-images.githubusercontent.com/54815905/220521074-6ab355c4-fe0b-42b5-88fa-cbd15de82b8a.png">
<img width="1125" alt="Screen Shot 2023-02-21 at 11 21 24 PM" src="https://user-images.githubusercontent.com/54815905/220521160-946a03db-cf80-4fbc-818c-3be337df3983.png">
<img width="1135" alt="Screen Shot 2023-02-21 at 11 21 43 PM" src="https://user-images.githubusercontent.com/54815905/220521210-981d6099-e093-4b8b-9913-2dafe5bef905.png">


@gante working on  XLNet