From 015d86b0e7b34b46850faf0453d41116c40aadaa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Roberto=20Di=20Remigio=20Eik=C3=A5s?= Date: Sun, 27 Sep 2026 13:34:50 +0200 Subject: [PATCH 1/3] =?UTF-8?q?docs(spec):=20=F0=9F=93=9D=20design=20for?= =?UTF-8?q?=20an=20SRC=20sweep=20over=20the=20GPUs=20of=20a=20node?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 2 of the multi-GPU plan: one process per GPU, NCCL collectives, the sketch index split among the GPUs, with a communicator passed in Resources. Targets strong scaling at the Phase 1 reference size. Assisted-by: Pi:claude-opus-5-5 --- .../specs/2026-09-25-src-multi-gpu-design.md | 327 ++++++++++++++++++ 1 file changed, 327 insertions(+) create mode 100644 docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md diff --git a/docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md b/docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md new file mode 100644 index 0000000..c281db9 --- /dev/null +++ b/docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md @@ -0,0 +1,327 @@ +# Design: SRC sweep over the GPUs of a node (Phase 2) + +- Status: approved in discussion, pending spec review +- Builds on: Phase 1, the out-of-core sweep + ([design](2026-09-25-src-out-of-core-design.md)). +- Follow-up: Phase 3, the bond split across nodes; see + [Forward compatibility](#forward-compatibility-with-phase-3). + +## Goal + +Run one `src` sweep on the GPUs of a node, splitting the sketch index among them, so +that the reference problem of Phase 1 finishes close to `G` times faster on `G` +GPUs. The mathematics, the random draws and the result do not change. + +## Background + +The Khatri-Rao sketch makes the `chi_out` (`l` below) sketch columns independent in +the environment pass and in the sketch step, so they can be split among GPUs: + +- GPU `g` owns a subset `J_g` of the sketch columns and builds only those columns + of every environment, with no communication. +- At every site of the right-to-left pass, the sketch columns are all-gathered, the + thin QR is repeated on every GPU, each GPU projects a block of rows of the new + projected environment `S`, and those rows are all-gathered. +- The cores of `M` are never split: every GPU streams the same cores. + +This is Algorithm 5 of the working notes (in `sandbox/`, not versioned), mirrored +to the direction of `_sweep.py`, which builds environments left-to-right and +compresses right-to-left. + +Because every GPU holds a full core of `M`, a full `S` and the full sketch, the +per-site working set does not shrink with `G`: the largest `D_M` stays near `10^4` +in complex128 at `l = 2000` on 80 GB GPUs, and only Phase 3 raises it. Phase 2 is +about time to solution. The environments do shrink, to `1/G` per GPU, so at the +reference size they fit in the combined device memory of a node and stop spilling. + +## Decisions + +| Question | Decision | +|---|---| +| Execution model | one process per GPU (SPMD), launched with `srun` or `mpirun` | +| Collectives | NCCL through `cupy.cuda.nccl`; `mpi4py` for bootstrap and small host messages | +| Public interface | an `mpi4py` communicator in `Resources(comm=...)`; nothing is collective without it | +| Output | rank 0 in memory by default; with `Resources(output_dir=...)`, rank 0 writes `.npy` files and every rank returns memmaps | +| Large inputs | must be file-backed (`np.memmap`, zarr, HDF5), so the page cache shares them among the ranks of a node | +| Success criterion | strong scaling at the reference size | + +## Requirements + +1. With `Resources(comm=...)`, every rank calls `src` (or `apply`, `compress`) with + the same arguments and the sweep runs on all of them. +2. The result equals the single-process result for the same seed, up to + floating-point rounding. +3. Without a communicator, or with a communicator of size one, the sweep is exactly + the Phase 1 sweep. +4. `mpi4py` and NCCL stay optional: a single-process run imports neither. +5. The distributed logic runs on the CPU backend too, with `mpi4py` on host arrays, + so that CI can test it without GPUs. +6. A failure must not leave ranks waiting in a collective. + +## Non-goals + +- The bond split and multiple nodes (Phase 3). +- Sharing in-memory inputs among ranks through MPI shared memory, or one reader per + node broadcasting cores; large inputs are file-backed instead. +- GPUDirect Storage and compressed inputs; see + [Future extensions](#future-extensions). +- Heterogeneous GPUs within a node. +- Raising the largest `D_M` or `l` a site can have. + +## Public interface + +`Resources` gains two optional fields: + +```python +Resources(gpu_memory=None, host_memory=None, scratch_dir=None, + comm=None, # an mpi4py communicator; None = single process + output_dir=None) # write the output cores here as .npy files +``` + +- `comm` is typed as `Any`, so `mpi4py` is only imported when a communicator is + given. +- `output_dir` also works without `comm`: the output is written and returned as + memmaps. It lives in `Resources` so that the signatures of `src`, `apply` and + `compress` do not change. +- `src`, `apply` and `compress` return `list[NDArray] | None`: rank 0 returns the + cores and the other ranks `None`; with `output_dir`, every rank returns read-only + memmaps of the same files, after a barrier. + +Contract of a distributed call: + +- Before any work, the ranks all-gather a fingerprint of the call: layer shapes and + dtypes, `chi_out`, `cutoff`, `dtype` and `seed`. On a mismatch, every rank raises + the same `ValueError`. +- With `seed=None`, rank 0 draws a seed and broadcasts it before the generator is + created. +- A rank that sees one device (the usual case under `srun`, which sets + `CUDA_VISIBLE_DEVICES`) uses it; a rank that sees several uses device + `local_rank mod count`. +- Two-site stacks take the exact path on every rank; the output is handled as above. + +## Communicator + +A new private module `_comm.py` defines the collectives the driver needs, and three +implementations of them: + +```python +class Communicator(Protocol): + rank: int; size: int # in the communicator + local_rank: int; local_size: int # within the node + def allgather(self, local: NDArray, axis: int) -> NDArray: ... + def allgather_objects(self, value: object) -> list[object]: ... + def bcast_int(self, value: int, root: int = 0) -> int: ... + def barrier(self) -> None: ... + +def make_communicator(comm: Any | None, xp: ModuleType) -> Communicator: ... +``` + +- `allgather` concatenates the blocks of every rank along `axis`, in rank order; + blocks may differ in size along `axis`. +- `SingleComm`: identity semantics, used when `comm` is `None` or has size one. +- `MpiHostComm`: `mpi4py` `Allgatherv` on NumPy arrays, for the CPU backend. +- `NcclComm`: `cupy.cuda.nccl` on CuPy arrays, on the compute stream so that the + collectives are ordered with the kernels. Data travels as bytes, because NCCL has + no complex type, and blocks are padded to a common size, because NCCL has no + all-gather with variable counts; the padding and unpadding are pure functions. + The NCCL unique id is broadcast over `mpi4py`, and the node layout comes from + `MPI.COMM_TYPE_SHARED`. + +The driver makes the same calls in all three cases, which keeps the backend +difference behind one seam, as `AGENTS.md` requires. Phase 3 adds rail and node +communicators as more instances of the same interface. + +## Driver + +`src` (in `stack.py`) prepares the call, so that the exact path benefits as well: + +1. Build the communicator with `make_communicator(resources.comm, xp)` and select + the device. +2. All-gather the fingerprint; raise the same `ValueError` on every rank on a + mismatch. +3. Agree on the seed, then create the generator. +4. Run the exact path or the sweep, passing the communicator to `sweep`. +5. Handle the output: rank 0 in memory, or `output_dir` followed by a barrier and + memmaps on every rank. + +In `sweep`: + +- Rank `g` owns the sketch columns `J_g = range(g, l, G)`, kept locally as columns + `0 .. n_g - 1`. Each `omega_j` is drawn in full on every rank, as in Phase 1, and + sliced as `omega[g::G]`. The environments have shape `(n_g, *left_bonds)`. +- The left-to-right pass is Phase 1's, on the local columns, with no communication. +- In the right-to-left pass, at each site: + - the local sketch `(eta, up, down, n_g)` is assembled from batches, then + all-gathered along the column axis into `(eta, up, down, l)`; the columns come + out permuted, identically on every rank, which leaves the range and hence `Q_k` + unchanged; + - the thin QR runs on every rank; with `cutoff > 0`, the rank is decided on rank 0 + and broadcast, and `truncated_qr` gains an optional `rank` argument to apply it; + - rank `g` projects a contiguous block of the rows of the new `S`, in batches, and + the blocks are all-gathered along the row axis; + - only rank 0 keeps the output core, in memory or written to `output_dir`. +- The first site's output core is computed on rank 0 only. + +If the GPU QR is not bitwise reproducible, the blocks of `P_k` (the new `S`) come +from copies of `Q_k` that differ at rounding level; the effect is at rounding level +too. + +## Planner + +Per rank: + +- Environment and sketch batches are planned for `n_g` columns and the local + environment shapes; projection batches for `ceil(l / G)` rows. +- The fixed memory gains the two gather buffers, each about the size of the full + sketch or the full new `S`. The full sketch, `Q` and `S` stay on every rank. +- The host budget is a share of the node: `(host - output) / local_size` on every + rank, where `output` is the output train when it is held in memory; rank 0 also + holds the output. The scratch disk is shared the same way. +- The GPU budget is unchanged: each rank has its own device. + +Plans may differ between ranks, by one column and by the output on rank 0; the +collectives do not depend on batch sizes. + +## Errors + +A collective blocks until every rank joins it, so a failure on one rank must not +leave the others waiting. + +- Up to the start of the sweep, errors are agreed: the fingerprint check raises on + every rank, and after planning the ranks all-gather either an ok or the text of + their `MemoryError`, and all raise the same error. +- During the sweep, an exception on one rank is logged with its rank, then + `MPI.Abort(1)` ends the job. NCCL and MPI collectives cannot be cancelled, and a + clear abort is better than a job that hangs until its time limit. + +## Dependencies + +- A new extra `mpi = ["mpi4py>=4"]`, which needs an MPI library: a system Open MPI or + MPICH, or an MPI wheel from PyPI. The implementation plan settles which one CI + uses. +- The `gpu-nvidia` extra gains `nvidia-nccl-cu12`. + +## Expected scaling + +Per-site budget at the reference size (50 sites, `D_M = 4000`, `l = 2000`, +complex128): + +| Item | Cost per site | +|---|---| +| Compute on one A100 | about 110 s | +| Compute on each of 8 A100 | about 14 s | +| The two all-gathers, about 9 GB over NVLink/NVSwitch | 30-40 ms on A100, less on H100 | +| The repeated QR (32000 x 2000, complex) | 0.1-0.2 s | +| Copies of the core of `M` to each GPU, twice per site, over shared PCIe switches | 0.3-0.6 s | +| Reading `M_k` from NVMe, once per node per pass thanks to the page cache | hidden by the prefetch thread | + +The parts that do not scale cost about 3-6% at 8 GPUs, so an efficiency near 90% is +expected. On small problems the repeated QR and the collectives weigh more and the +efficiency drops, as expected. + +## Testing + +Unit tests, CPU only, in the PR suite: + +- `SingleComm` semantics, and `make_communicator(None, xp)`. +- The padding and unpadding of uneven blocks for NCCL. +- The cyclic column split and the contiguous row blocks cover every index exactly + once, also when `l` is not a multiple of `G`. +- The planner plans for `n_g` columns and `ceil(l / G)` rows, includes the gather + buffers, and shares host memory and scratch disk; rank 0 holds the output unless + `output_dir` is set. +- `output_dir` in a single process: the memmaps equal the in-memory result and the + files exist. + +MPI tests on the CPU, in the normal suite, skipped without `mpi4py` or `mpirun`: + +- Each test launches `mpirun -n 2` or `-n 3` (an uneven split) on a small script in + `tests/mpi_scripts/`; each rank writes its result to `tmp_path`, and the test + compares. Every launch has a timeout, so a hang fails the test instead of + stalling the suite. +- The distributed result equals the single-process one for the same seed, in the + dense operator to `1e-10`, with and without `cutoff`; with `seed=None` the ranks + agree. +- Rank 0 returns the cores and the other ranks `None`; with `output_dir`, every rank + returns identical memmaps. +- A fingerprint mismatch raises `ValueError` on every rank; a planning failure on + one rank raises `MemoryError` on every rank. +- Tiny budgets force the disk tier on every rank, and every scratch directory is + removed. +- An exception injected during the sweep aborts the job: nonzero exit status, and a + log line naming the rank. +- CI installs an MPI library so that these tests run on every PR. + +GPU tests in `tests/test_gpu_backend.py`, skipped without two GPUs and `mpirun`: +`mpirun -n 2` with `NcclComm` gives the operator of a single-GPU run, with an odd +`l` to exercise the padding. + +## Acceptance on the cluster + +`benches/large/` gains an MPI mode, where every rank runs `run` and reports its plan +and times, and a `scaling` command, launched with `srun` for 1, 2, 4 and 8 GPUs, +that records the wall time, the parallel efficiency `T_1 / (G T_G)` and the distance +to the single-GPU result, computed from `output_dir` memmaps with TT inner products. + +Pass criteria: + +1. The distributed result matches the single-GPU result to `1e-10` in relative + Frobenius norm. +2. At the reference size, parallel efficiency at least 80% at 8 GPUs on 8xA100 and + on 8xH100, against Phase 1 on one GPU of the same node, with the 1/2/4/8 curve + recorded; at least 90% on 2xA100. +3. Without a communicator, or with one rank, the sweep is exactly Phase 1. + +## Documentation + +- A "Several GPUs" section in `docs/large-problems.md`: launching with `srun` or + `mpirun`, `Resources(comm=..., output_dir=...)`, file-backed large inputs, how the + ranks of a node share host memory and scratch disk, and the abort on errors + during the sweep. +- `README.md` points to it; `docs/developer-guide/dependencies.md` lists the `mpi` + extra; `docs/developer-guide/testing.md` explains how to run the MPI tests. + +## Forward compatibility with Phase 3 + +Phase 3 splits the bond of `M` across nodes and keeps the sketch split within a +node. The rail communicator (the GPUs with the same index on every node) and the +node communicator become two `Communicator` instances built from `mpi4py` +sub-communicators; the driver adds the reductions of the bond split along the rails +and keeps the Phase 2 all-gathers within the node. + +## Future extensions + +- **GPUDirect Storage.** cuFile, or KvikIO from Python, moves data between NVMe and + device memory without a host bounce buffer. It fits the environment spill files + first: they are private to each rank, written by the store, and their batch + offsets are multiples of the row size, hence aligned. The site source and the + disk tier of the store would read into a ring of device buffers instead of pinned + host buffers. For inputs it bypasses the page cache that the ranks of a node + share, so it would need one reader per node broadcasting the cores over NCCL; + `.npy` files would also want 4 KiB-aligned data. Zarr with uncompressed chunks + and KvikIO's zarr support, with nvCOMP for decompression on the GPU, is the route + if compressed inputs are ever wanted. +- **Compression.** Lossless compression is not planned: the mantissas of dense + cores, and in particular of the isometries SRC produces, compress by about + 1.0-1.2, and reading is not the bottleneck. It becomes worthwhile for structured + operators (MPOs of finite-state machines or of Pauli strings are mostly exact + zeros), or if a real core of `M` compresses by more than about 1.3 with Blosc + (Zstd, bitshuffle). Lossy compression would be a numerical decision, not an I/O + one. + +## Risks + +- NCCL or MPI availability and configuration differ between clusters; the NCCL + wheel reduces the risk on the GPU side, and the CPU backend with `mpi4py` keeps + the logic testable anywhere. +- `MPI.Abort` ends the whole job, including unrelated work in the same MPI program; + this is documented. +- With eight ranks streaming the same `M` through the page cache, an `M` larger + than the free host memory is read from disk once per pass per node rather than + once per pass; the prefetch still hides it at the reference size. +- Rank 0 carries extra host memory (the output) and a little extra work (the first + site and the output copies); the plan accounts for the memory, and the work is + small. +- The copies of the cores of `M` to each GPU run on the compute stream (a Phase 1 + refinement) and cost a few percent at 8 GPUs; they are the first optimisation to + try if the measured efficiency falls short. From 2776ad73bc9bcd5d946bfb6ceea80cc311ecf524 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Roberto=20Di=20Remigio=20Eik=C3=A5s?= Date: Tue, 29 Sep 2026 09:04:10 +0200 Subject: [PATCH 2/3] =?UTF-8?q?docs(spec):=20=F0=9F=93=9D=20fold=20the=20r?= =?UTF-8?q?efinements=20found=20while=20planning=20into=20the=20spec?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An abort method on the communicator, the truncation rank agreed through a callback, output sinks, a logged warning for large in-memory inputs, and MPICH in CI and in the dev shell. Assisted-by: Pi:claude-opus-5-5 --- .../specs/2026-09-25-src-multi-gpu-design.md | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md b/docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md index c281db9..c0b6b9f 100644 --- a/docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md +++ b/docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md @@ -1,6 +1,6 @@ # Design: SRC sweep over the GPUs of a node (Phase 2) -- Status: approved in discussion, pending spec review +- Status: approved; refined while planning (see the plan's "Refinements of the spec") - Builds on: Phase 1, the out-of-core sweep ([design](2026-09-25-src-out-of-core-design.md)). - Follow-up: Phase 3, the bond split across nodes; see @@ -98,6 +98,8 @@ Contract of a distributed call: `CUDA_VISIBLE_DEVICES`) uses it; a rank that sees several uses device `local_rank mod count`. - Two-site stacks take the exact path on every rank; the output is handled as above. +- With more than one rank per node, an in-memory input layer (not a memmap) above + 1 GiB is logged as a warning: every rank holds a copy. ## Communicator @@ -112,6 +114,7 @@ class Communicator(Protocol): def allgather_objects(self, value: object) -> list[object]: ... def bcast_int(self, value: int, root: int = 0) -> int: ... def barrier(self) -> None: ... + def abort(self, code: int) -> None: ... def make_communicator(comm: Any | None, xp: ModuleType) -> Communicator: ... ``` @@ -156,10 +159,12 @@ In `sweep`: out permuted, identically on every rank, which leaves the range and hence `Q_k` unchanged; - the thin QR runs on every rank; with `cutoff > 0`, the rank is decided on rank 0 - and broadcast, and `truncated_qr` gains an optional `rank` argument to apply it; + and broadcast: `truncated_qr` gains an `agree_rank` callback that receives the + local rank and returns the one to apply; - rank `g` projects a contiguous block of the rows of the new `S`, in batches, and the blocks are all-gathered along the row axis; - - only rank 0 keeps the output core, in memory or written to `output_dir`. + - only rank 0 keeps the output core, unpadded as it is produced, and hands it to an + output sink: in memory, or written to `output_dir`. - The first site's output core is computed on rank 0 only. If the GPU QR is not bitwise reproducible, the blocks of `P_k` (the new `S`) come @@ -197,8 +202,9 @@ leave the others waiting. ## Dependencies - A new extra `mpi = ["mpi4py>=4"]`, which needs an MPI library: a system Open MPI or - MPICH, or an MPI wheel from PyPI. The implementation plan settles which one CI - uses. + MPICH, the `mpich` or `openmpi` wheels of the mpi4py project, or `impi-rt`. CI + installs the system MPICH on its Ubuntu runner, and the Nix dev shell provides + MPICH. - The `gpu-nvidia` extra gains `nvidia-nccl-cu12`. ## Expected scaling From e173b0197612748dcf5b7fc4c263bcc015cf5400 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Roberto=20Di=20Remigio=20Eik=C3=A5s?= Date: Tue, 29 Sep 2026 09:04:15 +0200 Subject: [PATCH 3/3] =?UTF-8?q?docs(plan):=20=F0=9F=93=9D=20implementation?= =?UTF-8?q?=20plan=20for=20the=20multi-GPU=20SRC=20sweep?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eleven tasks, each test-first, from the communicator to the scaling runs. Written under the no-full-code rule of the design phase: interfaces, pseudocode and test specifications. Assisted-by: Pi:claude-opus-5-5 --- .../plans/2026-09-25-src-multi-gpu.md | 503 ++++++++++++++++++ 1 file changed, 503 insertions(+) create mode 100644 docs/superpowers/plans/2026-09-25-src-multi-gpu.md diff --git a/docs/superpowers/plans/2026-09-25-src-multi-gpu.md b/docs/superpowers/plans/2026-09-25-src-multi-gpu.md new file mode 100644 index 0000000..a22a7ba --- /dev/null +++ b/docs/superpowers/plans/2026-09-25-src-multi-gpu.md @@ -0,0 +1,503 @@ +# SRC Sweep over the GPUs of a Node (Phase 2): Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Run one `src` sweep on the GPUs of a node, one process per GPU, with the sketch index split among them, so that the Phase 1 reference problem finishes close to `G` times faster on `G` GPUs. + +**Architecture:** A small `Communicator` interface in `_comm.py` (single process, `mpi4py` on host arrays, NCCL on CuPy arrays) carries the few collectives the sweep needs. `src` builds the communicator from `Resources(comm=...)`, agrees on the call and the seed, and hands the output to a sink (rank 0 in memory, or `.npy` files and memmaps on every rank). The Phase 1 driver keeps one copy of the maths: each rank owns a cyclic subset of the sketch columns, and the right-to-left pass all-gathers the sketch columns and the rows of the projected environment at every site. The planner plans each rank's share. + +**Tech Stack:** Python 3.11-3.14, NumPy and CuPy (through `src_method.utils._backend`), `opt_einsum`, `structlog`, `mpi4py>=4` (new `mpi` extra), NCCL through `cupy.cuda.nccl` (`nvidia-nccl-cu12` in the `gpu-nvidia` extra), `pytest`, `cyclopts` (benchmark), `uv`, `ruff` through `prek`, GitButler (`but`). + +**Spec:** `docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md` (read it first; this plan argues from it). Phase 1 context: `docs/superpowers/specs/2026-09-25-src-out-of-core-design.md`. + +## How to read this plan + +This plan was written under a rule of the design phase: no code in full. Every task +gives the files, the exact interfaces (signatures, types, error messages), the +behaviour in prose or pseudocode, and the tests to write, each with its setup and +its assertions. The executor writes the code, test first, and runs the commands +given. Where a detail is left to the executor, the plan says so. + +## Global Constraints + +- Work on branch `feat/src-multi-gpu`, stacked on `feat/src-out-of-core`; the spec is committed there. Commit with `but commit -b feat/src-multi-gpu -m ... ` (IDs from `but diff`), never `git add`/`git commit`, and never the untracked `.codegraph/`. +- `from __future__ import annotations` at the top of every module under `src/`; type hints everywhere; Google-style docstrings without types on every public function, class and module; `ruff` is the source of truth for style. +- Backend-specific calls (CuPy, NCCL, device selection) live in `src_method.utils._backend` or in `_comm.py`; the algorithms get `xp` and a `Communicator`. `mpi4py` and `cupy.cuda.nccl` are imported lazily, only when a communicator is given: a single-process run imports neither. +- Log with `structlog`, never `print`, and never `warnings.warn` for run-time advice (`filterwarnings = ["error"]` turns warnings into test failures). +- `src`, `apply`, `compress` stay pure: never mutate inputs. +- Commits: `(): `, ending with `Assisted-by: Pi:claude-opus-5-5`, no `Co-authored-by`. +- Before pushing: `uv run prek run --all-files` and `uv run pytest` pass. +- Decisions of the spec, verbatim: one process per GPU (SPMD) launched with `srun` or `mpirun`; NCCL through `cupy.cuda.nccl` with `mpi4py` for bootstrap; `Resources(comm=...)`, nothing collective without it; output on rank 0 in memory, or `.npy` files and memmaps on every rank with `Resources(output_dir=...)`; large inputs file-backed; success is strong scaling (at least 80% parallel efficiency at 8 GPUs on 8xA100 and 8xH100 at the reference size, at least 90% on 2xA100). +- Without a communicator, or with one of size one, the sweep is exactly Phase 1: the existing tests must pass unchanged. + +## Refinements of the spec + +Settled while planning; the executor should not undo them. + +1. **`Communicator.abort(code)`** joins the interface: the driver calls it when an exception escapes during the sweep on more than one rank. `SingleComm.abort` is never called. +2. **The truncation rank is agreed through a callback**: `truncated_qr(matrix, cutoff, xp, agree_rank=None)` computes its local rank from the singular values of `R` and passes it through `agree_rank` (rank 0's value, broadcast) before truncating. The spec's "optional `rank` argument" cannot work, because the rank is only known inside the function. +3. **Output sinks and `unpad_site`**: the output goes through a small sink object, and each core is unpadded as it is produced, with a new `unpad_site` that mirrors `pad_site`. +4. **Large in-memory inputs are logged, not warned about**: when a distributed call on more than one rank per node gets an in-memory `np.ndarray` layer (not a memmap) above 1 GiB, `src` logs a warning through `structlog` that every rank holds a copy. +5. **CI uses the system MPICH** on the Ubuntu runner (`apt-get install mpich`) with the `mpi4py` wheel from PyPI; the MPI tests skip wherever `mpiexec` or a working `mpi4py` is missing. + +## File map + +| File | Status | Responsibility | +|---|---|---| +| `src/src_method/_comm.py` | create | `Communicator`, `SingleComm`, `MpiHostComm`, `NcclComm`, `make_communicator`; block and share helpers; `agree`, `collectively` | +| `src/src_method/_output.py` | create | output sinks: memory, directory, none | +| `src/src_method/_tensor_train.py` | modify | `unpad_site` | +| `src/src_method/utils/linalg.py` | modify | `truncated_qr(..., agree_rank=None)` | +| `src/src_method/utils/_backend.py` | modify | `select_device(local_rank, xp)`, `nccl_module(xp)` | +| `src/src_method/_plan.py` | modify | `Resources.comm`, `Resources.output_dir`; per-rank shares in `resolve_budgets`; `make_plan(..., columns=, rows=, gather=, holds_output=)` | +| `src/src_method/_sweep.py` | modify | column ownership, the two all-gathers, rank-0 output, collective error handling | +| `src/src_method/stack.py` | modify | communicator, fingerprint, seed, device, sinks | +| `pyproject.toml`, `uv.lock`, `flake.nix`, `.github/workflows/test.yml` | modify | `mpi` extra, NCCL wheel, MPICH in the dev shell and in CI | +| `tests/conftest.py` | create | `mpirun` fixture | +| `tests/mpi_scripts/*.py` | create | scripts run under `mpiexec` | +| `tests/test_comm.py`, `tests/test_output.py` | create | unit and MPI tests | +| `tests/test_plan.py`, `tests/test_stack.py`, `tests/test_tensor_train.py`, `tests/test_backend.py`, `tests/test_gpu_backend.py`, `tests/test_linalg.py` | modify / create | planner, integration, helpers, GPU | +| `benches/large/bench_large.py`, `benches/large/README.md` | modify | MPI mode, `scaling` and `efficiency` commands | +| `docs/large-problems.md`, `README.md`, `docs/developer-guide/dependencies.md`, `docs/developer-guide/testing.md` | modify | documentation | + +--- + +### Task 1: Work splitting and block helpers, and the single-process communicator + +**Files:** +- Create: `src/src_method/_comm.py` +- Test: `tests/test_comm.py` (create) + +**Interfaces (produced):** + +```python +class Communicator(Protocol): + rank: int + size: int + local_rank: int + local_size: int + def allgather(self, local: NDArray, axis: int) -> NDArray: ... + def allgather_objects(self, value: object) -> list[object]: ... + def bcast_int(self, value: int, root: int = 0) -> int: ... + def barrier(self) -> None: ... + def abort(self, code: int) -> None: ... + +class SingleComm: # rank 0 of 1, local rank 0 of 1 + ... # allgather returns `local` itself; allgather_objects -> [value]; + # bcast_int -> value; barrier no-op; abort raises SystemExit(code) + +def owned_columns(n: int, rank: int, size: int) -> slice # slice(rank, n, size) +def row_block(n: int, rank: int, size: int) -> slice # contiguous, sizes differ by <= 1, + # lower ranks take the larger blocks +def pad_blocks(block: NDArray, axis: int, rows: int) -> NDArray + # move `axis` to the front, make contiguous, pad with zeros to `rows` along it +def join_blocks(gathered: NDArray, counts: Sequence[int], axis: int) -> NDArray + # gathered: (size, rows_max, *rest), padded blocks in rank order; + # take counts[r] rows of block r, concatenate, move the axis back +def agree(comm: Communicator, fn: Callable[[], T]) -> T +def collectively(comm: Communicator, fn: Callable[[], T]) -> T +``` + +Behaviour: + +- `agree` runs `fn` on every rank, then all-gathers `None` or `(type name, message)` of the exception it raised. If any rank failed, every rank raises an exception of the same type (one of `ValueError`, `MemoryError`, `TypeError`, otherwise `RuntimeError`) with the message of the lowest failing rank, prefixed `"Rank {r}: "`. Without failures it returns `fn()`'s result. With `SingleComm` it is `fn()`. +- `collectively` runs `fn`; if an exception escapes and `comm.size > 1`, it logs it with `logger.exception("SRC failed; aborting", rank=comm.rank)` and calls `comm.abort(1)`; with one rank it re-raises. + +Tests to write (`tests/test_comm.py`): + +- [ ] **Step 1: Write the failing tests** + - `owned_columns`: for `n in (1, 7, 8, 2000)` and `size in (1, 2, 3, 8)`, the ranges of all ranks partition `range(n)` exactly once; rank `r` gets `ceil((n - r) / size)` columns. + - `row_block`: same partition property; blocks are contiguous and in rank order; block sizes differ by at most one. + - `pad_blocks` then `join_blocks` round-trips random blocks of uneven sizes along axis 0 and axis 3 of a rank-4 array, float64 and complex128: joining the padded blocks of three ranks equals `np.concatenate(blocks, axis)`. + - `SingleComm`: `allgather(x, axis)` is `x`; `allgather_objects(v) == [v]`; `bcast_int(5) == 5`; `rank, size, local_rank, local_size == 0, 1, 0, 1`. + - `agree` with `SingleComm`: returns the value; re-raises a `MemoryError` from `fn` unchanged in type. + - `collectively` with `SingleComm`: re-raises. +- [ ] **Step 2: Run** `uv run pytest tests/test_comm.py -o log_cli=false`; expected: `ModuleNotFoundError: No module named 'src_method._comm'`. +- [ ] **Step 3: Implement** `_comm.py` with the interfaces above. Keep `mpi4py` and CuPy imports out of this task. +- [ ] **Step 4: Run** the tests; expected: all pass. +- [ ] **Step 5: Lint** `uv run ruff format src/src_method/_comm.py tests/test_comm.py && uv run ruff check src tests`. +- [ ] **Step 6: Commit** `feat(comm): ✨ work splitting, block helpers and the single-process communicator`. + +--- + +### Task 2: MPI on host arrays, the `mpi` extra and the MPI test harness + +**Files:** +- Modify: `src/src_method/_comm.py` (add `MpiHostComm`, `make_communicator`) +- Modify: `pyproject.toml` (extra `mpi = ["mpi4py>=4"]`), `uv.lock` (`uv lock`), `flake.nix` (add `pkgs.mpich` to the dev shell packages and its `lib` to `LD_LIBRARY_PATH`), `.github/workflows/test.yml` (Ubuntu runner: `sudo apt-get install -y mpich` before syncing, and `--extra mpi` in `uv sync`; the GPU runner is unchanged) +- Create: `tests/conftest.py`, `tests/mpi_scripts/collectives.py` +- Test: `tests/test_comm.py` (append) + +**Interfaces (produced):** + +```python +class MpiHostComm: + def __init__(self, comm: Any) -> None: ... # an mpi4py Comm + # rank/size from comm; local_rank/local_size from comm.Split_type(MPI.COMM_TYPE_SHARED) + # allgather: moveaxis + contiguous uint8 view, Allgatherv with byte counts and + # displacements from an allgather of the block lengths, then join_blocks + # allgather_objects: comm.allgather; bcast_int: comm.bcast; barrier: comm.Barrier + # abort: comm.Abort(code) + +def make_communicator(comm: Any | None, xp: ModuleType) -> Communicator: + # None or comm.Get_size() == 1 -> SingleComm() + # host backend -> MpiHostComm(comm); GPU backend -> NcclComm(comm, xp) (Task 3) +``` + +Test harness (`tests/conftest.py`): + +```python +@pytest.fixture +def mpirun(tmp_path) -> Callable[..., MpiResult]: + # skip if shutil.which("mpiexec") is None, or if + # `python -c "from mpi4py import MPI"` fails in a subprocess + # run(script: str, n: int, *args: str, timeout: float = 120) -> MpiResult + # command: mpiexec [--oversubscribe if Open MPI] -n n sys.executable + # tests/mpi_scripts/