diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 7e07be1d..a1e23d66 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
@@ -35,7 +35,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
@@ -56,7 +56,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
@@ -77,8 +77,8 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
-        devices: ["1", "2"]
+        python-version: ["3.12"]
+        devices: ["1", "2", "3"]
 
     env:
       TEST_DEVICES: ${{ matrix.devices }}
@@ -101,7 +101,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
@@ -122,7 +122,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
         devices: ["1", "2", "3"]
 
     env:
@@ -146,7 +146,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
         devices: ["1", "2", "3"]
 
     env:
@@ -170,7 +170,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml
index 4e3249b3..747ab8af 100644
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -22,7 +22,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v3
       with:
-        python-version: "3.10"
+        python-version: "3.12"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/.github/workflows/testpypi.yml b/.github/workflows/testpypi.yml
index cfc99a27..920936d5 100644
--- a/.github/workflows/testpypi.yml
+++ b/.github/workflows/testpypi.yml
@@ -22,7 +22,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v3
       with:
-        python-version: "3.10"
+        python-version: "3.12"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 043a79c1..cc32a48e 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -20,7 +20,7 @@
 import torch
 import yaml
 from jsonargparse import Namespace, class_from_function
-from jsonargparse._loaders_dumpers import DefaultLoader
+from jsonargparse._loaders_dumpers import get_yaml_default_loader
 from jsonargparse._util import import_object
 from lightning.pytorch.cli import ArgsType, LightningArgumentParser, LightningCLI
 from torch.utils._pytree import tree_map
@@ -184,7 +184,12 @@ class CheckpointLoader(FileLoader):
     convert_fn: Callable[[Any], Any] | str | None = None
 
     def __new__(cls, file_path, attr=None, key=None, convert_fn=None):
-        return super().__new__(cls, file_path, CellariumModule.load_from_checkpoint, attr, key, convert_fn)
+        return super().__new__(cls, file_path, _load_cellarium_module, attr, key, convert_fn)
+
+
+def _load_cellarium_module(path: str) -> CellariumModule:
+    """Wraps load_from_checkpoint, injecting weights_only=False, due to pytorch change in 2.4.0."""
+    return CellariumModule.load_from_checkpoint(path, weights_only=False)
 
 
 def file_loader_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode) -> FileLoader:
@@ -202,7 +207,7 @@ def checkpoint_loader_constructor(loader: yaml.SafeLoader, node: yaml.nodes.Mapp
     return CheckpointLoader(**loader.construct_mapping(node))  # type: ignore[arg-type]
 
 
-loader = DefaultLoader
+loader = get_yaml_default_loader()
 loader.add_constructor("!FileLoader", file_loader_constructor)
 loader.add_constructor("!FileMultiLoader", file_multi_loader_constructor)
 loader.add_constructor("!CheckpointLoader", checkpoint_loader_constructor)
@@ -681,6 +686,38 @@ def _add_instantiators(self) -> None:
             # https://github.com/Lightning-AI/pytorch-lightning/pull/18105
             pass
 
+        def _prepare_subcommand_parser(self, klass, subcommand, **kwargs):
+            """Override the default checkpoint loading with weights_only=True in torch 2.4.0+"""
+            parser = super()._prepare_subcommand_parser(klass, subcommand, **kwargs)
+            parser.set_defaults({"weights_only": False})
+            return parser
+
+        def _parse_ckpt_path(self) -> None:
+            from pathlib import Path
+
+            if not self.config.get("subcommand"):
+                return
+            ckpt_path = self.config[self.config.subcommand].get("ckpt_path")
+            if ckpt_path and Path(ckpt_path).is_file():
+                ckpt = torch.load(ckpt_path, weights_only=False, map_location="cpu")
+                hparams = ckpt.get("hyper_parameters", {})
+                hparams.pop("_instantiator", None)
+                if not hparams:
+                    return
+                if "_class_path" in hparams:
+                    hparams = {
+                        "class_path": hparams.pop("_class_path"),
+                        "dict_kwargs": hparams,
+                    }
+                hparams = {self.config.subcommand: {"model": hparams}}
+                try:
+                    self.config = self.parser.parse_object(hparams, self.config)
+                except SystemExit:
+                    import sys
+
+                    sys.stderr.write("Parsing of ckpt_path hyperparameters failed!\n")
+                    raise
+
         def before_instantiate_classes(self):
             # issue a UserWarning if the subcommand is predict and return_predictions is not set to False
             if self.subcommand == "predict":
diff --git a/cellarium/ml/core/datamodule.py b/cellarium/ml/core/datamodule.py
index b5fb437d..f0aee7ce 100644
--- a/cellarium/ml/core/datamodule.py
+++ b/cellarium/ml/core/datamodule.py
@@ -126,9 +126,7 @@ def __init__(
         pin_memory: bool = False,
     ) -> None:
         super().__init__()
-        self.save_hyperparameters(logger=False)
-        # Don't save dadc to the checkpoint
-        self.hparams["dadc"] = None
+        self.save_hyperparameters(logger=False, ignore=["dadc"])
 
         self.dadc = dadc
         # IterableDistributedAnnDataCollectionDataset args
diff --git a/cellarium/ml/core/module.py b/cellarium/ml/core/module.py
index 5c1a4982..222accc0 100644
--- a/cellarium/ml/core/module.py
+++ b/cellarium/ml/core/module.py
@@ -79,6 +79,26 @@ def __init__(
             # Thus, we need to use manual optimization for the No Optimizer case.
             self.automatic_optimization = False
 
+    @classmethod
+    def load_from_checkpoint(
+        cls,
+        checkpoint_path,
+        map_location=None,
+        hparams_file=None,
+        strict=None,
+        weights_only=False,
+        **kwargs,
+    ):
+        """Replace the torch 2.4.0+ default weights_only=True with False"""
+        return super().load_from_checkpoint(
+            checkpoint_path,
+            map_location=map_location,
+            hparams_file=hparams_file,
+            strict=strict,
+            weights_only=weights_only,
+            **kwargs,
+        )
+
     def configure_model(self) -> None:
         """
         .. note::
diff --git a/cellarium/ml/data/distributed_anndata.py b/cellarium/ml/data/distributed_anndata.py
index 60eae53b..116c6810 100644
--- a/cellarium/ml/data/distributed_anndata.py
+++ b/cellarium/ml/data/distributed_anndata.py
@@ -9,7 +9,13 @@
 import pandas as pd
 from anndata import AnnData, concat
 from anndata._core.index import _normalize_indices
-from anndata.compat import Index, Index1D
+
+try:
+    from anndata.compat import Index, Index1D
+except ImportError:
+    from anndata.typing import Index, Index1D
+
+    dtype_deprecated = True
 from anndata.experimental.multi_files._anncollection import (
     AnnCollection,
     AnnCollectionView,
@@ -18,7 +24,7 @@
 from boltons.cacheutils import LRU
 from braceexpand import braceexpand
 
-from cellarium.ml.data.fileio import read_h5ad_file
+from cellarium.ml.data.fileio import GCS_CACHE_DIR, backed_mode_default, backed_mode_type, read_h5ad_file
 from cellarium.ml.data.schema import AnnDataSchema
 
 
@@ -28,6 +34,8 @@ class getattr_mode:
 
 _GETATTR_MODE = getattr_mode()
 
+allowed_backed_modes = [None, True, False, "r"]
+
 
 @contextmanager
 def lazy_getattr():
@@ -155,6 +163,19 @@ class DistributedAnnDataCollection(AnnCollection):
         obs_columns_to_validate:
             Subset of columns to validate in the :attr:`obs` attribute.
             If ``None``, all columns are validated.
+        backed:
+            Optional backing mode for the h5ad files. ``'r'`` will leave count matrices
+            on disk until specific cell indices are queried, enabling the use of very large
+            h5ad files, while ``None`` will load entire count matrices from individual h5ad files
+            into cached memory as needed: a strategy that necessitates smaller chunked h5ad files.
+            See :func:`anndata.read_h5ad` for details on backing modes.
+        cache_dir:
+            Directory for caching GCS-downloaded h5ad files on local disk. On first access each
+            shard is downloaded and written here; subsequent accesses (including after a checkpoint
+            restart on the same machine) read from disk instead of re-downloading. If a write fails
+            (e.g. disk full) the shard is streamed from GCS instead. Set to ``None`` to always
+            stream GCS files directly into memory with no local disk usage. Only affects
+            ``gs://`` paths; local and HTTP paths are unaffected.
     """
 
     def __init__(
@@ -171,6 +192,8 @@ def __init__(
         convert: ConvertType | None = None,
         indices_strict: bool = True,
         obs_columns_to_validate: Sequence[str] | None = None,
+        backed: backed_mode_type = backed_mode_default,
+        cache_dir: str | None = str(GCS_CACHE_DIR),
     ):
         self.filenames = list(braceexpand(filenames) if isinstance(filenames, str) else filenames)
         if (shard_size is None) and (last_shard_size is not None):
@@ -192,8 +215,14 @@ def __init__(
         self.cache = LRU(max_cache_size)
         self.max_cache_size = max_cache_size
         self.cache_size_strictly_enforced = cache_size_strictly_enforced
+        if backed not in allowed_backed_modes:
+            raise ValueError(f"Invalid backed mode: {backed}. Choose from {allowed_backed_modes}")
+        self.backed = backed
+        self.cache_dir = cache_dir
         # schema
-        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0])
+        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(
+            self.filenames[0], backed=backed, cache_dir=cache_dir
+        )
         if len(adata0) != limits[0]:
             raise ValueError(
                 f"The number of cells in the first anndata file ({len(adata0)}) "
@@ -203,7 +232,7 @@ def __init__(
         self.schema = AnnDataSchema(adata0, obs_columns_to_validate)
         # lazy anndatas
         lazy_adatas = [
-            LazyAnnData(filename, (start, end), self.schema, self.cache)
+            LazyAnnData(filename, (start, end), self.schema, self.cache, backed=backed, cache_dir=cache_dir)
             for start, end, filename in zip([0] + limits, limits, self.filenames)
         ]
         # use filenames as default keys
@@ -298,10 +327,14 @@ def __getstate__(self):
     def __setstate__(self, state):
         self.__dict__.update(state)
         self.cache = LRU(self.max_cache_size)
-        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0])
+        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(
+            self.filenames[0], backed=self.backed, cache_dir=self.cache_dir
+        )
         self.schema = AnnDataSchema(adata0, self.obs_columns_to_validate)
         self.adatas = [
-            LazyAnnData(filename, (start, end), self.schema, self.cache)
+            LazyAnnData(
+                filename, (start, end), self.schema, self.cache, backed=self.backed, cache_dir=self.cache_dir
+            )
             for start, end, filename in zip([0] + self.limits, self.limits, self.filenames)
         ]
         self.obs_names = pd.Index([f"cell_{i}" for i in range(self.limits[-1])])
@@ -323,6 +356,11 @@ class LazyAnnData:
             Schema used as a reference for lazy attributes.
         cache:
             Shared LRU cache storing buffered anndatas.
+        backed:
+            Optional backing mode for the anndata. ``'r'`` will leave count matrix
+            on disk, while ``None`` will load count matrix in memory (when the anndata is
+            cached by calling the `.adata` property).
+            See :func:`anndata.read_h5ad` for details on backing modes.
     """
 
     _lazy_attrs = ["obs", "obsm", "layers", "var", "varm", "varp", "var_names"]
@@ -343,10 +381,16 @@ def __init__(
         limits: tuple[int, int],
         schema: AnnDataSchema,
         cache: LRU | None = None,
+        backed: backed_mode_type = backed_mode_default,
+        cache_dir: str | None = str(GCS_CACHE_DIR),
     ):
         self.filename = filename
         self.limits = limits
         self.schema = schema
+        if backed not in allowed_backed_modes:
+            raise ValueError(f"Invalid backed mode: {backed}. Choose from {allowed_backed_modes}")
+        self.backed = backed
+        self.cache_dir = cache_dir
         if cache is None:
             cache = LRU()
         self.cache = cache
@@ -382,16 +426,16 @@ def cached(self) -> bool:
 
     @property
     def adata(self) -> AnnData:
-        """Return backed anndata from the filename"""
+        """Return anndata from the filename"""
         try:
             adata = self.cache[self.filename]
         except KeyError:
             # fetch anndata
-            adata = read_h5ad_file(self.filename)
+            adata = read_h5ad_file(self.filename, backed=self.backed, cache_dir=self.cache_dir)
             # validate anndata
             if self.n_obs != adata.n_obs:
                 raise ValueError(
-                    "Expected `n_obs` for LazyAnnData object and backed anndata to match "
+                    "Expected `n_obs` for LazyAnnData object and loaded anndata to match "
                     f"but found {self.n_obs} and {adata.n_obs}, respectively."
                 )
             self.schema.validate_anndata(adata)
@@ -426,8 +470,9 @@ def __repr__(self) -> str:
             buffered = "Cached "
         else:
             buffered = ""
-        backed_at = f" backed at {str(self.filename)!r}"
-        descr = f"{buffered}LazyAnnData object with n_obs × n_vars = {self.n_obs} × {self.n_vars}{backed_at}"
+        located_at = f" referencing {str(self.filename)!r}"
+        backed = " in backed mode" if (self.backed in [True, "r"]) else " in memory mode"
+        descr = f"{buffered}LazyAnnData object with n_obs × n_vars = {self.n_obs} × {self.n_vars}{located_at}{backed}"
         if self.cached:
             for attr in self._all_attrs:
                 keys = getattr(self, attr).keys()
diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index c1a3229e..888367f7 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -1,18 +1,36 @@
 # Copyright Contributors to the Cellarium project.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import contextlib
+import os
 import re
 import shutil
 import tempfile
 import urllib.request
+from pathlib import Path
+from typing import Literal
 
 from anndata import AnnData, read_h5ad
+from google.auth.exceptions import DefaultCredentialsError
 from google.cloud.storage import Client
 
 url_schemes = ("http:", "https:", "ftp:")
-
-
-def read_h5ad_gcs(filename: str, storage_client: Client | None = None) -> AnnData:
+backed_mode_type = Literal["r"] | bool | None
+backed_mode_default: backed_mode_type = False
+
+# Default persistent cache for GCS-downloaded h5ad files. Shards are written
+# here on first access and reused on subsequent accesses or after a checkpoint
+# restart on the same machine. Pass cache_dir=None to stream directly from GCS
+# with no local disk usage.
+GCS_CACHE_DIR: Path = Path.home() / ".cache" / "cellarium_gcs_cache"
+
+
+def read_h5ad_gcs(
+    filename: str,
+    storage_client: Client | None = None,
+    backed: backed_mode_type = backed_mode_default,
+    cache_dir: Path | str | None = GCS_CACHE_DIR,
+) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the Google Cloud Storage.
 
@@ -22,6 +40,13 @@ def read_h5ad_gcs(filename: str, storage_client: Client | None = None) -> AnnDat
 
     Args:
         filename: Path to the data file in Cloud Storage.
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
+        cache_dir: Directory for caching downloaded files on local disk. On first
+            access the shard is saved here; subsequent accesses read from disk instead
+            of re-downloading. If a write fails (e.g. disk full), falls back to
+            streaming from GCS. Set to ``None`` to always stream with no disk usage.
     """
     if not filename.startswith("gs:"):
         raise ValueError("The filename must start with 'gs:' protocol name.")
@@ -30,16 +55,57 @@ def read_h5ad_gcs(filename: str, storage_client: Client | None = None) -> AnnDat
     bucket_name, blob_name = filename.split("/", 1)
 
     if storage_client is None:
-        storage_client = Client()
+        try:
+            storage_client = Client()
+        except DefaultCredentialsError:
+            storage_client = Client.create_anonymous_client()
 
     bucket = storage_client.bucket(bucket_name)
     blob = bucket.blob(blob_name)
 
-    with blob.open("rb") as f:
-        return read_h5ad(f)
-
-
-def read_h5ad_url(filename: str) -> AnnData:
+    if cache_dir is not None:
+        local_path = Path(cache_dir) / bucket_name / blob_name
+        if local_path.exists():
+            return read_h5ad(str(local_path), backed=backed)
+        # Cache miss — download to a staging file then rename atomically so that
+        # a concurrent worker or an interrupted run never sees a partial file.
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+        staging = local_path.with_suffix(local_path.suffix + ".download")
+        try:
+            with open(staging, "wb") as f:
+                blob.download_to_file(f)
+                f.flush()
+            staging.rename(local_path)
+            return read_h5ad(str(local_path), backed=backed)
+        except OSError:
+            # Disk full or permission error — discard the staging file and fall
+            # through to the no-cache path below.
+            with contextlib.suppress(OSError):
+                staging.unlink()
+
+    # No cache (or cache write failed) — stream without leaving anything on disk.
+    if backed not in [True, "r"]:
+        with blob.open("rb") as f:
+            return read_h5ad(f)
+
+    # Backed mode without a persistent cache: download to an anonymous temp file.
+    # Flushed and closed before h5py opens it to avoid truncation from an
+    # unflushed write buffer. The unlink removes the directory entry immediately;
+    # on Linux/macOS the inode persists until h5py closes its fd (when the
+    # AnnData is GC'd or evicted from the LRU cache).
+    with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
+        temp_path = tmp_file.name
+        blob.download_to_file(tmp_file)
+        tmp_file.flush()
+
+    try:
+        return read_h5ad(temp_path, backed=backed)
+    finally:
+        with contextlib.suppress(OSError):
+            os.unlink(temp_path)
+
+
+def read_h5ad_url(filename: str, backed: backed_mode_type = backed_mode_default) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the URL.
 
@@ -48,45 +114,84 @@ def read_h5ad_url(filename: str) -> AnnData:
         >>> adata = read_h5ad_url(
         ...     "https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad"
         ... )
+        >>> adata = read_h5ad_url(
+        ...     "https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad",
+        ...     backed='r'
+        ... )
 
     Args:
         filename: URL of the data file.
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
     """
     if not any(filename.startswith(scheme) for scheme in url_schemes):
         raise ValueError("The filename must start with 'http:', 'https:', or 'ftp:' protocol name.")
-    with urllib.request.urlopen(filename) as response:
-        with tempfile.TemporaryFile() as tmp_file:
+
+    if backed not in [True, "r"]:
+        # Anonymous TemporaryFile: no path, no flush needed, deleted automatically.
+        with urllib.request.urlopen(filename) as response:
+            with tempfile.TemporaryFile() as tmp_file:
+                shutil.copyfileobj(response, tmp_file)
+                return read_h5ad(tmp_file)
+
+    # Backed mode needs a real path; flush and close before h5py opens it.
+    with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
+        temp_path = tmp_file.name
+        with urllib.request.urlopen(filename) as response:
             shutil.copyfileobj(response, tmp_file)
-            return read_h5ad(tmp_file)
+        tmp_file.flush()
+
+    try:
+        return read_h5ad(temp_path, backed=backed)
+    finally:
+        try:
+            os.unlink(temp_path)
+        except OSError:
+            pass
 
 
-def read_h5ad_local(filename: str) -> AnnData:
+def read_h5ad_local(filename: str, backed: backed_mode_type = backed_mode_default) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the local disk.
 
     Args:
         filename: Path to the local data file.
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
     """
     if not filename.startswith("file:"):
         raise ValueError("The filename must start with 'file:' protocol name.")
     filename = re.sub(r"^file://?", "", filename)
-    return read_h5ad(filename)
+    return read_h5ad(filename, backed=backed)
 
 
-def read_h5ad_file(filename: str, **kwargs) -> AnnData:
+def read_h5ad_file(
+    filename: str,
+    backed: backed_mode_type = backed_mode_default,
+    cache_dir: Path | str | None = GCS_CACHE_DIR,
+    **kwargs,
+) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from a filename.
 
     Args:
         filename: Path to the data file.
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
+        cache_dir: Directory for caching GCS-downloaded files on local disk.
+            Only used when ``filename`` starts with ``gs://``. Set to ``None`` to
+            stream GCS files directly into memory with no disk usage.
     """
     if filename.startswith("gs:"):
-        return read_h5ad_gcs(filename, **kwargs)
+        return read_h5ad_gcs(filename, backed=backed, cache_dir=cache_dir, **kwargs)
 
     if filename.startswith("file:"):
-        return read_h5ad_local(filename)
+        return read_h5ad_local(filename, backed=backed)
 
     if any(filename.startswith(scheme) for scheme in url_schemes):
-        return read_h5ad_url(filename)
+        return read_h5ad_url(filename, backed=backed)
 
-    return read_h5ad(filename)
+    return read_h5ad(filename, backed=backed)
diff --git a/pyproject.toml b/pyproject.toml
index 8ea9b06a..b549f22b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ authors = [
 ]
 description = "Machine learning library for single-cell data analysis"
 readme = "README.rst"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 license = {file = "LICENSE.md"}
 classifiers = [
   "Intended Audience :: Developers",
@@ -21,17 +21,18 @@ classifiers = [
 ]
 
 dependencies = [
-  "anndata",
+  "anndata>=0.12.4",
   "boltons",
   "braceexpand",
   "crick>=0.0.4",
   "google-cloud-storage",
-  "jsonargparse[signatures]==4.27.7",
-  "lightning>=2.2.0,<=2.5.2",
+  "jsonargparse[signatures]>=4.28.0",
+  "lightning>=2.2.0",
   "numpy<=1.26.4",
   "pyro-ppl>=1.9.1",
   "pytest",
   "scikit-misc==0.5.1",
+  "scverse-misc==0.1.1",
   "torch>=2.2.0",
   "transformers",
   "zuko==1.6.0",
@@ -50,7 +51,7 @@ test = [
   "tensorboard",
   "torchvision",
   "scvi-tools==1.3.3",
-  "zarr==2.18.3",
+  "zarr==3.1.5",
 ]
 docs = [
   "nbsphinx",
diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index 8923872f..0647a8ca 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from pathlib import Path
+from typing import Any
 
 import lightning.pytorch as pl
 import pytest
@@ -139,7 +140,7 @@ def _check_transform_lists_match(iterable_modules1, iterable_modules2, message):
     module = _new_module()
     trainer = pl.Trainer(accelerator=accelerator, devices=1, max_steps=2, default_root_dir=tmp_path)
     print("Training one more step starting from a checkpoint...")
-    trainer.fit(module, datamodule, ckpt_path=ckpt_path)
+    trainer.fit(module, datamodule, ckpt_path=ckpt_path, weights_only=False)
     print("    ... ✓")
 
     print("Ensuring the pipeline is correctly configured in the restarted Trainer... ")
@@ -176,14 +177,15 @@ def _check_transform_lists_match(iterable_modules1, iterable_modules2, message):
     "accelerator",
     ["cpu", pytest.param("gpu", marks=pytest.mark.skipif(not USE_CUDA, reason="requires_cuda"))],
 )
-@pytest.mark.parametrize("batch_size", [50, None])
-def test_datamodule(tmp_path: Path, batch_size: int | None, accelerator: str) -> None:
+@pytest.mark.parametrize("batch_size", [50, 100])
+def test_datamodule(tmp_path: Path, batch_size: int, accelerator: str) -> None:
+    dadc = DistributedAnnDataCollection(
+        filenames="https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad",
+        shard_size=100,
+    )
     datamodule = CellariumAnnDataDataModule(
-        DistributedAnnDataCollection(
-            filenames="https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad",
-            shard_size=100,
-        ),
-        batch_size=100,
+        dadc=dadc,
+        batch_size=batch_size,
         batch_keys={
             "x_ng": AnnDataField(attr="X", convert_fn=densify),
         },
@@ -193,15 +195,166 @@ def test_datamodule(tmp_path: Path, batch_size: int | None, accelerator: str) ->
     trainer.fit(module, datamodule)
 
     ckpt_path = str(tmp_path / "lightning_logs/version_0/checkpoints/epoch=0-step=1.ckpt")
-    adata = datamodule.dadc.adatas[0].adata
-    kwargs = {"dadc": adata}
+    kwargs: dict[str, Any] = {"dadc": dadc, "weights_only": False}
     if batch_size is not None:
         kwargs["batch_size"] = batch_size
     loaded_datamodule = CellariumAnnDataDataModule.load_from_checkpoint(ckpt_path, **kwargs)
 
     assert loaded_datamodule.batch_keys == datamodule.batch_keys
     assert loaded_datamodule.batch_size == batch_size or datamodule.batch_size
-    assert loaded_datamodule.dadc is adata
+    assert loaded_datamodule.dadc is dadc
+
+
+@pytest.fixture
+def fake_massive_dense_h5ad(tmp_path: Path) -> Path:
+    import h5py
+    import numpy as np
+
+    # Create a dataset that CLAIMS to be ~40GB but uses almost no disk space
+    n_obs = 2_000_000  # 2 million cells
+    n_vars = 5_000  # 5k genes
+
+    h5ad_path = tmp_path / "massive_fake.h5ad"
+
+    with h5py.File(h5ad_path, "w") as f:
+        # Create X dataset with claimed huge size but minimal actual storage
+        # Using fillvalue=0.0 with chunking - chunks are only allocated when written to
+        f.create_dataset(
+            "X",
+            shape=(n_obs, n_vars),
+            dtype=np.float32,
+            fillvalue=0.0,
+            chunks=True,  # Enable chunking so not all data needs to be stored
+            compression=None,  # No compression to keep it simple
+        )
+
+        # Create minimal obs metadata - just the index is required
+        obs_group = f.create_group("obs")
+        # Create a small obs index but tell HDF5 it could expand to n_obs
+        obs_index_data = np.array([f"CELL_{i:07d}".encode("utf-8") for i in range(n_obs)])
+        obs_group.create_dataset("_index", data=obs_index_data, maxshape=(n_obs,), dtype="S12")
+
+        # Create minimal var metadata - just the index is required
+        var_group = f.create_group("var")
+        var_index_data = np.array([f"GENE_{i:05d}".encode("utf-8") for i in range(n_vars)])
+        var_group.create_dataset("_index", data=var_index_data, dtype="S10")
+
+        # Set minimal h5ad format attributes that anndata expects
+        f.attrs["encoding-type"] = "anndata"
+        f.attrs["encoding-version"] = "0.1.0"
+
+    return h5ad_path
+
+
+def test_datamodule_massive_h5ad_backed(tmp_path: Path, fake_massive_dense_h5ad: Path) -> None:
+    # try training using a massive (faked) h5ad file which should only succeed if backed mode works
+    dadc = DistributedAnnDataCollection(
+        filenames=str(fake_massive_dense_h5ad),  # Use full path instead of just name
+        shard_size=2_000_000,
+        backed=True,
+    )
+    datamodule = CellariumAnnDataDataModule(
+        dadc=dadc,
+        batch_size=100,
+        batch_keys={
+            "x_ng": AnnDataField(attr="X", convert_fn=None),  # already dense
+        },
+    )
+    module = CellariumModule(model=BoringModel())
+    trainer = pl.Trainer(accelerator="cpu", devices=1, max_steps=1, default_root_dir=tmp_path)
+    trainer.fit(module, datamodule)
+    # the idea is if this can run without a memory overflow, backed mode is implemented correctly
+    # we have separately verified that backed=False will crash due to 40GB memory use
+
+
+@pytest.mark.parametrize("backed", [False, True], ids=["in_memory", "backed"])
+def test_datamodule_gcs_paths(backed: bool) -> None:
+    """End-to-end regression test: dataloader works with real GCS bucket paths.
+
+    Covers backed=False (in-memory) and backed=True so that both the write-buffer
+    flush fix and any backed-mode interaction with temporary files are exercised
+    against the actual GCS storage layer.
+    """
+    gcs_filenames = [
+        "gs://dsp-cellarium-cas-public/test-data/test_0.h5ad",
+        "gs://dsp-cellarium-cas-public/test-data/test_1.h5ad",
+        "gs://dsp-cellarium-cas-public/test-data/test_2.h5ad",
+    ]
+    dadc = DistributedAnnDataCollection(
+        filenames=gcs_filenames,
+        shard_size=100,
+        backed=backed,
+    )
+    datamodule = CellariumAnnDataDataModule(
+        dadc=dadc,
+        batch_size=50,
+        batch_keys={
+            "x_ng": AnnDataField(attr="X", convert_fn=densify),
+            "var_names_g": AnnDataField(attr="var_names"),
+        },
+        num_workers=0,
+    )
+    datamodule.setup("fit")
+    batch = next(iter(datamodule.train_dataloader()))
+
+    assert "x_ng" in batch
+    assert "var_names_g" in batch
+    x_ng = batch["x_ng"]
+    var_names_g = batch["var_names_g"]
+    assert x_ng.shape == (50, len(var_names_g)), (
+        f"Unexpected batch shape {x_ng.shape}; expected (50, {len(var_names_g)})"
+    )
+
+
+def test_read_h5ad_gcs_flushes_write_buffer(tmp_path: Path) -> None:
+    """Regression test: read_h5ad_gcs must flush the write buffer before reading.
+
+    The real GCS client streams data in chunks. Python's BufferedWriter accumulates
+    those chunks; when the buffer fills (every ~8 KB) it is auto-flushed to the OS,
+    but the final partial buffer fill stays in userspace memory until flush()/close()
+    is called. If read_h5ad opens the same path before that flush, h5py sees a
+    truncated file and raises OSError ("truncated file: eof = X, stored_eof = Y").
+
+    This flush risk only applies to the backed-mode code path, which downloads to a
+    named temp file. Non-backed mode streams directly via blob.open() with no temp
+    file involved. The test therefore uses backed='r'.
+
+    A single large write would bypass the buffer entirely, so this test writes in
+    1 KB chunks without flushing — exactly replicating the real client's behavior.
+    """
+    from unittest.mock import MagicMock
+
+    import anndata as ad
+    import numpy as np
+
+    from cellarium.ml.data.fileio import read_h5ad_gcs
+
+    rng = np.random.default_rng(42)
+    adata = ad.AnnData(X=rng.random((4, 3)).astype(np.float32))
+    small_h5ad = tmp_path / "small.h5ad"
+    adata.write_h5ad(small_h5ad)
+    h5ad_bytes = small_h5ad.read_bytes()
+
+    # The real GCS client streams in chunks; a single large write bypasses the buffer
+    # entirely and reaches disk immediately (no flush needed). Writing in 1 KB chunks
+    # means the last partial buffer fill (<8 KB) stays in Python's write buffer, so
+    # h5py would see a truncated file without an explicit flush().
+    chunk_size = 1024
+
+    def download_to_file_chunked(f: object) -> None:
+        for i in range(0, len(h5ad_bytes), chunk_size):
+            f.write(h5ad_bytes[i : i + chunk_size])  # type: ignore[attr-defined]
+
+    mock_blob = MagicMock()
+    mock_blob.download_to_file.side_effect = download_to_file_chunked
+    mock_bucket = MagicMock()
+    mock_bucket.blob.return_value = mock_blob
+    mock_client = MagicMock()
+    mock_client.bucket.return_value = mock_bucket
+
+    result = read_h5ad_gcs("gs://fake-bucket/fake.h5ad", storage_client=mock_client, backed="r")
+
+    assert result.shape == adata.shape
 
 
 @pytest.mark.parametrize(
diff --git a/tests/dataloader/test_dataset.py b/tests/dataloader/test_dataset.py
index 007989ef..3e0d3f51 100644
--- a/tests/dataloader/test_dataset.py
+++ b/tests/dataloader/test_dataset.py
@@ -370,10 +370,10 @@ def test_load_from_checkpoint(
     )
     try:
         ckpt_path = tmp_path / f"checkpoints/epoch=0-step={resume_step}.ckpt"
-        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path)
+        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path, weights_only=False)
     except FileNotFoundError:
         ckpt_path = tmp_path / f"checkpoints/epoch=1-step={resume_step}.ckpt"
-        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path)
+        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path, weights_only=False)
 
     iter_data1 = collate_fn(module1.model.iter_data)
     iter_data2 = collate_fn(module2.model.iter_data)
diff --git a/tests/dataloader/test_distributed_anndata.py b/tests/dataloader/test_distributed_anndata.py
index dad39235..d8ada2a6 100644
--- a/tests/dataloader/test_distributed_anndata.py
+++ b/tests/dataloader/test_distributed_anndata.py
@@ -1,6 +1,7 @@
 # Copyright Contributors to the Cellarium project.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import importlib.util
 import os
 import pickle
 from pathlib import Path
@@ -10,6 +11,12 @@
 import pytest
 from anndata import AnnData
 
+if importlib.util.find_spec("anndata.compat.Index") is None:
+    # anndata 0.13.x +
+    dtype_deprecated = True
+else:
+    dtype_deprecated = False
+
 from cellarium.ml.data import (
     DistributedAnnDataCollection,
     IterableDistributedAnnDataCollectionDataset,
@@ -49,9 +56,10 @@ def adatas_path(tmp_path: Path):
         },
         index=[f"gene{i:03d}" for i in range(n_gene)],
     )
+    dtype_kwargs = {"dtype": X.dtype} if not dtype_deprecated else {}
     adata = AnnData(
         X,
-        dtype=X.dtype,
+        **dtype_kwargs,
         obs=obs,
         var=var,
         layers={"L": L},
@@ -69,22 +77,31 @@ def adatas_path(tmp_path: Path):
 
 @pytest.fixture
 def adt(adatas_path: Path):
-    # single anndata
-    adt = read_h5ad_file(str(os.path.join(adatas_path, "adata.h5ad")))
+    # single anndata in memory
+    adt = read_h5ad_file(str(os.path.join(adatas_path, "adata.h5ad")), backed=None)
     return adt
 
 
-@pytest.fixture(params=[(i, j) for i in (1, 2, 3) for j in (True, False)])
+@pytest.fixture(
+    params=[(i, j, k) for i in (1, 2, 3) for j in (True, False) for k in (None, "r")],
+    ids=[
+        f"cache{i}-enforce{j}-{'memory' if k is None else 'backed'}"
+        for i in (1, 2, 3)
+        for j in (True, False)
+        for k in (None, "r")
+    ],
+)
 def dat(adatas_path: Path, request: pytest.FixtureRequest):
     # distributed anndata
     filenames = str(os.path.join(adatas_path, "adata.{000..002}.h5ad"))
     limits = [2, 5, 10]
-    max_cache_size, cache_size_strictly_enforced = request.param
+    max_cache_size, cache_size_strictly_enforced, backed_mode = request.param
     dat = DistributedAnnDataCollection(
         filenames,
         limits,
         max_cache_size=max_cache_size,
         cache_size_strictly_enforced=cache_size_strictly_enforced,
+        backed=backed_mode,
     )
     return dat
 
diff --git a/tests/test_scvi.py b/tests/test_scvi.py
index 68388e8b..6e6e4bb7 100644
--- a/tests/test_scvi.py
+++ b/tests/test_scvi.py
@@ -2375,8 +2375,8 @@ def test_predict_reconstructed_counts_realdata(train_cellarium_model, train_scvi
     print(f"training data measured counts normalized to {reconstructed_library_size}")
     df = pd.DataFrame(
         (
-            np.array(train_adata_cellarium.X[:, :n_genes_to_reconstruct].todense()).squeeze()
-            / np.array(train_adata_cellarium.X.sum(axis=1)).squeeze()[:, None]
+            np.array(train_adata_cellarium.X[:, :n_genes_to_reconstruct].toarray()).squeeze()
+            / np.asarray(train_adata_cellarium.X[:].sum(axis=1)).squeeze()[:, None]  # materialize for backed case
             * reconstructed_library_size
         ),
         index=train_adata_cellarium.obs_names,