feat: Sequential beam search(a.k.a Low-memory beam search)

[Saibo-creator]

What does this PR do?

This is to address issue in #22639
This PR is based on the idea from
#22639 (comment)

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

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?

Fixes #22639

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@gante @ArthurZucker

[gante]

Hey @Saibo-creator 👋 My apologies for the delayed response, I was out of office 🌴

Thank you for opening this PR! Looking at the code, in principle, it looks good to me. The only thing I would change is the flag to trigger this low-memory behavior. We have a flag for that in contrastive search, so we could repurpose it -- then it is True, we would run one beam at a time. That way we avoid adding a new argument to generate (which already has too many arguments).

If you update your code to reuse this existing flag, I'd be happy to approve the PR :) I suspect the code could also be simplified as a result of this change.

P.S.: a test is also missing, to confirm that low-memory code path yields the same results. LMK if you need a hand with that :)

[github-actions]

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 are likely to be ignored.

[Saibo-creator]

repurpose

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

[Saibo-creator]

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:

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

[gante]

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 :)

[gante]

Happy to approve the PR as soon as we replace the added test by a mixin test :)

[gante]

Instead of adding an integration test, I'd like to request adding a mixin -- like this one for contrastive search with low memory.

That way, we can ensure it is fully compatible with all models :) Like in the test linked above, it is okay to skip some models if it is failing.

[Saibo-creator]

Sure, sounds good, I will do it ;)

One question, should I add @slow for this mixin test ?

[gante]

@Saibo-creator no need to add the @slow decorator -- we reserve that decorator for tests that use large models. In the mixing tests, we are using small and randomly initialized models :)

[Saibo-creator]

Hello @gante , I would like to have your opinions over these failure cases

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!

[Saibo-creator]

Relevant discussions:

• https://discuss.huggingface.co/t/results-of-model-generate-are-different-for-different-batch-sizes-of-the-decode-only-model/34878
• llama : understand why GPU results are different for different batch sizes ggerganov/llama.cpp#3014
• model generate with different batch size but get different results #23017, seems to be related to FP16!

[gante]

Hey @Saibo-creator 👋 We might be seeing the same effect as described here

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 :)

[Saibo-creator]

Hello @gante , thanks for your feedback, I tried runningpytest-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")

[gante]

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 :)

[Saibo-creator]

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) 

[gante]

@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

[Saibo-creator]

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.

[Saibo-creator]

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

[Saibo-creator]

Hello @gante ，
After a whole afternoon debugging, I found the bug, it's a tiny bug in this line. 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]

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

[JulesGM]

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

[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.

[JulesGM]

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

[Saibo-creator]

@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 ?

[ArthurZucker]

LGTM, thanks for you patience!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[gante]

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 🤗

[gante]

@Saibo-creator thank you for iterating with us 💛

[Saibo-creator]

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.

[gante]

@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]

@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 ?

[JulesGM]

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.

[JulesGM]

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.

[gante]

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

[Saibo-creator]

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

👍🏻 I will make a PR shortly

[JulesGM]

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.

[Saibo-creator]

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!

[Saibo-creator]

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 :)

[thanhlecongg]

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:

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

[Saibo-creator]

Hello @thanhlecongg , Thanks for the question!

You are looking at the right place.

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

[thanhlecongg]

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.