feat: ⚡ device SVDs and distributed CholeskyQR for large GPU sweeps

docs/content/docs/contributing/testing.mdx

@@ -95,9 +95,13 @@ alone, and `tests/test_gpu_backend.py` derives GPU budgets from a plan made with
 the host, through the same code as on GPUs: the split of the sketch columns, the
 collectives and the shared reading of the cores. `tests/test_stack.py` compares
 such runs with a single device, also with uneven blocks (`chi_out` not a
-multiple of `n`, or smaller than it), a cutoff and spilling. Only the peer copies
-and device selection are GPU-specific; their tests in `tests/test_gpu_backend.py`
-skip unless at least two GPUs are visible.
+multiple of `n`, or smaller than it), a cutoff and spilling. The distributed
+CholeskyQR only engages from `chi_out = DISTRIBUTED_QR_MIN` (1024), so those tests
+run both QR paths through the `cholesky` fixture, which lowers it to 0; the
+decomposition itself, including the sketches it must decline, is tested over
+simulated devices in `tests/test_linalg.py`. Only the peer copies and device
+selection are GPU-specific; their tests in `tests/test_gpu_backend.py` skip unless
+at least two GPUs are visible.
 
 ## Logging in tests
 

docs/content/docs/features/large-problems.mdx

@@ -93,8 +93,8 @@ out = src(
 )
 ```
 
-`devices` is a count, for the first `n` visible GPUs, or a list of device ids;
-the first one gathers the sketches and runs the QR. Everything runs in the calling
+`devices` is a count, for the first `n` visible GPUs, or a list of device ids,
+the first one leading. Everything runs in the calling
 process, one thread per GPU, so nothing changes for the caller: the call returns
 the output cores as with one GPU, errors are raised as usual, and the inputs are
 read once and shared. Each thread runs on a non-blocking CUDA stream of its own GPU,
@@ -106,11 +106,20 @@ current stream, so inputs given as CuPy arrays may still be in flight.
 The sketch index is split: each GPU owns a contiguous block of the `chi_out`
 sketch columns and computes only those columns of every environment and of every
 sketch, so the environments, the dominant memory, are split among the GPUs too.
-At every site of the right-to-left pass, the first GPU gathers the sketch, runs
-the QR and gives each GPU a block of rows of the output core; each GPU projects
-its rows of the next projected environment, and those rows are gathered on every
-GPU. The random draws are those of one GPU, so the result is that of one GPU up to
-rounding.
+At every site of the right-to-left pass, the sketch is orthonormalized where it
+lies when `chi_out` is at least 1024: the GPUs swap their blocks of columns for
+blocks of rows, and a shifted CholeskyQR (at most three passes, each exchanging
+one `chi_out x chi_out` Gram matrix per GPU) builds their rows of the orthonormal
+basis without moving the sketch again. Below that size, and at the sites whose
+sketch is too ill-conditioned for CholeskyQR (numerically rank-deficient, as at
+small bonds or early in an evolution), the first GPU gathers the sketch and runs a
+Householder QR instead. Either way, each GPU then has a block of rows of the output
+core, projects its rows of the next projected environment, and those rows are
+gathered on every GPU. The random draws are those of one GPU, so the result is that
+of one GPU up to rounding and a gauge: the two QRs span the same basis but may
+pick a different orthonormal frame of it, which changes the cores, not the
+operator they represent. With `DEBUG` logging, each sweep reports
+at how many sites it tried CholeskyQR and how many of those fell back.
 
 What is not split:
 
@@ -120,7 +129,8 @@ What is not split:
   GPUs reach each other directly (NVLink, or PCIe peer-to-peer), each uploads a
   slice of the core and gathers the rest from the others; otherwise each uploads
   the whole core;
-- the QR runs on the first GPU while the others wait.
+- below `chi_out = 1024`, and where CholeskyQR falls back, the QR runs on the
+  first GPU while the others wait.
 
 These costs grow more slowly with the problem than the contractions do, so large
 problems, the ones this is for, scale close to linearly, while small ones may run

pyproject.toml

@@ -268,7 +268,7 @@ python-version = "3.11"
 
 [tool.ty.analysis]
 # Optional GPU dependency, imported lazily in `utils._backend`.
-allowed-unresolved-imports = ["cupy", "cupyx"]
+allowed-unresolved-imports = ["cupy", "cupyx", "cupyx.**"]
 
 [tool.ty.src]
 include = ["src/src_method", "tests"]

src/src_method/_group.py

@@ -264,6 +264,31 @@ def scatter(self, rank: int, full: Array | None, root: int = 0) -> Array:
         self._settle(rank, parity)
         return local
 
+    def alltoall(self, rank: int, sends: Sequence[Array]) -> Array:
+        """Send ``sends[s]`` to rank ``s``; concatenate what every rank sent here.
+
+        Args:
+            rank: The calling rank.
+            sends: One C-contiguous block per rank, on this rank's device; the
+                blocks sent to one rank may differ along the first axis only.
+
+        Returns:
+            The blocks sent to this rank, concatenated along the first axis in rank
+            order, on this rank's device.
+        """
+        if self.size == 1:
+            return sends[0]
+        parity, parts = self._publish(rank, list(sends))
+        mine = [_Part(part.array[rank], part.ready) for part in parts]
+        full = self._empty_like_parts(mine)
+        lo = 0
+        for part in mine:
+            hi = lo + part.array.shape[0]
+            self._pull(part, full[lo:hi])
+            lo = hi
+        self._settle(rank, parity)
+        return full
+
     def _empty_like_parts(self, parts: Sequence[_Part]) -> Array:
         first = parts[0].array
         rows = sum(part.array.shape[0] for part in parts)

src/src_method/_sweep.py

@@ -18,17 +18,21 @@
 
 With several devices (`DeviceGroup`), each owns a contiguous block of the sketch
 columns. The left-to-right pass needs no communication: the columns are
-independent. At every site of the right-to-left pass, the blocks of the sketch are
-gathered on rank 0, which runs the QR and scatters blocks of rows of the output
-core; each device projects its rows of the new ``S``, and the rows are gathered on
-every device. The random draws, and hence the result, are those of one device.
+independent. At every site of the right-to-left pass, the sketch is factored by
+`cholesky_qr` where it is large and tall enough: an all-to-all turns the blocks of
+columns into blocks of rows, the devices exchange only Gram matrices, and a second
+all-to-all hands each device its block of rows of the output core. Elsewhere, and
+where CholeskyQR finds the sketch too ill-conditioned, the blocks are gathered on
+rank 0, which runs a Householder QR and scatters the rows. Each device then
+projects its rows of the new ``S``, and the rows are gathered on every device. The
+random draws, and hence the result, are those of one device.
 """
 
 from __future__ import annotations
 
 import logging
 import threading
-from math import ceil
+from math import ceil, prod
 from time import perf_counter_ns
 from typing import TYPE_CHECKING
 
@@ -42,11 +46,15 @@
 from ._store import EnvironmentStore
 from ._tensor_train import unpad
 from .utils import (
+    cholesky_qr,
+    current_stream,
     device_pool_bytes,
     device_pool_limit,
     gaussian_sketch,
     new_stream,
     release_cached_blocks,
+    thin_svd,
+    to_host_async,
     to_numpy,
     truncated_qr,
     use_device,
@@ -64,6 +72,11 @@
 
 logger = logging.getLogger(__name__)
 
+# Smallest sketch size factored by `cholesky_qr` on several devices
+DISTRIBUTED_QR_MIN = 1024
+# Numerical rank deficiency follows the bond profile of the operator
+CHOLESKY_QR_MAX_SKIP = 16
+
 
 def _ranges(n: int, batch: int) -> Iterator[tuple[int, int]]:
     """Split ``range(n)`` into consecutive ``(lo, hi)`` batches."""
@@ -264,20 +277,19 @@ def _truncate(
     for j in range(len(sites) - 1):
         left, right, up, down = centre.shape
         matrix = centre.transpose(0, 2, 3, 1).reshape(left * up * down, right)
-        # A tall matrix is reduced by QR first, so the SVD sees a small square. It
-        # runs on the host, as in `truncated_qr`: device SVDs are slow at this size.
+        # A tall matrix is reduced by QR first, so the SVD sees a small square.
         Q, R = xp.linalg.qr(matrix) if matrix.shape[0] > right else (None, matrix)
-        U, S, Vh = np.linalg.svd(to_numpy(R), full_matrices=False)
+        U, S, Vh = thin_svd(R)
         rank = min(chi_out, S.size)
         if cutoff > 0:
             rank = min(rank, max(1, int((cutoff * S[0] <= S).sum())))
-        U = xp.asarray(U[:, :rank])
+        U = U[:, :rank]
         U = U if Q is None else Q @ U
         sites[j] = to_numpy(
             xp.ascontiguousarray(U.reshape(left, up, down, rank).transpose(0, 3, 1, 2))
         )
         centre = xp.tensordot(
-            xp.asarray(S[:rank, None] * Vh[:rank]), xp.asarray(sites[j + 1]), axes=1
+            xp.asarray(S[:rank, None]) * Vh[:rank], xp.asarray(sites[j + 1]), axes=1
         )
     sites[-1] = to_numpy(centre)
     return sites
@@ -311,6 +323,33 @@ def get(self, j: int, up: int, down: int) -> np.ndarray:
         return omega
 
 
+class _Backoff:
+    """Whether to try `cholesky_qr` at the next site (see `CHOLESKY_QR_MAX_SKIP`).
+
+    Every rank sees the same outcomes, so every rank takes the same decisions.
+    """
+
+    def __init__(self) -> None:
+        self.tried = self.fallbacks = self.skipped = 0
+        self._failed = self._skip = 0
+
+    def attempt(self) -> bool:
+        """Whether to try at this site; counts the attempt or the skip."""
+        if self._skip:
+            self._skip -= 1
+            self.skipped += 1
+            return False
+        self.tried += 1
+        return True
+
+    def record(self, *, succeeded: bool) -> None:
+        """Record the outcome of an attempt."""
+        self._failed = 0 if succeeded else self._failed + 1
+        self.fallbacks += not succeeded
+        if self._failed > 1:
+            self._skip = min(2 ** (self._failed - 1), CHOLESKY_QR_MAX_SKIP)
+
+
 class _Sweep:
     """The per-device part of a sweep; `run` executes it on one rank."""
 
@@ -332,6 +371,10 @@ def __init__(
         self.chi_out, self.cutoff, self.sketches = chi_out, cutoff, sketches
         self.work, self.depth, self.optimize = work, depth, optimize
         self.columns = blocks(chi_out, group.size)
+        # The host output cores of the right-to-left pass, by site, written by rank
+        # 0 or, after `cholesky_qr`, by every rank in its own block of rows.
+        self._cores: dict[int, np.ndarray] = {}
+        self._lock = threading.Lock()
 
     def _tag(self, rank: int) -> str:
         if self.group.size == 1:
@@ -408,7 +451,7 @@ def _right_to_left(
         root = rank == 0
         col_lo, col_hi = self.columns[rank]
         n_local = col_hi - col_lo
-        eta_reversed: list[np.ndarray] = []
+        backoff = _Backoff()
         proj = xp.ones((1,) * (self.depth + 1), dtype=dtype)
         for j in range(len(source) - 1, 0, -1):
             cores = source.next(rank)
@@ -422,24 +465,15 @@ def _right_to_left(
                 part = kernels.sketch(store.get(j, lo, hi), cores, proj)
                 local[lo:hi] = xp.moveaxis(part, -1, 0)
             store.drop(j)
-            sketch = group.gather(rank, local)
+            shape = (proj.shape[0], up, down)
+            eta_local = None
+            tall = prod(shape) >= self.chi_out >= DISTRIBUTED_QR_MIN
+            if group.size > 1 and tall and backoff.attempt():
+                eta_local = self._cholesky_qr(rank, j, local, shape)
+                backoff.record(succeeded=eta_local is not None)
+            if eta_local is None:
+                eta_local = self._householder_qr(rank, j, local, shape)
             del local
-            eta_full = None
-            if root:
-                assert sketch is not None  # noqa: S101  (gathered on rank 0)
-                rows = proj.shape[0] * up * down
-                Q = truncated_qr(
-                    xp.reshape(sketch, (self.chi_out, rows)).T, self.cutoff
-                )
-                del sketch
-                eta_full = xp.reshape(
-                    xp.ascontiguousarray(Q.T), (-1, proj.shape[0], up, down)
-                )
-                del Q
-            eta_local = group.scatter(rank, eta_full)
-            if eta_full is not None:
-                eta_reversed.append(to_numpy(eta_full))
-                del eta_full
             new_local = xp.empty(
                 (eta_local.shape[0], *(c.shape[0] for c in cores)), dtype=dtype
             )
@@ -454,10 +488,102 @@ def _right_to_left(
         cores = source.next(rank)
         if not root:
             return None
+        if backoff.tried:
+            logger.debug(
+                "%sCholeskyQR at %d sites, %d of them too ill-conditioned for it; "
+                "%d sites skipped",
+                self._tag(rank),
+                backoff.tried,
+                backoff.fallbacks,
+                backoff.skipped,
+            )
         up, down = cores[0].shape[2], cores[-1].shape[3]
         first = xp.empty((proj.shape[0], up, down), dtype=dtype)
         for lo, hi in _ranges(proj.shape[0], plan.sites[0].project_batch):
             first[lo:hi] = xp.reshape(
                 kernels.first(cores, proj[lo:hi]), (hi - lo, up, down)
             )
-        return [to_numpy(first)[None], *reversed(eta_reversed)]
+        # The other ranks' blocks of the cores are complete once `run` returns.
+        return [to_numpy(first)[None], *(self._cores[j] for j in range(1, len(source)))]
+
+    def _householder_qr(
+        self, rank: int, j: int, local: Array, shape: tuple[int, ...]
+    ) -> Array:
+        """Factor the sketch of site ``j`` on rank 0; return this rank's rows of eta.
+
+        Args:
+            rank: The calling rank.
+            j: The site.
+            local: This rank's columns of the sketch, ``(columns, *shape)``.
+            shape: The bonds of the output core but its left one.
+
+        Returns:
+            This rank's block of rows of the output core, ``(rows, *shape)``.
+        """
+        xp, group = self.group.xp, self.group
+        sketch = group.gather(rank, local)
+        eta_full = None
+        if rank == 0:
+            assert sketch is not None  # noqa: S101  (gathered on rank 0)
+            Q = truncated_qr(
+                xp.reshape(sketch, (self.chi_out, prod(shape))).T, self.cutoff
+            )
+            del sketch
+            eta_full = xp.reshape(xp.ascontiguousarray(Q.T), (-1, *shape))
+            del Q
+        eta_local = group.scatter(rank, eta_full)
+        if eta_full is not None:
+            self._cores[j] = to_numpy(eta_full)
+        return eta_local
+
+    def _cholesky_qr(
+        self, rank: int, j: int, local: Array, shape: tuple[int, ...]
+    ) -> Array | None:
+        """Factor the sketch of site ``j`` by `cholesky_qr`, on every rank.
+
+        Args:
+            rank: The calling rank.
+            j: The site.
+            local: This rank's columns of the sketch, ``(columns, *shape)``.
+            shape: The bonds of the output core but its left one.
+
+        Returns:
+            This rank's block of rows of the output core, ``(rows, *shape)``, which
+            it also copied into the host core. ``None`` on every rank if the sketch
+            is too ill-conditioned for CholeskyQR.
+        """
+        xp, group = self.group.xp, self.group
+        rows = prod(shape)
+        flat = xp.reshape(local, (local.shape[0], rows))
+        block = group.alltoall(
+            rank,
+            [
+                xp.ascontiguousarray(flat[:, lo:hi])
+                for lo, hi in blocks(rows, group.size)
+            ],
+        )
+        del flat
+
+        def gram(part: Array) -> Array:
+            parts = group.allgather(rank, part)
+            # Summed in rank order on every rank, so the ranks see the same bits.
+            return xp.sum(xp.reshape(parts, (group.size, *part.shape)), axis=0)
+
+        Qt = cholesky_qr(block, rows, self.cutoff, gram)
+        del block
+        if Qt is None:
+            return None
+        out_rows = blocks(Qt.shape[0], group.size)
+        mine = group.alltoall(
+            rank, [xp.ascontiguousarray(Qt[lo:hi].T) for lo, hi in out_rows]
+        )
+        del Qt
+        eta_local = xp.reshape(xp.ascontiguousarray(mine.T), (-1, *shape))
+        del mine
+        with self._lock:
+            if j not in self._cores:
+                self._cores[j] = np.empty((out_rows[-1][1], *shape), dtype=self.work)
+            core = self._cores[j]
+        lo, hi = out_rows[rank]
+        to_host_async(eta_local, core[lo:hi], current_stream(xp))
+        return eta_local