From 4349740270a93c1da60364bd7b5a5100eaf5ec5d Mon Sep 17 00:00:00 2001 From: Panadestein Date: Wed, 28 Oct 2026 17:12:39 +0100 Subject: [PATCH 1/5] =?UTF-8?q?feat(linalg):=20=E2=9A=A1=20run=20large=20t?= =?UTF-8?q?hin=20SVDs=20on=20the=20device?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit truncated_qr and the oversampled truncation took every SVD of the sketch's R factor on the host. From n = 512 on, cuSOLVER's gesvd on an A100 beats numpy's OpenBLAS gesdd (0.63 s against 2.2 s at n = 2048, complex128), and it converges on the ill-conditioned factors where gesdd fails. thin_svd sends factors with a side of at least 512 to the device and keeps smaller ones on the host. An oversampled kicked Ising compression at chi = 1024 drops from 148 s to 96 s. Assisted-by: Claude Code:claude-opus-5-5 --- src/src_method/_sweep.py | 10 +++++----- src/src_method/utils/__init__.py | 3 ++- src/src_method/utils/linalg.py | 34 +++++++++++++++++++++++++++----- tests/test_gpu_backend.py | 16 +++++++++++++++ 4 files changed, 52 insertions(+), 11 deletions(-) diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index f4b92de..8d08cb5 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -47,6 +47,7 @@ gaussian_sketch, new_stream, release_cached_blocks, + thin_svd, to_numpy, truncated_qr, use_device, @@ -264,20 +265,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 diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index bed085c..dd7b325 100644 --- a/src/src_method/utils/__init__.py +++ b/src/src_method/utils/__init__.py @@ -29,7 +29,7 @@ use_device, use_stream, ) -from .linalg import truncated_qr +from .linalg import thin_svd, truncated_qr __all__ = [ "Array", @@ -52,6 +52,7 @@ "pinned_empty", "release_cached_blocks", "sketch_dtype", + "thin_svd", "to_device_async", "to_host_async", "to_numpy", diff --git a/src/src_method/utils/linalg.py b/src/src_method/utils/linalg.py index b8d00f0..9382fb3 100644 --- a/src/src_method/utils/linalg.py +++ b/src/src_method/utils/linalg.py @@ -7,11 +7,17 @@ import numpy as np from array_api_compat import array_namespace -from ._backend import to_numpy +from ._backend import is_host, to_numpy if TYPE_CHECKING: from ._backend import Array +# Smallest side for which a device SVD (cuSOLVER gesvd) beats numpy's host gesdd: +# one complex128 n x n SVD on an A100 against OpenBLAS takes 0.023 s against +# 0.015 s at n = 256, 0.058 s against 0.071 s at 512 and 0.63 s against 2.2 s at +# 2048. gesvd also converges on the ill-conditioned factors where gesdd fails. +DEVICE_SVD_MIN = 512 + def truncated_qr(matrix: Array, cutoff: float) -> Array: """QR with SVD-based rank truncation, returning only the isometry. @@ -36,9 +42,27 @@ def truncated_qr(matrix: Array, cutoff: float) -> Array: transpose = matrix.shape[0] < matrix.shape[1] Q, R = xp.linalg.qr(matrix.T if transpose else matrix) - # R is small: its SVD runs on the host. - R = to_numpy(R) - U, S, _ = np.linalg.svd(R.T if transpose else R, full_matrices=False) + U, S, _ = thin_svd(R.T if transpose else R) rank = max(1, int((cutoff * S[0] <= S).sum())) - U = xp.asarray(U[:, :rank]) + U = U[:, :rank] return U if transpose else Q @ U + + +def thin_svd(matrix: Array) -> tuple[Array, np.ndarray, Array]: + """Thin SVD of a small square-ish factor, where it runs fastest. + + On the host below `DEVICE_SVD_MIN`, on the device of ``matrix`` above it. + + Args: + matrix: Input matrix of shape (m, n), on any backend. + + Returns: + ``U`` and ``Vh`` on the backend of ``matrix``, and the singular values + ``S`` as a host array, for the rank decisions. + """ + xp = array_namespace(matrix) + if is_host(xp) or min(matrix.shape) < DEVICE_SVD_MIN: + U, S, Vh = np.linalg.svd(to_numpy(matrix), full_matrices=False) + return xp.asarray(U), S, xp.asarray(Vh) + U, S, Vh = xp.linalg.svd(matrix, full_matrices=False) + return U, to_numpy(S), Vh diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 9788162..58412e0 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -23,6 +23,8 @@ make_plan, ) from src_method._sites import padded_shapes, site_bytes +from src_method.utils import thin_svd +from src_method.utils.linalg import DEVICE_SVD_MIN cupy = pytest.importorskip("cupy") @@ -508,3 +510,17 @@ def fn(rank: int) -> tuple[int, int]: return rank, stream.device_id assert group.run(fn) == [(0, 0), (1, 1)] + + +@pytest.mark.parametrize("n", [DEVICE_SVD_MIN // 2, DEVICE_SVD_MIN + 8]) +def test_thin_svd_matches_host_on_either_side_of_the_threshold(n: int) -> None: + """Below the threshold the SVD runs on the host, above it on the device.""" + rng = np.random.default_rng(3) + host = rng.standard_normal((n + 5, n)) + 1j * rng.standard_normal((n + 5, n)) + U, S, Vh = thin_svd(cupy.asarray(host)) + assert isinstance(U, cupy.ndarray) + assert isinstance(Vh, cupy.ndarray) + assert isinstance(S, np.ndarray) + np.testing.assert_allclose(S, np.linalg.svd(host, compute_uv=False), rtol=1e-10) + rebuilt = cupy.asnumpy((U * cupy.asarray(S)) @ Vh) + np.testing.assert_allclose(rebuilt, host, atol=1e-10 * np.abs(host).max()) From 4e41f81a588f93be55c78e8eb3481faabae2e666 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Sun, 1 Nov 2026 12:46:59 +0100 Subject: [PATCH 2/5] =?UTF-8?q?feat(sweep):=20=E2=9A=A1=20factor=20multi-G?= =?UTF-8?q?PU=20sketches=20by=20distributed=20CholeskyQR?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With several devices, every site of the right-to-left pass gathered the sketch on rank 0, ran a Householder QR there and scattered the rows of the output core, while the other devices waited: the serial part that kept 4-GPU runs at 1.5x of one GPU. From chi_out = 1024, the sketch is now factored where it lies: an all-to-all turns the devices' blocks of sketch columns into blocks of rows, a shifted CholeskyQR3 (Fukaya et al. 2020) orthonormalizes them exchanging only chi_out x chi_out Gram matrices, and a second all-to-all hands every device its rows of the output core, which each device also copies to the host. Sketches too ill-conditioned for it (numerically rank-deficient, as early in an evolution at cutoff 0) fall back to the Householder QR on rank 0; after repeated fallbacks the next sites skip the attempt. On four A100s, one compression of a 127-site complex128 kicked Ising stack takes 0.72x as long at chi = 1024, 0.68x at 2048 and 0.76x at 4096 (1.25x at 512, hence the threshold). Outputs are orthonormal to 2e-14, as with Householder, and agree with it to rounding. Assisted-by: Claude Code:claude-opus-5-5 --- docs/content/docs/contributing/testing.mdx | 10 +- docs/content/docs/features/large-problems.mdx | 26 ++- src/src_method/_group.py | 25 +++ src/src_method/_sweep.py | 182 +++++++++++++++--- src/src_method/utils/__init__.py | 5 +- src/src_method/utils/_backend.py | 21 ++ src/src_method/utils/linalg.py | 75 +++++++- tests/test_gpu_backend.py | 34 ++++ tests/test_group.py | 17 ++ tests/test_linalg.py | 86 +++++++++ tests/test_stack.py | 54 ++++++ 11 files changed, 498 insertions(+), 37 deletions(-) create mode 100644 tests/test_linalg.py diff --git a/docs/content/docs/contributing/testing.mdx b/docs/content/docs/contributing/testing.mdx index ce1aa0c..8c29b4a 100644 --- a/docs/content/docs/contributing/testing.mdx +++ b/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 diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx index dc6cca2..f0221f9 100644 --- a/docs/content/docs/features/large-problems.mdx +++ b/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 diff --git a/src/src_method/_group.py b/src/src_method/_group.py index 413a11f..1928b8f 100644 --- a/src/src_method/_group.py +++ b/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) diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index 8d08cb5..fcd5547 100644 --- a/src/src_method/_sweep.py +++ b/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,12 +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, @@ -65,6 +72,19 @@ logger = logging.getLogger(__name__) +# Smallest sketch size factored by `cholesky_qr` on several devices; below it the +# Householder QR on rank 0 is cheaper than the extra exchanges. On four A100s +# (NVLink), compressing a 127-site complex128 MPO stack with CholeskyQR took 1.25x +# as long at chi_out = 512, and 0.72x, 0.68x and 0.76x as long at 1024, 2048 and +# 4096. +DISTRIBUTED_QR_MIN = 1024 +# Numerical rank deficiency follows the bond profile of the operator, so it spans +# runs of neighbouring sites (all of them early in an evolution at cutoff 0, where +# the bonds are padded to chi_out), and a CholeskyQR that falls back costs almost as +# much as one that succeeds. From the second consecutive fallback on, the next 2, +# 4, 8, ... sites, at most this many, go straight to the Householder QR. +CHOLESKY_QR_MAX_SKIP = 16 + def _ranges(n: int, batch: int) -> Iterator[tuple[int, int]]: """Split ``range(n)`` into consecutive ``(lo, hi)`` batches.""" @@ -311,6 +331,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 +379,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 +459,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 +473,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 +496,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 diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index dd7b325..3f015b4 100644 --- a/src/src_method/utils/__init__.py +++ b/src/src_method/utils/__init__.py @@ -23,18 +23,20 @@ pinned_empty, release_cached_blocks, sketch_dtype, + solve_triangular, to_device_async, to_host_async, to_numpy, use_device, use_stream, ) -from .linalg import thin_svd, truncated_qr +from .linalg import cholesky_qr, thin_svd, truncated_qr __all__ = [ "Array", "NullEvent", "NullStream", + "cholesky_qr", "copy_into", "current_device", "current_stream", @@ -52,6 +54,7 @@ "pinned_empty", "release_cached_blocks", "sketch_dtype", + "solve_triangular", "thin_svd", "to_device_async", "to_host_async", diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index 3124d57..d320d85 100644 --- a/src/src_method/utils/_backend.py +++ b/src/src_method/utils/_backend.py @@ -333,6 +333,27 @@ def copy_into(dst: Array, src: Array, stream: Any) -> None: # noqa: ANN401 ) +def solve_triangular(a: Array, b: Array, *, lower: bool) -> Array: + """Solve ``a @ x = b`` for a triangular ``a``, which the array API standard lacks. + + On the device this is cuBLAS's ``trsm``. The host only runs it with simulated + devices, so a general solve stands in there. + + Args: + a: The ``(n, n)`` triangular matrix. + b: The right-hand sides, ``(n, k)``. + lower: Whether ``a`` is lower triangular. + + Returns: + ``x``, ``(n, k)``, on the backend of ``a``. + """ + if isinstance(a, np.ndarray): + return np.linalg.solve(a, b) + from cupyx.scipy.linalg import solve_triangular # noqa: PLC0415 (optional) + + return solve_triangular(a, b, lower=lower) + + def device_memory(xp: ModuleType) -> tuple[int, int]: """Return the device bytes available to the sweep and the device total. diff --git a/src/src_method/utils/linalg.py b/src/src_method/utils/linalg.py index 9382fb3..d9c2a7a 100644 --- a/src/src_method/utils/linalg.py +++ b/src/src_method/utils/linalg.py @@ -7,9 +7,11 @@ import numpy as np from array_api_compat import array_namespace -from ._backend import is_host, to_numpy +from ._backend import is_host, solve_triangular, to_numpy if TYPE_CHECKING: + from collections.abc import Callable + from ._backend import Array # Smallest side for which a device SVD (cuSOLVER gesvd) beats numpy's host gesdd: @@ -18,6 +20,11 @@ # 2048. gesvd also converges on the ill-conditioned factors where gesdd fails. DEVICE_SVD_MIN = 512 +# A plain CholeskyQR pass leaves its output orthonormal to working precision once +# the Gram matrix it factors lies within this Frobenius distance of the identity: +# the condition number of its input is then at most sqrt(1.1 / 0.9). +CHOLESKY_QR_TOL = 0.1 + def truncated_qr(matrix: Array, cutoff: float) -> Array: """QR with SVD-based rank truncation, returning only the isometry. @@ -48,6 +55,72 @@ def truncated_qr(matrix: Array, cutoff: float) -> Array: return U if transpose else Q @ U +def cholesky_qr( + block: Array, rows: int, cutoff: float, gram: Callable[[Array], Array] +) -> Array | None: + """`truncated_qr` of a tall matrix split by rows, by shifted CholeskyQR. + + The matrix ``Y`` (``rows x n``, ``rows >= n``) is held in blocks of rows, each + transposed, possibly on several devices: ``block`` is one of them. The blocks + only exchange ``n x n`` Gram matrices, through ``gram``. This is shifted + CholeskyQR3 (Fukaya et al., SIAM J. Sci. Comput. 42, A477 (2020)): a first, + shifted pass makes ``Y`` well conditioned whenever its condition number is + below about ``1 / eps``, and at most two plain passes then orthonormalize it. + The second plain pass is skipped when the first already sees an orthonormal + input. + + Args: + block: Some rows of ``Y``, transposed: ``(n, rows_block)``. + rows: The number of rows of ``Y``. + cutoff: As in `truncated_qr`. + gram: Sums the ``(n, n)`` partial Gram matrices of every block. Every + block's call must return the same bits, so that they all take the same + decisions. + + Returns: + The same rows of the isometry ``Q``, transposed: ``(rank, rows_block)``. + ``None`` if ``Y`` is too ill-conditioned for CholeskyQR, typically + numerically rank-deficient; `truncated_qr` handles that case. + """ + xp = array_namespace(block) + n = block.shape[0] + eye = xp.eye(n, dtype=block.dtype) + factors = [] + for p in range(3): + G = gram(xp.conj(block) @ block.T) + G = (G + xp.conj(G.T)) / 2 + if p == 0: + # trace(G) bounds ||Y||_2^2 from above, so the shift is safe. + unit = xp.finfo(block.dtype).eps / 2 + trace = float(xp.real(xp.linalg.trace(G))) + G = G + 11 * (rows * n + n * (n + 1)) * unit * trace * eye + last = False + else: + last = float(xp.linalg.vector_norm(G - eye)) <= CHOLESKY_QR_TOL + if p == 2 and not last: + return None + try: + L = xp.linalg.cholesky(G) + except np.linalg.LinAlgError: + return None + # A failed factorization (CuPy does not raise) leaves a pivot <= 0 or NaN. + if not float(xp.min(xp.real(xp.linalg.diagonal(L)))) > 0: + return None + # Y = Q L^H, so the blocks of Q^T are conj(L)^-1 times those of Y^T. + block = solve_triangular(xp.conj(L), block, lower=True) + factors.append(L) + if last: + break + if cutoff <= 0: + return block + R = xp.conj(factors[0].T) + for L in factors[1:]: + R = xp.conj(L.T) @ R + U, S, _ = thin_svd(R) + rank = max(1, int((cutoff * S[0] <= S).sum())) + return U[:, :rank].T @ block + + def thin_svd(matrix: Array) -> tuple[Array, np.ndarray, Array]: """Thin SVD of a small square-ish factor, where it runs fastest. diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 58412e0..d4118cd 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -5,6 +5,7 @@ from __future__ import annotations +import logging from typing import Any import numpy as np @@ -406,6 +407,39 @@ def test_two_gpus_match_one( assert list(tmp_path.iterdir()) == [] +@pytest.mark.usefixtures("two_gpus") +@pytest.mark.parametrize("cutoff", [0.0, 1e-3]) +def test_two_gpus_factor_by_cholesky_qr( + cutoff: float, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """The distributed CholeskyQR on two GPUs gives the operator of one GPU.""" + rng = np.random.default_rng(11) + stack = [random_mpo_arrays([4, 8, 8, 8, 8, 4], rng) for _ in range(3)] + one = src(*stack, chi_out=40, cutoff=cutoff, seed=2, device="gpu") + + monkeypatch.setattr(sweep_module, "DISTRIBUTED_QR_MIN", 0) + with caplog.at_level(logging.DEBUG, logger="src_method"): + two = src( + *stack, + chi_out=40, + cutoff=cutoff, + seed=2, + device="gpu", + resources=Resources(devices=2), + ) + + assert "CholeskyQR at" in caplog.text + assert [t.shape for t in two] == [t.shape for t in one] + reference = dense_mpo(one) + error = np.linalg.norm(dense_mpo(two) - reference) / np.linalg.norm(reference) + assert error < 1e-10 + for core in two[1:]: + matrix = core.reshape(core.shape[0], -1) + np.testing.assert_allclose( + matrix @ matrix.conj().T, np.eye(core.shape[0]), atol=1e-13 + ) + + @pytest.mark.usefixtures("two_gpus") def test_two_gpus_take_device_inputs() -> None: """Cores given on one GPU reach the other through the staging buffer.""" diff --git a/tests/test_group.py b/tests/test_group.py index bb3dc0a..800678c 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -29,6 +29,7 @@ def fn(rank): assert group.allgather(rank, local) is local assert group.gather(rank, local) is local assert group.scatter(rank, local) is local + assert group.alltoall(rank, [local]) is local return rank assert group.run(fn) == [0] @@ -64,6 +65,22 @@ def fn(rank): np.testing.assert_array_equal(results[0][1], expected) +@pytest.mark.parametrize("size", [2, 3]) +def test_alltoall_concatenates_what_each_rank_sent(size): + group = DeviceGroup("cpu", size) + # Rank r sends rank s an (r + 1, s + 2) block filled with 10 r + s. + sends = [ + [np.full((r + 1, s + 2), 10.0 * r + s) for s in range(size)] + for r in range(size) + ] + + results = group.run(lambda rank: group.alltoall(rank, sends[rank])) + + for rank, received in enumerate(results): + expected = np.concatenate([sends[r][rank] for r in range(size)]) + np.testing.assert_array_equal(received, expected) + + def test_many_collectives_in_a_row(): group = DeviceGroup("cpu", [0, 1, 2]) diff --git a/tests/test_linalg.py b/tests/test_linalg.py new file mode 100644 index 0000000..58d8f09 --- /dev/null +++ b/tests/test_linalg.py @@ -0,0 +1,86 @@ +"""Test the decompositions: CholeskyQR over row blocks against Householder QR.""" + +import numpy as np +import pytest + +from src_method._group import DeviceGroup, blocks +from src_method.utils import cholesky_qr, truncated_qr + + +@pytest.fixture +def rng(): + return np.random.default_rng(1234) + + +def conditioned(m, n, singular_values, rng): + """A complex ``m x n`` matrix with the given singular values.""" + + def unitary(k): + z = rng.standard_normal((m, k)) + 1j * rng.standard_normal((m, k)) + return np.linalg.qr(z)[0] + + U = unitary(n) + V = np.linalg.qr(rng.standard_normal((n, n)) + 1j * rng.standard_normal((n, n)))[0] + return (U * singular_values) @ V.conj().T + + +def split_qr(Y, size, cutoff=0.0): + """`cholesky_qr` of ``Y`` over ``size`` simulated devices; Q or ``None``.""" + group = DeviceGroup("cpu", size) + rows = blocks(Y.shape[0], size) + + def fn(rank): + lo, hi = rows[rank] + + def gram(part): + parts = group.allgather(rank, part) + return np.sum(parts.reshape(size, *part.shape), axis=0) + + return cholesky_qr(np.ascontiguousarray(Y[lo:hi].T), Y.shape[0], cutoff, gram) + + results = group.run(fn) + if any(r is None for r in results): + assert all(r is None for r in results) + return None + return np.concatenate([r.T for r in results]) + + +def projector(Q): + return Q @ Q.conj().T + + +@pytest.mark.parametrize("size", [1, 3]) +@pytest.mark.parametrize("condition", [1.0, 1e4, 1e8, 1e12]) +def test_cholesky_qr_spans_the_householder_range(size, condition, rng): + S = np.logspace(0, -np.log10(condition), 40) + Y = conditioned(160, 40, S, rng) + + Q = split_qr(Y, size) + + assert Q.shape == (160, 40) + np.testing.assert_allclose(Q.conj().T @ Q, np.eye(40), atol=1e-13) + residual = Y - projector(Q) @ Y + assert np.linalg.norm(residual) < 1e-13 * np.linalg.norm(Y) + Q_ref = truncated_qr(Y, 0.0) + np.testing.assert_allclose(projector(Q), projector(Q_ref), atol=1e-6 * condition) + + +@pytest.mark.parametrize("size", [1, 2]) +def test_cholesky_qr_truncates_like_truncated_qr(size, rng): + S = np.concatenate([np.logspace(0, -2, 25), np.logspace(-6, -9, 15)]) + Y = conditioned(120, 40, S, rng) + + Q = split_qr(Y, size, cutoff=1e-4) + + Q_ref = truncated_qr(Y, 1e-4) + assert Q.shape == Q_ref.shape == (120, 25) + np.testing.assert_allclose(Q.conj().T @ Q, np.eye(25), atol=1e-13) + np.testing.assert_allclose(projector(Q), projector(Q_ref), atol=1e-10) + + +@pytest.mark.parametrize("size", [1, 3]) +def test_cholesky_qr_declines_a_rank_deficient_matrix(size, rng): + S = np.concatenate([np.ones(30), np.zeros(10)]) + Y = conditioned(160, 40, S, rng) + + assert split_qr(Y, size) is None diff --git a/tests/test_stack.py b/tests/test_stack.py index 0a6f52e..6b9ff10 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -702,9 +702,18 @@ def test_out_of_memory_on_one_row_raises(rng, monkeypatch): # ------------------------------- +@pytest.fixture(params=[False, True], ids=["householder", "cholesky"]) +def cholesky(request, monkeypatch): + """Factor the sketches of several devices by Householder QR or by CholeskyQR.""" + if request.param: + monkeypatch.setattr(sweep_module, "DISTRIBUTED_QR_MIN", 0) + return request.param + + @pytest.mark.parametrize( ("devices", "chi"), [(2, 16), (3, 16), (3, 17), ([2, 0, 1], 16), (4, 3)] ) +@pytest.mark.usefixtures("cholesky") def test_devices_match_one_device(devices, chi, rng): stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] @@ -730,6 +739,7 @@ def test_devices_are_exact_without_truncation(spec, rng): assert rel_error(out, dense_stack(*stack)) < 1e-10 +@pytest.mark.usefixtures("cholesky") def test_devices_with_cutoff_keep_the_bonds(rng): stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(3)] @@ -740,6 +750,50 @@ def test_devices_with_cutoff_keep_the_bonds(rng): assert rel_error(many, dense(one)) < 1e-10 +@pytest.mark.parametrize("cutoff", [0.0, 0.05]) +def test_devices_factor_tall_sketches_by_cholesky_qr(cutoff, rng, monkeypatch, caplog): + stack = [random_mpo([3, 4, 4, 4, 4, 3], rng) for _ in range(3)] + one = src(*stack, chi_out=24, cutoff=cutoff, seed=6, dtype=np.complex128) + + monkeypatch.setattr(sweep_module, "DISTRIBUTED_QR_MIN", 0) + with caplog.at_level(logging.DEBUG, logger="src_method"): + many = src( + *stack, + chi_out=24, + cutoff=cutoff, + seed=6, + dtype=np.complex128, + resources=Resources(devices=3), + ) + + assert "CholeskyQR at 4 sites" in caplog.text + assert [t.shape for t in many] == [t.shape for t in one] + assert rel_error(many, dense(one)) < 1e-10 + for core in many[1:]: + matrix = core.reshape(core.shape[0], -1) + np.testing.assert_allclose( + matrix @ matrix.conj().T, np.eye(core.shape[0]), atol=1e-13 + ) + + +def test_cholesky_qr_backs_off_after_repeated_fallbacks(): + backoff = sweep_module._Backoff() + outcomes = iter([False, False, False, True, False, True]) + + tried = [] + for _ in range(10): + if backoff.attempt(): + tried.append(True) + backoff.record(succeeded=next(outcomes)) + else: + tried.append(False) + + # Fail, fail (skip 2), fail (skip 4), ...: only the second fallback in a row + # starts skipping. + assert tried[:6] == [True, True, False, False, True, False] + assert (backoff.tried, backoff.fallbacks) == (4, 3) + + def test_devices_spill_and_clean_up(rng, tmp_path, plans): stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] tight = Resources(host_memory="1MB", scratch_dir=tmp_path, devices=2) From 7e8f37d859b97027226595fd97d9ddbf52198c02 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Mon, 2 Nov 2026 10:32:42 +0100 Subject: [PATCH 3/5] =?UTF-8?q?chore(lint):=20=F0=9F=9A=A8=20let=20ty=20sk?= =?UTF-8?q?ip=20unresolved=20cupyx=20submodules?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `solve_triangular` imports `cupyx.scipy.linalg`, which the CuPy-less lint job cannot resolve; `allowed-unresolved-imports` only covered the top-level `cupyx`. Assisted-by: Claude Code:claude-opus-5-5 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f4b1098..e2dc9c4 100644 --- a/pyproject.toml +++ b/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"] From dbe645b5f62a5beb0353b073a6e44196b5560079 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Mon, 2 Nov 2026 10:32:42 +0100 Subject: [PATCH 4/5] =?UTF-8?q?test(linalg):=20=E2=9C=85=20accept=20exact?= =?UTF-8?q?=20CholeskyQR=20factorizations=20of=20rank-deficient=20sketches?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Whether rounding leaves the null space of a numerically rank-deficient sketch resolvable depends on the BLAS: with the lowest supported NumPy, CholeskyQR factored the test matrix instead of declining it (28 of 400 seeds locally too), and the result was exact (orthonormal to 7e-16, residual 4e-16). The test now accepts either outcome but checks the factorization when there is one, and a matrix with an exactly zero column covers the decline deterministically. Assisted-by: Claude Code:claude-opus-5-5 --- src/src_method/utils/linalg.py | 4 +++- tests/test_linalg.py | 22 +++++++++++++++++++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/src/src_method/utils/linalg.py b/src/src_method/utils/linalg.py index d9c2a7a..8337eee 100644 --- a/src/src_method/utils/linalg.py +++ b/src/src_method/utils/linalg.py @@ -80,7 +80,9 @@ def cholesky_qr( Returns: The same rows of the isometry ``Q``, transposed: ``(rank, rows_block)``. ``None`` if ``Y`` is too ill-conditioned for CholeskyQR, typically - numerically rank-deficient; `truncated_qr` handles that case. + numerically rank-deficient; `truncated_qr` handles that case. Whether a + rank-deficient ``Y`` is declined depends on the rounding; a returned + isometry is accurate either way. """ xp = array_namespace(block) n = block.shape[0] diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 58d8f09..4e33b3d 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -79,8 +79,24 @@ def test_cholesky_qr_truncates_like_truncated_qr(size, rng): @pytest.mark.parametrize("size", [1, 3]) -def test_cholesky_qr_declines_a_rank_deficient_matrix(size, rng): - S = np.concatenate([np.ones(30), np.zeros(10)]) - Y = conditioned(160, 40, S, rng) +def test_cholesky_qr_declines_an_exactly_singular_matrix(size, rng): + Y = conditioned(160, 40, np.ones(40), rng) + Y[:, 7] = 0 assert split_qr(Y, size) is None + + +@pytest.mark.parametrize("seed", range(8)) +@pytest.mark.parametrize("size", [1, 3]) +def test_cholesky_qr_declines_or_factors_a_rank_deficient_matrix(size, seed): + # Whether rounding leaves the null space resolvable depends on the BLAS; either + # way the result must be exact: declined, or an orthonormal basis of Y. + S = np.concatenate([np.ones(30), np.zeros(10)]) + Y = conditioned(160, 40, S, np.random.default_rng(seed)) + + Q = split_qr(Y, size) + + if Q is not None: + np.testing.assert_allclose(Q.conj().T @ Q, np.eye(40), atol=1e-13) + residual = Y - projector(Q) @ Y + assert np.linalg.norm(residual) < 1e-13 * np.linalg.norm(Y) From e2f010caeb5c96116d65e12f3058aa5c187637f9 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Mon, 2 Nov 2026 10:54:05 +0100 Subject: [PATCH 5/5] chore: clean-up verbose comments. --- src/src_method/_sweep.py | 12 ++---------- src/src_method/utils/linalg.py | 7 +------ 2 files changed, 3 insertions(+), 16 deletions(-) diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index fcd5547..f9cfbeb 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -72,17 +72,9 @@ logger = logging.getLogger(__name__) -# Smallest sketch size factored by `cholesky_qr` on several devices; below it the -# Householder QR on rank 0 is cheaper than the extra exchanges. On four A100s -# (NVLink), compressing a 127-site complex128 MPO stack with CholeskyQR took 1.25x -# as long at chi_out = 512, and 0.72x, 0.68x and 0.76x as long at 1024, 2048 and -# 4096. +# Smallest sketch size factored by `cholesky_qr` on several devices DISTRIBUTED_QR_MIN = 1024 -# Numerical rank deficiency follows the bond profile of the operator, so it spans -# runs of neighbouring sites (all of them early in an evolution at cutoff 0, where -# the bonds are padded to chi_out), and a CholeskyQR that falls back costs almost as -# much as one that succeeds. From the second consecutive fallback on, the next 2, -# 4, 8, ... sites, at most this many, go straight to the Householder QR. +# Numerical rank deficiency follows the bond profile of the operator CHOLESKY_QR_MAX_SKIP = 16 diff --git a/src/src_method/utils/linalg.py b/src/src_method/utils/linalg.py index 8337eee..509c1af 100644 --- a/src/src_method/utils/linalg.py +++ b/src/src_method/utils/linalg.py @@ -15,14 +15,9 @@ from ._backend import Array # Smallest side for which a device SVD (cuSOLVER gesvd) beats numpy's host gesdd: -# one complex128 n x n SVD on an A100 against OpenBLAS takes 0.023 s against -# 0.015 s at n = 256, 0.058 s against 0.071 s at 512 and 0.63 s against 2.2 s at -# 2048. gesvd also converges on the ill-conditioned factors where gesdd fails. DEVICE_SVD_MIN = 512 -# A plain CholeskyQR pass leaves its output orthonormal to working precision once -# the Gram matrix it factors lies within this Frobenius distance of the identity: -# the condition number of its input is then at most sqrt(1.1 / 0.9). +# A plain CholeskyQR pass leaves its output orthonormal to working precision CHOLESKY_QR_TOL = 0.1