speed up PrefixConstrainedLogitsProcessor
Author: vukasin
Pull request: yes

# Speed up PrefixConstrainedLogitsProcessor

The `__call__` method creates an empty tensor and then updates it repeatedly. This causes a significant slowdown when executing on a GPU since it requires updating the GPU memory over and over. This pull request creates the mask in CPU memory and then creates a tensor on GPU with one call. The tests I ran (the test case from https://github.com/worldbank/REaLTabFormer/tree/main/colab) has a speed increase of over 200% up to depending on the batch size, 300%. 




## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
cc @gante @zucchini-nlp 

@gante @zucchini-nlp Hi, is there anything else you need me to do before this can be reviewed? 

no particular reason, I can change it to use torch instead

updated the code so it keeps to torch functions and only manipulates the device

@zucchini-nlp good call, this just became a pretty trivial change :-) but still has the same speedup (in my tests at least)

@zucchini-nlp anything else I need to do to get this moving forward?

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_35275). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@gante anything else I need to do?

Hi @vukasin 

Thank you for your PR! For future reference, please share a stand-alone benchmarking script so that everyone can immediately assess the impact of the change :)

In this particular case, the proposed change actually results in a slowdown if everything is done as efficiently as possible. The best speed is achievable when `prefix_allowed_tokens_fn` consists of tensor manipulations, where the tensors never leave the GPU -- the slowdown you observe is because we are slicing `mask`, a tensor on the GPU, with a tensor located on the CPU. If `prefix_allowed_tokens_fn` relies on CPU operations, then your change is slightly faster (as measured on my machine).

I would accept two changes:
1 - document this important property that `prefix_allowed_tokens_fn` should have for optimal performance 🤗
2 - vectorize the `mask`  update: collect all `prefix_allowed_tokens` to then update `mask` once. Because the update is done once, it should be faster when `prefix_allowed_tokens_fn` relies on CPU ops (at no cost of performance on the GPU case)

________________________________________________________
For full reference, here's a benchmarking script

```py
import torch
import math
import time
from tqdm import tqdm

NUM_BEAMS = 4
BATCH_SIZE = 2
VOCAB_SIZE = 50256

def prefix_allowed_tokens_fn_cpu(batch_id, input_ids):
    return input_ids.tolist()

def prefix_allowed_tokens_fn_gpu(batch_id, input_ids):
    return input_ids


def cpu_mask(input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
    mask = torch.full_like(scores, -math.inf, device=torch.cpu.current_device())
    for batch_id, beam_sent in enumerate(input_ids.view(-1, NUM_BEAMS, input_ids.shape[-1])):
        for beam_id, sent in enumerate(beam_sent):
            prefix_allowed_tokens = prefix_allowed_tokens_fn_cpu(batch_id, sent)
            if len(prefix_allowed_tokens) == 0:
                raise ValueError(
                    f"`prefix_allowed_tokens_fn` returned an empty list for batch ID {batch_id}."
                    f"This means that the constraint is unsatisfiable. Please check your implementation"
                    f"of `prefix_allowed_tokens_fn` "
                )
            mask[batch_id * NUM_BEAMS + beam_id, prefix_allowed_tokens] = 0

    scores_processed = scores + mask.to(scores.device)
    return scores_processed


def gpu_mask(input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
    mask = torch.full_like(scores, -math.inf)
    for batch_id, beam_sent in enumerate(input_ids.view(-1, NUM_BEAMS, input_ids.shape[-1])):
        for beam_id, sent in enumerate(beam_sent):
            prefix_allowed_tokens = prefix_allowed_tokens_fn_gpu(batch_id, sent)
            if len(prefix_allowed_tokens) == 0:
                raise ValueError(
                    f"`prefix_allowed_tokens_fn` returned an empty list for batch ID {batch_id}."
                    f"This means that the constraint is unsatisfiable. Please check your implementation"
                    f"of `prefix_allowed_tokens_fn` "
                )
            mask[batch_id * NUM_BEAMS + beam_id, prefix_allowed_tokens] = 0

    scores_processed = scores + mask
    return scores_processed


cpu_times = []
gpu_times = []
for _ in tqdm(range(1000)):
    input_ids = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE * NUM_BEAMS, 5)).to("cuda")
    scores = torch.randn(BATCH_SIZE * NUM_BEAMS, VOCAB_SIZE).to("cuda")

    start = time.time()
    cpu_mask(input_ids, scores)
    end = time.time()
    cpu_times.append(end - start)

    start = time.time()
    gpu_mask(input_ids, scores)
    end = time.time()
    gpu_times.append(end - start)

print("CPU time:", sum(cpu_times) / len(cpu_times))
print("GPU time:", sum(gpu_times) / len(gpu_times))
```