bug when Bert train on multi gpus
Author: z379035389
Pull request: no

### System Info

- `transformers` version: 4.28.1
- Platform: Linux-3.10.0-957.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.16
- Huggingface_hub version: 0.15.1
- Safetensors version: not installed
- PyTorch version (GPU?): 1.12.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

 @ArthurZucker  @younesbelkada @sgugger

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I had alter the BertEncoder defined in modeling_bert.py, like below:
```
class BertEncoder(nn.Module):
    def __init__(self, config, meta_layer_index=None, scale=1):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False
# added by me 
        self.meta_layer = BertLayer(config) 
        self.meta_layer_index = meta_layer_index 
        self.scale = scale 
        self.optimizer_for_meta_layer = torch.optim.SGD(self.meta_layer.parameters(), lr=1e-5, weight_decay=0.005)
        self.inputs_for_metalayer = None
        self.outputs_for_metalayer = None
        self.meta_layer_outputs = None
        self.st_loss = None
        
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
        self.inputs_for_metalayer = (hidden_states.clone().detach(), head_mask[0] if head_mask is not None else None, past_key_values[0] if past_key_values is not None else None)
        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        next_decoder_cache = () if use_cache else None
        
        # 
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None
            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, past_key_value, output_attentions)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(layer_module),
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )
            # layer_outputs : Tuple[torch.Tensor]
            hidden_states = layer_outputs[0]
            #added by me
            
            if i == self.meta_layer_index - 1:
                self.inputs_for_metalayer = (hidden_states.clone().detach(), layer_head_mask, past_key_value)
            if i == self.meta_layer_index + self.scale - 1:
                self.outputs_for_metalayer = hidden_states.clone().detach()
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
        if self.inputs_for_metalayer is not None and self.outputs_for_metalayer is not None:
            self.meta_layer_outputs = self.meta_layer(
                    self.inputs_for_metalayer[0],
                    attention_mask,
                    self.inputs_for_metalayer[1],
                    encoder_hidden_states,
                    encoder_attention_mask,
                    self.inputs_for_metalayer[2],
                    output_attentions,
                )[0]
            self.st_loss = torch.mean((self.meta_layer_outputs - self.outputs_for_metalayer) ** 2)
            if self.st_loss.requires_grad is True:
                self.optimizer_for_meta_layer.zero_grad()
                self.st_loss.backward()
                self.optimizer_for_meta_layer.step()
            

            
            
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )
```


then I just train a Bert Model on multi gpus, and that didn't work because optimizer_for_meta_layer is None. But it works using only one gpu. The train code like below:
```
from transformers import BertLayer,BertConfig,BertModel,BertForMaskedLM
from transformers import BertForMaskedLM
from transformers import BertConfig
from transformers import BertTokenizer
import datasets
import json
import sys
import copy
from datasets import load_dataset
BertBaseconfig = BertConfig()
BertBase = BertForMaskedLM(BertBaseconfig)
layer = BertLayer(BertBaseconfig)
MetaModel = BertForMaskedLM.from_pretrained('/home/wanzhipeng/deepincubation/MetaModel_bert_wiki/checkpoint-36500')
MetaModelEncoderBertLayer = MetaModel.bert.encoder.layer
BaseModelEncoderBertLayer = BertBase.bert.encoder.layer
BaseLayerNums = BaseModelEncoderBertLayer.__len__()
MetaLayerNums = MetaModelEncoderBertLayer.__len__()
Submodules = []



# 
def initSubmodules():
    global Submodules
    Submodules = []
    scale = BaseLayerNums // MetaLayerNums   
    for i in range(MetaLayerNums):
        layers = [BertLayer(BertBaseconfig) for _ in range(scale)]
        Submodule = copy.deepcopy(MetaModel)
        Submodule.bert.encoder.layer = Submodule.bert.encoder.layer[0:i+1] + layers + Submodule.bert.encoder.layer[i+1:]
        del Submodule.bert.encoder.layer[i]
        Submodules.append(Submodule)


def tokenize_function(examples):
    return tokenizer(examples["text"])
initSubmodules()
model=Submodules[0]
model.config.num_hidden_layers = 6
model.bert.encoder.meta_layer_index = 0

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased',use_fast=True) 
datasets = load_dataset('wikitext', 'wikitext-2-raw-v1')
tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=["text"])

block_size = 128
def group_texts(examples):
    # Concatenate all texts.
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
        # customize this part to your needs.
    total_length = (total_length // block_size) * block_size
    # Split by chunks of max_len.
    result = {
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
    result["labels"] = result["input_ids"].copy()
    return result
lm_datasets = tokenized_datasets.map(
    group_texts,
    batched=True,
    batch_size=1000,
    num_proc=4,
)

from transformers import Trainer, TrainingArguments
from transformers import DataCollatorForLanguageModeling
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)
training_args = TrainingArguments(
    output_dir="sub1",          # output directory to where save model checkpoint
    evaluation_strategy="steps",    # evaluate each `logging_steps` steps
    logging_strategy="steps",
    overwrite_output_dir=True,      
    num_train_epochs=10,            # number of training epochs, feel free to tweak
    logging_steps=10,             # evaluate, log and save model checkpoints every 1000 step
    save_steps=10,
    load_best_model_at_end=True,  # whether to load the best model (in terms of loss) at the end of training
    save_total_limit=3,           # whether you don't have much space so you let only 3 model weights saved in the disk
    learning_rate=1e-5,
    weight_decay=0.01,
    warmup_steps=10000,
    per_device_train_batch_size=64,
    per_gpu_eval_batch_size=64,
)
from transformers import Trainer, TrainingArguments,EarlyStoppingCallback,TrainerCallback
from transformers import DataCollatorForLanguageModeling
import torch
## 不同卡的情况下会出问题
class CallbackForMetaLayer(TrainerCallback):
    def __init__(self):
        super().__init__()
        self.meta_layer_outputs = None

    def on_step_begin(self, args, state, control, model=None, **kwargs):
        self.meta_layer_outputs = [model.bert.encoder.meta_layer_outputs]
        print("step_begin:")
        print(id(self.meta_layer_outputs[0]))
           
        
    
    def on_step_end(self, args, state, control, model=None, **kwargs):
        
        print("step_end:")
        print(id(self.meta_layer_outputs[0]))
        # print("*********************************************************")
        # print(model.bert.encoder.meta_layer_outputs)
        # print("*********************************************************")
        # print(model.bert.encoder.outputs_for_metalayer)
        # print("*********************************************************")
        
        # model.bert.encoder.st_loss = torch.mean((model.bert.encoder.meta_layer_outputs - model.bert.encoder.outputs_for_metalayer) ** 2)
        # model.bert.encoder.optimizer_for_meta_layer.zero_grad()
        # model.bert.encoder.st_loss.backward()
        # model.bert.encoder.optimizer_for_meta_layer.step()
        
                          
trainer = Trainer(
    model=model.to("cuda"),
    args=training_args,
    train_dataset=lm_datasets["train"],
    eval_dataset=lm_datasets["validation"],
    data_collator=data_collator,
    callbacks = [EarlyStoppingCallback(early_stopping_patience=5)],
)        

# trainer.train(resume_from_checkpoint=True)
trainer.train()

```
error : 
Traceback (most recent call last):
  File "test2/sub1.py", line 130, in <module>
    trainer.train()
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/trainer.py", line 1662, in train
    return inner_training_loop(
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/trainer.py", line 1929, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/trainer.py", line 2699, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/trainer.py", line 2731, in compute_loss
    outputs = model(**inputs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 168, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 178, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/_utils.py", line 461, in reraise
    raise exception
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 1393, in forward
    outputs = self.bert(
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 1055, in forward
    encoder_outputs = self.encoder(
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 643, in forward
    self.meta_layer_outputs = self.meta_layer(
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 495, in forward
    self_attention_outputs = self.attention(
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 425, in forward
    self_outputs = self.self(
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 284, in forward
    mixed_query_layer = self.query(hidden_states)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/wanzhipeng/miniconda3/envs/dl/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Output 101 of BroadcastBackward is a view and its base or another view of its base has been modified inplace. This view is the output of a function that returns multiple views. Such functions do not allow the output views to be modified inplace. You should replace the inplace operation by an out-of-place one.


### Expected behavior

I think it shoule be trained on multi gpus.

Comments:
The normal BERT model can train on multiple GPUs, the bug is thus likely due to the modifications you made. You should ask no the [forums](https://discuss.huggingface.co/) to help debug your code.

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/main/CONTRIBUTING.md) are likely to be ignored.