Make beam sample more robust
Author: ydshieh
Pull request: yes

# What does this PR do?

Make beam sample more robust.

The probability (more generally, scores) in `torch.multinomial` could not contain `nan`. However, the computation
```python
probs = nn.functional.softmax(next_token_scores, dim=-1)
```
could have `next_token_scores` being all `-inf` (along some batch dim.) due to the processing in logit warpers and beam scorers, 
 and this leads to `probs`  being all `nan` along that batch dim, and we sometimes get the following error in the sampling-style generation methods
```
RuntimeError: probability tensor contains either inf, nan or element < 0
```

This PR makes sure the probability to `torch.multinomial` is a valid input while not affecting the existing generation process.

Related PRs:

#17972

#18053

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

Uhmmm... if all tokens in a batch have `-inf` scores, something has gone very wrong. `-inf` can only happen when some tokens are forbidden (e.g. through `NoBadWordsLogitsProcessor`), and if they are all forbidden then it means `.generate` is incorrectly parameterized. 

I think the runtime error is adequate in that situation -- perhaps we could make it an informative exception explaining the problem?

@gante Nothing is wrong, at least in the situation I observed:

- The `logits_warper` I checked contains:
  - TemperatureLogitsWarper 
  - TopKLogitsWarper 
  - TopPLogitsWarper 

Here, 
  - `TopKLogitsWarper`: keeps top 10 logits, set others to `-inf`
  - `TopPLogitsWarper `: keeps top 1 ([`<eos>` token]) logits, which has probability > top_p (= 0.7), and  set others to `-inf`.
  
Then in beam scores, it can't choose `<eos>`, so other tokens are chosen with scores `-inf`.
In the next iteration, the line
```
next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores)
```
introduce the `all -inf` for some batch dimension, as `beam_scores` contains some `-inf`.

The `all -inf` comes from the fact the scores go through `TopPLogitsWarper` while `eos` has a higher probability > `top_p` and `min_tokens_to_keep = 1`.

WDYT? 

@ydshieh I understand what causes it, but I still think we shouldn't change the code -- if this situation happens in production, it means that the user has selected a bad combination of model + processors/warpers. Allowing this behavior results in a silent failure, which is much worse than a crash :)

The solution should be a fix that would be applicable in real usage, i.e. fixing parameterization. For instance, changing `TopPLogitsWarper` to have `min_tokens_to_keep=2` would fix this issue (and is a potential solution if this problem was happening in an actual use case)

Well, I do agree some arguments, but I also don't think this is a real problem: Given a set of parametrization, the algorithm is to give the result which makes its own sense. In this case, the result (at some batch dimension) are with scores `-inf`, which totally makes sense: the users can verify the scores and decide what to do on their own.

**IMO, as long as we define clearly the behavior, and document it, it is fine.**

# Examples

Taking one example (not 100% relevant): when we want to sample without replacement, but the places with positive probabilities > 0 are fewer than the sample size request.

### PyTorch gives what you want, even those elements with probability 0 
```python
import torch
device = "cpu"

probs = torch.tensor([[1, 1, 0, 0, 0, 0]], device=device, dtype=torch.float)
o = torch.multinomial(probs, num_samples=4)
print(o)

# tensor([[0, 1, 5, 4]])
# The last 2 elements are meaningless
# On CPU/GPU, the behavior seems different too!
```
While

### NumPy throws an error
```python
import numpy as np

o = np.random.choice(6, 4, replace=False, p=[1.0/3, 1.0/3, 1.0/3, 0, 0, 0])

# ValueError: Fewer non-zero entries in p than size
```

Also, thinking in batch mode, what if a user really want to use a fixed parametrization, but with it, some example gives error while all other examples could generate successfully? Force them to change the parametrization doesn't seem really good, and it is also not easy to determine beforehand which example will fail with this kind of error.

And yet another argument:

> The solution should be a fix that would be applicable in real usage, i.e. fixing parameterization. For instance, changing TopPLogitsWarper to have min_tokens_to_keep=2 would fix this issue (and is a potential solution if this problem was happening in an actual use case)

Well, how about a user have extra `NoBadWordsLogitsProcessor`, which have a bad word, and that one has the higher probability together with the `eos` token for some example(s) in the dataset? Then the generation will fail again, and to make it work, users have to change to `min_tokens_to_keep=2`, while all other examples (in validation/test datasets) all work with previous parametrization? Such failure in the middle of the process will be really annoying IMO.

@ydshieh The thing is, all the examples you pointed out won't happen unless the user has made a mistake. Beam search methods require at least two tokens per round to operate correctly, so `min_tokens_to_keep=2` should always be set in `beam_sample` (perhaps we can modify `.generate()` to set it by default when `num_beams>1`). 

Again, if we merge this behavior, we expose ourselves to many silent failure modes, where the user will say "the model is bad"/"HF's generate is wrong" instead of being pointed at the root cause. Here are a few examples:
1. logits processors that force tokens in a certain position with unfeasible constraints (e.g.`ForcedBOSTokenLogitsProcessor` + `ForceTokensLogitsProcessor`)
2. logits processors that prevent tokens in a certain position / all positions with unfeasible constraints (e.g. `NoBadWordsLogitsProcessor`, `WhisperTimeStampLogitsProcessor`)
3. A combination of the above

I'm sorry, I'll be very stubborn against this change :)  

> won't happen unless the user has made a mistake

I don't agree with this. A parametrization may work very well for all examples in a dataset but will fail with a single one
  - it's really arguable what is right and wrong here
  - to reiterate again, this is much more annoying if one has to figure out what parametrization to change for a single/few examples, especially in the prod environment, where a crash should really be avoided
    - current implementation doesn't allow one to deal with such error with try/except(they can, but doing this in batch mode while the failing case may just be a single example in that batch is annoying)

> Again, if we merge this behavior, we expose ourselves to many silent failure modes, where the user will say "the model is bad"/"HF's generate is wrong" instead of being pointed at the root cause.

- For users that are not developers + with no motivation to dive into + just want complain:
  - I am not sure if a crush will make them to change their mind and motivate them to dive into
- For other users who are willing to debug
  - Make sure we communicate well/clearly the returned scores should be checked (at least when debugging/analyze logs) should be already super good
  
> Beam search methods require at least two tokens per round to operate correctly, so `min_tokens_to_keep=2` should always be set in `beam_sample` (perhaps we can modify `.generate()` to set it by default when `num_beams>1`). 

- Hmm, the test I checked has `logits_warper_kwargs, logits_warper = self._get_warper_and_kwargs(num_beams=1)`, in a test named `test_beam_sample_generate_dict_output` 😕 . I can re-work the test, but this is the least point in our discussion here.

> I'm sorry, I'll be very stubborn against this change :)  

I understand, but nice to have a discussion anyway. Let's bring @sgugger and @patrickvonplaten into discussion and see what they think :-)

Would adding warning around the changes made in this PR will change your mind, @gante ?

> Would adding warning around the changes made in this PR will change your mind, @gante ?

Haha no, users ignore warnings (and documentation) most of the time 🙃 

> > Would adding warning around the changes made in this PR will change your mind, @gante ?
> 
> Haha no, users ignore warnings (and documentation) most of the time 🙃

Well, I do agree with you (with this) 🙃🙃

I'd also be more on the error side, with a clearer error message and actionable reactions printed to the user. If we get to an all nan line in the batch, there is nothing we can really generate.

I also think that a runtime error should be thrown here ideally. **However** checking all values for `-inf` can seriously slow down generation, so let's maybe make sure to first test for a potential slow down. 

Also to me it seems this PR is just to make the tests less flaky. Can we maybe instead try to relax parameters like `top_k` (instead of using default 50, we disable it) to minimize flakiness ? 

> I also think that a runtime error should be thrown here ideally. **However** checking all values for `-inf` can seriously slow down generation, so let's maybe make sure to first test for a potential slow down.
> 

OK, I understand better now what you mean, but I left a comment above.

> Also to me it seems this PR is just to make the tests less flaky. Can we maybe instead try to relax parameters like `top_k` (instead of using default 50, we disable it) to minimize flakiness ?

The flaky tests have been addressed in #21445



As we all agree that an error should be thrown with clear message, I am going to close this PR. The work on make more clear message should be done in another PR :-)