Failing to train on multiple GPUs in jupyter notebook

[chem1kal1]

After running SCVI.setup_anndata and creating the model, when training the model with devices != 1 (ex devices = 2) an error occurs:

RuntimeError: CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

import your h5ad data,

scvi.model.SCVI.setup_anndata(adata, layer = "counts", categorical_covariate_keys=['Sample'],
                             continuous_covariate_keys=['pct_counts_mt', 'total_counts', 'pct_counts_ribo'])
model = scvi.model.SCVI(adata)
model.train(devices = 2)

RuntimeError: CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Versions:

scvi-tools '1.1.2'
torch '2.2.1+cu118'

OS is ubuntu 22.04

I have tested ddp_notebook_find_unused_parameters_true
and it does not work either.
In fact, providing a strategy parameter causes it to fail.

Running on 1 gpu works fine. Sometime last year I was able to run on two gpus, I don't remember the exact time.

[martinkim0]

Hi, sorry you're running into this issue. Did you happen to try running multi-GPU training outside of the notebook? Does it work then or is it the same error?

[chem1kal1]

It is the same error when running it in a .py file.

[martinkim0]

Apologies for the delay. Is there a different error when passing in a strategy parameter? Could you try passing in the DDPStrategy using the spawn start method?

[Bo-UT]

Got the exactly same error here