beam_sample throws a nan error on long generations
Author: fpgaminer
Pull request: no

### System Info

- `transformers` version: 4.29.0.dev0
- Platform: Linux-5.15.0-67-generic-x86_64-with-glibc2.35
- Python version: 3.10.10
- Huggingface_hub version: 0.13.4
- Safetensors version: 0.3.0
- PyTorch version (GPU?): 2.0.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- 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?

@gante

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

It seems that `beam_sample` throws a NaN exception when generating long sequences.  Specifically the call `next_tokens = torch.multinomial(probs, num_samples=2 * num_beams)`.  Example generate call that causes the bug:

```
output_sequences = model.generate(
	input_ids=encoded_prompt,
	max_length=512 + len(encoded_prompt[0]),
	temperature=0.7,
	num_return_sequences=1,
	num_beams=2,
	do_sample=True,
)
```

Reliably throws a NaN on my system and @diegomontoya 's system.  In my testing this occurs when the requested number of new tokens is roughly >=256.  In the example above I use 512 just to be sure.

Based on the debugging I've done so far, what's happening is `beam_scores` increases exponentially with each iteration of the inner beam search loop.  It does this until it reaches a very large negative number, causing `next_token_scores` to contain all `-inf`, which causes `probs` to be all `nan` and then `multinomial` throws.

As for why this occurs, a rough summary of the inner loop elucidates:

```
while
    next_token_scores = ...
    next_token_scores = next_token_scores + beam_scores
    next_token_scores = logits_warper(..., next_token_scores)
    
    beam_scores = beam_scorer.process(..., beam_scores, next_token_scores)
```

Specifically, beam_scores feeds back into itself with every iteration.  If the inner loop was additive only, this would be fine, and `beam_scores` would increase linearly with length.  But this is not the case.  `logits_warper` makes the loop non-additive.  In the example above it behaves as approximately multiplying `next_token_scores` by 1.5.  Hence `beam_scores` goes exponential and the function eventually throws.

I don't know enough about how `beam_sample` is meant to function to analyze further.  It does seem odd to me, though, that the sampling is dependent on the current beam score.  Since the beam score is a scalar value, it affects the probabilities of all tokens equally, so ... it shouldn't have any effect at all?  So why apply it to the sampling logic?  It seems more reasonable to me, and would indeed fix this bug, if it were added after sampling and before handing the scores off to the BeamScorer for processing.

### Expected behavior

`generate` shouldn't throw a `nan` error under reasonable circumstances.

Comments:
Hey @fpgaminer 👋 

My first recommendation would be to use "normal" `sample`, perhaps with a slightly lower temperature. If you think about it, `beam_sample` is a sample-based strategy that greedily picks the best scores among the drawn sequences, which is similar to `sample` with a lower temperature (which also favors high-scoring tokens). `sample` is also faster (no beam-related operations), and subject to much more maintenance :)

If you still want to use `beam_sample`, my recommendation would be to add the `remove_invalid_values` flag ([docs](https://huggingface.co/docs/transformers/v4.28.1/en/main_classes/text_generation#transformers.GenerationConfig.remove_invalid_values)).

Hello @gante,

Thanks for the response.  I have no intention of using beam sampling myself.  I'm bubbling up a bug report by @diegomontoya from my GPTQ-triton repo, that turned out to just be a bug in `transformers` itself.  It was a curious enough bug that I got nerd-sniped by it...

> If you still want to use beam_sample, my recommendation would be to add the remove_invalid_values flag ([docs](https://huggingface.co/docs/transformers/v4.28.1/en/main_classes/text_generation#transformers.GenerationConfig.remove_invalid_values)).

I don't think that would work.  The bug results from `beam_scores` exploding, which drives all the scores down to `-inf`.  Invalid tokens are removed in the `logits_processor` pass, before `beam_scores` is added.  Even if it were applied after, it would just set _all_ tokens to `max` which I think would cause softmax->multinomial to just throw anyway.

------

I've looked at the code more, and read up on beam search more.  I think my initial take is correct.  I see no reason to feed the beam_scores to the logit processors.  It's a scalar value added to all the logits/probs, so what effect could it possibly have?  Temperature, for example, is completely unaffected as proven like so:

```
Suppose we have a vector `x`
Softmax is `e**x / sum(e**x)`

Suppose we add a scalar `b`: `x + b`
Softmax is now: `e**(x + b) / sum(e**(x + b))`
Exponential law: `e**x * e**b / sum(e**x * e**b)`
Simplify: `e**x * e**b / (sum(e**x) * e**b)`
Simplify: `e**x / sum(e**x)`
Q.E.D.
```
It's possible that `b`, aka the beam score, has an effect on other logit processors, but I can't fathom what effect one would _want_ it to have on things like top p, top k, typical, etc.  I'd have to go through each in more detail to have a stronger opinion here.  It just feels wrong, since I think all those logit processors were introduced in the context of greedy sampling.  They weren't designed to take a global scalar like beam score into account.

So I argue that `beam_sample` should be modified to _not_ include the `beam_scores` when calling `logits_warper`, and when doing multinomial sampling.  It should be added after the tokens have been sampled.

-------

I also think there is other oddness to the way `beam_sample` samples.  Consider the simplified forms of `sample` vs `beam_sample`:

sample:
```
next_token_logits = outputs.logits[:, -1, :]
next_token_scores = logits_processor(input_ids, next_token_logits)
next_token_scores = logits_warper(input_ids, next_token_scores)
probs = nn.functional.softmax(next_token_scores, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
```

beam_sample:
```
next_token_logits = outputs.logits[:, -1, :]
next_token_scores = log_softmax(next_token_logits, dim=-1)
next_token_scores_processed = logits_processor(input_ids, next_token_scores)
next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores)
next_token_scores = logits_warper(input_ids, next_token_scores)
probs = nn.functional.softmax(next_token_scores, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=2 * num_beams)
... beam search stuff ...
```

Why does `beam_sample` apply a `log_softmax` to the logits before feeding them to `logits_processor` when the sample method doesn't?  That seems odd, especially when all the logit processors are expecting, well, logits, not the log softmax of logits.

The same goes for `logits_warper`, which also applies a sequence of LogitProcessors.  They aren't likely to be expecting log softmaxed values.

And then `softmax` gets applied afterwards to values in the log softmax domain... very confusing.

----

So I propose for beam_sample (simplified/pseudo):

```
next_token_logits = outputs.logits[:, -1, :]
next_token_scores = logits_processor(input_ids, next_token_logits)
next_token_scores = logits_warper(input_ids, next_token_scores)
probs = nn.functional.softmax(next_token_scores, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=2 * num_beams)
... gather tokens, scores ...
... add beam_scores to respective scores ...
... beam processing ...
```

---

> If you think about it, beam_sample is a sample-based strategy that greedily picks the best scores among the drawn sequences, which is similar to sample with a lower temperature (which also favors high-scoring tokens). sample is also faster (no beam-related operations), and subject to much more maintenance :)

My quick take: sure, maybe.  But in theory beam search and beam sampling still provide potential value over low temp sampling.  They can explore the landscape more thoroughly and potentially find more globally optimal sequences that a greedy sampling method usually won't.  I dunno.

I'm personally in the "better logit processors" and "better models" camp than futzing with beam search.  But since HF includes beam sampling, might as well make it work as well as possible?

@gante I am not qualified to comment on the internal code itself so I will only report from a user level perspective:

1.  Adding `remove_invalid_values=True` does not resolve the issue. I am still getting the exact same nan/inf exceptions with num_beams = 2 on input+output (expected) total token values > 256.  I added it to both generate_config and directly to generate() method and it still threw exceptions. Am I using it correctly?

```probability tensor contains either `inf`, `nan` or element < 0```

2. Having read the naive concepts of beam search and also huggingface's own interpretations of the beam search,  I don't understand why user have to care about a `remove_invalid_values` toggle. Isn't it implied that generate wrapper, which most user and external libs use, should auto remove and bypass any invalid values during gen stages?  This add another chicken and egg problem, if we don't add `remove_invalid_values`, only a runtime generate will find out that inf/nan tokens are generated and then we apply a `remove_invalid_values` pass which negates any performance. As result, as an end-user, I will always set `remove_invalid_values` with `num_beams` >1, but if the both options are symbiotic, they should be done internally by the library and not exposed to user. 

3. I am using beam search because I believe it may resolve an issue that is outlined by the beam search principle. I can lower the the temperature but that requires that:

* I can detect my result from higher temperature is wrong, very difficult for my problem set. 
* Even if I can detect error due to higher temp, I need re-run pass in lower temp which is basically beams in operation. 
* Not possible to predetermine whether lower/higher temp result in better answer. In my test case use of beam-search. I am relying on the idea that `num_beams=2` select two paths, and only until the end, compare the prob score of the result and give me the best one.  

@fpgaminer @diegomontoya Let me split my comment in three: `remove_invalid_values`, how beam sample is implemented, and a suggestion based on @diegomontoya 3rd point in the last comment :)

___________________________________________________________________________________________

`remove_invalid_values` was created to avoid errors with extreme numbers, as a last resort. When it needs to be used, it means that there is something unstable in the process. I was double-checking it and it is missing the `-inf` case, which is probably why it didn't immediately solve your case (I'll open a PR). However, it should still be avoided, and the cases where you actually need it are very very uncommon.

> Isn't it implied that generate wrapper, which most user and external libs use, should auto remove and bypass any invalid values during gen stages?

Definitely not. Our guiding principles for building blocks like `.generate()`, sorted by priority, are 1. keep retrocompatibility (unless it is to fix bugs) and 2. build a default behavior that works in most cases and minimizes black-box behavior. Having `remove_invalid_values` on by default would go against 2 -- if there is something wrong in the generation strategy, we'd rather show it up to the user.
___________________________________________________________________________________________

The same discussion and arguments you wrote about `beam_sample` were also written in the past, by myself included :) (a few examples: [1](https://github.com/huggingface/transformers/pull/5420#discussion_r449779867) [2](https://github.com/huggingface/transformers/pull/21341#discussion_r1089223478)). 

TL;DR: I agree with your point of view, but a) `beam_sample` is not an official implementation so the order of operations is not right or wrong, it is a matter of taste of its creator b) because of the principles I wrote above, ensuring retrocompatibility > individual opinion. 

Our codebase is fully open, so feel free to monkey patch on your end any different perspective 🤗 And my apologies for the nerd snipe, beam methods are indeed a strong magnet!

__________________________________________________________________________________________

@diegomontoya if beam sample keeps failing after I add the `-inf` case and monkey patching is not an option, try the following:
1. Use `sample`
2. Set `num_return_sequences` to an integer, which will make `generate` return these many sequences per input
3. Set `output_scores` and `return_dict_in_generate` to `True`, so you have access to the scores
4. Pick the output with the highest score ([this function may help](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores))

This is essentially a poor man's version of beam sample. While beam sample greedily optimizes the score in the intermediary steps, this will retain full randomness.

___________________________________________________________________________________________

I hope this (long) comment helps understanding why we make certain decisions, even if you don't agree with them :)



@gante Thank you. Got much more info than I had hoped in return and not only did it clarify it for me but your poor-man's beam really opened up my mind about how I should properly use and approach my future usage of generate as a whole. 

btw, the error you've seen is very likely related to this one: https://github.com/huggingface/transformers/issues/22979

TL;DR -- pytorch's sampling function is buggy atm, being able to pick tokens with 0 probability 👀 

Just adding that it could be CUDA, bitsandbytes and pytorch related.

The same error happens for me as well on `torch==1.13.1` with model call:
`tokens = model.generate(**inputs,  max_new_tokens=500, do_sample=True, temperature=0.9, streamer=streamer)`

This call does not throw the error, but returns gibberish:
`tokens = model.generate(**inputs,  max_new_tokens=25, do_sample=True, num_beams=1, temperature=0.9, streamer=streamer, remove_invalid_values=True)`
returns for example:
`ováBit}")VAjem ubuntu米 alwaysicago connectingselection Rewrite perceMillBLoll Forschavano economic pygindi Pent öss fs file`

For me the issue happens on my multi gpu ubuntu 22.04 system with CUDA 12.0 (python detects 11.8 interestingly).
It does not happen on my single gpu ubuntu 20.04 system with CUDA 11.6.

Also, this only happens when I load the model in 8-bit with `bitsandbytes`. Loading the model without `load_in_8bit=True` is very slow (5-10 seconds per token), but returns text that makes sense and does not throw any error.

Further testing shows that after downgrading from CUDA 11.8 to CUDA 11.6, I no longer receive this error when using `load_in_8bit=True` and `tokens = model.generate(**inputs,  max_new_tokens=25, do_sample=True, temperature=0.9, streamer=streamer)`. However, I still get gibberish results:
`ток hastICEyk char sunny少 hardwareington chi GraphSecondsesser引 conser conformygieneOriuvimplughtub`.
The winning combo for 'no error and words that make sense' seems to be either:
- CUDA 11.6, `load_in_8bit=True` and a single GPU system.
- or CUDA 11.6, `load_in_8bit=False` and a multi GPU system.

**Update: ** it's not pytorch related, happens for both 2.0.1 and 1.13.1. See https://github.com/huggingface/transformers/issues/23989

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.