CUDA Error with Typical Sampling mass of 1.0
Author: harubaru
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.18.0.dev0
- Platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.12
- Huggingface_hub version: 0.4.0
- PyTorch version (GPU?): 1.10.0+cu111 (True)
- Tensorflow version (GPU?): 2.8.0 (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

@cimeister - Original author of TypicalLogitsWarper code

Models:

- GPT-2, GPT: @patrickvonplaten, @LysandreJik

Library:

- Text generation: @patrickvonplaten @narsil

## Information

Model I am using (Bert, XLNet ...): distilgpt2

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Initialize any Causal Language Model and load it on a GPU.
2. Initialize a ``TypicalLogitsWarper`` with a ``mass`` of ``1.0``
3. Attempt to sample the model using the ``sample()`` function with the ``TypicalLogitsWarper``.

I have a Google Colab notebook setup to reproduce this issue, which can be found here:

https://colab.research.google.com/drive/1AX62JAl_yXNXPsTZvUevUyNfcS_eyAXm?usp=sharing

Here is the actual code from the notebook used to reproduce the issue:

```python
import torch
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          LogitsProcessorList, MaxLengthCriteria,
                          StoppingCriteriaList, LogitsWarper)
from transformers.generation_logits_process import TypicalLogitsWarper

model_name = 'distilgpt2'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

outputs = model.sample(
    input_ids = tokenizer.encode('test', return_tensors='pt').to(device),
    logits_warper=LogitsProcessorList([TypicalLogitsWarper(1.0)]),
    stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(10)]),
    pad_token_id=tokenizer.eos_token_id
)
```

And, this is the resulting stack trace when running the above code on a GPU:

```
---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

[<ipython-input-3-3aaa15148b4a>](https://localhost:8080/#) in <module>()
     16     logits_warper=LogitsProcessorList([TypicalLogitsWarper(1.0)]),
     17     stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(10)]),
---> 18     pad_token_id=tokenizer.eos_token_id
     19 )

[/usr/local/lib/python3.7/dist-packages/transformers/generation_utils.py](https://localhost:8080/#) in sample(self, input_ids, logits_processor, stopping_criteria, logits_warper, max_length, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, **model_kwargs)
   1906             # sample
   1907             probs = nn.functional.softmax(next_token_scores, dim=-1)
-> 1908             next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
   1909 
   1910             # finished sentences should have their next token be a padding token

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
```

## Expected behavior

I highly suspect that the behavior above is closely related to this issue on PyTorch, https://github.com/pytorch/pytorch/issues/9062, as the script above runs perfectly fine on a CPU. I believe the expected behavior would be to raise an exception if a ``mass`` of 1.0 or above is used.



Comments:
Using `CUDA_LAUNCH_BLOCKING=1` I was able to pintpoint the issue into TypicalLogitsWarper, not `multinomial` (you need to set the flag otherwise cuda can report incorrect issues since all computations are asynchronous).


I have pinpointed tentatively that the issue was with the  `last_ind` variable.

```python
        last_ind = (cumulative_probs < self.mass).sum(dim=1)
        last_ind[last_ind < 0] = 0
        # Last in can be equal to 50257, which is the vocab size, which crashes in the gather because it cannot gather
        # into a non existing index.
        # Essentially `sorted_indices_to_remove` should be True.
        sorted_indices_to_remove = sorted_scores > sorted_scores.gather(1, last_ind.view(-1, 1))
```

One tentative fix would be to check if the `last_ind` was the vocab size, and then use something like

```python
min_score = sorted_scores.gather(1, last_ind.view(-1, 1)) if last_ind < vocab_size else -1
sorted_indices_to_remove = sorted_scores > True
```

Pinging @patrickvonplaten which has more overview on the soundness of Mass=1.0 anyway and might have better ideas to fix.

Hey @harubaru,

I don't think it makes sense to set the value to 1.0. You can see here: https://github.com/huggingface/transformers/pull/15504/files#r825865764 that we disregard the typical sampler in this case as it probably has no effect. 

Could you try if values < 1.0 work instead? 

Also cc @cimeister in case you have more details here :-)

Thanks! It works with values less than 1.0, it didn't come to my mind that it would be an invalid value unless it threw an exception, as other ``LogitsWarper`` objects throw an exception if it comes across invalid values such as in the ``TopPLogitsWarper`` object here:
https://github.com/huggingface/transformers/blob/8f3ea7a1e1a85e80210b3d4423b674d9a61016ed/src/transformers/generation_logits_process.py#L186
But, I am not sure if it was intended to not raise an exception during the ``__init__`` call of ``TypicalLogitsWarper`` for an invalid ``mass`` value and I think it'd be a good idea to have it raise an exception of sorts if a ``mass`` value greater than or equal to ``1.0`` is passed.

@harubaru I agree that it would make sense to include something similar in the `__init__` for the logits warper! Should probably be the exact same as for top_p. I can start a PR later today unless you'd like to take care of that