Resolvi.load overflows memory

[jairaj-mathur]

When I load a saved model via RESOLVI.load, the memory consumption to just load the model is higher than while the model was being trained. I am using a 256gb ram and 22 GB GPU. While loading, the model, the GPU only uses about 200 MB VRAM while the RAM overshoots the 256 GB limit. During training, model used about 180 GB RAM and 2-5 GB VRAM.

Xenium dataset, 200k cells 5k genes

RESOLVI.load(/path/to/model, data, accelerator='gpu')

The python process exited with exit code 137 (SIGKILLED: Killed). This may have been caused by an OOM error. Check your commands memory usage.

Versions:

1.4.1

[canergen]

I actually assumed this was fixed. Did you run with the most recent scvi-tools version?

[ori-kron-wis]

@jairaj-mathur please use the branch I suggested in the other issue, see if this issue is solved as well.

[LizEve]

@ori-kron-wis I'm a coworker of @jairaj-mathur. I think our issue is that Databricks upgraded to scipy==1.17.0 recently. Setting scipy==1.16.3 seems to solve my OOM issue when loading models.

[ori-kron-wis]

@LizEve Thanks for letting us know.
I can confirm the memory footprint increased by 10% on my test of loading resolvi model with scipy 1.17.