gallantlab/himalaya

Specifying Specific Cuda Device (PyTorch backend)

MShinkle opened this issue · 5 comments

MShinkle commented

Currently, specifying 'torch_cuda' as the backend appears to select the first CUDA device visible to PyTorch (cuda:0). However, in multi-gpu systems, it would be useful to specify a specific CUDA device through something like:
set_backend("torch_cuda:3")
which would tell Himalaya to use CUDA device 3. set_backend("torch_cuda") would still function equivalently to how it currently does.

Are there any interest or plans for this feature? From glancing through the Himalaya PyTorch backend I don't think implementing this would be too involved, but I could be mistaken.

mvdoc commented
  • 👍1
MShinkle commented

PyTorch generally recommends against setting cuda device via this method, in favor of device('cuda:3') or device('cuda', index=3) syntax. Though I think the benefits of the latter (better support for using multiple CUDA devices within the same process) are unlikely to impact 99% of use cases, so maybe not worth the time to incorporate.

Thanks!

mvdoc commented

I think that adding this option to himalaya will complicate the code unnecessarily.
But what happens if you manually push the features and data to the gpu that you want to use, and then those tensors are passed to the himalaya solvers? I wonder whether if this can be a workaround to avoid using the environment variable.

MShinkle commented

That's an interesting idea, my expectation is that it would convert it to whatever cuda:0 is in the current environment, but I'll give it a test.

MShinkle commented

Looks like backend.asarray moves tensors to cuda:0 regardless of the original device of the tensor. For example:

import torch
from himalaya.ridge import Ridge
from himalaya.backend import set_backend, get_backend

set_backend('torch_cuda')
backend = get_backend()

print(backend.asarray(torch.zeros(10, device='cuda:1')).device)

Prints cuda:0

  • 😕1