Beam search uses large amounts of VRAM even with depth of 1
Author: monsieurpooh
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.15.0
- Platform: Windows-10-10.0.19041-SP0
- Python version: 3.8.5
- PyTorch version (GPU?): 1.10.2+cu113 (True)
- Tensorflow version (GPU?): 2.5.1 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
@patrickvonplaten @narsil

## Information

Model I am using (KoboldAI/GPT-Neo-2.7B-AID):

The problem arises when using:
* A simple script that calls model.generate() with beam search of 50 beams and max_new_token of 1. When the prompt gets long enough, e.g. a few hundred tokens, CUDA will run out of memory with 8gb vram


The tasks I am working on is:
* I am simply attempting to generate all the probabilities of the next predicted token, similar to OpenAI's playground with the option "display probabilites: Full spectrum" enabled. For example, given "Jane and I went to the", the program should output something like "park: 0.2, store: 0.15, party: 0.12, etc..". According to https://github.com/huggingface/transformers/issues/10012 this is only possible to do via "beam search". 

## To reproduce

Steps to reproduce the behavior:

1. Call model.generate with beam search of 50 beams, max_new_token of 1, and a prompt of about 300 tokens. 

## Expected behavior
I do not understand why this task requires so much VRAM. It should be no more demanding than generating 50 new tokens serially. In fact, it should be even less demanding than that, because each iteration we stay with the same token size, whereas under normal operation, the input will keep growing. 

Comments:
Hi @monsieurpooh ,

I think it's because you're using `num_beams=50` quite simply.
We don't include optimizations to get only 1 beam when `max_new_token=1` so even though for this specifically we don't need the `50 copies` there are actually created (this actually speeds up later when generating more tokens since there's less data moving around later).

So `100 tokens x 50 beams x Vocab_size` logits, it gets fast quick.

If you are only interested in the `top_k` logits on a single step, then I think you shouldn't use `num_beams`.

Two suggestions:

- Using a custom logits processor to intercept the logits while using generate.
```python
class LogitsInspector(LogitsProcessor):
    def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.FloatTensor:
        # Do what you want, `scores` should contain everything you need. they are the logits returned by the model.
        
model.generate(... Old arguments...., logits_processor=LogitsProcessorList(LogitsInspector())
```

- Using a simpler loop
```python
model = AutoModelForConditionalGeneration(... ) # Maybe ForCausalLM, I don't remember if gpt-neo2.7B is decoder-only or not)
tokenizer = AutoTokenizer.from_pretrained(...)

inputs = tokenizer(sentence, return_tensors="pt")
outputs = model(**inputs)
# outputs.logits should contain what you need
```

Depending on your context either of these solutions might suit better.
The first one is more involved but allow you finer control and maybe modifying the logits on the fly if you want.
The second one is simpler, but you might need to readd some of the stuff `.generate()` takes care for your (all information coming from the config, and selecting the logits if you want an interative way to see the logits at each step).


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.

Thank you @Narsil ; forgive me for the noob question but in the 2nd solution how do I get the corresponding tokens of the outputs.logits? It seems to be just a 2D tensor of floats.

Same question about the scores in the 1st solution; how do I get which tokens each score corresponds to?

Also I'm unable to get the code to work for the 1st solution. What should the LogitsInspector return?

I found out the answer to my first question. It's a 2d tensor of size 115, and 50000+. The 50000+ corresponds to the tokens, and the position in the list corresponds to the index of the token as determined by the tokenizer.

However, after being able to sort these by score, I'm not sure how to convert the score into the actual probabilities or logprobs.

I found the answer to my question; just call .softmax(-1) on logits to get the probabilities. 

Oddly after finishing this implementation I got very slightly differing scores from the beam search implementation. 

Beam search:
```
" she": 0.826171875
" it": 0.0556640625
" that": 0.04119873046875
" this": 0.0152435302734375
" her": 0.0082855224609375
" the": 0.00806427001953125
" your": 0.006633758544921875
" there": 0.0038547515869140625
" you": 0.003482818603515625
" one": 0.003108978271484375
```

Simple logits style: 
```
" she": 0.849609375
" it": 0.057830810546875
" that": 0.042236328125
" this": 0.0157470703125
" her": 0.00847625732421875
" the": 0.00844573974609375
" your": 0.00679779052734375
" there": 0.0039825439453125
" you": 0.0035495758056640625
" one": 0.0031948089599609375
```

I'm not sure which one is more accurate; I presume the latter, because the former required a 1-line hack from https://github.com/huggingface/transformers/issues/10012. 

In any case the 2nd version is wayyyyyy faster, so that's good. 

EDIT: Actually, the 1st version is the correct one. I had to call .softmax(-1) _before_ filtering to topk, not after. Now the outputs are identical, and the 2nd version is still faster.

I'm still curious to know if there's any tutorial on how to use LogitsInspector, in case I need to do some more complex logic.

> I'm still curious to know if there's any tutorial on how to use LogitsInspector, in case I need to do some more complex logic.

Not really that's intended for power users.
But if you feel that displaying the entire probability distribution is worth it, feel free to open a PR, we'll take a look and could add it to the lib so that it's easier for others to use maybe.
When doing so, try to share as early as possible your design and intent so the community can give feedback before investing too much work and effort which might have some flaw.

Cheers,
Nicolas