Run on certain GPUs

[quancs]

🚀 Feature

Can we add a cudavd arg in Trainer to select the gpus to use for the whole program? This feature is important for the user to select the gpus to use (for memory reseans or for usage reasons, or....).

Below is my code for doing this.

  parser = ArgumentParser()
  parser = Trainer.add_argparse_args(parser)
  parser.add_argument('--cudavd', type=str, default='')
  temp_args, _ = parser.parse_known_args()
  if len(temp_args.cudavd) > 0:
      import os
      os.environ["CUDA_VISIBLE_DEVICES"] = temp_args.cudavd

[ananthsub]

Does the trainer gpus flag work for your use case? https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html

[quancs]

The gpus is good, but not perfect. Here is an example where gpus not working: run on the second gpu without editing the code. If I pass --gpus 2 to the program, it will select 2 gpus for me.

[quancs]

Without editing code, there is no simple way provided by lightning to make the code run on one specified gpu. That's the shortcoming of the current design of gpus.

[Borda]

the same is already implemented with Trainer argument gpus if you pass a list with selected indexes (the same as you would set CUDA_VISIBLE_DEVICES); see docs example: https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html?highlight=gpus#gpus
To have it easily exposed you can use LightningCLI witch automatically expose all Traner arguments, see: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_cli.html?highlight=lightning%20cli

[quancs]

Yes, I know it... The problem is that a single value for the gpus will be treated as the number of gpus not the id of gpus. So, if I want to run my model on a single certain gpu, I could not choose which gpu to run if I don't edit the gpus in the code.

[Borda]

yes, passing 1, will be interpreted as list
https://github.com/PyTorchLightning/pytorch-lightning/blob/e54180363663b690367f0051484650fcd12e5de2/pytorch_lightning/utilities/device_parser.py#L124-L131

[quancs]

Solutions from #7461: --gpus x, can solve this....

[tchaton]

Trainer(gpus=[3, 4, 6]) would select the gpus 3, 4, 6 automatically.

[marsggbo]

Trainer(gpus=[3, 4, 6]) would select the gpus 3, 4, 6 automatically.

work on lightning v1.5.8, not work on lightning v1.6.5. That's weird and took me a long time to debug....