Donut model.generate is extremely slow when run inference
Author: minhqngo
Pull request: no

I trained a Donut model for document classification using my custom dataset (format similar to RVL-CDIP). However, when I run inference, the model.generate() run extremely slow (5.9s ~ 7s).

Inference device: NVIDIA A100 40GB.

Requirements:
CUDA 11.7
torch==1.13.1+cu117
torchvision==0.14.1+cu117
datasets==2.10.1
transformers==4.26.1
sentencepiece==0.1.97
onnx==1.12.0
protobuf==3.20.0

Here is the GPU when I run inference:

![image](https://user-images.githubusercontent.com/38152758/230818381-d2489865-78b2-410f-b2d7-1a487a98f3eb.png)

This is my code for inference

```
model = VisionEncoderDecoderModel.from_pretrained(CKPT_PATH, config=config)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

accs = []
model.eval()
for i, sample in tqdm(enumerate(val_ds), total=len(val_ds)):
    pixel_values = sample["pixel_values"]
    pixel_values = torch.unsqueeze(pixel_values, 0)
    pixel_values = pixel_values.to(device)

    start = time.time()
    task_prompt = "<s_fci>"
    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
    decoder_input_ids = decoder_input_ids.to(device)
    print(f"Tokenize time: {time.time() - start:.4f}s")

    start = time.time()
    outputs = model.generate(
        pixel_values,
        decoder_input_ids=decoder_input_ids,
        max_length=model.decoder.config.max_position_embeddings,
        early_stopping=True,
        pad_token_id=processor.tokenizer.pad_token_id,
        eos_token_id=processor.tokenizer.eos_token_id,
        use_cache=True,
        num_beams=1,
        bad_words_ids=[[processor.tokenizer.unk_token_id]],
        return_dict_in_generate=True,
    )
    print(f"Inference time: {time.time() - start:.4f}s")

    # turn into JSON
    start = time.time()
    seq = processor.batch_decode(outputs.sequences)[0]
    seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
    seq = re.sub(r"<.*?>", "", seq, count=1).strip()  # remove first task start token
    seq = processor.token2json(seq)
    if "class" not in seq.keys():
        seq["class"] = "other"
    print(f"Decoding time: {time.time() - start:.4f}s")

    gt = sample["labels"]
    score = float(seq["class"] == gt["class"])

    accs.append(score)

acc_score = np.mean(accs)

print(f"Accuracy: {acc_score * 100:.4f}%")
```

Can someone look into this issue? Thank you very much.

Comments:
cc @gante and @younesbelkada 

@gante and @younesbelkada can you guys look into this issue? Thank you very much.

Hey @MS1908 👋 

Text generation can be quite slow. 6-7s is within the expected time for large-ish models, and the generation time greatly depends on the length of the output. As a reference, [the example in the documentation](https://huggingface.co/docs/transformers/main/en/model_doc/donut#inference) takes ~1s on `.generate()` on an nvidia 3090, and the output is very short.

My top recommendation would be to batch your inputs, as opposed to calling `.generate()` with one example at a time. The execution time of generate grows very slowly with the batch size -- the biggest limitation is GPU memory, which you have plenty :D

On the `generate` side, we have some speedup tricks like using smaller variable representation. Sadly, as far as I know, most of them don't work out of the box with multimodal models like Donut (is this correct, @younesbelkada?). The only option that I see is to use [PT2.0+dynamo](https://pytorch.org/docs/stable/dynamo/index.html) to compile `.generate()`.

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.