Compute logits and past_key_values once for the initial context rather than `num_beams` times?
Author: tom-p-reichel
Pull request: no

### Feature request

It looks like (at least) the following methods:
- sample
- beam search
- beam sample
- group beam search

all immediately expand their input batch size to either `num_beams` or `num_return_sequences`. This means that the initial context, which can be very long, will be run through the model `num_beams`/`num_return_sequences` times, which seems to be unnecessary and causes a large overhead, especially when the initial context is long.

It seems to be possible to run the model once on the initial context within `beam_search` and related functions and broadcast the model's outputs to the desired batch size rather than expand the batch beforehand and run a full batch of the same context.


### Motivation

I wrote a short hack in the contribution section of the report that fixes this redundancy, at least for beam search, and at least for decoder only models, etc-- it's not a complete PR, but it demonstrates the kind of change that should be made.

I also tested the performance using this script: 

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.set_default_device("cuda")
model = AutoModelForCausalLM.from_pretrained("PY007/TinyLlama-1.1B-step-50K-105b", load_in_4bit=True)
tokenizer = AutoTokenizer.from_pretrained("PY007/TinyLlama-1.1B-step-50K-105b")
inputs = tokenizer('''The unanimous Declaration of the thirteen united States of America, When in the Course of human events, it becomes necessary for one people to dissolve the political bands which have connected them with another, and to assume among the powers of the earth, the separate and equal station to which the Laws of Nature and of Nature's God entitle them, a decent respect to the opinions of mankind requires that they should declare the causes which impel them to the separation.

We hold these truths to be self-evident, that all men are created equal, that they are endowed by their Creator with certain unalienable Rights, that among these are Life, Liberty and the pursuit of Happiness.--That to secure these rights, Governments are instituted among Men, deriving their just powers from the consent of the governed, --That whenever any Form of Government becomes destructive of these ends, it is the Right of the People to alter or to abolish it, and to institute new Government, laying its foundation on such principles and organizing its powers in such form, as to them shall seem most likely to effect their Safety and Happiness. Prudence, indeed, will dictate that Governments long established should not be changed for light and transient causes; and accordingly all experience hath shewn, that mankind are more disposed to suffer, while evils are sufferable, than to right themselves by abolishing the forms to which they are accustomed. But when a long train of abuses and usurpations, pursuing invariably the same Object evinces a design to reduce them under absolute Despotism, it is their right, it is their duty, to throw off such Government, and to provide new Guards for their future security.--Such has been the patient sufferance of these Colonies; and such is now the necessity which constrains them to alter their former Systems of Government. The history of the present King of Great Britain is a history of repeated injuries and usurpations, all having in direct object the establishment of an absolute Tyranny over these States. To prove this, let Facts be submitted to a candid world.

''', return_tensors="pt", return_attention_mask=False)

import time

start = time.time()
for _ in range(10):
	outputs = model.generate(**inputs, max_new_tokens=15, do_sample=False, num_beams=6, num_return_sequences=6)
print("elapsed",(time.time()-start)/10)


text = tokenizer.batch_decode(outputs[:,-15:])
print(text)
```

Before the diff below was applied, this was the output:
```
elapsed 4.5888751745224
['We hold these truths to be self-evident, that all men', 'We hold these truths to be self evident, that all men are created', 'We hold these truths to be self-evident, that all Men', 'He has refused his Assent to Laws for establishing Judiciary', 'We hold these truths to be self-evident, that allmen', 'He has refused his Assent to Laws for establishing Judicial Systems']
```

After the diff below was applied, this was the output:
```
elapsed 2.614666700363159
['We hold these truths to be self-evident, that all men', 'We hold these truths to be self evident, that all men are created', 'We hold these truths to be self-evident, that all Men', 'He has refused his Assent to Laws for establishing Judiciary', 'We hold these truths to be self-evident, that allmen', 'He has refused his Assent to Laws for establishing Judicial Systems']
```

So, the change seems to give a ~1.7x speedup for this specific case while leaving the results of `beam_search` unchanged.

The speedup should be better for higher `num_beams` or longer contexts.

### Your contribution


As an example, I wrote a short hack that changes the behavior of `beam_search` on the main branch of huggingface so that it does not recompute `logits` and `past_key_values` `num_beams` times:
```diff
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
index 4dbfc3670..63f863c80 100644
--- a/src/transformers/generation/utils.py
+++ b/src/transformers/generation/utils.py
@@ -1821,13 +1821,6 @@ class GenerationMixin:
                 num_beam_hyps_to_keep=generation_config.num_return_sequences,
                 max_length=generation_config.max_length,
             )
-            # 12. interleave input_ids with `num_beams` additional sequences per batch
-            input_ids, model_kwargs = self._expand_inputs_for_generation(
-                input_ids=input_ids,
-                expand_size=generation_config.num_beams,
-                is_encoder_decoder=self.config.is_encoder_decoder,
-                **model_kwargs,
-            )
             # 13. run beam search
             return self.beam_search(
                 input_ids,
@@ -3142,12 +3135,7 @@ class GenerationMixin:
         batch_size = len(beam_scorer._beam_hyps)
         num_beams = beam_scorer.num_beams
 
-        batch_beam_size, cur_len = input_ids.shape
-
-        if num_beams * batch_size != batch_beam_size:
-            raise ValueError(
-                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
-            )
+        batch_beam_size, cur_len = batch_size*num_beams, input_ids.shape[1]
 
         # init attention / hidden states / scores tuples
         scores = () if (return_dict_in_generate and output_scores) else None
@@ -3171,6 +3159,7 @@ class GenerationMixin:
         beam_scores[:, 1:] = -1e9
         beam_scores = beam_scores.view((batch_size * num_beams,))
 
+        expander = lambda x : x.expand([num_beams,*x.shape[1:]])
         this_peer_finished = False  # used by synced_gpus only
         while True:
             if synced_gpus:
@@ -3192,6 +3181,13 @@ class GenerationMixin:
                 output_hidden_states=output_hidden_states,
             )
 
+
+            if num_beams * batch_size != input_ids.shape[0]:
+                # ok, NOW expand the batch dimension
+                # AFTER the hard part is done
+                outputs["logits"] = expander(outputs["logits"])
+                outputs["past_key_values"] = [[expander(y) for y in x] for x in outputs["past_key_values"]]
+
             if synced_gpus and this_peer_finished:
                 cur_len = cur_len + 1
                 continue  # don't waste resources running the code we don't need
@@ -3237,6 +3233,11 @@ class GenerationMixin:
             next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
             next_tokens = next_tokens % vocab_size
 
+            if input_ids.shape[0] != num_beams:
+                    input_ids = expander(input_ids)
+                    if "attention_mask" in model_kwargs:
+                            model_kwargs["attention_mask"] = expander(model_kwargs["attention_mask"])
+
             # stateless
             beam_outputs = beam_scorer.process(
                 input_ids,
```

a change would have to be made to each one of the affected functions.

I may be able to expand this into a PR if desired.

Comments:
cc @gante 

Hi @tom-p-reichel 👋 Thank you for opening this issue. It is something I really want to bring to `generate`! In fact, a few weeks ago at the PyTorch conference, they mentioned that isolating this stage (also known as prefill) is a significant source of speedups.

The devil is, as always, in the details. We can't move the input expansion inside the decoding functions (like `sample`), as the functions are part of our public API and it would be a breaking change. We can, however, create a `prefill` function in `GenerationMixin`, and call it before `# 10. go into different generation modes`. The prepared inputs in `prefill` would then be expanded on a need basis, before entering the decoding functions. This would preserve the API while enabling a faster prefill stage!

Would you be interested in working on this? 🤗 I'd be happy to provide pointers and reviews along the way!

@gante Sure, I can give this a try. Will send some initial questions soon.

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.

> The devil is, as always, in the details. We can't move the input expansion inside the decoding functions (like `sample`), as the functions are part of our public API and it would be a breaking change. We can, however, create a `prefill` function in `GenerationMixin`, and call it before `# 10. go into different generation modes`. The prepared inputs in `prefill` would then be expanded on a need basis, before entering the decoding functions. This would preserve the API while enabling a faster prefill stage!

Hi Joao! I think nowadays this idea is even more relevant as prefixes of thousands of tokens are common (tools/safety/context "help" and this will probably become even worse in the future). I got a noticeable speedup when a prefix is prefilled once before beam expansion (sometimes seconds of GPU time -- depending on the number of beams and length of course).

So since `_prefill()` now exists as a separate stage at GenerationMixin exactly like you foresaw, and that caches now support batch expansion, this may be easier to implement now cleanly.

If you agree, I'll be happy to work on this and add some tests to show it's a meaningful and non-breaking improvement. I already contributed some minor stuff to group beam search that was merged and am eager to improve the core transformers too!

Would you prefer `_prefill()` itself to also handle beam-related input and cache expansion, with the beam parameters passed as arguments? Or should `_prefill()` remain a generic helper, while the broader prefill stage around it handles those beam-specific (expansion-related) steps?

cc @Cyrilvallez : this issue was not completed, however the generation code has changed since then: `_prefill()` now exists as a separate helper and caches support batch expansion. Therefore, now it may be possible to implement the original idea cleanly exactly as @gante suggested: prefill the original batch once, then decently expand the inputs and cache for beam search. This can save quite a GPU time with longer prefixes (becoming an everyday reality) and many beams. Would you mind taking a look at my previous comment: https://github.com/huggingface/transformers/issues/27449#issuecomment-4959486740