Distributed Training with Triplet Loss and DistilRoberta encoder
Author: rafaljanwojcik
Pull request: no

I wasn't sure if this issue is suitable as a Bug Report, so I copied the template only, so that its readable

## 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.11.3
- Platform: Linux
- Python version: 3.6.2
- PyTorch version (GPU?): 1.3.1
- Tensorflow version (GPU?): --
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help

Hello @LysandreJik - I'm using 'distilroberta-base', so I wasn't sure whom I should ping, thanks for help! I was also using Trainer library, so I allowed myself to tag @sgugger 

## Information

Model I am using: DistilRoberta

The problem arises when using:
* distributed training along with triplet loss training scheme

The tasks I am working on is:
* Semantic Search using Triplet Loss and Transformer Encoder

Note: when using single GPU without Distributed training, the training normally progresses and converges, the model is just training very slowly, hence I'm trying to implement distributed training 

There were two errors I experienced, using Trainer API:
First, related to flagging in autograd one variable ready two times:
```
RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons: 
1) Use of a module parameter outside the `forward` function. Please make sure model parameters are not shared across multiple concurrent forward-backward passes. or try to use _set_static_graph() as a workaround if this module graph does not change during training loop.
2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple `checkpoint` functions to wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases in default. You can try to use _set_static_graph() as a workaround if your module graph does not change over iterations.
Parameter at index 99 with name backbone.encoder.encoder.layer.5.output.LayerNorm.weight has been marked as ready twice. This means that multiple autograd engine  hooks have fired for this particular parameter during this iteration.
```

Second, related to in place modification of tensor:
```
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.LongTensor [4, 447]] is at version 3; expected version 2 instead. 
Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

backtrace:
  File "[...]/lib/python3.6/site-packages/transformers/models/roberta/modeling_roberta.py", line 132, in forward
    token_type_embeddings = self.token_type_embeddings(token_type_ids)
  File "[...]/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "[...]/lib/python3.6/site-packages/torch/nn/modules/sparse.py", line 160, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "[...]/lib/python3.6/site-packages/torch/nn/functional.py", line 2043, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
 ```

## To reproduce

Currently I was able to reproduce the second error outside of Trainer API, still working on the first one (sorry for some hacky parts, like dataset definition and iteration, I hope it will still be understandable - just wanted to create something reproducible quickly):

```python
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

from transformers import AutoConfig, AutoModel

def run_training(rank):
    torch.autograd.set_detect_anomaly(True)
    torch.distributed.init_process_group(backend="nccl", rank=rank)
    
    device = torch.device('cuda', rank)
    
    data = [{'sent1': {'input_ids': torch.randint(0, 200, (16, 256)), 'attention_mask': torch.ones(16, 256)}, 
             'sent2': {'input_ids': torch.randint(0, 200, (16, 256)), 'attention_mask': torch.ones(16, 256)},
             'sent3': {'input_ids': torch.randint(0, 200, (16, 256)), 'attention_mask': torch.ones(16, 256)}
            }]*8
    
    train_sampler = DistributedSampler(
                    data,
                    num_replicas=torch.cuda.device_count(),
                    rank=rank,
                    seed=44,
                )
    
    dll = DataLoader(
            data,
            batch_size=1,
            sampler=train_sampler,
            drop_last=False,
            num_workers=2,
            pin_memory=True,
        )

    config = AutoConfig.from_pretrained('distilroberta-base')
    model = AutoModel.from_pretrained('distilroberta-base', config=config, add_pooling_layer=False)
    model.to(device)
    model.gradient_checkpointing_enable()
    
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank], output_device=rank, find_unused_parameters=False)
    optimizer = torch.optim.Adam(model.parameters())
    metric = lambda x,y: 1.0 - F.cosine_similarity(x, y)
    criterion = nn.TripletMarginWithDistanceLoss(distance_function=metric, margin=0.2, reduction='none')
    
    for n, b in enumerate(dll):
        print(n)
        model.zero_grad()
        sent1 = {k:v.squeeze(0) for k,v in b['sent1'].items()}
        sent2 = {k:v.squeeze(0) for k,v in b['sent2'].items()}
        sent3 = {k:v.squeeze(0) for k,v in b['sent3'].items()}
        
        emb1 = model(**sent1)[0][:, 0, :]
        emb2 = model(**sent2)[0][:, 0, :]
        emb3 = model(**sent3)[0][:, 0, :]
        
        losses = criterion(emb1, emb2, emb3)
        loss = losses.mean()
        loss.backward()
        
        optimizer.step()
        print('Model device: {}, loss device: {}, loss: {}'.format(model.device, loss.device, loss))


def main():
    world_size = torch.cuda.device_count()
    os.environ["MASTER_PORT"] = '1234'
    os.environ["MASTER_ADDR"] = '127.0.0.1'
    os.environ["WORLD_SIZE"] = str(world_size)
    mp.spawn(run_training,
        nprocs=world_size,
        join=True)

if __name__ == "__main__":
    main()

```


## Expected behavior
Training without errors


Comments:
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.