Contrastive Search peak memory reduction
Author: blbadger
Pull request: yes

# Contrastive Search memory reduction via sequential topk embedding recovery

This PR describes a new feature for contrastive search generation that may be of interest to the community.

Problem: 
Contrastive search is an effective method for LLM text generation, but as currently implemented requires far more maximum memory (tested with VRAM) than comparable methods like nucleus search.

Solution:
Extra memory (ie more than greedy generation) required for contrastive search primarily comes from two sources: storing of the last hidden layers per token and the parallelized computation of the last hidden layer embeddings for `top_k` tokens.  This PR addresses the second source by providing a switch to sequential last hidden layer computation.  

The result is that far less maximum memory is required during generation: for example, generation max memory usage for Llama 13b using Int4 quantization (@qlora) for 1k tokens reduces from >15GB to ~9GB.  The generation process necessarily becomes somewhat slower as well.

Code Snippet:
Pass the kwarg `low_memory` to the generation as follows:
```python
model.generate(input_ids, top_k=4, penalty_alpha=0.6, low_memory=True)
```

No additional modules are required for use.

## 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.
- [ x] 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
Anyone who would like to review is welcome!



Comments:
To make our CI go green, you may need to:
1. rebase with `main` (LMK if you need instructions)
2. run `make fixup` and commit the changes

Thanks very much for the review and good comments and edits @gante ! I have committed your changes and am adding the tests to `transformers/tests/generation/test_utils.py` and will commit that when done.  I might need some pointers for rebasing the PR after:)

Instructions to rebase:
1. get the latest `main`. If your fork is not synched with upstream, you may need to follow [these instructions](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork#syncing-a-fork-branch-from-the-web-ui) before running the commands below.
```
git checkout main
git pull
```

2. rebase your branch
```
git checkout your_branch
git rebase origin/main
```

3. force push your changes
```
git push origin your_branch -f
```

Thanks for the rebase instructions! I have rebased and added the required tests. 

While testing, I found and removed a bug that caused the contrastive search to degenerate into greedy generation.  After fixing this issue, however, the low-memory contrastive search still does not yield exactly the same output as batched (normal) contrastive search.  

After looking into this today, it seems that this issue is caused by numerical errors in forward passes between batched versus unbatched inputs.  These errors seem to propegate in the generation process (via saved hidden layers) such that a dozen tokens or so after generation starts,  one starts to get different tokens.

I am also not sure how the memory test should be performed with many models sequentially.  I typically have been measuring the footprint on GPU, and with one model the low-memory version has a smaller footprint for any given model tested.  But footprint typically does not decrease predictably between model initializations, so this test will likely fail as written.

I'd be happy to go into the former issue in more detail, or I can simply substitute a test to show that the output is approximately the same for batched versus unbatched contrastive search.  And would be happy to change the memory test to include only one model too

Although with some more testing it seems that the numerical errors are not propegating, but there is a problem with the past key value cache in the low memory approach.  Will work on fixing this!

OK all is fixed and ready to go. Selecting `low_memory` does not change the output tokens for contrastive search but reduces the memory footprint for longer (>1k tokens) sequence generation by the amount mentioned above.  Generally the longer the sequence length, the larger the difference between low memory and normal contrastive search.

I am not sure how to write the memory test in the style of the other tests which iterate through a set of models, as unless the models are sorted by increasing size then the test will fail (memory measurement via `torch.cuda.max_memory_allocated()` does not reliably decrease between model re-initializations or even cache cleaning).

Final note, the added code is a bit messy and could be reduced to around 20 lines if https://github.com/huggingface/transformers/issues/17016 were to be implemented.   I can clean it somewhat if required.

> I am not sure how to write the memory test in the style of the other tests which iterate through a set of models, as unless the models are sorted by increasing size then the test will fail (memory measurement via torch.cuda.max_memory_allocated() does not reliably decrease between model re-initializations or even cache cleaning).

Could you share a memory measurement script then, so we can keep it here in the PR for future reference? :)

No problem, for a newly spun up single GPU the following can be used to check the memory reduction.

```python
!pip install -q -U git+https://github.com/blbadger/transformers.git

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
# any model compatible with contrastive search
tokenizer = AutoTokenizer.from_pretrained("gpt2") 
model = AutoModelForCausalLM.from_pretrained("gpt2")

input_ids = tokenizer.encode('This is a new prompt', return_tensors='pt').to('cuda')
model = model.to('cuda')

low_output = model.generate(input_ids, top_k=4, penalty_alpha=0.8, low_memory=True, max_new_tokens=600)
low_mem = torch.cuda.max_memory_allocated()

high_output = model.generate(input_ids, top_k=4, penalty_alpha=0.8, low_memory=False, max_new_tokens=600)
high_mem = torch.cuda.max_memory_allocated()

print (torch.all(low_output == high_output))
print (low_mem, high_mem)
```

I went ahead and removed the memory portion of the contrastive search test module, but we can add it back if necessary.

@blbadger Thank you for the script! (I've confirmed that it works as expected on my end 👍 )

To get the CI green you'll need to run `make fixup` and then commit the changes. You also have errors on the test that's being added on this PR :)

As soon as our CI becomes green, I'll tag a core maintainer so we can merge the PR 🙌 

@gante you are most welcome, thanks for flagging the test failing too.  I have moved the `low_memory` kwarg to the config, cleaned up the low memory code, and fixed up the test module.  

It looks like we are all set with respect to CI requirements except for the code formatting.  Is there a way you recommend fixing this so that `black` does not throw an error?

@blbadger try running `make fixup` and then committing :) If that fails, then try rebasing.

@gante thanks!  `make fixup` completes but there were no changes to commit, same goes for rebasing after syncing the branch but unfortunately `black` still fails with
```
Oh no! 💥 💔 💥
2 files would be reformatted, 2482 files would be left unchanged.

Exited with code exit status 1
```


Just heads up: with (much) more testing it appears that the batch versus no-batch numerical errors mentioned above do indeed propegate such that choosing `low_memory` results in divergence from normal contrastive search for long outputs.  It seems this divergence is not picked up by the existing test suite because these outputs are limited in length, and because the models tested are relatively small.

For an example of this divergence, for Guanaco 33b loaded in `torch.bfloat16`  with an input prompt `Write a full Shakespearean sonnet about the virtues of apples` with `low_memory=False` gives

```
When like a ripe and ruddy cherub's face,
The apple doth in greenness shine so bright,
Its luscious juice and crispness doth embrace,
A feast for taste, both sweet and mild of might.

Yet more than this, it doth our health impart,
With vitamins rich, and fiber to sustain,
Our bodies nourished, and our spirits cheered,
In vigor and in joy, we're once again regained.

Oh, let us praise the apple, in each part,
For its delights so pure, so fresh and fair,
A gift from Nature's bounty, without art,
A treasure true, that ne'er shall disappear.
```

but `low_memory=True` gives

```
When like a ripe and ruddy cherub's face,
The apple doth in greenness shine so bright,
Its luscious juice and crispness doth embrace,
A feast for taste, both sweet and mild of might.

Yet more than this, it doth our health impart,
Its vitals rich, a cure for many ills,
In olden days, 'twas deemed a symbol art,
Of love and beauty, a fruit most fulfills.

O then, let us rejoice in its bounty great,
And praise the Lord who doth such gifts bestow,
Whose handiwork doth manifest such treat,
A blessed fruit, the apple, let us bow.
```

for the same penalty_alpha and top_k.

I would not really classify this as an issue (as the low memory output is likely more numerically accurate than standard contrastive search) but we might want to include this in the documentation for to avoid confusion.


@gante OK looks like the CI is green, running `make quality` after `make fixup` was able to reformat the code properly.

_The documentation is not available anymore as the PR was closed or merged._

@amyeroberts quick TL;DR -- This PR adds a new flag that, when turned on, reduces the memory requirements of contrastive search. There are small numerical, which is expected -- the order of operations with floating point may produce minor variations, as usual in this type of changes

@gante Happy to do so, I appreciate all your help with this PR! 

Would it be better to keep the low memory flag optional and then cast to a bool, or to enforce the low memory flag as a bool? I can see benefits of both options.
 
@amyeroberts Any time, thanks for the detailed review!
1.  Perhaps a better name would be `low_cache_memory` or `single_pkv` as we are here avoiding the generation of past key value caches for all `top_k` tokens.  Is that on the right track?
2. Agreed that the code can be refactored, I will work on integrating your suggestions and will try to simplify the logic.  



@amyeroberts I'd rather keep the generic name (`low_memory`) with an unequivocally unset state by default (`None`) unless you strongly oppose :) 

Here's my reasoning:
- Naming: `.generate()` already has a very big number of configuration parameters, to the point of being one of the most challenging problems to manage. It's very hard to add new options without flags, making the flag discoverability process even harder. I'll gladly take any chance to consolidate flag names, even if it comes at the cost of possible extra code logic in the future 🤗 If some new `.generate()`-level memory-reduction technique comes out, and if it is mutually exclusive with the existing techniques, then we can bump the flag complexity by allowing it to be a string.
- Default: In the recent transition from `model.config` to `model.generate_config`, all non-`None` defaults were a massive pain -- we can't distinguish a default value from an intentionally set value that matches the default. A `None` default protects us from that :)

_____________________________

Regarding complex logic, it won't be a problem as soon as I get my hands on refactoring generate 💪 

@gante OK, understood :) 

* For `low_memory` I agree for the generation config that generic is better. For the `contrastive_search` method, I'd still prefer the kwarg to be something clearer e.g. `sequential`. WDYT?
* Completely agree - I'd prefer `None` as default! 
* Complex logic - We're sweeping things under the rug a bit but OK if a refactor is happening soon. Only thing I'd mention is that the code at the moment makes reviewing hard and longer. The sooner it's tidied up, the quicker new features can (safely) be added! 

> For low_memory I agree for the generation config that generic is better. For the contrastive_search method, I'd still prefer the kwarg to be something clearer e.g. sequential. WDYT?

Sounds good 👍 

Sounds great to me too! @gante would you like me to go ahead and change the kwarg in contrastive_search and assign it the value of `low_memory` in the generation config?

@blbadger yes 👍 

@gante @amyeroberts just renamed the flag, looks like we are ready to go!

Awesome, merging the PR 🔥 Thank you for the cool contribution @blbadger! 

My pleasure, and thanks so much for all your help @gante @amyeroberts!

Thanks for the PR! Do we know if set `low_memory=True` will make the generation how much slower than before? And if `top_k` is small enough like 5, do you still suggest using this method? Thanks!

That is a good question, thanks @yuchenlin! In the worst case the generation time per token is `top_k` times what it would be with `low_memory=False`, and the memory required is asymptotically 1/top_k (as the number of generated tokens increases to infinity).   

In practice for generation around the 1k-2k token length for a medium sized 4-bit quantized (~10b parameter) model I have found that choosing `top_k=3` or `top_k=4` leads to approximately 2x the time required per token, and the peak memory required is around 60% of what is needed with `low_memory=False`. This is somewhat model-dependent, however.

Generally the smaller the `top_k` the faster generation will be, but the less memory savings you will see:)

Hi @blbadger,

Thank you sooo much for your help! It is really insightful! I have another quick question about contrastive searching in general, and hope you would please help! 

From the description of this contrastive searching method, it should generate a deterministic output by setting `top_k=x` and `penalty_alpha=y`. Do you need to set `do_sample=True`? I initially thought we should not because this is a `searching` method and the penalty helps us do `argmax` for selecting from the top_k, which should not involve any randomness via sampling. But there will be a warning message telling me that I should enable `do_sample` when I set `top_k` with something. 

I tried both enabling `do_sample` and not, and I found that if you set `do_sample=True`, there will be randomness in multiple runs, while if I set `do_sample=False`, the results are the same as the greedy decoding results. I'm not sure if I did this correctly. Thank you so much in advance! :D 


No worries @yuchenlin!  You can leave the default value `do_sample=False` as setting this to True will over-ride other arguments and contrastive search will not be used. Make sure you provide argument values for both `penalty_alpha` and `top_k` as giving only a `top_k` while setting `do_sample=True` will raise the warning you mention and generation will revert to greedy decoding.  You are correct that the contrastive search outputs should be deterministic, or very nearly so.

As an example, below is a typical call to `model.generate` that activates the low-memory version of contrastive search:

```python
model.generate(input_ids, top_k=4, penalty_alpha=0.6, low_memory=True, max_new_tokens=200)
```

Hope that is helpful!

hi @blbadger  thank you sooo much for clarifying this! I really appreciate it! :D 

Hi @blbadger, sorry to bother you again. I have another quick question about contrastive searching and I guess you're the best one who can answer it. LOL. 

My understanding is that `model.generate` with contrastive searching will consider the similarity between next top_k and _all_ previous tokens (including both the previous generated tokens and the given prompt, i.e., input_ids). Is this correct? If so, is it possible to make it ignore the tokens that are in the given prompt? In other words, can we make the penalty computation only about the generated tokens and ignore the tokens within `input_ids`.

Thank you so much! :D 


Happy to help @yuchenlin 

Another interesting question! I think that you are correct in that contrastive search will  calculate the degeneration penalty-based ranking using hidden states from all previous tokens, not only tokens generated by the model.  

Specifically the hidden states the model first returns (while generating the first new token) will contain the last hidden states for all previous input tokens (ie if there are 26 input tokens then there are 26 last hidden states), and subsequently single last hidden states are concatenated for each model output token generated.

Just to check we can explicitly return the hidden states in the model and find the token dimension size. For Llama models this can be done as follows:

```python
generation_kwargs = dict(
    output_hidden_states = True,
    return_dict_in_generate = True
)
output = model.generate(input_ids, top_k=4, penalty_alpha=0.6, **generation_kwargs, max_new_tokens=10)
print (len(output.hidden_states[0][0][0])) # n tokens per hidden state
```

and it can be confirmed that the contrastive search ranking uses the last hidden state of the same token dimension: in the codebase the ranking is performed by `ranking_fast(context_hidden, next_hidden, top_k_probs, penalty_alpha, top_k)` where the token dimension of `context_hidden` matches what we found above.

A contrastive search which ignores the input's hidden states would have to be implemented in the future as far as I know.

@gante please let me know if I am in the wrong here! I think that an option to ignore input tokens while performing contrastive search would not be too difficult to add and I would be happy to work on this if there is a desire for this kind of feature?















@blbadger everything you wrote in your comments above is correct! Thank you for helping out here 💛

@yuchenlin The modified version you'd like to experiment with would require modifying `contrastive_search` itself, so my suggestion would be to load a model, monkey patch `contrastive_search` to your liking, and experiment with it :)