Have a beam search sub batch size to limit memory use
Author: JulesGM
Pull request: no

### Feature request

Currently, beam search effectively multiplies the batch size memory-wise and compute-wise by the batch size. If you have a batch size of 1 and a beam search of 8, `model.forward` sees 8 samples at once. This becomes an unnecessary problem for situations where needs that beam size but is only able to fit a smaller quantity of samples into memory.

That's why I suggest that a "beam search sub-batch" size be added to limit the number of samples (including beams) that are seen at once by the model. Eg, if the beam search batch size is 4, even if the main batch size is 1 and the beam search is 8, `model.foward` would only be called on 4 samples at once. Same for if the main batch size is 2 and the beam size is 8.

### Motivation

With the size of LLM & their popularity, people are often stretching the capacities of their hardware, & in some cases (like majority voting for chain-of-thoughts), need beam search with some beam size.

If that beam size can't fit all at once in the memory, they are currently doomed, which is unnecessary, as the beams could be computed in a smaller quantity at once.


Comments:
@sgugger 

cc @gante 

Hey @JulesGM! Thank you for suggesting sub-batch size beam search. It is somewhat related to #22340 (add some option to stop spending resources on beams that have already been finished).

Currently, we are unable to fulfill all requests, and I haven't seen demand for this particular feature. As such, I'll offer my standard pact: if/when this comment reaches 10 reactions, I'll put the feature on my todo list :) (and whoever does the 10th react, plz ping me)

Alternatively, if you'd like to implement this feature yourself, I'd be happy to guide you 🙌 

I was thinking something like taking https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L2833

```python
outputs = self(
    **model_inputs,
    return_dict=True,
    output_attentions=output_attentions,
    output_hidden_states=output_hidden_states,
)
```
Into, fairly naively splitting into sub-batches then concatenating should be fine:

```python

# Make it more explicit that all kwargs are shared by the two forward modes
# & make the code shorter
forward_kwargs = dict(                
    return_dict=True,
    output_attentions=output_attentions,
    output_hidden_states=output_hidden_states,
)

if beam_search_batch_size:
    bsbs = beam_search_batch_size  # Repeated a bunch of times, makes code too long
    outputs_per_sub_batch = []

    for i in range(0, seq_len, bsbs):
        model_inputs_sub_batch = {
             k: v[i * bsbs: (i + 1) * bsbs] 
             for k, v in model_inputs
         }
         sub_batch_outputs = self(
             **model_inputs_sub_batch,
             **forward_kwargs,
         )
         outputs_per_sub_batch.append(sub_batch_outputs)

    keys = outputs_per_sub_batch[0].keys()
    outputs = {
        k: torch.cat([sub_batch[k] for sub_batch in outputs_per_sub_batch] f, dim=0) 
        for k in keys
    }

else: # Unchanged original behavior
    outputs = self(
        **model_inputs,
        **forward_kwargs,
    )
```

This could likely be moved to a new function called "forward on beams" that could be called on `model_inputs` in all beam search variants.

I guess the one thing I'm not confident about is how to add the "beam_search_sub_batch" parameter to the `generate` call in a way that respects huggingface's .. interface objectives

Ideally, no parameterization would be needed (e.g. split the beam search batch size in two every time it hits a memory exception). But a try/except on the main code path would be ugly 🤔 

I am working on a plan to increase the flexibility of `.generate()`, hopefully this sort of features will become trivial to integrate 🙏 

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.