(Apologies if there is an existing similar topic. The search doesn’t seem to be working properly. The topics I could find via searching on Google did not seem to answer my question.)
I am parallelizing across multiple GPUs. I want to parallelize across non-contiguous device IDs (e.g., 0 and 2), as I am running another process on device 1. I can set CUDA_VISIBLE_DEVICES=0,2 and I can wrap my model in nn.DataParallel (torch.cuda.device_count() correctly returns 2). However, the DataParellel constructor by default assigns to list(range(# available devices) (see https://pytorch.org/docs/stable/_modules/torch/nn/parallel/data_parallel.html 1), which means that it will try to assign using IDs [0, 1]. This means it will never try to allocate to device 2 (and consequently run out of memory).
If I pass the DataParellel constructor the list of devices [0, 2], it throws an error as device ID 2 == device_count() (this is in cuda/__init__.py:292). Thus, it won’t let me assign to any GPU IDs >= the device count – if I want to use all available GPUs, their IDs must be contiguous.
Is there any way for me to force DataParellel to use device 2?
Thanks!