Contrastive Search peak memory reduction
# 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!


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!