Speed up repetition penalty logits processor
Author: bdalal
Pull request: no

# 🚀 Feature request
Hey Team, Thanks for the great work on the project to date! This is more of an enhancement so putting it here instead of as a bug.

<!-- A clear and concise description of the feature proposal.
     Please provide a link to the paper and code in case they exist. -->
The [`RepetitionPenaltyLogitsProcessor`](https://github.com/huggingface/transformers/blob/a1bbcf3f6c20e15fe799a8659d6b7bd36fdf11ed/src/transformers/generation_logits_process.py#L147-L155) which is used to enforce the repetition penalty when generating tokens from a Seq2Seq head is extremely slow for long sequences due to the nested loops.

A vectorized implementation will be much much faster.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request
     related to a problem? e.g., I'm always frustrated when [...]. If this is related
     to another GitHub issue, please link here too. -->
Here's a minimal example to reproduce the slow behavior:
```
import torch
from transformers import RepetitionPenaltyLogitsProcessor

import timeit

def vectorized(input_ids, scores, penalty):
    score_range = torch.arange(scores.shape[0])
    score = scores[score_range[:, None], input_ids]
    score[score >= 0] = score[score >= 0] / penalty
    score[score < 0] = score[score < 0] * penalty
    scores[score_range[:, None], input_ids] = score

input_ids = torch.randint(0, 10000, (256, 256))
scores = torch.randn(256, 10000)

rep_proc = RepetitionPenaltyLogitsProcessor(2.0)

print(f"Existing impl time for 10 iterations on CPU = {timeit.timeit(lambda: rep_proc(input_ids, scores), number=10)}")
print(f"Vectorized impl time for 10 iterations on CPU = {timeit.timeit(lambda: vectorized(input_ids, scores, 2.0), number=10)}")

if torch.cuda.is_available():
    input_ids = input_ids.cuda()
    scores = scores.cuda()

    print(f"Existing impl time for 10 iterations on GPU = {timeit.timeit(lambda: rep_proc(input_ids, scores), number=10)}")
    print(f"Vectorized impl time for 10 iterations on GPU = {timeit.timeit(lambda: vectorized(input_ids, scores, 2.0), number=10)}")
```
Here's the speedups on CPU and GPU with the vectorized version:
```
Existing impl time for 10 iterations on CPU = 23.23520456800179
Vectorized impl time for 10 iterations on CPU = 0.035849231004249305
```
```
Existing impl time for 10 iterations on GPU = 42.0977192690043
Vectorized impl time for 10 iterations on GPU = 0.008036320999963209
```
These numbers are from a machine with 16vCPUs@2.3GHz and an Nvidia T4 GPU.

## Your contribution

<!-- Is there any way that you could help, e.g. by submitting a PR?
     Make sure to read the CONTRIBUTING.MD readme:
     https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md -->
I'll have a PR up for this shortly.

@patrickvonplaten pinging you on this for your thoughts because I saw your last few commits on this code.

Thanks!

Comments:
PR is up [8598](https://github.com/huggingface/transformers/pull/8598)

Can you please check speed of this implementation also:
```
# if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability
scores = torch.where(scores < 0, scores * penalty, scores / penalty)
```

@LSinev Thanks for the suggestion! This is slightly faster than my implementation and much more elegant. Also makes sense because we don't really need to look at previous tokens for modifying the score. 
I'll replace my implementation with this. Credit to you!

Oh, no. I didn't mean that this should be the only code used. Of course, only input ids should be penalized, not everything. And of course, one should ensure (and probably add tests) that this solution works per row for input batches (batches as input and batches after the expansion due to num_return_sequences > 1). This was just a hypothesis that such code may work faster.

Hmm, that makes sense. I hadn't considered that earlier. All - greedy and beam search and sampling may produce incorrect tokens because only the `torch.where` approach will alter scores for all other tokens.
Having accounted for only the input_ids, the `torch.where` approach is still marginally faster than performing the score modification in 2 steps.
I'll update the code accordingly. Thanks for your help!