inf/nan in generate (beam_sample) with small temperature values
## Environment info

- transformers` version: transformers version: '4.6.0.dev0'
- Platform: Linux
- Python version: 3.6.9
- PyTorch version (GPU?): '1.8.0' (yes)


## Information
The `generate` function (`beam_sample`) throws error when passing small temperature values.

## To reproduce
```
from transformers import (
    AutoModelForSeq2SeqLM, 
    AutoTokenizer
  )

model_name = "sshleifer/distilbart-xsum-12-3"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "New York City (NYC), often simply called New York, is the most populous city in the United States"
input_ids = tokenizer.encode(text, return_tensors='pt')

sample_outputs = model.generate(input_ids, 
                                num_beams=3,
                                do_sample=True,
                                temperature=0.2
                                )
```
```
Traceback (most recent call last):
  File "test.py", line 16, in <module>
    temperature=0.2
  File "/opt/anaconda3/envs/tensorflow2/lib/python3.6/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/anaconda3/envs/tensorflow2/lib/python3.6/site-packages/transformers/generation_utils.py", line 1113, in generate
    **model_kwargs,
  File "/opt/anaconda3/envs/tensorflow2/lib/python3.6/site-packages/transformers/generation_utils.py", line 2134, in beam_sample
    next_tokens = torch.multinomial(probs, num_samples=2 * num_beams)
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0
```
Another way to reproduce this error is using higher temperatures and more iterations (generate a longer output).

It looks like this error is caused by `next_token_scores` growing to -inf and `probs` becoming nan.
Apparently, large absolute values accumulate over iterations because `next_token_scores` are no longer normalized after adding unnormalized `beam_scores`. 
`beam_scores` are calculated form the output of  `logits_warper(input_ids, next_token_scores)` ,
and can grow fast with low temperatures (warper does: `scores = scores / self.temperature`).

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
Is the increase of unscaled values a desired behaviour and should one just implement their own `logits_warper` handling float overflow?

If not, a quick fix, just for demonstration, is scaling the values of `beam_scores` added to `next_token_scores` by replacing: 
`next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores)` 
with: 
`beam_scores_softmax = F.softmax(beam_scores, dim=-1) `
`next_token_scores = next_token_scores + beam_scores_softmax[:, None].expand_as(next_token_scores)` 
It works fine but changes absolute values of scores users may rely on.


Hi @elsanns 

Great catch, thanks for the detailed explanation!

Your observation is right and I can re-produce this error.

Re, your question
>  should one just implement their own logits_warper handling float overflow?

Actually there's an `InfNanRemoveLogitsProcessor ` (#10769) which does just that,  and can be enabled by passing `remove_invalid_values=True` to `generate`. But the problem is that it replaces the `inf`   values by the maximum value for the current `dtype` which is still quite large and ends up becoming `inf` again after adding the `beam_scores`.

Also if you use `InfNanRemoveLogitsProcessor` as `logits_warper` (so that it gets applied after adding the `beam_scores`) then it no longer gives this error but seems to be shifting the distribution and the generated output doesn't make sense. 

I tried your fix of normalizing `beam_scores` and it seems to be working.

One possible solution would be to add `normalize_beam_scores` argument and when it is `True`, `BeamScorer` would return the normalized `beam_scores`.

What do you think @patrickvonplaten?<|||||>Hi @patil-suraj,

Thanks for replying!
I think there are several ways of scaling `beam_scores`, e.g. using `beam_sample` with a custom `beam_scorer` scaling the input before processing. Pros: no changes to the code, cons: not available through `generate`.
Another approach would be applying logits processors and warper before softmax but it could be a breaking change for users writing custom processors.<|||||>Hey @elsanns,

Sorry for answering so late! My answer here: https://github.com/huggingface/transformers/issues/14993#issuecomment-1003945387 might also be of relevance. 

In short, I think there are a couple of things here:
- `beam_sample()` is quite an edge-case because encoder-decoder models are usually evaluated with `beam_search` instead? Could I ask why you chose to use `beam_sample()` here? Did it give better results for you? 
- Distilled models, like `distilbart` tend to have more extreme output logits as lots of knowledge is compressed into comparably little capacity
- Lastly, as said in the linked answer above, I don't know of an "official" beam sample algorithm which is the reason `transformers` `beam_sample()` algorithm is not implemented according to an official paper or any mathematically sound algorithm. 

IMO a better solution than having the beam score normalize it's outputs would be to maybe add a `Normalizer` to the logits warper so that before the logits are sampled they are being normalized.

In case we see more and more of issues like this one or https://github.com/huggingface/transformers/issues/14993 we might also consider change the `beam_sample()` algorithm to follow the approach proposed in Algorithm 2 in https://github.com/huggingface/transformers/issues/14993 . This would however be a big breaking change and I am currently really not sure that it is worth it
<|||||>Hi @patrickvonplaten,

Thank you for a detailed answer. 

I noticed this behaviour testing various decoding methods, and I don't recall seeing a significant advantage of `beam_sample` in any particular use case.

Since the new approach would be a breaking change, it seems a right solution to keep it the way it is for now.

Thanks again for your answer<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.