feat: Sequential beam search(a.k.a Low-memory beam search)
Author: Saibo-creator
Pull request: yes

# What does this PR do?
This is to address issue in https://github.com/huggingface/transformers/issues/22639 
This PR is based on the idea from
https://github.com/huggingface/transformers/issues/22639#issuecomment-1507525155

The original implementation of 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 as effective batch size. This implementation is not necessary per se and can consume a lot of memory.

The new implementation split the full_batch(num_beam x batch size) inputs into a list of reduced_batch(beam_search_batch_size), run them sequentially and concat them back to a single model_output object. It involves two helper function:
-  `def concat_model_outputs(objs: List[ModelOutput]) -> ModelOutput` 
-  `def split_model_inputs( obj: Union[ModelOutput, Dict], split_size: int, full_batch_size: int ) -> List[Union[ModelOutput, Dict]]`

The new implementation can be used in 4 decoding methods: `beam_search`, `beam_sample`, `group_beam_search` and `constrained_beam_search`

The expected behavior is that it produces exactly the same output(logits) as the original implementation.

I tested it with the following quick test and it works regardless of input

```python

from
transformers import GPT2Tokenizer, AutoModelForCausalLM
import numpy as
np
from transformers import (
    AutoTokenizer,

AutoModelForSeq2SeqLM,
    LogitsProcessorList,

MinLengthLogitsProcessor,
    BeamSearchScorer,
)

tokenizer =
GPT2Tokenizer.from_pretrained("gpt2")
model =
AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer.pad_token_id =
tokenizer.eos_token_id
model_inputs = tokenizer('I enjoy walking with my
cute dog', return_tensors='pt')

# activate beam search and
early_stopping
beam_output = model.generate(
    **model_inputs,

max_new_tokens=40,
    num_beams=5,

early_stopping=True
)

print("Output:\n" + 100 *
'-')
print(tokenizer.decode(beam_output[0],
skip_special_tokens=True))

beam_output_w_subbatch = model.generate(

**model_inputs,
    max_new_tokens=40,
    num_beams=5,

early_stopping=True,
    beam_search_batch_size=1
)

print("Output:\n" +
100 * '-')
print(tokenizer.decode(beam_output_w_subbatch[0],
skip_special_tokens=True))

assert (beam_output ==
beam_output_w_subbatch).all(), "Beam search results from sub batch and
full batch are different"
```

TODO: only did it for pytorch models, if you think this PR is promising, I can do it for tensorflow too.

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->


Fixes #22639 


## 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#start-contributing-pull-requests),
      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.
- [ ] 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?

@gante @ArthurZucker 


Comments:
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.

> repurpose

Hey @gante , sorry I just saw your message. I agree with your suggestion and I will update the code and add a test!


Hello @gante ,
I have reused the flag `low_memory` in the code and try to unify part of the code used in sequential beam search and sequential contrastive search. 
A test is also added.

What makes me confused is the test case where I failed:
```terminal
self = <tests.models.xglm.test_modeling_xglm.XGLMModelTest testMethod=test_tf_from_pt_safetensors>

    @is_pt_tf_cross_test
    def test_tf_from_pt_safetensors(self):
        for model_class in self.all_model_classes:
            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
    
            tf_model_class_name = "TF" + model_class.__name__  # Add the "TF" at the beginning
            if not hasattr(transformers, tf_model_class_name):
                # transformers does not have this model in TF version yet
                return
    
            tf_model_class = getattr(transformers, tf_model_class_name)
    
            pt_model = model_class(config)
    
            with tempfile.TemporaryDirectory() as tmpdirname:
                pt_model.save_pretrained(tmpdirname, safe_serialization=True)
                tf_model_1 = tf_model_class.from_pretrained(tmpdirname, from_pt=True)
    
                pt_model.save_pretrained(tmpdirname, safe_serialization=False)
                tf_model_2 = tf_model_class.from_pretrained(tmpdirname, from_pt=True)
    
                # Check models are equal
                for p1, p2 in zip(tf_model_1.weights, tf_model_2.weights):
>                   self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
E                   AssertionError: False is not true

tests/test_modeling_common.py:3246: AssertionError


FAILED tests/models/speech_to_text/test_modeling_speech_to_text.py::Speech2TextModelTest::test_tf_from_pt_safetensors - AssertionError: False is not true
FAILED tests/models/transfo_xl/test_modeling_transfo_xl.py::TransfoXLModelTest::test_tf_from_pt_safetensors - AssertionError: False is not true
FAILED tests/models/xglm/test_modeling_xglm.py::XGLMModelTest::test_tf_from_pt_safetensors - AssertionError: False is not true

Exited with code exit status 255
```

This doesn't seem to be directly related to my changes. Maybe you know what's wrong here ;D




Hey @Saibo-creator!

The failing test is indeed unrelated :) We have skipped the test while we work on it, rebasing this branch should yield a green CI :)

Hello @gante , I would like to have your opinions over these failure cases

```terminal
FAILED tests/models/blenderbot/test_modeling_blenderbot.py::BlenderbotModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[1, 51, 51, 51, 51], [1, 37, 37, 43, 43]] != [[1, 37, 37, 43, 43], [1, 74, 74, 74, 74]]
FAILED tests/models/bloom/test_modeling_bloom.py::BloomModelTest::test_beam_search_low_memory - RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 2 but got size 8 for tensor number 1 in the list.
FAILED tests/models/codegen/test_modeling_codegen.py::CodeGenModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[135, 102, 167, 94, 172, 103, 209], [233, 32, 101, 213, 234, 172, 103]] != [[135, 102, 167, 94, 172, 103, 209], [233, 32, 101, 203, 147, 47, 164]]
FAILED tests/models/clvp/test_modeling_clvp.py::ClvpDecoderTest::test_beam_search_low_memory - AssertionError: Lists differ: [[38, 235, 110, 144, 159], [93, 177, 267, 144, 55]] != [[38, 235, 110, 177, 144], [93, 177, 67, 235, 116]]
FAILED tests/models/git/test_modeling_git.py::GitModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[56, 96, 46, 38, 92, 42, 1], [25, 63, 12, 22, 49, 91, 19]] != [[56, 96, 46, 38, 92, 42, 1], [25, 63, 12, 45, 72, 77, 69]]
FAILED
FAILED tests/models/gpt_neox/test_modeling_gpt_neox.py::GPTNeoXModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[6, 34, 87, 20, 93, 20, 93], [82, 10, 76, 68, 62, 48, 10]] != [[6, 34, 87, 20, 93, 20, 93], [82, 10, 76, 68, 62, 68, 62]]
FAILED tests/models/gptj/test_modeling_gptj.py::GPTJModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[84, 2, 29, 55, 59, 47, 88], [50, 23, 76, 50, 76, 16, 20]] != [[84, 2, 29, 94, 76, 16, 20], [50, 23, 76, 50, 94, 76, 50]]
FAILED
FAILED tests/models/imagegpt/test_modeling_imagegpt.py::ImageGPTModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[57, 97, 6, 36, 77, 62, 18], [11, 23, 94, 85, 88, 81, 61]] != [[57, 97, 6, 36, 77, 62, 18], [11, 23, 94, 54, 9, 72, 18]]
FAILED tests/models/llama/test_modeling_llama.py::LlamaModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[83, 48, 56, 54, 56, 54, 56], [91, 84, 33, 72, 56, 54, 56]] != [[83, 48, 56, 54, 56, 54, 56], [91, 84, 33, 72, 56, 70, 78]]
FAILED tests/models/mega/test_modeling_mega.py::MegaModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[53, 16, 33, 21, 21, 34, 70], [82, 45, 21, 21, 34, 34, 70]] != [[53, 16, 33, 21, 34, 34, 70], [82, 45, 21, 21, 34, 34, 70]]
FAILED tests/models/mistral/test_modeling_mistral.py::MistralModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[52, 24, 42, 54, 4, 34, 21], [31, 71, 18, 5, 27, 28, 52]] != [[52, 24, 42, 54, 4, 34, 21], [31, 71, 18, 5, 73, 73, 73]]
FAILED tests/models/mpt/test_modeling_mpt.py::MptModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[8, 13, 85, 19, 19, 19, 19], [35, 61, 88, 19, 19, 19, 19]] != [[8, 13, 85, 19, 19, 19, 19], [35, 61, 88, 92, 92, 92, 92]]
FAILED tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[81, 83, 11, 58, 81, 28, 26], [54, 81, 18, 38, 18, 38, 18]] != [[81, 83, 11, 85, 58, 13, 38], [54, 81, 18, 38, 18, 38, 18]]
FAILED tests/models/phi/test_modeling_phi.py::PhiModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[28, 14, 49, 19, 56, 62, 50], [9, 30, 60, 19, 19, 19, 57]] != [[28, 14, 49, 37, 88, 35, 80], [9, 30, 60, 19, 57, 9, 57]]
```

The mixin tests reveal that a subset of models will generate different outputs over `sequential beam search` and `beam search`, though the difference is not very big.

It seems all these models are recent llm, such as `phi`, `llama`, `mistral` etc. While for other models, the outputs are identical.

I have tried debugging locally, and I confirm that:
- the beam search is deterministic across multi-runs
- the sequential beam search is also deterministic across multi-runs

So indeed these models's outputs differ in the SBS vs BS

Do you have any insights into what may have caused this difference?

On interesting question is : do these models give same outputs when running with different batch sizes ? Because SBS is just like unbatched BS.

Thank you!



Relevant discussions:
- https://discuss.huggingface.co/t/results-of-model-generate-are-different-for-different-batch-sizes-of-the-decode-only-model/34878
- https://github.com/ggerganov/llama.cpp/issues/3014
- https://github.com/huggingface/transformers/issues/23017, seems to be related to FP16!



Hey @Saibo-creator 👋 We might be seeing the same effect as described [here](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535)

If you run the test multiple times on a single model, from the list of failing tests above, do you always get a failure? In other words, is the test stable or flaky? 

To see whether a test is flaky, install
```
pip install pytest-flakefinder
```

and then run
```
pytest --flake-finder --flake-runs=100 tests/test_failing_test.py
```

👉 Since the models in the mixin are randomly initialized, and the inputs are random as well, all runs are different. If the test is flaky, then the mismatch is probably due to the effect described in the threat at the beginning of this comment, and there are quick solutions for the test :)

Hello @gante , thanks for your feedback, I tried running`pytest-flakefinder` and indeed the tests pass 7 out of 100 times and fails the other 93 times. 

I think it's indeed due to the effect you explained in the thread(btw, bravo for the detailed explanation!) 

What is your solution for it ? :)

In addition to the flaky failure case, there are some other failure cases, mostly due to model specific implementations, I dive into some of them:
- bloom: special implementation of past key,value tensor shape
- ctrl: TODO
- ~fsmt~: old model with different cache format, won't fix
- gpt_bigcode: due to non-standard implementation of past key values , easy to fix
- reformer: old model with different cache format, won't fix
- transfo_xl: TODO
- xlnet: TODO
- cpm: TODO

I would like to know what is the general principle about these specific models? On the one hand, we could try to add if-else in the code to handle their specificity, but this also decreases the code quality a bit and our energy.

Let me know which models do you think we must have them supported. :)

(just saw this from https://github.com/huggingface/transformers/blob/main/tests/generation/test_utils.py#L1979)
>  3. TODO (joao): A few models have different formats, skipping those until the cache refactor is complete
>  models_without_standard_cache = ("bloom", "ctrl", "fsmt", "gptbigcode", "mega", "reformer")
        






Hey @Saibo-creator!

Regarding the model-specific implementations, feel free to ignore them for now (including skipping the tests, as you pasted at the end of your comment) :)

Regarding flakiness: you mentioned "the tests pass 7 out of 100 times and fails the other 93 times". Is this per model, or when you test all models at once? 
👉 If it is the latter, for all models at once, then it means the per model failure rate is low -- we can simply add the `@is_flaky` test decorator, adding a comment pointing to the comment I linked. 
👉 If it is the former, per model, then the failure rate is super high! There may be something else that we must uncover before merging :)

Hey @gante 

I think it's the former case, per model... 

I need to dive deep into the problem to understand why..

If you are interested,  here is what is happening, If I run `pytest --flake-finder --flake-runs=100 tests/models/llama/test_modeling_llama.py`,

 I get  ` 93 failed, 8307 passed, 4700 skipped, 1613 warnings in 930.30s (0:15:30) ` 

and the failing cases are from `test_beam_search_low_memory` like:
```
FAILED tests/models/llama/test_modeling_llama.py::LlamaModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[7, 68, 64, 28, 26, 37, 49], [52, 58, 92, 50, 7, 37, 49]] != [[7, 68, 64, 7, 37, 49, 37], [52, 58, 92, 71, 41, 7, 41]]
FAILED tests/models/llama/test_modeling_llama.py::LlamaModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[32, 79, 18, 50, 94, 48, 93], [52, 51, 3, 56, 24, 12, 24]] != [[32, 79, 18, 50, 94, 5, 25], [52, 51, 3, 56, 24, 12, 24]]
FAILED tests/models/llama/test_modeling_llama.py::LlamaModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[39, 97, 28, 57, 0, 0, 0], [34, 69, 91, 34, 63, 0, 0]] != [[39, 97, 28, 57, 0, 0, 0], [34, 69, 91, 5, 58, 18, 14]]
```

So this is 93/ 100 where the two outputs are different


For the other models such as `GPT2`, I also get som failure cases, though not frequent
with `pytest --flake-finder --flake-runs=10 tests/models/gpt2/test_modeling_gpt2.py`

```
FAILED tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[29, 58, 3, 68, 68, 68, 68], [90, 65, 45, 45, 68, 68, 68]] != [[29, 58, 3, 68, 68, 68, 68], [90,...
FAILED tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[18, 32, 29, 29, 87, 87, 87], [40, 85, 64, 64, 30, 30, 30]] != [[18, 32, 29, 87, 87, 87, 87], [4...
FAILED tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_low_memory - AssertionError: Lists differ: [[98, 79, 92, 9, 9, 9, 9], [46, 16, 93, 9, 9, 9, 9]] != [[98, 79, 92, 9, 9, 9, 9], [46, 16, 93, 9...
=============================== 3 failed, 907 passed, 510 skipped, 197 warnings in 470.60s (0:07:50) 
```





@Saibo-creator sadly I have no easy advice here :( it could be a kernel-related numerical issue due to different shapes, like when using `past_key_values`, or it can be a subtle bug.

I see two paths forward:
1 - you are able to pin the source of the mismatch, and we can safely confirm that it is something like a numerical issue
2 - you run a few benchmarks over your PR -- if the resulting metrics are similar with the low memory mode, then there shouldn't be a bug

I spent some time investigating the reason for mismatch. I notice that if batch-size = 1, all tests pass without any failure(as shown in the latest commit)

So it seems this is not a numerical issue but more like something to do with batch processing. 

I will try to figure out if my implementation has a bug with batched input.

I confirm that if we disable `use_cache`, the output are identical regardless of `batch_size`
So now we identified the origin was from the `key_value_cache`

Hello @gante ，
     After a whole afternoon debugging, I found the bug, it's a tiny bug [in this line](https://github.com/epfl-dlab/transformers-GCD-PR/blob/fix_issue_22639/src/transformers/generation/utils.py#L4955). I somehow wrote `for i in range(0, full_batch_size, 1)` instead of `for i in range(0, full_batch_size, split_size)`. The former makes no sense at all so I guess it was just a typo, but hidden very well....

I'm happy that it works now : )

(The failure cases are irrelevant)

@Saibo-creator any reason this hasn't been merged?

or @gante ?

Thanks for implementing it btw @Saibo-creator 

I could really use this

> @Saibo-creator any reason this hasn't been merged?

Hello Jules! Sorry I was a bit out of bandwidth last week, I'm fixing the small things pointed by Arthur now and will push immediately : )

@ArthurZucker  I updated  code to incorporate your feedbacks. Let me know any other things to do 

Is there a way the batch size of the beam-search could be made separate from `num_return_sequences`?

For example, for RL, I would like to generate let's say 64 different outputs with a single input sequence. However, my hardware (multiple gpu or not) hardware can't handle more than 16 beams at once.

From trying to run the code, it seems like I couldn't get those 64 outputs? Just 16?

@JulesGM 


> For example, for RL, I would like to generate let's say 64 different outputs with a single input sequence. However, my hardware (multiple gpu or not) hardware can't handle more than 16 beams at once.

In your case, say you have a batch size = 1(single input), you want to do beam search of k=64 with num_return_sequences=64 and the maximum number of sequences your gpu can handle in parallel is 16.
If you do beam_search with `low_memory=True` then this method will divide your beam into 64 batches(each of 1 sequence), and run them sequentially. 
Currently there is no argument to specify the `parallel_size`, it's set to be identical to the  `batch_size`, this is something we could discuss if we want to add. It involves adding a new argument to the generation function.



> From trying to run the code, it seems like I couldn't get those 64 outputs? Just 16?

That is something I don't expect to happen. If you set beam search k =64 and `num_return_seq`=64, you should get 64 outputs, regardless you have `low_memory` option True or False.  Am I wrong ? Could you provide a reproducible example ?





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

Merging as it's approved as it solves the original target use case.

@JulesGM don't refrain from commenting if it doesn't solve your use case: we may still open a follow-up PR to further refine this 🤗 

@Saibo-creator thank you for iterating with us 💛 

> Merging as it's approved as it solves the original target use case.
> 
> @JulesGM don't refrain from commenting if it doesn't solve your use case: we may still open a follow-up PR to further refine this 🤗

@JulesGM If you think it's necessary,  we could consider adding more flexibility by allowing users to specify the `parallel_size` for example. It would be an easy add-on. 

@Saibo-creator alternatively, a try/except loop could be tried (try with `batch_size` -> if it fails, try with `batch_size/2` -> ... -> try with batch_size = 1). I'd prefer that to an additional flag -- we already have too many flags in `generate` :)

> @Saibo-creator alternatively, a try/except loop could be tried (try with `batch_size` -> if it fails, try with `batch_size/2` -> ... -> try with batch_size = 1). I'd prefer that to an additional flag -- we already have too many flags in `generate` :)

Thanks for the suggestion @gante  ! I think this would be a good improvement to reduce work from user side. Should I open another PR to proceed ?

> > Merging as it's approved as it solves the original target use case.
> > @JulesGM don't refrain from commenting if it doesn't solve your use case: we may still open a follow-up PR to further refine this 🤗
> 
> @JulesGM If you think it's necessary, we could consider adding more flexibility by allowing users to specify the `parallel_size` for example. It would be an easy add-on.

I think that this would be great. In my personal fork, I currently look at if low_memory is an int, and use it as a sub-batch size if it is the case ahah. Ofc I can see arguments as to why a separate argument would be better, I'm not attached to this solution at all.

I also extended the code in my personal fork to beam_sample, basically copy-pasting the beam_search sequential section. I think that that would be a good thing to add. 

@Saibo-creator yeah, open a new PR, I'd be happy to merge it 🤗 

> @Saibo-creator yeah, open a new PR, I'd be happy to merge it 🤗 

👍🏻 I will make a PR shortly


So I noticed that in practice this didn't allow me to increase the beam size by that much at all if the sequence size is long. So I'm experimenting with offloading the key value cache to cpu & only loading the part of the cache that is needed, & it seems to really help.

> So I noticed that in practice this didn't allow me to increase the beam size by that much at all if the sequence size is long. So I'm experimenting with offloading the key value cache to cpu & only loading the part of the cache that is needed, & it seems to really help.

Oh, thanks for pointing out @JulesGM . I think your observation makes sense. Offloading the kv-cache should indeed be more efficient!

> So I noticed that in practice this didn't allow me to increase the beam size by that much at all if the sequence size is long. So I'm experimenting with offloading the key value cache to cpu & only loading the part of the cache that is needed, & it seems to really help.

Hey @JulesGM , this approach is really cool,  would you like to share a link to your implementation of this feature? I could benefit from it too :)

Hi @Saibo-creator, I have tried to run the model with `low_memory=True` with high. However, it is quite counter-intuitive that memory usage in `low-memory beam search` is higher than the original one. 

```
### Low_memory = True
Memory usage:  19.74514389038086
Time taken:  63.430684089660645


### Normal
Memory usage:  10.743547439575195
Time taken:  2.9829838275909424`
```

This is my code:

```python
import torch
from transformers import GPT2Tokenizer, AutoModelForCausalLM
import numpy as np
import time
from transformers import (
    AutoTokenizer,

AutoModelForSeq2SeqLM,
    LogitsProcessorList,

MinLengthLogitsProcessor,
    BeamSearchScorer,
)

def get_memory_usage():
    return torch.cuda.max_memory_allocated(device=None)/1024/1024/1024

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")
model.to('cuda')
tokenizer.pad_token_id = tokenizer.eos_token_id
model_inputs = tokenizer('I enjoy walking with my cute dog' * 50, return_tensors='pt')
model_inputs.to('cuda')

start = time.time()

beam_output = model.generate(**model_inputs,
                            max_new_tokens=40,
                            num_beams=100,
                            early_stopping=True
)

print("Output:\n" + 100 *'-')
print(tokenizer.decode(beam_output[0], skip_special_tokens=True))

# Uncomment the following lines to run the code with low_memory = True
# beam_output_w_subbatch = model.generate(
#     **model_inputs,
#     max_new_tokens=40,
#     num_beams = 100,
#     early_stopping = True,
#     low_memory = True
# )

# print("Output:\n" + 100 * '-')
# print(tokenizer.decode(beam_output_w_subbatch[0], skip_special_tokens=True))

print("Memory usage: ", get_memory_usage())
print("Time taken: ", time.time() - start)
```

I'm thinking that it is because you are stacking outputs on these [lines](https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L3231-L3235) before releasing memory. So it consume double memory (both in outputs_per_sub_batch and outputs) than expected. Do you have any thoughts on this?

Thanks. 

Hello @thanhlecongg , Thanks for the question!

You are looking at the right [place](https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L3231-L3235). 


## Why is the memory usage higher in `low_memory=True` than the vanilla beam search?
The operation of stacking `outputs_per_sub_batch` effectively doubles the memory usage of the model output.

In cases where the model output is small, this is not a problem. But in cases where the model output is long ( your case), the memory usage can overwhelm the memory usage for computing the model output.

For example, with beam size = 20, sequence length = 350, the `model_output.logits` is 20 x 350 x vocab_size(50K) x 4 bytes ~ 1.4GB. 


## This is counter-intuitive, when `low_memory=True` is supposed to reduce memory usage ?
This is not a big problem per se because depending on the use case, the output may not be huge and the memory bottleneck is the model size and the intermediate variables in forward pass.

One example of a use case where `low_memory=True` is useful is when the model is very large and the input is relatively short,
for example, use
-  `gpt2-xl`(1.2 Trillion params) instead of `gpt2` (150 Million Params)
-  `I enjoy walking with my cute dog * 2` instead of `I enjoy walking with my cute dog * 50`
- beam_size=4

Check this [gist](https://gist.github.com/Saibo-creator/f8eb37ef7ca4a61b8a9105bbbdc7fe4d) to verify.

In this case, the vanilla beam search will OOM (my GPU is 8GB) but `low_memory=True` will work.


## Fix 
One simple fix is to do the **stacking** operation in CPU memory and then move it back to GPU memory.

This will make sure the memory usage in case of `low_memory=True` is always less than the vanilla beam search.



@gante , do you think this is a valid issue? If yes, I can create a PR for this. 


Thanks for your quick response and clear explanation. I wonder if we really need to stack all model outputs or only top-k is enough? If only top-k is enough, the memory usage could be lower.