Enhancing auto_select_GPU

[sebastienwood]

🚀 Feature

Automatically pick a GPU that has enough VRAM for the workload.

Motivation

Auto_select_GPU #1426 offers potential to automatically adjust the choice of GPU when for example they may be in use by someone else.
However the current implementation only check if it's possible to use the GPU not if the GPU has enough VRAM for the usage.
The proposed feature would have the following workflow :

• check which GPU can be used
• on train start :
  • check if one batch can be sent through the network on current GPU
  • else change current GPU
  • repeat until there are no more GPU or a capable GPU has been found
  • (possibly ?) wait a bit then retry

Pitch

Keep the flag auto_select_GPU but add checks in the lifecycle to ensure you can actually use the GPU selected.

Alternatives

Add a user try/catch in on_train_start.

Additional context

The use case for this feature is when you want to run multiple jobs in parallel/on a shared cluster without manually tuning the appropriate GPU.

[cgerum]

Does one want to sort the list of GPUs by the average load or ram usage over the last n minutes?

[Borda]

@SkafteNicki mind have look? or @justusschock

[SkafteNicki]

I think this would be a great addition, I sometimes have the same problem.
@sebastienwood would you be up for trying to implement this?

On a side-note:
I had a hard time finding the auto_select_gpus flag in the documentation. Seems it is missing https://pytorch-lightning.readthedocs.io/en/latest/trainer.html#trainer-flags and should probably also be mentioned in https://pytorch-lightning.readthedocs.io/en/latest/single_gpu.html and/or https://pytorch-lightning.readthedocs.io/en/latest/multi_gpu.html#select-gpu-devices

[sebastienwood]

I can give it a go !

[Borda]

@SkafteNicki is it solved with #2836 ?

[SkafteNicki]

@Borda no, I only added the missing docs to auto_select_gpus. What this PR is actually proposing is an enhancement of the feature, which @sebastienwood would do.

[justusschock]

@sebastienwood please give it a try. If you got stuck, please ping one of us :)

[sebastienwood]

PR submitted !

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!