AssertionError with model_parallel in run_clm.py
## 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.1.1
- Platform: AWS Sagemaker
- Python version:  3.6
- PyTorch version (GPU?): 1.6
- Tensorflow version (GPU?):
- Using GPU in script?: YES
- Using distributed or parallel set-up in script?: YES

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
 albert, bert, GPT2, XLM: @LysandreJik 
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Speed and Memory Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 TextGeneration: @TevenLeScao 
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten @TevenLeScao
 Blenderbot: @patrickvonplaten
 Bart: @patrickvonplaten
 Marian: @patrickvonplaten
 Pegasus: @patrickvonplaten
 mBART: @patrickvonplaten
 T5: @patrickvonplaten
 Longformer/Reformer: @patrickvonplaten
 TransfoXL/XLNet: @TevenLeScao
 RAG: @patrickvonplaten, @lhoestq
 FSMT: @stas00
 examples/seq2seq: @patil-suraj
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->
@LysandreJik  @sgugger @alexorona 


## Information

Model I am using (Bert, XLNet ...): GPT2

The problem arises when using:
* [X] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

I am fine-tuning GPT2 using my own dataset, using the examples/language-modelling/run_clm.py script, and I want to use the new model_parallel feature in v4.1.1. I am using a multi-gpu instance (AWS p3.8xlarge - with 4 gpus). 

But I get this error:
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [2], output_device 2, and module parameters {device(type='cpu')}.

## To reproduce

Steps to reproduce the behavior:

1. Run run_clm.py with the following params:
```
python -m torch.distributed.launch \
    --nproc_per_node 4 run_clm.py \
    --do_train \
    --do_eval \
    --fp16 \
    --logging_first_step \
    --model_parallel \
    --evaluation_strategy epoch \
    --logging_steps 50 \
    --model_name_or_path gpt2 \
    --model_type gpt2 \
    --num_train_epochs 1 \
    --output_dir /opt/ml/model/ \
    --per_device_eval_batch_size 2 \
    --per_device_train_batch_size 2 \
    --save_steps 50 \
    --save_total_limit 1 \
    --train_file /opt/ml/input/data/data/train.txt \
    --validation_file /opt/ml/input/data/data/val.txt
```

2. I get this error when training starts:
```
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
#015 95%|ââââââââââ| 62/65 [00:04<00:00, 12.58ba/s]    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [1], output_device 1, and module parameters {device(type='cpu')}.
#015 98%|ââââââââââ| 64/65 [00:05<00:00, 12.26ba/s]#015100%|ââââââââââ| 65/65 [00:05<00:00, 12.61ba/s]
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [2], output_device 2, and module parameters {device(type='cpu')}.
#015 95%|ââââââââââ| 62/65 [00:05<00:00, 11.84ba/s]#015 98%|ââââââââââ| 64/65 [00:05<00:00, 12.18ba/s]#015100%|ââââââââââ| 65/65 [00:05<00:00, 12.60ba/s]
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [3], output_device 3, and module parameters {device(type='cpu')}.
#015 98%|ââââââââââ| 64/65 [00:05<00:00, 11.50ba/s]#015100%|ââââââââââ| 65/65 [00:05<00:00, 12.48ba/s]
[INFO|trainer.py:388] 2020-12-22 00:54:34,892 >> The following columns in the training set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: .
[INFO|trainer.py:388] 2020-12-22 00:54:34,892 >> The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: .
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [0], output_device 0, and module parameters {device(type='cpu')}.
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/conda/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/launch.py", line 261, in <module>
    main()
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/launch.py", line 257, in main
    cmd=cmd)
subprocess.CalledProcessError: Command '['/opt/conda/bin/python', '-u', 'run_clm.py', '--local_rank=3', '--do_train', '--do_eval', '--fp16', '--logging_first_step', '--model_parallel', '--evaluation_strategy', 'epoch', '--logging_steps', '50', '--model_name_or_path', 'gpt2', '--model_type', 'gpt2', '--num_train_epochs', '1', '--output_dir', '/opt/ml/model/', '--per_device_eval_batch_size', '2', '--per_device_train_batch_size', '2', '--save_steps', '50', '--save_total_limit', '1', '--train_file', '/opt/ml/input/data/data/train.txt', '--validation_file', '/opt/ml/input/data/data/val.txt']' returned non-zero exit status 1.
```


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
I expect the fine-tuning script to run successfully. 

If I remove the --model_parallel in the args, then it does run successfully in distributed mode. But I want to use this new feature to reduce memory usage, and increase batch_size

I'm not the expert on the model parallel feature, but I think it's not supposed to be launched with `torch.distributed` as it will only use one process, then split the layers of your model on several GPUs.<|||||>Thanks for the fast response @sgugger 

I tried your suggestion and ran the following (removing the torch.distributed.launch):
```
python run_clm.py \
    --do_train \
    --do_eval \
    --fp16 \
    --logging_first_step \
    --model_parallel \
    --evaluation_strategy epoch \
    --logging_steps 50 \
    --model_name_or_path gpt2 \
    --model_type gpt2 \
    --num_train_epochs 1 \
    --output_dir /opt/ml/model/ \
    --per_device_eval_batch_size 2 \
    --per_device_train_batch_size 2 \
    --save_steps 50 \
    --save_total_limit 1 \
    --train_file /opt/ml/input/data/data/train.txt \
    --validation_file /opt/ml/input/data/data/val.txt
```

I then see this error:
```
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 799, in train
    tr_loss += self.training_step(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1137, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1163, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 732, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 895, in forward
    return_dict=return_dict,
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 732, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 681, in forward
    inputs_embeds = self.wte(input_ids)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 732, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/sparse.py", line 126, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/functional.py", line 1814, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_index_select
```<|||||>Like I said, not the parallel expert, you should try tagging the person who added the functionality :-)<|||||>@laphang Somewhere **data parallelism** is being triggered in run_clm.py. I highly suspect it's because you set `per_device_train_batch_size` and `per_device_eval_batch_size` to a value greater than 1, so the script is probably confused. You indicated you wanted to use **model parallelism** -- i.e. _split a single model into pieces and distribute those pieces across several devices_ so that, for example, the embedding layers and the first several attention blocks are only on the first GPU. A sample starts on the first GPU and is automatically handed off to another GPU as it goes through the mode. Then you indicated that you want to assign different batches to different GPUs when they all have to start on the first GPU. This might not be a problem though because it actually sounds like you want data parallelism, which _duplicates the model to each device to train a larger batch_. You don't need model parallelism for that. Data parallelism is the default behavior of Trainer. So to summarize:

- **Model parallelism**: let's you train bigger models (e.g. gpt2-xl). Set the `per_device_eval_batch_size  `and `per_device_train_batch_size `to 1.
- **Data parallelism**: let's you train bigger batch sizes by duplicating the model to several GPUs and training on more samples at the same time. Set `model_parallel `to false and the trainer will automatically default to data parallelism when you have more than one GPU.<|||||>@sgugger Short term, we need to add this to `TrainingArguments`:
```
if self.model_parallel:
    assert self.per_device_train_batch_size == 1, "Model is parallelized, but per_device_train_batch_size is not 1. Model parallelism only supports a batch size of one at this time."
    assert self.per_device_eval_batch_size == 1, "Model is parallelized, but per_device_eval_batch_size is not 1. Model parallelism only supports a batch size of one at this time."
```

In the long-term, we need to figure out how to enable batches for model parallelism. Batches aren't assigned to devices, so the current arguments `per_device`... only makes sense for data parallelism. 

<|||||>Hi @alexorona, thanks for the quick response. I've just gotten back from the Christmas / New year break, and getting back into things. 

I tried setting the batch_sizes to 1, but I still seem to get basically the same errors. (I also switched to gpt2-large)

A) when running this:
```
python -m torch.distributed.launch \
    --nproc_per_node 4 run_clm.py \
    --do_train \
    --do_eval \
    --fp16 \
    --logging_first_step \
    --model_parallel \
    --evaluation_strategy epoch \
    --logging_steps 50 \
    --model_name_or_path gpt2-large \
    --model_type gpt2 \
    --num_train_epochs 1 \
    --output_dir /opt/ml/model/ \
    --per_device_eval_batch_size 1 \
    --per_device_train_batch_size 1 \
    --save_steps 50 \
    --save_total_limit 1 \
    --train_file /opt/ml/input/data/data/train.txt \
    --validation_file /opt/ml/input/data/data/val.txt
```

I get this:
```
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
#015 86%|âââââââââ | 56/65 [00:04<00:00, 12.03ba/s]    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [1], output_device 1, and module parameters {device(type='cpu')}.
#015 86%|âââââââââ | 56/65 [00:04<00:00, 11.81ba/s]#015 89%|âââââââââ | 58/65 [00:04<00:00, 12.31ba/s]#015 89%|âââââââââ | 58/65 [00:04<00:00, 12.26ba/s]#015 89%|âââââââââ | 58/65 [00:04<00:00, 12.08ba/s]#015 92%|ââââââââââ| 60/65 [00:04<00:00, 12.52ba/s]#015 92%|ââââââââââ| 60/65 [00:04<00:00, 12.40ba/s]#015 92%|ââââââââââ| 60/65 [00:04<00:00, 12.21ba/s]#015 95%|ââââââââââ| 62/65 [00:04<00:00, 12.58ba/s]#015 95%|ââââââââââ| 62/65 [00:04<00:00, 12.47ba/s]#015 95%|ââââââââââ| 62/65 [00:05<00:00, 12.22ba/s]#015 98%|ââââââââââ| 64/65 [00:05<00:00, 12.20ba/s]#015100%|ââââââââââ| 65/65 [00:05<00:00, 12.64ba/s]
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [2], output_device 2, and module parameters {device(type='cpu')}.
#015 98%|ââââââââââ| 64/65 [00:05<00:00, 12.09ba/s]#015100%|ââââââââââ| 65/65 [00:05<00:00, 12.56ba/s]
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [3], output_device 3, and module parameters {device(type='cpu')}.
#015 98%|ââââââââââ| 64/65 [00:05<00:00, 11.61ba/s]#015100%|ââââââââââ| 65/65 [00:05<00:00, 12.42ba/s]
[INFO|trainer.py:388] 2021-01-04 23:04:51,956 >> The following columns in the training set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: .
[INFO|trainer.py:388] 2021-01-04 23:04:51,957 >> The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: .
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 681, in train
    else True
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 282, in __init__
    ).format(device_ids, output_device, {p.device for p in module.parameters()})
AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device CUDA modules, but got device_ids [0], output_device 0, and module parameters {device(type='cpu')}.
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/conda/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/launch.py", line 261, in <module>
    main()
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/launch.py", line 257, in main
    cmd=cmd)
subprocess.CalledProcessError: Command '['/opt/conda/bin/python', '-u', 'run_clm.py', '--local_rank=3', '--do_train', '--do_eval', '--fp16', '--logging_first_step', '--model_parallel', '--evaluation_strategy', 'epoch', '--logging_steps', '50', '--model_name_or_path', 'gpt2-large', '--model_type', 'gpt2', '--num_train_epochs', '1', '--output_dir', '/opt/ml/model/', '--per_device_eval_batch_size', '1', '--per_device_train_batch_size', '1', '--save_steps', '50', '--save_total_limit', '1', '--train_file', '/opt/ml/input/data/data/train.txt', '--validation_file', '/opt/ml/input/data/data/val.txt']' returned non-zero exit status 1.
```

B) when running this:
```
python run_clm.py \
    --do_train \
    --do_eval \
    --fp16 \
    --logging_first_step \
    --model_parallel \
    --evaluation_strategy epoch \
    --logging_steps 50 \
    --model_name_or_path gpt2-large \
    --model_type gpt2 \
    --num_train_epochs 1 \
    --output_dir /opt/ml/model/ \
    --per_device_eval_batch_size 1 \
    --per_device_train_batch_size 1 \
    --save_steps 50 \
    --save_total_limit 1 \
    --train_file /opt/ml/input/data/data/train.txt \
    --validation_file /opt/ml/input/data/data/val.txt
```
I get this:
```
Traceback (most recent call last):
  File "run_clm.py", line 374, in <module>
    main()
  File "run_clm.py", line 344, in main
    trainer.train(model_path=model_path)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 799, in train
    tr_loss += self.training_step(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1137, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1163, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 732, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 895, in forward
    return_dict=return_dict,
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 732, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 681, in forward
    inputs_embeds = self.wte(input_ids)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 732, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/sparse.py", line 126, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/functional.py", line 1814, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_index_select
#015  0%|          | 0/813 [00:00<?, ?it/s]
```

I had a few questions / comments:
1) For model_parallel, am I supposed to use torch.distributed.launch or not? I wasn't 100% clear on that. 
2) Seems like I'm still getting similar errors with batch_sizes 1 as I was getting previously, any other thoughts on what the issue is?
3) For some use cases, I would be interested in using model parallelism and data parallelism together (e.g. for models that currently just fit on the gpu memory with batch size 1 or 2 - I am presuming that splitting the model with model parallelism would allow space in the gpu memory for larger batch sizes and increase speed?). So would definitely be interested in any future changes that allow for that. (per your last comment)<|||||>FYI I noticed that this was included in v4.2.0, removing model_parallel arg from trainer. It hadn't been made to work yet.
https://github.com/huggingface/transformers/pull/9451

I'll wait for it to be included in the trainer. 
<|||||>It is included, there is just no need for the flag that wasn't doing anything special (parallelizing the model was the user's responsibility and still is). We just automatically detect if the model is parallelized now, without needing the flag.<|||||>This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.