clip_grad_norm on Multiple GPUs: (CUDA error: device-side assert triggered)
Author: apteryxlabs
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.0.0
- Platform: Linux-5.4.0-53-generic-x86_64-with-debian-buster-sid
- Python version: 3.7.9
- PyTorch version (GPU?): 1.7.0 (True)
- Tensorflow version (GPU?): 2.3.1 (True)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes


### Who can help
@LysandreJik 
@sgugger


## Information

Model I am using (Bert, XLNet ...):
RoBERTa

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)
trainer.train() runs for a bit, then fails with the following output:
```
RuntimeError                              Traceback (most recent call last)
<ipython-input-11-3435b262f1ae> in <module>
----> 1 trainer.train()

~/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py in train(self, model_path, trial)
    759                         torch.nn.utils.clip_grad_norm_(amp.master_params(self.optimizer), self.args.max_grad_norm)
    760                     else:
--> 761                         torch.nn.utils.clip_grad_norm_(model.parameters(), self.args.max_grad_norm)
    762 
    763                     if is_torch_tpu_available():

~/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/utils/clip_grad.py in clip_grad_norm_(parameters, max_norm, norm_type)
     33         total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
     34     clip_coef = max_norm / (total_norm + 1e-6)
---> 35     if clip_coef < 1:
     36         for p in parameters:
     37             p.grad.detach().mul_(clip_coef.to(p.grad.device))

RuntimeError: CUDA error: device-side assert triggered
```

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)
Training RoBERTa for sequence classification from text to binary.

## To reproduce
Steps to reproduce the behavior:
1. Load pre-processed dataset from disk using datasets.Dataset.load_from_disk()
2. Instantiate RoBERTa from pretrained (roberta-base) with config mods (num_labels = 2)
3. Create and run trainer. See full code below (most imports omitted).

```
from transformers import (RobertaTokenizerFast)

BLOCK_SIZE = 512
tok = RobertaTokenizerFast.from_pretrained("./art_tok_onefile_roberta_tuned/")

ds_root = '/media/b/My Passport/datasets/'
tokenized = datasets.Dataset.load_from_disk(os.path.join(ds_root, 'art_unit_tokenized_balanced'))

columns_to_return = ['input_ids', 'attention_mask', 'labels']
tokenized.set_format(type='torch', columns=columns_to_return)

from transformers import RobertaForSequenceClassification  

config = RobertaConfig(
    vocab_size=tok.vocab_size,
    max_position_embeddings=514,
    num_labels = 2
)

config = RobertaConfig.from_pretrained("roberta-base", 
                                        vocab_size=tok.vocab_size,
                                        max_position_embeddings=514,
                                        num_labels = 2)

model = RobertaForSequenceClassification.from_pretrained('roberta-base', config=config)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

for param in model.base_model.parameters():
    param.requires_grad = False

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./roberta_train_test",
    overwrite_output_dir=True,
    num_train_epochs=5,
    per_device_train_batch_size=128,
    save_steps=50,
    save_total_limit=2,
    logging_steps=10,
    #fp16 = True #Enable low-precision via AMP - omitted for now.
)

train_test_bal = tokenized.train_test_split(test_size=0.1)

trainer = Trainer(
    model=model,
    args=training_args,
    #data_collator=collate_fn,
    train_dataset=train_test_bal['train']
)

trainer.train()
```

## Expected behavior

The model trains for the duration of the training cycle.


Comments:
I would guess this is a memory error. Have you tried monitoring the memory available on your GPUs while the training is running?

A CUDA device-side assert triggered means a bad index error somewhere, and persists until you restart your kernel. The code you provide does not allow us to reproduce the bug because it uses a tokenizer and a datasets we don't have access to. There are thus multiple reasons for a bad index error. If you want us to help, you'll need to give a reproducer using a pretrained tokenizer of the hub and a dataset on the hub (for instance GLUE MRPC is great since it's tiny).

To debug your problem locally:
```
for batch in trainer.get_train_dataloader():
    break
model.cpu()(**batch)
```
as on the CPU you will get a clear indication of where the index error is.

@sgugger I'll work on making the datasets public and will post here. In the meantime, I'll run your snippet. 

@LysandreJik , it's not a memory issue - all four GPUs are at ~87% volatile util for the duration.

@sgugger @LysandreJik I've made our bucket public, and the relevant material is in gs://bao-ai/transfer; you should be able to pull stuff down in jupyter via:
`!gsutil -m cp -r gs://bao-ai/transfer .`
... though I haven't tried that command on an unauthenticated computer.

Also, note I'm having the same issue in Colab ([notebook here](https://colab.research.google.com/drive/1y-Tgl_zPJzjrzsq3WeYeUFL9A6hGhLhD?usp=sharing)), so I suspect it's an issue with the dataset as suggested above. @sgugger could you elaborate on what sort of issues you mean by 'bad index'? Would rebuilding the dataset from our source files help? If so, are there steps I can take to make sure indexing issues don't arise?

Note I also periodically get the following error messages at train time:
```
/home/b/anaconda3/envs/transformers_3/lib/python3.7/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
  warnings.warn('Was asked to gather along dimension 0, but all '
```
That's copied directly - it cuts off like that in Jupyter.

I rebuilt the dataset using the following script, reloaded, and trained - still breaking with the same error message. Here's the rebuild code:
```
balanced = pickle.load(open('./balanced_ds.pickle', 'rb'))
bal_df = datasets.Dataset.from_pandas(pd.DataFrame.from_records(balanced, columns = ['txt', 'labels']))

BLOCK_SIZE = 512
tok = RobertaTokenizerFast.from_pretrained("./art_tok_onefile_roberta_tuned/")

ds_tokenized_no_special = bal_df.map(lambda example: tok(example['txt'], 
                                                padding='max_length', 
                                                max_length=BLOCK_SIZE, 
                                                truncation=True,
                                                add_special_tokens = False), batched=True)

ds_tokenized_no_special.save_to_disk('./art_unit_tokenized_balanced_rebuild')
```

This uses the same imports (probably redundantly) as the main script. You can access all the data using `!gsutil -m cp -r gs://bao-ai/transfer .`, same as above.


I'm going to loop through all the data in the dataloader and see if it's returning anything janky. We're expecting tensors of size BATCHxBLOCK_SIZE for attention_mask and input_ids, and a tensor of size BATCHx1 for labels, right? Our labels are currently just 0 or 1, depending on whether a tokenized json document falls within a certain document class (art unit 3600, to be precise - this work is for patent law analysis).

Note, I found experimentally that special tokens have to be removed from the tokenizer in order to be properly passed through the RobertaForSequenceClassification model; otherwise, we get an index error in the torch.nn.embedding step, due to the vocabulary exceeding the Roberta vocab size by 2.

On Colab, before the trainer crashes, I get lots of these messages in the runtime logs:
```/pytorch/aten/src/ATen/native/cuda/Indexing.cu:658: indexSelectLargeIndex: block: [374,0,0], thread: [19,0,0] Assertion `srcIndex < srcSelectDimSize` failed.```

Ran the following checks on the entire dataset - it passed. Training still fails. This would suggest to me that the dataset is not actually the issue:
```
def no_nans(t):
  return bool(t.flatten().isnan().sum() == 0)

def check_batch(ex):
  masks = ex['attention_mask']
  masks_expected_shape = torch.Size([BATCH_SIZE, BLOCK_SIZE])
  masks_are_expected_shape = (masks.shape == masks_expected_shape)

  allowed_masks = [0,1]
  only_allowed_masks = all(i in allowed_masks for i in masks.flatten())

  masks_not_nan = no_nans(masks)

  masks_valid = (masks_are_expected_shape and only_allowed_masks and masks_not_nan)

  ids = ex['input_ids']
  ids_expected_shape = torch.Size([BATCH_SIZE, BLOCK_SIZE])
  ids_are_expected_shape = (ids.shape == ids_expected_shape)

  ids_within_vocab_range = ((ids.max() < tok.vocab_size + 4) and (ids.min() >= 0))

  ids_not_nan = no_nans(ids)

  ids_valid = (ids_are_expected_shape and ids_within_vocab_range)


  labels = ex['labels']
  allowed_labels = [0,1]
  only_allowed_labels = all(i in allowed_labels for i in labels.flatten())

  labels_are_expected_shape = labels.shape == torch.Size([BATCH_SIZE])

  labels_not_nan = no_nans(labels)

  labels_valid = (only_allowed_labels and labels_are_expected_shape and labels_not_nan)

  failures = {
      'masks_are_expected_shape': masks_are_expected_shape,
      'only_allowed_masks': only_allowed_masks,
      'masks_not_nan': masks_not_nan,
      #'masks_valid': masks_valid,
      'ids_are_expected_shape': ids_are_expected_shape,
      'ids_within_vocab_range': ids_within_vocab_range,
      'ids_not_nan': ids_not_nan,
      #'ids_valid': ids_valid,
      'only_allowed_labels': only_allowed_labels,
      'labels_are_expected_shape': labels_are_expected_shape,
      'labels_not_nan': labels_not_nan,
      #'labels_valid': labels_valid
  }

  
  return ((masks_valid and ids_valid and labels_valid), ex, failures)

failed = []
for idx, ex in tqdm(enumerate(iter(loader)), total=len(loader)):
  passed, ex, fail = check_batch(ex)
  if not passed:
    print(f'{idx} Failed!')
    failed.append((idx, ex, fail))
```

(This, again, uses all the above code as a base).

I also deactivated the train test split and deleted the cache files in the dataset - also fails.

Your check for ids in the proper change seems incorrect:
```
ids_within_vocab_range = ((ids.max() < tok.vocab_size + 4) and (ids.min() >= 0))
```
will allow for the indices `tok.vocab_size` to `tok.vocab_size+3` which are all going to generate an index error given the fact your model as `vocab_size = tok.vocab_size`.


@sgugger vocab indices +1 thru +4 are the special tokens, though, right? And the model should be able to accept them, right?

*or +0 thru +3 if we're being pythonic with our indexing

The highest token index in the entire dataset is the pad token. That shouldn't throw an indexing error. Can you reproduce on your end?

> Your check for ids in the proper change seems incorrect:
> 
> ```
> ids_within_vocab_range = ((ids.max() < tok.vocab_size + 4) and (ids.min() >= 0))
> ```
> 
> will allow for the indices `tok.vocab_size` to `tok.vocab_size+3` which are all going to generate an index error given the fact your model as `vocab_size = tok.vocab_size`.

What should the proper check for this step be?

And regardless of the checks - what would token indices have to do with the ultimate error, which is in clip_grad.py? 

Truncated from above:
```
RuntimeError                              Traceback (most recent call last)
<ipython-input-11-3435b262f1ae> in <module>()
----> 1 trainer.train()

1 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/utils/clip_grad.py in clip_grad_norm_(parameters, max_norm, norm_type)
     36         total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
     37     clip_coef = max_norm / (total_norm + 1e-6)
---> 38     if clip_coef < 1:
     39         for p in parameters:
     40             p.grad.detach().mul_(clip_coef.to(p.grad.device))

RuntimeError: CUDA error: device-side assert triggered
```

I faced similar issue while running on colab with Linux OS . Ttried restarting and resetting the kernal error disappeared . 

@shivaraj1994 can you define 'kernel'? Do you mean the Jupyter Kernel, the python Kernel, or the Linux/Mac/Windows kernel? 

This problem was replicated both on Linux and on a Colab account; I don't think it's an issue with the operating system of a given computer. 

**NOTE:** I was able to fix the problem by ditching the datasets library and using the older pytorch dataset paradigm. Full code below (split between three files - note, I'm not importing the HF datasets library; rather, I'm importing a custom module called datasets.py, from the same directory in which I'm running train.py):

# train.py
```
from transformers import (RobertaTokenizerFast,
                          RobertaForSequenceClassification,
                          RobertaConfig)

from transformers import Trainer, TrainingArguments

import pickle

from datasets import BalancedDataset
from collators import DataCollatorForDocumentClassificationBATCH


'''
TODO:
Use custom tok (w/ special tokens removed)
Train more layers
Use longformer as drop-in for Roberta
Make nicer dataset (with train/test split, etc...)
'''

BLOCK_SIZE = 512
BATCH_SIZE = 32

balanced = pickle.load(open('./balanced_ds.pickle', 'rb'))


tok = RobertaTokenizerFast.from_pretrained('roberta-base') #"./art_tok_onefile_roberta_tuned/")

bal_ds = BalancedDataset(tok, balanced, BLOCK_SIZE)

collator = DataCollatorForDocumentClassificationBATCH()


config = RobertaConfig.from_pretrained("roberta-base",
                                      vocab_size=tok.vocab_size,
                                      max_position_embeddings=514,
                                      num_labels = 2)

model = RobertaForSequenceClassification.from_pretrained('roberta-base',
                                                         config=config)


#Disable training on all but the Classification Head!
for param in model.base_model.parameters():
    param.requires_grad = False

training_args = TrainingArguments(
    output_dir="./roberta_train_test",
    overwrite_output_dir=True,
    num_train_epochs=5,
    per_device_train_batch_size=BATCH_SIZE,
    save_steps=150,
    save_total_limit=2,
    logging_steps=20,
    max_grad_norm = 5,
    dataloader_num_workers = 15,
    #fp16 = True #Enable low-precision via AMP
)


trainer = Trainer(
    model = model,
    args = training_args,
    data_collator = collator,
    train_dataset = bal_ds
)

trainer.train()
```

# datasets.py
```
import torch
from torch.utils.data.dataset import Dataset
from tqdm import tqdm



class BalancedDataset(Dataset):
    def __init__(self, tokenizer, data, block_size: int, limit=None):
        self.block_size = block_size
        self.tok = tokenizer

        print('Ingesting data!')
        # Load Data
        self.txt = [i[0] for i in tqdm(data[:limit])]
        self.labels = torch.tensor([i[1] for i in tqdm(data[:limit])])

    def __len__(self):
        return len(self.txt)

    def __getitem__(self, item):
        d = self.tok(self.txt[item], padding='max_length',
                     truncation=True, max_length=self.block_size,
                     return_tensors='pt')
        d['labels'] = self.labels[item]
        return d
```

# collators.py
```
from dataclasses import dataclass
from typing import Dict, List, Union

import torch


@dataclass
class DataCollatorForDocumentClassificationBATCH:
    def __call__(
            self, examples: List[Union[List[int], torch.Tensor, Dict[str, torch.Tensor]]]
    ) -> Dict[str, torch.Tensor]:
        return {
            'input_ids': torch.stack([e['input_ids'] for e in examples]).squeeze(),
            'attention_mask': torch.stack([e['attention_mask'] for e in examples]).squeeze(),
            'labels': torch.stack([e['labels'] for e in examples]),
        }

```

I will keep the bao-ai bucket open to the public for a bit longer so y'all can attempt to replicate the original issue. We still need to figure out what was causing the clip_grad_norm issue in the first place.



> Also, note I'm having the same issue in Colab ([notebook here](https://colab.research.google.com/drive/1y-Tgl_zPJzjrzsq3WeYeUFL9A6hGhLhD?usp=sharing)), so I suspect it's an issue with the dataset as suggested above. @sgugger could you elaborate on what sort of issues you mean by 'bad index'? Would rebuilding the dataset from our source files help? If so, are there steps I can take to make sure indexing issues don't arise?

**Replication note:** I continued to debug in that Colab notebook; if you want to replicate the original issue, you'll need to use the old code, not the code that currently exists at the end of that link.

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.