Scoring translations is unacceptably slow
Author: erip
Pull request: no

### System Info

- `transformers` version: 4.29.0
- Platform: Linux-3.10.0-862.11.6.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.9.16
- Huggingface_hub version: 0.12.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No


### Who can help?

@ArthurZucker and @younesbelkada

### 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. Install transformers, pytorch, tqdm
2. Create `forced_decode.py` [^1] 
3. Create `repro.sh` [^2]
4. Run `bash repro.sh` and observe extremely slow scoring speeds.



[^1]: 
    ```python
    #!/usr/bin/env python3
    
    import itertools
    from argparse import ArgumentParser, FileType
    from tqdm import tqdm
    
    import torch
    
    from transformers import PrefixConstrainedLogitsProcessor, AutoTokenizer, AutoModelForSeq2SeqLM
    
    def setup_argparse():
        parser = ArgumentParser()
        parser.add_argument("-t", "--tokenizer", type=str, required=True)
        parser.add_argument("-m", "--model", type=str, required=True)
        parser.add_argument("-bs", "--batch-size", type=int, default=16)
        parser.add_argument("-i", "--input", type=FileType("r"), default="-")
        parser.add_argument("-o", "--output", type=FileType("w"), default="-")
        parser.add_argument("-d", "--delimiter", type=str, default="\t")
        parser.add_argument("--device", type=str, default="cpu")
        return parser
    
    
    def create_processor_fn(ref_tokens_by_segment):
        def inner(batch_id, _):
            return ref_tokens_by_segment[batch_id]
        return inner
    
    def tokenize(src, tgt, tokenizer):
        inputs = tokenizer(src, text_target=tgt, padding=True, return_tensors="pt")
        return inputs
    
    def forced_decode(inputs, model, num_beams=5):
        inputs = inputs.to(model.device)
        logit_processor = PrefixConstrainedLogitsProcessor(create_processor_fn(inputs["labels"]), num_beams=num_beams)
        output = model.generate(**inputs, num_beams=num_beams, logits_processor=[logit_processor], return_dict_in_generate=True, output_scores=True)
        return output.sequences_scores.tolist()
    
    def batch_lines(it, batch_size):
        it = iter(it)
        item = list(itertools.islice(it, batch_size))
        while item:
            yield item
            item = list(itertools.islice(it, batch_size))
    
    if __name__ == "__main__":
        args = setup_argparse().parse_args()
        f_tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
        f_model = torch.compile(AutoModelForSeq2SeqLM.from_pretrained(args.model).to(args.device))
    
        with args.input as fin:
            inputs = list(batch_lines(map(str.strip, fin), args.batch_size))
    
        inputs_logits = []
        for batch in tqdm(inputs):
            src, tgt = zip(*[line.split(args.delimiter) for line in batch])
            inputs_logits.append(tokenize(src, tgt, f_tokenizer))
    
        with args.output as fout, torch.no_grad():
            for input in tqdm(inputs_logits):
                scores = forced_decode(input, f_model)
                print(*scores, sep="\n", file=fout)
    ```

[^2]:

    ```bash
    #!/usr/bin/env bash
    
    function get_input {
      curl -s https://gist.githubusercontent.com/erip/e37283b8f51d4e2c16996fc8a6a01aa7/raw/f5a3daffb04dad76464188c2a6949649f5cf3f9c/en-de.tsv
    }
    
    python forced_decode.py \
       -t Helsinki-NLP/opus-mt-en-de -m Helsinki-NLP/opus-mt-en-de \
       -i <(get_input) \
       --device cuda:0 \
       -bs 16
    ```


### Expected behavior

Scoring should be _very fast_ since the beam doesn't actually need to be searched, but I'm finding speeds on the order of seconds per batch which is far slower than generating.

Comments:
cc @gante 

Hey @erip 👋 

Sadly, I'm out of bandwidth to dive into the performance of very specific generation modes (in this case, beam search with `PrefixConstrainedLogitsProcessor`). If you'd like to explore the issue and pinpoint the cause of the performance issue, I may be able to help, depending on the complexity of the fix.

Meanwhile, I've noticed that you use `torch.compile`. I would advise you not to use it with text generation, as your model observes different shapes at each forward pass call, resulting in potential slowdowns :) 

Cheers, @gante. I'll try removing the compilation to see how far that moves the needle. I'm trying to score ~17m translations which tqdm is reporting will take ~50 days so we'll see what the delta is without `torch.compile`. I'll post updates here as well.

Edit: 96 days w/o `torch.compile` :-)

@erip have you considered applying 4-bit quantization ([docs](https://huggingface.co/docs/transformers/v4.31.0/en/main_classes/quantization#load-a-large-model-in-4bit), reduces the GPU ram requirements to ~1/6 of the original size AND should result in speedups) and then increasing the batch size as much as possible?  

You may be able to get it <1 week this way, and the noise introduced by 4 bit quantization is small.

I guess I'm more concerned that this is going to take a lot of time at all. Fairseq, Marian, and Sockeye can score translations extremely quickly (17m would probably take ~1-2 days on similar hardware). Transformers can translate in that amount of time, so I'm lead to conclude that logits processors are just performance killers.

@erip some of them are performance killers (e.g. `PrefixConstrainedLogitsProcessor ` seems to need vectorization). Our Pytorch beam search implementation is not optimized either, compared to our TF/FLAX implementation.

We focus on breadth of techniques and models, but welcome optimization contributions 🤗 

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.