fix(resolvi): reduce peak memory during RESOLVI.load() (#3887)

Author: ori-kron-wis

CHANGELOG.md

@@ -14,19 +14,20 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
     enabling memory-efficient training on large-scale datasets stored as sharded Zarr collections,
     with support for batch covariates, {pr}`3620`.
 - Add support for rapids-singlecell, {pr}`3811`.
-- Add {class}`scvi.external.CytoVI` KNN imputation backend option to be cuML, {pr}`3821`.
+- Add {class}`scvi.external.CYTOVI` KNN imputation backend option to be cuML, {pr}`3821`.
 
 #### Fixed
 
 - Fix list of metrics to be recorded in {class}`scvi.autotune.AutotuneExperiment`, {pr}`3816`.
+- Fix {class}`scvi.external.RESOLVI` preparing data for every load, {pr}`3887`.
 
 #### Changed
 
 - Changed {class}`scvi.external.Tangram` backend to be in Pytorch, {pr}`3786`.
 - Consolidate parts of the training and data loading between {class}`~scvi.external.GIMVI`
     and {class}`scvi.external.DIAGVI`, {pr}`3830`.
-- support validation set in DestVI training and raise clear errors for unsupported validation
-    in RESOLVI, {pr}`3881`.
+- support validation set in {class}`scvi.model.DestVI` training and raise clear errors for
+    unsupported validation in {class}`scvi.external.RESOLVI`, {pr}`3881`.
 
 #### Removed
 

docs/api/developer.md

@@ -72,6 +72,7 @@ DataLoaders for loading tensors from AnnData objects. DataSplitters for splittin
 
    dataloaders.AnnDataLoader
    dataloaders.AnnTorchDataset
+   dataloaders.AnnbatchDataModule
    dataloaders.CollectionAdapter
    dataloaders.ConcatDataLoader
    dataloaders.DataSplitter

src/scvi/external/decipher/_model.py

@@ -10,7 +10,6 @@
 from scvi.data import AnnDataManager
 from scvi.data.fields import LayerField
 from scvi.model.base import BaseModelClass, PyroSviTrainMixin
-from scvi.train import PyroTrainingPlan
 from scvi.utils import setup_anndata_dsp
 
 from ._module import DecipherPyroModule
@@ -88,7 +87,7 @@ def train(
         shuffle_set_split: bool = True,
         batch_size: int = 128,
         early_stopping: bool = False,
-        training_plan: PyroTrainingPlan | None = None,
+        training_plan: DecipherTrainingPlan | None = None,
         datasplitter_kwargs: dict | None = None,
         plan_kwargs: dict | None = None,
         **trainer_kwargs,
@@ -119,12 +118,12 @@ def train(
         early_stopping
             Perform early stopping. Additional arguments can be passed in ``**trainer_kwargs``.
         training_plan
-            Training plan instance. If ``None``, a default :class:`~scvi.train.PyroTrainingPlan`
-            is used.
+            Training plan instance. If ``None``,
+            a default :class:`~scvi.train.DecipherTrainingPlan` is used.
         datasplitter_kwargs
             Additional keyword arguments passed into :class:`~scvi.dataloaders.DataSplitter`.
         plan_kwargs
-            Keyword arguments for :class:`~scvi.train.PyroTrainingPlan`.
+            Keyword arguments for :class:`~scvi.train.DecipherTrainingPlan`.
         **trainer_kwargs
             Additional keyword arguments passed to :class:`~scvi.train.Trainer`.
         """

src/scvi/external/decipher/_trainingplan.py

@@ -4,10 +4,10 @@
 from scvi.module.base import (
     PyroBaseModuleClass,
 )
-from scvi.train import LowLevelPyroTrainingPlan
+from scvi.train import PyroTrainingPlan
 
 
-class DecipherTrainingPlan(LowLevelPyroTrainingPlan):
+class DecipherTrainingPlan(PyroTrainingPlan):
     """Lightning module task to train the Decipher Pyro module.
 
     Parameters

src/scvi/external/resolvi/_model.py

@@ -355,7 +355,21 @@ def setup_anndata(
         if prepare_data:
             if prepare_data_kwargs is None:
                 prepare_data_kwargs = {}
-            RESOLVI._prepare_data(adata, batch_key=batch_key, **prepare_data_kwargs)
+            effective_config = {"batch_key": batch_key, **prepare_data_kwargs}
+            stored_config = adata.uns.get("_resolvi_prepare_data_config", None)
+            neighbors_valid = (
+                stored_config == effective_config
+                and "index_neighbor" in adata.obsm
+                and "distance_neighbor" in adata.obsm
+                and int(adata.obsm["index_neighbor"].max()) < adata.n_obs
+            )
+            if not neighbors_valid:
+                print(
+                    "Preparing data for training. This may take a while. "
+                    "RAPIDS SingleCell will be used if installed."
+                )
+                RESOLVI._prepare_data(adata, batch_key=batch_key, **prepare_data_kwargs)
+                adata.uns["_resolvi_prepare_data_config"] = effective_config
 
         anndata_fields = [
             LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True),

src/scvi/module/base/_base_module.py

@@ -375,17 +375,31 @@ def list_obs_plate_vars(self):
         return {"name": "", "in": [], "sites": {}}
 
     def on_load(self, model, **kwargs):
-        """Callback function run in :meth:`~scvi.model.base.BaseModelClass.load`.
-
-        For some Pyro modules with AutoGuides, run one training step prior to loading state dict.
-        """
+        """Callback function run in :meth:`~scvi.model.base.BaseModelClass.load`."""
         pyro.clear_param_store()
         old_history = model.history_.copy() if model.history_ is not None else None
-        model.train(max_steps=1, **self.on_load_kwargs)
         model.history_ = old_history
-        if "pyro_param_store" in kwargs:
-            # For scArches shapes are changed, and we don't want to overwrite these changed shapes.
-            pyro.get_param_store().set_state(kwargs["pyro_param_store"])
+        if "pyro_param_store" in kwargs and kwargs["pyro_param_store"] is not None:
+            store_state = kwargs["pyro_param_store"]
+            # AutoGuide parameters are lazy — they only exist as nn.Module attributes
+            # after the first forward call. Run one no-grad guide pass BEFORE restoring
+            # the saved store so the CPU module always sees a device-neutral empty store.
+            # Restoring first would expose CUDA constraint tensors (e.g. AffineTransform
+            # bounds derived from CUDA buffers) that clash with the CPU batch tensors.
+            if model.adata is not None:
+                if hasattr(self.model, "n_obs"):
+                    self.model.n_obs = model.adata.n_obs
+                if hasattr(self.guide, "n_obs"):
+                    self.guide.n_obs = model.adata.n_obs
+                dl = model._make_data_loader(model.adata, batch_size=2)
+                batch = next(iter(dl))
+                args, guide_kwargs = self._get_fn_args_from_batch(batch)
+                with torch.no_grad():
+                    self.guide(*args, **guide_kwargs)
+            # Restore the saved store after warmup. For scArches the saved store
+            # carries old reference shapes — those are intentionally kept here.
+            # load_state_dict + model.to_device() handle the final device placement.
+            pyro.get_param_store().set_state(store_state)
 
     def create_predictive(
         self,

tests/external/decipher/test_decipher.py

@@ -68,6 +68,27 @@ def test_decipher_model_methods(adata):
     assert z_cov.shape == (adata.n_vars, model.module.dim_z)
 
 
+def test_decipher_save_load(adata, save_path):
+    import os
+
+    Decipher.setup_anndata(adata)
+    model = Decipher(adata)
+    model.train(max_epochs=2, train_size=0.5)
+
+    v_before = model.get_latent_representation(give_z=False)
+    z_before = model.get_latent_representation(give_z=True)
+
+    path = os.path.join(save_path, "test_decipher")
+    model.save(path, overwrite=True, save_anndata=True)
+    model2 = Decipher.load(path)
+
+    v_after = model2.get_latent_representation(give_z=False)
+    z_after = model2.get_latent_representation(give_z=True)
+
+    np.testing.assert_array_almost_equal(v_before, v_after, decimal=4)
+    np.testing.assert_array_almost_equal(z_before, z_after, decimal=4)
+
+
 def test_decipher_trajectory(adata):
     Decipher.setup_anndata(adata)
     model = Decipher(adata)

tests/model/test_amortizedlda.py

@@ -96,7 +96,7 @@ def test_lda_model(n_topics: int):
 
 
 @pytest.mark.parametrize("n_topics", [5])
-def test_lda_model_save_load(save_path: str, n_topics: int):
+def test_lda_model_save_load(n_topics: int, save_path: str):
     adata = synthetic_iid()
     AmortizedLDA.setup_anndata(adata)
     mod = AmortizedLDA(adata, n_topics=n_topics)

tests/model/test_pyro.py

@@ -339,7 +339,7 @@ def test_pyro_bayesian_save_load(save_path):
     linear_median = quants["linear.weight"][0].detach().cpu().numpy()
 
     model_save_path = os.path.join(save_path, "test_pyro_bayesian/")
-    mod.save(model_save_path)
+    mod.save(model_save_path, overwrite=True, save_anndata=True)
 
     # Test setting `on_load_kwargs`
     mod.module.on_load_kwargs = {"batch_size": 8}