Generate text with `model.generate` on TPU does not work
Author: stekiri
Pull request: no
Labels: name: WIP

## Environment info
- `transformers` version: 4.7.0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29 (Ubuntu 20.04.2 LTS)
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.1+cu102 (False)
- PyTorch XLA version: 1.8.1
- Using GPU in script?: No, using TPU
- Using distributed or parallel set-up in script?: No, using a single TPU core

### Who can help
@patrickvonplaten

## Information

Model I am using (Bert, XLNet ...): `facebook/m2m100_1.2B`, but other text generating models have the same problem.

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

On a machine with a TPU run:

```python
import torch_xla.core.xla_model as xm
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

model_name = 'facebook/m2m100_1.2B'
source_lang = 'en'
target_lang = 'de'

docs = [
    "This is some document to translate.",
    "And another document to translate."
]

device = xm.xla_device()

model = M2M100ForConditionalGeneration.from_pretrained(model_name).to(device)

tokenizer = M2M100Tokenizer.from_pretrained(model_name, src_lang=source_lang)
encoded_docs = tokenizer(docs, return_tensors='pt', padding=True).to(device)

generated_tokens = model.generate(**encoded_docs, forced_bos_token_id=tokenizer.get_lang_id(target_lang))
```

The call to `model.generate()` runs without ever terminating. It seems to be stuck somewhere in the beam search.

The same code runs perfectly fine on CPUs and GPUs.

## Expected behavior

I'd expect that the generation of text works in the same way as for CPUs and GPUs.


Comments:
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.

The issue still seems to be unresolved. Maybe the generation is not supposed to be supported on TPUs? In that case, a short note in the documentation could be helpful 😄 

@stekiri - I think we indeed don't support PyTorch-XLA generation yet and yes a comment in the docs would be good! Would you be interested in making such a PR? I've put "enabling TPU generation for PyTorch-XLA" on my TODO list now so hope to tackle this sometime in 3-4 weeks

I have the same situation on GCP Cloud TPU v2-8 (summarization pipeline with T5ForConditionalGeneration).
I'm eagerly waiting for the support.

This is still broken

@patil-suraj it worked for you on TPU no?

It's possible to use PT `generate` on TPU with `accelerate`, here's a colab which uses GPT2 as an example
https://colab.research.google.com/drive/1OqCLWuEbWLp4fLLWcT-vteEJZHHZ3SZ5?usp=sharing

> It's possible to use PT `generate` on TPU with `accelerate`, here's a colab which uses GPT2 as an example https://colab.research.google.com/drive/1OqCLWuEbWLp4fLLWcT-vteEJZHHZ3SZ5?usp=sharing

This doesn't seem faster than GPU. I tried to run it on 400 examples and it was slower than GPU. Seems like you are just unloading the model back to CPU, which is actually slower, and we can't really leverage the power of TPU.

Is there any update on this?

> Is there any update on this?

I had an exchange with @gante about it and it seems like the code will need major refactoring for this. https://huggingface.co/spaces/joaogante/tf_xla_generate_benchmarks/discussions/1#62eb9350985a691200cf2921

@mikcnt @divyanshuaggarwal The previous TF generate function was almost a (reduced) copy of the current PT generate function. We had to do a major rework of the TF generate function to make it compatible with XLA, so yeah... PT needs the same treatment if we want to use it with XLA :D 

I've shared a twitter thread today about the subject: https://twitter.com/joao_gante/status/1555527603716444160

@gante thanks a lot for the super exhaustive explanation. Do you think we can expect a refactoring for PT some time soon?

Otherwise, do you know of a temporary workaround to use the generate method on TPU?

@mikcnt we don't have refactoring PT's generate in our short-term plans -- it is a very labor-intensive refactor whose main benefit is to enable TPU usage (i.e. a niche usage :) ). For context, the TF refactor took me >2 months of dedicated effort (contrarily to PT, the old TF implementation was slow on GPUs, and it was much smaller than PT's generate).

There are no immediate alternatives -- not being fully compatible with XLA implies that it can't get on a TPU effectively. Maybe the model you want to use exists on FLAX/TF, whose generate is compatible with TPUs.

I don't want to clip your wings, so here's my suggestion: our efforts go towards what the majority of the community wants. If you open an issue in `transformers` and attract some attention to PT generation on TPU, the odds of it happening soon increase significantly!