[generate] ✨ vectorized beam search ✨
Author: gante
Pull request: yes

# What does this PR do?

✨ Vectorizes beam search ✨ 

Fixes #34843 
Fixes #35618 
Fixes #34574
Groundwork for: #35561 #35451 #35436 #32097 #30810
Other related issues: #35023

⚠️ Potential friction points:
1. This refactor removes the use of the `BeamScorer` inside `_beam_search()`. Advanced users relying on their custom `BeamScorer` instances should pin transformers to `<=4.48` for now. Custom beam scoring metrics will be added in a future PR;
2. The `low_memory=True` path [which runs the forward pass for each beam in sequence rather than in parallel] was not coded in the refactor. Since this flag was added, several memory-reduction techniques have been introduced, so I'm assuming it's niche now. It would also bloat the code significantly. An informative error message was added.

### Note to reviewers

Don't look at the diff for beam search, but rather read the new `_beam_search` as if it was a new function. One of goals of the refactor is improving code readability for future hacking (beam search is a form of test-time compute), and thus the  new code is heavily commented -- please add a review comment if a certain part of the code isn't immediately understandable, so we can improve it.

### Benefits of the vectorization:
1. 🚤 Much better scalability with respect to `num_beams` (see benchmarks below). Note that we can further improve these numbers -- `beam_search` should now be compatible with `torch.compile` (to be explored in a future PR);
3. 🔍 Simpler to read and understand -- it's still quite complex, but now we don't have custom state variables nor a large number of nested calls;
4. 🐛 Bugfix: Beam search now accepts arbitrary stopping criteria.

### Benchmarks 

<details>
  <summary>Test script (small model, beam search, no compilation)</summary>
  
  ```py
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

DEVICE = "cuda"
MODEL_IDS = [
    "Qwen/Qwen2.5-0.5B-Instruct",
]
NUM_BEAMS = [2, 8, 32, 128, 512]
NUM_RUNS = 5
NUM_WARMUP_RUNS = 2
MAX_NEW_TOKENS = 128

generation_kwargs = {
    "do_sample": False,
    "max_new_tokens": MAX_NEW_TOKENS,
    "min_new_tokens": MAX_NEW_TOKENS,  # forces the generation of `max_new_tokens`
}
memory = {}  # Max memory, MB
throughput = {}  # Throughput, tokens/sec

for model_id in MODEL_IDS:
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)

    inputs = tokenizer(["The quick brown"], return_tensors="pt").to(model.device)

    memory[model_id] = {}
    throughput[model_id] = {}

    for num_beams in NUM_BEAMS:
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        torch.cuda.reset_peak_memory_stats(DEVICE)
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

        # warmup
        for _ in range(NUM_WARMUP_RUNS):
            _ = model.generate(**inputs, **generation_kwargs, num_beams=num_beams)

        # measure
        start_event.record()
        for _ in range(NUM_RUNS):
            _ = model.generate(**inputs, **generation_kwargs, num_beams=num_beams)
        end_event.record()

        torch.cuda.synchronize()
        max_memory = torch.cuda.max_memory_allocated(DEVICE) * 1e-6
        tok_per_sec = (NUM_RUNS * MAX_NEW_TOKENS) / (start_event.elapsed_time(end_event) * 1.0e-3)
        memory[model_id][num_beams] = max_memory
        throughput[model_id][num_beams] = tok_per_sec

        print("\nNum Beams: ", num_beams)
        print("Max memory (MB): ", max_memory)
        print("Throughput (tokens/sec): ", tok_per_sec)

    print("\nModel: ", model_id)
    print("Num Beams: ", NUM_BEAMS)
    print("Max memory (MB): ", memory[model_id])
    print("Throughput (tokens/sec): ", throughput[model_id])
  ```
</details>

Measured on an RTX4090

Throughput (tokens/sec):
| num_beams | old | new | ratio (new/old, higher throughput is better) |
|---|---|---|---|
| 2 | 62.98 | 62.82 | 1.00 |
| 8 | 60.10 | 62.67 | 1.04 |
| 32 | 54.25 | 61.24 | 1.13 |
| 128 | 38.15 | 56.51 | 1.48 |
| 512 | 15.95 | 36.65 | 2.30 |

Memory (MB):
| num_beams | old | new | ratio (new/old, lower memory is better) |
|---|---|---|---|
| 2 | 1017.32 | 1016.15 | 1.00 |
| 8 | 1050.93 | 1045.92 | 1.00 |
| 32 | 1206.99 | 1188.17 | 0.98 |
| 128 | 1717.07 | 1641.76 | 0.96 |
| 512 | 3848.44 | 3547.72 | 0.92 |

### Tests run
checked box = no regressions vs `main`
- [x] `RUN_SLOW=1 py.test tests/models/gpt2/test_modeling_gpt2.py -vv` -- slow gpt2 tests (reference model for older implementations)
- [x] `RUN_SLOW=1 py.test tests/models/llama/test_modeling_llama.py -vv` -- slow llama tests (key model)
- [x] `RUN_SLOW=1 py.test tests/models/whisper/test_modeling_whisper.py -vv` -- slow whisper tests (key model, has hard beam search tests) [note: 13 failures on `main`]
- [x] `RUN_SLOW=1 py.test tests/models/t5/test_modeling_t5.py -vv` -- slow t5 tests (has hard beam search tests)
- [x] `RUN_SLOW=1 py.test tests/models/bart/test_modeling_bart.py -vv` -- slow bart tests (has hard beam search tests)
- [x] `RUN_SLOW=1 py.test tests/models/rag/test_modeling_rag.py -vv` -- slow rag tests (has hard beam search tests) [note: 4 failures on `main`]
- [x] `RUN_SLOW=1 py.test tests/utils/test_cache_utils.py -vv` -- slow cache tests (sanity check)

### Follow-up
- Deprecate `BeamScorer`
- Deprecate other niche beam methods (e.g. `constrained_beam_search`) -- maybe we can expose the scoring method, and write the other beam methods as a scoring modification?
- Remove documentation from deprecated methods and classes
- explore torch.compile (probably requires isolating prefil first)



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

@ArthurZucker @zucchini-nlp ready for a re-review, all comments addressed and/or answered.

Major changes:
1. (Arthur's suggestion) Major logic blocks were abstracted as functions;
2. (Arthur's suggestion) Use static shapes, as opposed to tensors that grow as we generate. This resulted in a small speedup, due to the removal of `torch.cat` ops;

After the changes, we get the same memory footprint and ~1% speedup (vs version you reviewed).

❓ Open question: What's more readable, functions as major abstraction blocks (current version) or no functions (version you reviewed before)? The performance is the same, readability is the key criteria. I have a very small preference for the "no functions" version, but I understand it might be too much to look at the whole beam search method at once -- I've spent dozens of hours debugging beam search, so it no longer looks daunting to me :D Your opinion has more weight here.

@zucchini-nlp @ArthurZucker PR comments addressed (namely, fns moved out `_beam_search` and re-introduced one set of `torch.cat`) 🤗