Choosing the device for lightning.gpu

#1

Hello everyone,

I am currently running some tests with lightning.gpu on a server with 4 gpus. I am wondering if there is a way to select which gpu to run on… perhaps something aking to torch.device("cuda:x") where x corresponds to the number of the gpu you want to run on.

From my tests, I can see that the default behaviour is to run on gpu0, even though that might be at 100% usage or the vram is full.

Thanks in advance,
Patrick.

#2

HI @podagiu, welcome to the forum!

You should set the CUDA_VISIBLE_DEVICES environment variable before running the script or setting it inside Python (using os.setenv ).

$ export CUDA_VISIBLE_DEVICES=1 # uses the second card as the number starts from 0 
$ python myscript.py

Or inside Python

import os 
os.environ['CUDA_VISIBLE_DEVICES'] = 1 # safer to use before loading lightning.gpu

Let me know if this solves your problem!

#3

Hey Catalina,

Thanks for the reply! What you proposed works!

However, this really needs to be set at the start of the script, before importing any of my modules, and therefore I cannot automatise something like finding the least busy GPU and running on that. It would be nice if in the future you could select the specific device like one does in pytorch.

Nevertheless, this works for now. Thanks again!

#4

Hi @podagiu,

Handling this inside our plugin is not the best way to do it. You can probably make a shell script parsing the result of nvidia-smi and choose the best GPU.

We may however explore the possibility of programatically choosing a GPU in future pennylane-lightning releases. So thank you for bringing this up!