Does cell2location support multi-gpu ?

[AtlasMaroc]

Please use the template below to post a question to https://discourse.scverse.org/c/ecosytem/cell2location/.

Problem

i have a large visium HD dataset, i am trying to train the model, however i run into CudaOutOfMemory issue, i tried adding trainer_kwargs argument from scvi tools, but i run into 'devices getting multiple argument' error

Description of the data input and hyperparameters

model_B408.train(max_epochs=30000,
batch_size=None,
train_size=1,
datasplitter_kwargs = {
"num_workers": 7,
"pin_memory": True, # optional (GPU)
"persistent_workers": True },
)

Single cell reference data: number of cells, number of cell types, number of genes

after filtering, 69k cells * 8588 genes

Single cell reference data: technology type (e.g. mix of 10X 3' and 5')

Spatial data: number of locations numbers, technology type (e.g. Visium, ISS, Nanostring WTA)

Visium HD, 678771 cells * 8588 genes
...

[vitkl]

Please check other multi GPU and Visium HDissues for suggested alternative solutions.

Cell2location has obs_names specific cell abundance parameters so conceptualising what multi GPU training is less straightforward. You need to do both data parallelism without minibatches (data-load once like the current default) and also put obs_names specific cell abundance parameters on specific GPUs in a paired way - AFAIK this is very hard to implement and it's not like putting all parameters and random minibatches on different GPUs.

Modelling Visium HD at 2um resolution likely requires more thinking and adding eg a CNN that sees histology and counts data and constructs cell abundance from that - while accounting for row and columns technical effects. I tried this is hires model branch but looks like the CNN architecture needs more work and optimisation - likely starting with a pre trained histology CNN for H&E and pre trained scVI encoder on scRNA to embed observed counts and Conv2D summed counts (eg 10um bins), combining the outputs and fine tuning on Visium HD. This approach would require training in spatially adjacent minibatches like the branch suggests and on multiple GPUs to simultaneously cover many tissue areas (likely 10+) - but it requires substantial extension of the model.

[AtlasMaroc]

Thank you for the detailed response @vitkl, do you advise to use GPU with higher unified memory on cloud instances ?(e.g. B300-B200, H200), in my case i don't have high resolution images, so segmentation is out of the equation, and most benchmarks i read point to cell2location having the highest score compared to ground truth.

My fear is that i would run into compatibility issue since they have much recent CUDA version (=>13.00)