Avoid nan during sampling in generate()

[ydshieh]

What does this PR do?

Fix CI test error

            # sample
            probs = nn.functional.softmax(next_token_scores, dim=-1)
>           next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
E           RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

in
https://github.com/huggingface/transformers/runs/6959698965?check_suite_focus=true

The test test_sample_generate may still fail at

transformers/tests/generation/test_generation_utils.py

Line 711 in 8f40077

self.assertListEqual(output_sample.tolist(), output_generate.tolist())

for some unknown reason. I think it is better to investigate this in another PR.

[ydshieh]

I have some doubts here, as this will make all tokens having equal probability to be sampled. But with all -inf, nothing could be sampled which leads to error. I feel there is no well-defined expected results in such edge cases.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

@ydshieh, softmax should be able to handle -inf correctly actually.
You can try:

torch.nn.functional.softmax(torch.tensor([0, float("-inf")]))

which works as mathematically expected.

It's only when all values are -inf that it doesn't work in which case this fix won't help because the generation is broken.

[ydshieh]

This will fix the nan issue actually. The concern is that it doesn't really make sense, as it changes the probability to uniform distribution along vocab dim, while in the broken cases, it is nothing can't be sampled (all probability 0 , mathematically)

[patrickvonplaten]

As explained here: https://github.com/huggingface/transformers/pull/17937/files#r910424357
this won't fix the problem. Also note that generation is used a lot so it's every additional operation (torch.max(...)) leads to a tiny slow down.

Usually if you get nan's after the softmax it means that the generation is broken anyways which can happen and I think there is little we can do against it

[ydshieh]

Yes, that happens only when all -inf along the vocab dim. I will close this PR, and we have to maybe create a doc with all possible flaky tests :-)