In a multi-gpu setup, use a fraction of GPUs for training and the other fraction for some other computations 
### Feature request

When running pytorch examples with transformers library, it seems that the number of GPUs is inferred at https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/training_args.py#L1091 
and it assumes that all `CUDA_VISIBLE_DEVICES` will be used for training. In some cases, one would like to use other GPUs at the machine to perform some other computations during the training but still keep the training run on a single GPU.
For example, with `CUDA_VISIBLE_DEVICES=0,1` one would like to run training on `cuda:0` only and use `cuda:1` for some additional custom computations running in parallel with the training process.

### Motivation

Useful when one wants to utilize other GPUs on the machine to perform some heavy computations from the data that is collected during the training process. Allowing the model to use all devices for training would unnecessarily slow down a run that can easily fit on a single GPU device.

### Your contribution

A hacky solution that might work here would be to override `_n_gpu` in the main script, but I am not sure if this would have some negative consequences on other parts of the library (i.e. things initialized before main is reached and the new value to `_n_gpu` is set). Also, dynamically modifying the `_n_gpu` probably wouldn't be considered as the cleanest solution.
(not sure if I should tag someone for this, a wild guess: @LysandreJik @patil-suraj )

Edit by Lysandre: @sgugger :)<|||||>This is only used for `DataParallel`, which is not the recommended way to run training on several GPUs (as recommended by the [PyTorch doc](https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html) itself) so we won't invest any time in making it more customizable than it is.<|||||>This is also used in single-GPU runs, which is why I've opened to ask about it. For example, I think it's currently impossible to run finetuning of bert-base-uncased on SQuAD on 1-GPU, and to use the second GPU to do something else while the finetuning is running on the first one.<|||||>i.e. this line https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/training_args.py#L1091 will "force" usage of DataParallel just because there are two GPUs available on the machine (even though I want to run on the first one only). Constraining it to the first GPU only via `CUDA_VISIBLE_DEVICES` will block usage of the second GPU for some additional stuff during the finetuning run.<|||||>I don't know how to run a training in parallel to some other computation without launching two different scripts, and you can then leverage `CUDA_VISIBLE_DEVICES` for the first. Could you give me an explicit example of what you're thinking of?<|||||>Oh sorry for not clarifying this properly. For example, let's say we want to run the default example of finetuning the bert-base-uncased model on SQuAD dataset ([squad pytorch example](https://github.com/huggingface/transformers/tree/main/examples/pytorch/question-answering#:~:text=python%20run_qa.py%20%5C%0A%20%20%2D%2Dmodel_name_or_path%20bert%2Dbase%2Duncased%20%5C%0A%20%20%2D%2Ddataset_name%20squad%20%5C%0A%20%20%2D%2Ddo_train%20%5C%0A%20%20%2D%2Ddo_eval%20%5C%0A%20%20%2D%2Dper_device_train_batch_size%2012%20%5C%0A%20%20%2D%2Dlearning_rate%203e%2D5%20%5C%0A%20%20%2D%2Dnum_train_epochs%202%20%5C%0A%20%20%2D%2Dmax_seq_length%20384%20%5C%0A%20%20%2D%2Ddoc_stride%20128%20%5C%0A%20%20%2D%2Doutput_dir%20/tmp/debug_squad/)). And let's say that during this finetuning run, we want to collect gradients from each step and calculate some interesting statistics from them. Given that our GPU is already at limits with memory, we would like to utilize the second GPU available at the machine. To be able to use it, we would have to make it visible to pytorch via `CUDA_VISIBLE_DEVICES=0,1`. But rerunning the finetuning example with `0,1` will trigger `DataParallel` which we don't want to use (too slow, not recommended, and so on). We would still like to run the example on the first GPU, as it would normally run without `DataParallel` being triggered, and to use the second GPU to calculate some custom stuff that can't fit on the first GPU. <|||||>I'm still very confused about how you are going to achieve this with the `Trainer`. I'd really love to see some code to see how we can design an API that solves the problem.<|||||>For example, a simple solution to the problem could be something like using a training argument to specify which GPUs to use, instead of taking all available GPUs at the machine. I could provide some code, but I don't want to waste your time on digging through this. Basically, this could be a solution to this problem:
changing this line https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/training_args.py#L1091
to
```python
self._n_gpu = len(self.devices)
```
where `self.devices` is an argument of the `TrainingArgument` class, that could look like this:
```python
gpus: List[str] = field(default=['cuda:0'], metadata={"help": "Available GPUs for model's training."})
```
I.e. there are use-cases where not all available GPUs at the machine should be used for training.
With the aforementioned fix, one could run finetuning of the bert-base-uncased on the first GPU only (via `--gpus ['cuda:0']`) and still use the second GPU for some custom computations (for example attaching gradient hooks to the model and dumping them on the `cuda:1` device for additional processing).<|||||>I think we're just going in circles, as I still don't see how you are going to be able to your extra computations on the second GPU with the Trainer.<|||||>For example like this:
```python
for name, param in trainer.model.named_parameters():
    param.register_hook(lambda grad: grad.to('cuda:1'))
```
Edit: normally, this `lambda` would do something like `deepcopy` of gradients to `cuda:1`, and then for example sort them.<|||||>I am still not seeing what you are doing with those during the training. What piece of code executed by the `Trainer` will then give you anything useful from those? The reason I'm asking is that whatever you want to do probably requires subclassing the `Trainer` with some custom code, so instead of adding a new argument for everyone that will only work in DP setting and thus confuse users (the Trainer supporting way more setups) just for something that can't work out of the box anyway is not a good idea. If you end up subclassing the `Trainer` anyway, you should then be able to also implement the changes to only use GPU 0.<|||||>More specifically, I am trying to identify less important weights during the training and to compress them. Identifying them requires doing some computations with gradients, and that's the reason why I need to use the additional GPU (but this is probably not relevant to the problem). And yes, I have subclassed the `Trainer` class to implement these additional custom things. I am not using anything from the `Trainer` class except the default training pipeline. 

So, back to the original question and to add more context:
1. when this line is invoked: https://github.com/huggingface/transformers/blob/0645b07daf9a78ad64472a36288710b9fa32087a/examples/pytorch/question-answering/run_qa.py#L212
the `training_args` will grab all available gpus at the machine and set them as `self._n_gpu` attribute
https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/training_args.py#L1091
2. then, `trainer.train()` is invoked to begin the training phase. At this line
https://github.com/huggingface/transformers/blob/6bc6797e04811176f4244a42c86f8a65a1e1c455/src/transformers/trainer.py#L1366 the model gets wrapped into `DataParallel` class because `self._n_gpu` grabbed all available GPUs at the machine. And this is the core of the problem: the model is `DataParallel` just because there are multiple GPUs on the machine detected with `torch.cuda.device_count()`. If there were 100 GPUs, then the model would be trained on 100 GPUs just because they are detected with `torch.cuda.device_count()`. As I have mentioned earlier, forcing the number of visible GPUs with `CUDA_VISIBLE_DEVICES=0` is not a solution because then all other GPUs become invisible to pytorch and these custom computations during the training. 
3. so a possible solution to avoid adding new arguments to the `Trainer`: I could for example override the `self._n_gpu` attribute of the `training_args` object before the `Trainer` is initialized and model is wrapped (so from their point of view, there would be only one available GPU). What I wanted to ask is, if there are some other things that I should be careful about if I override dynamically the `self._n_gpu` attribute. I don't feel comfortable overriding it just because it's seems to be a "protected" attributed of the `TrainingArguments` class. Probably another possible solution could be to subclass the `TrainingArguments` class and override the behaviour of https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/training_args.py#L1034 by setting `self._n_gpu` attributed based on the given argument `--gpus` instead of grabbing all visible GPUs with `torch.cuda.device_count()`.
<|||||>TLDR for a possible enhancement of the existing pipeline: 
- Is it better to give users a choice to set which devices they want to use for training via an argument compared to always grabbing all visible devices via `torch.cuda.device_count()`?
- In any case, effort for end-users will be the same, as they would just use an argument (for example `--gpus ['cuda:0']`) instead of `CUDA_VISIBLE_DEVICES` to specify which devices should be used for their experiment. The former gives more flexibility and enables use-cases as described above.<|||||>Thanks for humoring me and detailing a bit more your use case :-).
I'm convinced, we can add a new training argument that will accept a list (if used outside of the parser) or a comma-separated string (for use as CLI argument) with elements being what `torch.device` accepts (so ints or strings like `'cuda:0'), which will default to `None` (in which case we take everything available).

Does that sound like an acceptable solution? Would you like to open a PR with that?<|||||>Yes sure, that solution would be perfect (providing both functionalities, and also backward compatible). Thanks a lot for your time and very helpful suggestions, I will open a PR for this. 
And sorry for spamming you with probably unnecessary details of my exact use case. I could have explained the issue in a much simpler way 😄<|||||>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.