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/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"] 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 f4b92de..f9cfbeb 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,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 diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index bed085c..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 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,8 @@ "pinned_empty", "release_cached_blocks", "sketch_dtype", + "solve_triangular", + "thin_svd", "to_device_async", "to_host_async", "to_numpy", 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 b8d00f0..509c1af 100644 --- a/src/src_method/utils/linalg.py +++ b/src/src_method/utils/linalg.py @@ -7,11 +7,19 @@ import numpy as np from array_api_compat import array_namespace -from ._backend import 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: +DEVICE_SVD_MIN = 512 + +# A plain CholeskyQR pass leaves its output orthonormal to working precision +CHOLESKY_QR_TOL = 0.1 + def truncated_qr(matrix: Array, cutoff: float) -> Array: """QR with SVD-based rank truncation, returning only the isometry. @@ -36,9 +44,95 @@ 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 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. 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] + 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. + + 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..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 @@ -23,6 +24,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") @@ -404,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.""" @@ -508,3 +544,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()) 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..4e33b3d --- /dev/null +++ b/tests/test_linalg.py @@ -0,0 +1,102 @@ +"""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_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) 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)