Repository: transformers
Issue: an inplace operation preventing TorchDistributor training
Labels: WIP

### System Info

databricks

### Who can help?

@ArthurZucker @younesbelkada

Hi team, 

I got an error message by using TorchDistributor. 


I have checked in the class BertEmbeddings (url as below), line 238, embeddings += position_embeddings is an inplace operation, would you be able to change to embeddings = embeddings + position_embeddings, to allow TOrchDistributor? 

BertEmbeddings url:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py

TorchDistributor sample code: 
https://docs.databricks.com/_extras/notebooks/source/deep-learning/torch-distributor-notebook.html

Thank you very much! 
Ling

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
single_node_single_gpu_dir = create_log_dir()
print("Data is located at: ", single_node_single_gpu_dir)

def train_one_epoch(model, device, data_loader, optimizer, epoch):
  torch.autograd.set_detect_anomaly(True)
  model.train()
  for batch_idx, (data, labels) in enumerate(data_loader):
    inputs1, inputs2 = data[0], data[1]
    inputs1 = {key: val.to(device) for key, val in inputs1.items()}
    inputs2 = {key: val.to(device) for key, val in inputs2.items()}
    # labels = labels.float().to(device)
    labels = labels.to(device)

    optimizer.zero_grad()
    # Compute embeddings
    embeddings1 = model(inputs1)['sentence_embedding']
    embeddings2 = model(inputs2)['sentence_embedding']
    
    # Compute loss
    loss = cosine_similarity_loss(embeddings1, embeddings2, labels)

    loss.backward()
    optimizer.step()
    if batch_idx % log_interval == 0:
      print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
          epoch, batch_idx * len(data), len(data_loader) * len(data),
          100. * batch_idx / len(data_loader), loss.item()))
      
      if int(os.environ["RANK"]) == 0:
        mlflow.log_metric('train_loss', loss.item())

def save_checkpoint(log_dir, model, optimizer, epoch):
  filepath = log_dir + '/checkpoint-{epoch}.pth.tar'.format(epoch=epoch)
  state = {
    'model': model.module.state_dict(),
    'optimizer': optimizer.state_dict(),
  }
  torch.save(state, filepath)

# For distributed training we will merge the train and test steps into 1 main function
def main_fn(directory):
  
  #### Added imports here ####
  import mlflow
  import torch.distributed as dist
  from torch.nn.parallel import DistributedDataParallel as DDP
  from torch.utils.data.distributed import DistributedSampler
  
  ############################

  ##### Setting up MLflow ####
  # We need to do this so that different processes that will be able to find mlflow
  os.environ['DATABRICKS_HOST'] = db_host
  os.environ['DATABRICKS_TOKEN'] = db_token

  # We set the experiment details here
  experiment = mlflow.set_experiment(experiment_path)
  ############################
  
  print("Running distributed training")
  dist.init_process_group("nccl")
  
  local_rank = int(os.environ["LOCAL_RANK"])
  global_rank = int(os.environ["RANK"])
  
  if global_rank == 0:
    train_parameters = {'batch_size': batch_size, 'epochs': num_epochs, 'trainer': 'TorchDistributor'}
    mlflow.log_params(train_parameters)
  
  model = SentenceTransformer(modelname)

  filepath = "../../dbfs/mnt/path2data/"
  df_train = readData('train', filepath)
  df_train = df_train.head(10000)

  train_text = df_train[['sentA', 'sentB', 'score']].values.tolist()
  train_examples = [InputExample(texts=[a, b], label=s) for [a, b, s] in train_text]
  train_dataset = SentencesDataset(train_examples, model)
  #### Added Distributed Dataloader ####
  train_sampler = DistributedSampler(dataset=train_dataset)
  data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler)
  ######################################
  data_loader.collate_fn = model.smart_batching_collate

  model = model.to(local_rank)
  #### Added Distributed Model ####
  ddp_model = DDP(model, device_ids=[local_rank], output_device=local_rank)
  #################################

  optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

  for epoch in range(1, num_epochs + 1):
    train_one_epoch(ddp_model, local_rank, data_loader, optimizer, epoch)
    
    if global_rank == 0: 
      save_checkpoint(directory, ddp_model, optimizer, epoch)
  
  dist.destroy_process_group()
  
  return "finished" # can return any picklable object

# single node distributed run to quickly test that the whole process is working
with mlflow.start_run():
  mlflow.log_param('run_type', 'test_dist_code')
  main_fn(single_node_single_gpu_dir)
```

### Expected behavior

below error disappear. 

![image](https://github.com/huggingface/transformers/assets/48280760/54232ecc-29da-4d37-a0b5-7ca4b8a0d22a)