NLLB-200-1.3B multi-GPU finetuning generates RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1
Author: molokanov50
Pull request: no

### System Info

- `transformers` version: 4.32.1
- Platform: Linux-5.15.0-78-generic-x86_64-with-glibc2.31
- Python version: 3.9.7
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 1.10.1+cu113 (True)

The versions of the following packages are not specified and, therefore, are the latest:
- sentencepiece
- sacrebleu
- sacremoses
- psutil
- nltk
- evaluate
- scikit-learn

### Who can help?

@SunMarc @younesbelkada 

[data.csv](https://github.com/huggingface/transformers/files/12472315/data.csv)


### 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

This code fails on 2 and more GPUs, obviously no matter what version of NLLB is in my `modelPath` (I checked `NLLB-200-1.3B` and `NLLB-200-distilled-1.3B`, the model is pre-downloaded and saved on my local machine):

```
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer
import pandas as pd
from sklearn.model_selection import train_test_split
import torch
import torch.utils.data
from transformers import DataCollatorForSeq2Seq
import evaluate
import numpy as np
from argparse import ArgumentParser

modelPath = "initialmodel"

tokenizer = AutoTokenizer.from_pretrained(modelPath)
model = AutoModelForSeq2SeqLM.from_pretrained(modelPath, device_map="auto")

parser = ArgumentParser()
parser.add_argument('--source-lang', type=str, default='eng_Latn')
parser.add_argument('--target-lang', type=str, default='rus_Cyrl')
parser.add_argument('--delimiter', type=str, default=';')
args = parser.parse_args()

dff = pd.read_csv('dataset/data.csv', sep=args.delimiter)

source = dff[args.source_lang].values.tolist()
target = dff[args.target_lang].values.tolist()

max = 512
X_train, X_val, y_train, y_val = train_test_split(source, target, test_size=0.2)
X_train_tokenized = tokenizer(X_train, padding=True, truncation=True, max_length=max, return_tensors="pt")
y_train_tokenized = tokenizer(y_train, padding=True, truncation=True, max_length=max, return_tensors="pt")
X_val_tokenized = tokenizer(X_val, padding=True, truncation=True, max_length=max, return_tensors="pt")
y_val_tokenized = tokenizer(y_val, padding=True, truncation=True, max_length=max, return_tensors="pt")

class ForDataset(torch.utils.data.Dataset):
    def __init__(self, inputs, targets):
        self.inputs = inputs
        self.targets = targets

    def __len__(self):
        return len(self.targets)

    def __getitem__(self, index):
        input_ids = torch.tensor(self.inputs["input_ids"][index]).squeeze()
        target_ids = torch.tensor(self.targets["input_ids"][index]).squeeze()

        return {"input_ids": input_ids, "labels": target_ids}

train_dataset = ForDataset(X_train_tokenized, y_train_tokenized)
test_dataset = ForDataset(X_val_tokenized, y_val_tokenized)

data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, return_tensors="pt")

metric = evaluate.load("sacrebleu")

def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [[label.strip()] for label in labels]

    return preds, labels

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    if isinstance(preds, tuple):
        preds = preds[0]
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

    result = metric.compute(predictions=decoded_preds, references=decoded_labels)
    result = {"bleu": result["score"]}

    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
    result["gen_len"] = np.mean(prediction_lens)
    result = {k: round(v, 4) for k, v in result.items()}
    return result

training_args = Seq2SeqTrainingArguments(
    output_dir="mymodel",
    evaluation_strategy="epoch",
    save_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    weight_decay=0.01,
    save_total_limit=3, 
    num_train_epochs=20, 
    predict_with_generate=True,
    load_best_model_at_end=True
)

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

trainer.save_model('finalmodel')
```

Text of the shell file used to run my code:
`python3 finetune.py --source-lang eng_Latn --target-lang rus_Cyrl --delimiter ';'`

My finetuning data (placed in the file `dataset/data.csv`): see attached file.

Error:

```
Traceback (most recent call last):
  File "/app/finetune.py", line 106, in <module>
    trainer.train()
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1662, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1929, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 2699, in training_step
    loss = self.compute_loss(model, inputs)
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 2731, in compute_loss
    outputs = model(**inputs)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 1335, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 1208, in forward
    encoder_outputs = self.encoder(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 837, in forward
    layer_outputs = encoder_layer(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 405, in forward
    hidden_states = residual + hidden_states
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!
```

I tried [this solution](https://discuss.huggingface.co/t/runtimeerror-expected-all-tensors-to-be-on-the-same-device-but-found-at-least-two-devices-cuda-1-and-cuda-0/39548/14), but the error repeated again. My code block to construct a custom device_map is as follows:
```
max_memory = get_balanced_memory(
    model,
    max_memory=None,
    no_split_module_classes=["M2M100EncoderLayer", "M2M100DecoderLayer", "M2M100Attention", "LayerNorm", "Linear"],
    dtype='float16',
    low_zero=False,
)
device_map = infer_auto_device_map(
    model,
    max_memory=max_memory,
    no_split_module_classes=["M2M100EncoderLayer", "M2M100DecoderLayer", "M2M100Attention", "LayerNorm", "Linear"],
    dtype='float16'
)
model = dispatch_model(model, device_map=device_map)
```

Based on the idea that the manipulations with input data between devices performed by `accelerate` should be essentially the same in _inference_ and in _finetuning_, I tried a [fix encoder hook](https://github.com/huggingface/transformers/pull/25735) proposed for [multi-GPU NLLB-moe inference](https://github.com/huggingface/transformers/issues/23385), but it also didn't help me.

### Expected behavior

Successful completion of `trainer.train()` and the presence of a finetuned model (in a huggingface format) in my `finalmodel` path.

Comments:
Hi @molokanov50,

Could you try to add `"M2M100EncoderLayer"`
to 
https://github.com/huggingface/transformers/blob/07998ef39926b76d3f6667025535d0859eed61c3/src/transformers/models/m2m_100/modeling_m2m_100.py#L542

and see if this works? If not, could you post the new traceback log after this change? Thanks!

@ydshieh I have done it but it is not working. 
Here I retained my both custom device_map and fix encoder hook enabled. Should I maybe disable any of them?
New traceback:
```
Traceback (most recent call last):
  File "/app/finetune.py", line 158, in <module>
    trainer.train()
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1555, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1837, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 2682, in training_step
    loss = self.compute_loss(model, inputs)
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 2707, in compute_loss
    outputs = model(**inputs)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 1314, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 1201, in forward
    encoder_outputs = self.encoder(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 837, in forward
    layer_outputs = encoder_layer(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 388, in forward
    hidden_states = self.self_attn_layer_norm(hidden_states)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/normalization.py", line 190, in forward
    return F.layer_norm(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/functional.py", line 2515, in layer_norm
    return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument weight in method wrapper_CUDA__native_layer_norm)
```

Hi @molokanov50 , what is the output of the following script ? 
```py
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-1.3B")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-1.3B", device_map="auto")
print(model.hf_device_map)
input = 'We now have 4-month-old mice that are non-diabetic that used to be diabetic," he added.'

input = tokenizer(input, return_tensors="pt")

translated_tokens = model.generate(
    **input, forced_bos_token_id=tokenizer.lang_code_to_id["fra_Latn"]
)
print(tokenizer.decode(translated_tokens[0], skip_special_tokens=True))
```



Hi @SunMarc, 
I've received the following output:
```
{'': 0}
Traceback (most recent call last):
  File "/home/molokanov/tf_new/testtest.py", line 10, in <module>
    translated_tokens = model.generate(
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/transformers/generation/utils.py", line 1492, in generate
    model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/transformers/generation/utils.py", line 661, in _prepare_encoder_decoder_kwargs_for_generation
    model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/transformers/models/m2m_100/modeling_m2m_100.py", line 786, in forward
    inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/torch/nn/modules/sparse.py", line 162, in forward
    return F.embedding(
  File "/home/molokanov/tf_new/env/lib/python3.9/site-packages/torch/nn/functional.py", line 2210, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)
```


Looks like you only have access to one GPU. The error appears because in this case (only one device in your device_map), you need to move the input to the right device :` input = tokenizer(input, return_tensors="pt").to(0)`. Can you try with multiple gpus as it was the initial issue ? 

@SunMarc Yeah, of course, I accidentally ran it on a different node with one GPU.
Now I repeated your script on a 2-GPU machine. Here is the output:
```
{'model.shared': 0, 'lm_head': 0, 'model.encoder.embed_tokens': 0, 'model.encoder.embed_positions': 0, 'model.encoder.layers.0': 0, 'model.encoder.layers.1': 0, 'model.encoder.layers.2': 0, 'model.encoder.layers.3': 0, 'model.encoder.layers.4': 0, 'model.encoder.layers.5': 0, 'model.encoder.layers.6': 0, 'model.encoder.layers.7': 0, 'model.encoder.layers.8': 0, 'model.encoder.layers.9': 0, 'model.encoder.layers.10': 0, 'model.encoder.layers.11': 0, 'model.encoder.layers.12': 0, 'model.encoder.layers.13': 0, 'model.encoder.layers.14': 0, 'model.encoder.layers.15': 1, 'model.encoder.layers.16': 1, 'model.encoder.layers.17': 1, 'model.encoder.layers.18': 1, 'model.encoder.layers.19': 1, 'model.encoder.layers.20': 1, 'model.encoder.layers.21': 1, 'model.encoder.layers.22': 1, 'model.encoder.layers.23': 1, 'model.encoder.layer_norm': 1, 'model.decoder': 1}
/usr/local/lib/python3.9/site-packages/transformers/generation/utils.py:1535: UserWarning: You are calling .generate() with the `input_ids` being on a device type different than your model's device. `input_ids` is on cpu, whereas the model is on cuda. You may experience unexpected behaviors or slower generation. Please make sure that you have put `input_ids` to the correct device by calling for example input_ids = input_ids.to('cuda') before running `.generate()`.
  warnings.warn(
Nous avons maintenant des souris de 4 mois qui ne sont pas diabétiques et qui étaient diabétiques", a-t-il ajouté.
```

The `generate`, hence the `forward` is working fine as the text was translated correctly without an error. It is strange that for your training script, the error happens during the `forward` of the model. Can you try one last time with your fine-tuning script ? 

@SunMarc I tried it again but I replaced an `nllb-200-1.3B`, which is saved on my disk, with `facebook/nllb-200-distilled-1.3B`.
New traceback:

```
Traceback (most recent call last):
File "/app/finetune.py", line 158, in <module>
    trainer.train()
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1555, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1916, in _inner_training_loop
    self.optimizer.step()
  File "/usr/local/lib/python3.9/site-packages/accelerate/optimizer.py", line 145, in step
    self.optimizer.step(closure)
  File "/usr/local/lib/python3.9/site-packages/torch/optim/lr_scheduler.py", line 69, in wrapper
    return wrapped(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/optim/optimizer.py", line 280, in wrapper
    out = func(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/optim/adamw.py", line 171, in step
    adamw(
  File "/usr/local/lib/python3.9/site-packages/torch/optim/adamw.py", line 321, in adamw
    func(
  File "/usr/local/lib/python3.9/site-packages/torch/optim/adamw.py", line 566, in _multi_tensor_adamw
    denom = torch._foreach_add(exp_avg_sq_sqrt, eps)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 32.00 MiB (GPU 1; 23.70 GiB total capacity; 22.51 GiB already allocated; 18.81 MiB free; 22.68 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```
My device_map:
`{'model.shared': 0, 'lm_head': 0, 'model.encoder.embed_tokens': 0, 'model.encoder.embed_positions': 0, 'model.encoder.layers.0': 0, 'model.encoder.layers.1': 0, 'model.encoder.layers.2': 0, 'model.encoder.layers.3': 0, 'model.encoder.layers.4': 0, 'model.encoder.layers.5': 0, 'model.encoder.layers.6': 0, 'model.encoder.layers.7': 0, 'model.encoder.layers.8': 0, 'model.encoder.layers.9': 0, 'model.encoder.layers.10': 0, 'model.encoder.layers.11': 0, 'model.encoder.layers.12': 0, 'model.encoder.layers.13': 0, 'model.encoder.layers.14': 0, 'model.encoder.layers.15': 1, 'model.encoder.layers.16': 1, 'model.encoder.layers.17': 1, 'model.encoder.layers.18': 1, 'model.encoder.layers.19': 1, 'model.encoder.layers.20': 1, 'model.encoder.layers.21': 1, 'model.encoder.layers.22': 1, 'model.encoder.layers.23': 1, 'model.encoder.layer_norm': 1, 'model.decoder': 1}`

Seems like `forward` is fine but at some stage I run out of memory. Strange. The model takes 5 GB, and my video memory is 48 GB. What can I change in my `training_args` in order to satisfy my resources? I'm on 2x 24 GB Nvidia RTX 3090 Ti. By the way, I nicely finetuned this model without parallelism on one 40 GB A100 GPU.
Now I will try `facebook/nllb-200-1.3B`. I will report you.

@SunMarc I tried `facebook/nllb-200-1.3B`.
The result is the same `torch.cuda.OutOfMemoryError` as for `facebook/nllb-200-distilled-1.3B`. Indeed, to reduce memory consumption in training is a separate issue, I also need a recommendation in this. But, as for now, it's obvious that `facebook/nllb-200-1.3B` which is downloaded in script performance process and the same pre-downloaded model from a folder act differently. The model, if it is taken from a local folder, produces `RuntimeError: Expected all tensors to be on the same device`.
The files [here](https://huggingface.co/facebook/nllb-200-1.3B/tree/main) are probably not fresh enough.

Amazing that we were finally able to find the issue ! For memory consumption, try the usual technique like gradient checkpointing, gradient accumulation with lower batch_size ... Check this [documentation](https://huggingface.co/docs/transformers/perf_train_gpu_one) for more details !