Fail to accerlerate scvi analysis using gpu on mac M1 about scvi-tools HOT 7 CLOSED

This is a problem on the MPS PyTorch side and not much we can do about it. While lgamma is implemented now (this is the most likely source as e.g. PeakVI, which uses a reconstruction loss without lgamma, worked in the past with MPS). The current lgamma implementation is in beta state and in our hands we also yield None errors.

from scvi-tools.

This is a problem on the MPS PyTorch side and not much we can do about it. While lgamma is implemented now (this is the most likely source as e.g. PeakVI, which uses a reconstruction loss without lgamma, worked in the past with MPS). The current lgamma implementation is in beta state and in our hands we also yield None errors.

I was thinking whether it might be possible to support TensorFlow for GPU invocation... Although the speed and performance may not be as good as PyTorch, it is still significantly faster than computing with only a CPU.

from scvi-tools.

If you want to implement tensorflow scVI, I would point to https://github.com/scverse/simple-scvi. It will require major engineering and there's no intention from our side to develop this, especially as it is easy running code on Google Colab to get hardware acceleration with Nvidia GPUs (for larger datasets Mac hardware will also not be sufficient, it's still considerably slower compared to dedicated GPUs). I think when looking at macOS - MLX sounds like the way to go. I would be curious to see an implementation of scVI in MLX. Closing here as not planned.

from scvi-tools.

as I have successfully used TensorFlow for acceleration in VeloVI

Could you elaborate on this please? VeloVI is written using scvi-tools and only supports PyTorch. Did you rewrite the module to use TensorFlow?

from scvi-tools.

Just to reiterate what @canergen mentioned, this is fundamentally an issue on the PyTorch side. It does look like they are working on adding better support for the Metal backend, so it's likely that we will see stable support for this in an upcoming PyTorch release. Please hold tight until then!

Alternatively, you may wish to use our experimental JAX implementation with scvi.model.JaxSCVI and pairing that with a jax-metal installation. I can't comment on how stable this is vs. PyTorch, but it may be worth trying out if you'd like to take full advantage of the Metal backend.

from scvi-tools.

as I have successfully used TensorFlow for acceleration in VeloVI

Could you elaborate on this please? VeloVI is written using scvi-tools and only supports PyTorch. Did you rewrite the module to use TensorFlow?

I apologize for the confusion. It is true that Velovi did not utilize GPU acceleration. However, I employed another velocity analysis software called UNITvelo, which successfully invoked the GPU. The code and output are as follows:

import os
os.environ["TF_USE_LEGACY_KERAS"] = "1" 
os.chdir("/Users/myapple/Documents/RStudio/Single_cell/myproject-1209") 
import numpy as np
np.random.seed(520)
import unitvelo as utv
import tensorflow as tf
import scvelo as scv
print ("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
velo_config = utv.config.Configuration()
velo_config.R2_ADJUST = True
velo_config.IROOT = None
velo_config.FIT_OPTION = '1'
velo_config.AGENES_R2 = 1
velo_config.GPU = 0
adata = utv.run_model('tmp/seuSMC.h5ad', "celltype", config_file=velo_config)

------> Manully Specified Parameters <------
------> Model Configuration Settings <------
N_TOP_GENES:	2000
LEARNING_RATE:	0.01
FIT_OPTION:	1
DENSITY:	SVD
REORDER_CELL:	Soft_Reorder
AGGREGATE_T:	True
R2_ADJUST:	True
GENE_PRIOR:	None
VGENES:	basic
IROOT:	None
--------------------------------------------

Filtered out 37 genes that are detected 20 counts (shared).
WARNING: Did not normalize X as it looks processed already. To enforce normalization, set `enforce=True`.
WARNING: Did not normalize spliced as it looks processed already. To enforce normalization, set `enforce=True`.
WARNING: Did not normalize unspliced as it looks processed already. To enforce normalization, set `enforce=True`.
Skip filtering by dispersion since number of variables are less than `n_top_genes`.
WARNING: adata.X seems to be already log-transformed.
[/opt/miniconda3/envs/velovi/lib/python3.10/site-packages/scvelo/preprocessing/utils.py:705](https://file+.vscode-resource.vscode-cdn.net/opt/miniconda3/envs/velovi/lib/python3.10/site-packages/scvelo/preprocessing/utils.py:705): DeprecationWarning: `log1p` is deprecated since scVelo v0.3.0 and will be removed in a future version. Please use `log1p` from `scanpy.pp` instead.
  log1p(adata)
Logarithmized X.
Extracted 1963 highly variable genes.
Computing moments for 1963 genes with n_neighbors: 30 and n_pcs: 30
computing moments based on connectivities
    finished (0:00:01) --> added 
    'Ms' and 'Mu', moments of un/spliced abundances (adata.layers)

# of velocity genes 762 (Criterion: positive regression coefficient between un/spliced counts)
# of velocity genes 682 (Criterion: std of un/spliced reads should be moderate, w/o extreme values)
# of velocity genes 679 (Criterion: genes have reads in more than 5% of total cells)
Using GPU card: 0
Loss (Total): 928198.315, (Spliced): 440206.892, (Unspliced): 487991.423: 100%|█████████▉| 11999/12000 [35:39<00:00,  6.48it/s]

from scvi-tools.

Just to reiterate what @canergen mentioned, this is fundamentally an issue on the PyTorch side. It does look like they are working on adding better support for the Metal backend, so it's likely that we will see stable support for this in an upcoming PyTorch release. Please hold tight until then!

Alternatively, you may wish to use our experimental JAX implementation with scvi.model.JaxSCVI and pairing that with a jax-metal installation. I can't comment on how stable this is vs. PyTorch, but it may be worth trying out if you'd like to take full advantage of the Metal backend.

By the way, I used the function scvi.model.JaxSCVI and still encountered an error：

XlaRuntimeError: UNKNOWN: [/var/folders/ww/75_k3cyx02j9mydt1qpsk5_r0000gn/T/ipykernel_78760/4044802829.py:4:0](https://file+.vscode-resource.vscode-cdn.net/var/folders/ww/75_k3cyx02j9mydt1qpsk5_r0000gn/T/ipykernel_78760/4044802829.py:4:0): error: failed to legalize operation 'chlo.lgamma'
[/var/folders/ww/75_k3cyx02j9mydt1qpsk5_r0000gn/T/ipykernel_78760/4044802829.py:4:0](https://file+.vscode-resource.vscode-cdn.net/var/folders/ww/75_k3cyx02j9mydt1qpsk5_r0000gn/T/ipykernel_78760/4044802829.py:4:0): note: see current operation: %0 = "chlo.lgamma"(%arg0) : (tensor<128x8000xf32>) -> tensor<128x8000xf32>

from scvi-tools.