Repository: huggingface/transformers

Issue:
Trainer is using DataParallel on parallelized models 
## Environment info

- `transformers` version: 4.2.0
- Platform: Ubuntu 20.04
- Python version:  3.8.5
- PyTorch version (GPU?): 1.7.1 / CUDA 11.2
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

### Who can help
@sgugger @stas00 


## Information

I'm trying out the 4.2.0 release with a training script that had been working in 4.1.1. 

I'm parallelizing my model over two GPUs, and I had been using the `--model_parallel` training arg in the previous version. Now that it's no longer used, I removed the arg from my training command, but I'm getting an error as though the DataParallel is being used and the model isn't being detected as parallelized:
`RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cuda:1`

I did some debugging, and everything seems okay with my model (`trainer. is_model_parallel` returns True). But the `trainer. args.n_gpu` is still 2.

I admit that I don't totally understand what's happening in the trainer code, but it might be an error on line 289?
[`self.args._n_gpu = 1`](https://github.com/huggingface/transformers/blob/126fd281bc309ec29caef99e982640265c8a4fba/src/transformers/trainer.py#L289)

Should that be `self.args.n_gpu = 1`, without the leading underscore?


## To reproduce

Steps to reproduce the behavior:

1. Parallelize a model
2. Train on a machine with multiple GPUs






Discussion:
The `self.args._n_gpu = 1` is to avoid parallelizing the data so it has nothing to do with your problem (and it is right, we can't set `self.args.n_gpu` which is a property but that's a whole different story!)

How is your model parallelized? Without that piece of code we can't reproduce the bug and help you.
Thanks @sgugger.

In my test, I'm using some code originally derived from the run_clm.py example. I'm trying to fine-tune a GPT2 model I've trained from scratch. The model was parallelized with the following lines, and this exact fine-tuning script ran successfully yesterday in 4.1.1, using the `--model_parallel` training arg. 

```
    device_map = {0: range(0, 15),
                  1: range(15, 32)}
    model.parallelize(device_map)
```

The error I'm getting now looks a lot like what would happen if I left out the `--model_parallel` flag in 4.1.1.


Resolving patch:
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -16,7 +16,7 @@
 import os
 from dataclasses import asdict, dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 from .file_utils import cached_property, is_torch_available, is_torch_tpu_available, torch_required
 from .trainer_utils import EvaluationStrategy, SchedulerType
@@ -426,7 +426,6 @@ def __post_init__(self):
 
         if is_torch_available() and self.device.type != "cuda" and self.fp16:
             raise ValueError("Mixed precision training with AMP or APEX (`--fp16`) can only be used on CUDA devices.")
-        self._n_gpu = torch.cuda.device_count()
 
     def __repr__(self):
         # We override the default repr to remove deprecated arguments from the repr. This method should be removed once
@@ -467,14 +466,14 @@ def eval_batch_size(self) -> int:
 
     @cached_property
     @torch_required
-    def _setup_devices(self) -> Tuple["torch.device", int]:
+    def _setup_devices(self) -> "torch.device":
         logger.info("PyTorch: setting up devices")
         if self.no_cuda:
             device = torch.device("cpu")
-            n_gpu = 0
+            self._n_gpu = 0
         elif is_torch_tpu_available():
             device = xm.xla_device()
-            n_gpu = 0
+            self._n_gpu = 0
         elif self.local_rank == -1:
             # if n_gpu is > 1 we'll use nn.DataParallel.
             # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
@@ -485,9 +484,7 @@ def _setup_devices(self) -> Tuple["torch.device", int]:
             device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
             # Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
             # the default value.
-            if self._n_gpu == -1:
-                self._n_gpu = torch.cuda.device_count()
-            n_gpu = self._n_gpu
+            self._n_gpu = torch.cuda.device_count()
         else:
             # Here, we'll use torch.distributed.
             # Initializes the distributed backend which will take care of synchronizing nodes/GPUs
@@ -507,12 +504,12 @@ def _setup_devices(self) -> Tuple["torch.device", int]:
             else:
                 torch.distributed.init_process_group(backend="nccl")
             device = torch.device("cuda", self.local_rank)
-            n_gpu = 1
+            self._n_gpu = 1
 
         if device.type == "cuda":
             torch.cuda.set_device(device)
 
-        return device, n_gpu
+        return device
 
     @property
     @torch_required
@@ -520,7 +517,7 @@ def device(self) -> "torch.device":
         """
         The device used by this process.
         """
-        return self._setup_devices[0]
+        return self._setup_devices
 
     @property
     @torch_required
@@ -532,7 +529,9 @@ def n_gpu(self):
             This will only be greater than one when you have multiple GPUs available but are not using distributed
             training. For distributed training, it will always be 1.
         """
-        return self._setup_devices[1]
+        # Make sure `self._n_gpu` is properly setup.
+        _ = self._setup_devices
+        return self._n_gpu
 
     @property
     @torch_required


Test patch:
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -381,9 +381,11 @@ def test_data_is_not_parallelized_when_model_is_parallel(self):
         # Make the Trainer believe it's a parallelized model
         model.is_parallelizable = True
         model.model_parallel = True
-        trainer = Trainer(model=model, train_dataset=RegressionDataset(), eval_dataset=RegressionDataset())
+        args = TrainingArguments("./regression", per_device_train_batch_size=16, per_device_eval_batch_size=16)
+        trainer = Trainer(model, args, train_dataset=RegressionDataset(), eval_dataset=RegressionDataset())
         # Check the Trainer was fooled
         self.assertTrue(trainer.is_model_parallel)
+        self.assertEqual(trainer.args.n_gpu, 1)
 
         # The batch size of the training and evaluation dataloaders should be 16, not 16 * n_gpu
         self.assertEqual(trainer.get_train_dataloader().batch_size, 16)