cuda devices

[kullrich]

Dear Mesih,

I get an memory allocation error, this is because on our GPU server multiple devices are running and it seems that the code always only uses the first GPU:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 26.00 MiB (GPU 0; 10.75 GiB total capacity; 25.19 MiB already allocated; 11.62 MiB free; 42.00 MiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Could you please provide a @click option to select the cuda devices for pytorch, e.g. cuda:1

Since on other devices the memory is not full and I would like to be able to select the device or even cuda:all

|   0  NVIDIA GeForce ...  On   | 00000000:1A:00.0 Off |                  N/A |
|  0%   22C    P8     1W / 250W |  10421MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  On   | 00000000:1B:00.0 Off |                  N/A |
|  0%   22C    P8    10W / 250W |    343MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |

Thank you in anticipation

Best regards

Kristian

[kullrich]

This fixed it for me, anyhow an option to select would be nice.

#esmts25_13.py

torch.set_num_threads(int(os.popen('nproc').read().strip()))
torch.cuda.set_device('cuda:1')

...

def _embed(seq):
    _, _, toks = batch_converter([("prot",seq)])
    if torch.cuda.is_available():
        toks = toks.to(device="cuda:1", non_blocking=True)
    results = esm1b(toks)
    for i in range(len(results)):
        results[i] = results[i].to(device="cpu")[0].detach().numpy()

[MesihK]

Hello @kullrich,

Thanks for providing the code that worked for you. I'll be implementing this feature on the next releases.