fix(sweep): 🐛 split a batch that runs out of GPU memory

.gitignore

@@ -31,7 +31,6 @@ build*/
 # Executables
 *.exe
 *.out
-!benches/*/leonardo/logs/*.out
 *.app
 
 ### Python ###

benches/README.md

@@ -12,38 +12,18 @@ different SRC variants. They are grouped by the kind of workload exercised:
   wall time on small chains, and quench dynamics of 50 and 100 qubits on the GPU.
   See [`stack/README.md`](stack/README.md).
 - [`large/`](large/) — Out-of-core SRC of `N . V . M . U` with a large `M` read
-  from disk, on one GPU: plan, wall time, memory peaks and stall time. See
-  [`large/README.md`](large/README.md).
+  from disk, on one GPU or several: plan, wall time, memory peaks, stall time and
+  parallel efficiency. See [`large/README.md`](large/README.md).
 
 The scripts need the `bench` dependency group (included in `dev`):
 
 ```bash
 uv sync --group bench
 ```
 
-## Leonardo
+## HPC clusters
 
-Each folder has a `leonardo/` subfolder with a Slurm script for the Booster
-partition of [Leonardo](https://docs.hpc.cineca.it/general/getting_started.html)
-at CINECA (one node: 32 Xeon 8358 cores, 4 A100-64GB, 512 GB of memory), and the
-logs of the recorded runs. Set up the environment once, from the repository
-root on a login node:
-
-```bash
-# Keep the uv cache and Python installs out of $HOME.
-export UV_CACHE_DIR=$PWD/.uv-cache UV_PYTHON_INSTALL_DIR=$PWD/.uv-python
-uv sync --all-groups --extra gpu-nvidia
-
-# The Booster driver (535) predates CUDA 13: unpack NVIDIA's forward-compat libcuda.
-mkdir -p .cuda-compat && cd .cuda-compat
-curl -sSfO https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-compat-13-4-615.71.09-1.el8.x86_64.rpm
-rpm2cpio cuda-compat-13-4-*.rpm | cpio -idm && cd ..
-```
-
-The scripts prepend `.cuda-compat/usr/local/cuda-13.4/compat` to
-`LD_LIBRARY_PATH` (override with `CUDA_COMPAT`) and keep CuPy's kernel cache in
-`.cupy-cache/`. Submit from the `leonardo/` folder, passing the benchmark's own
-arguments, e.g. `sbatch run.sh --n-sites 25 --chi-out 1000 --device gpu`. The
-default QoS allows long runs but can queue for hours; for jobs under 30 minutes
-add `--qos=boost_qos_dbg --time=00:30:00`, which starts quickly but takes at
-most two jobs per user.
+The recorded results come from the Booster partition of Leonardo at CINECA. The
+[HPC clusters](https://algorithmiq.github.io/src-method/hpc/leonardo) section of
+the documentation covers the setup and has an example job script; pass it the
+benchmark command, e.g. `python benches/large/bench_large.py run ...`.

benches/large/README.md

@@ -46,25 +46,21 @@ size on 8xH100 or 8xH200 (and 4xA100 on Leonardo):
 1. Every distance is below `1e-10`.
 2. The parallel efficiency is at least 80% on all the GPUs of the node.
 
+`scaling` reports two distances. `distance` comes from inner products and cancels
+down to about `sqrt(eps)`, near `1.5e-8` in complex128, so it cannot show `1e-10`.
+`QR distance` sweeps a QR over the MPO `a - b` (the direct sum of the two outputs)
+and resolves differences down to rounding; the acceptance applies to it.
+
 ## Leonardo
 
-The [`leonardo/`](leonardo/) folder holds the Slurm script and the logs; see the
-[Leonardo section](../README.md#leonardo) for the environment. Booster nodes have
-no local disk and a 10 GB `/tmp`, so the script points `TMPDIR` at
-`$CINECA_SCRATCH`, and the stacks live there too (Lustre, not NVMe):
+Recorded on Booster nodes of Leonardo at CINECA; see
+[HPC clusters](https://algorithmiq.github.io/src-method/hpc/leonardo) for the setup
+and a job script. The nodes have no local disk, so the stacks and the spilled
+environments live on Lustre scratch rather than NVMe.
 
-```bash
-D=$CINECA_SCRATCH/src-large
-sbatch run.sh generate $D/m4000 --bond-m 4000
-sbatch run.sh --debug run $D/m4000 --chi-out 2000 --scratch-dir $D/scratch
-sbatch run.sh generate $D/m1000 --bond-m 1000
-sbatch run.sh --debug compare $D/m1000 --chi-out 500 --small 4GB --large 60GB \
-    --scratch-dir $D/scratch
-```
+### One GPU: D_M = 1000, chi_out = 500
 
-`generate` writes the `D_M = 1000` stack (12 GB) in 51 s.
-
-`compare` at `D_M = 1000`, `chi_out = 500`, complex128, on one A100-64GB:
+`compare --small 4GB --large 60GB`, complex128, on one A100-64GB:
 
 | GPU budget | Wall time (s) | Pool (GB) | Planned device peak (GB) | Host peak (GB) | Tiers (device/host/disk) |
 |---|---|---|---|---|---|
@@ -74,3 +70,49 @@ sbatch run.sh --debug compare $D/m1000 --chi-out 500 --small 4GB --large 60GB \
 The two plans keep opposite tiers yet give bit-identical outputs (relative distance
 `0.000e+00`). Spilling 49 of 50 environments to the host costs 18 % of the wall
 time. Both pools stay within their caps (budget plus 10 % of the card, 6.9 GB).
+
+### Several GPUs: D_M = 1000, chi_out = 500
+
+One node, 4 A100-SXM-64GB fully connected by NVLink, with peer access on every
+pair, so each GPU uploads a slice of every core and gathers the rest from its
+peers. `scaling --devices 1 2 4`, complex128, every environment on the device:
+
+| GPUs | Time (s) | Speed-up | Efficiency | QR distance |
+|---|---|---|---|---|
+| 1 | 109.9 | 1.00 | 100% | |
+| 2 | 43.6 | 2.52 | 126% | 2.8e-13 |
+| 4 | 32.7 | 3.36 | 84% | 4.4e-13 |
+
+- The single-GPU time is that of a second SRC run in its job; the first SRC run of
+  a job takes about 22 s longer, for a reason not isolated.
+- The 2-GPU point is superlinear. The plan and tiers are the same on every count,
+  and the sketch batch drops from 480 to 224 columns per GPU, which may run the
+  contractions more efficiently; this was not isolated.
+- This problem fits on one GPU, so it says little about the reference size.
+
+### Several GPUs: D_M = 4000, chi_out = 2000 (reference size)
+
+50 sites, complex128, 184 GB of input: 49 environments of 8.19 GB, 401 GB in all.
+Speed-up and efficiency are against the single-GPU run:
+
+| GPUs | Time (s) | Speed-up | Efficiency | Left-to-right (s) | Right-to-left (s) | Environments device / host / disk | QR distance |
+|---|---|---|---|---|---|---|---|
+| 1 | 3925.3 | 1.00 | 100% | 794.8 | 3112.1 | 0 / 45 / 4 | |
+| 2 | 2348.8 | 1.67 | 84% | 730.6 | 1606.3 | 1 / 35 / 14 | 7.5e-13 (vs 4 GPUs) |
+| 4 | 1099.9 | 3.57 | 89% | 240.4 | 844.1 | 1 / 45 / 3 | |
+
+- Both efficiencies are above the 80% target.
+- Before the out-of-memory split, every run at this size failed. The live memory
+  stayed within the plan, but CuPy's pool held up to 66.7 GB in split blocks and a
+  12.6-14.7 GB allocation did not fit under its 67.6 GB cap. Each GPU of the
+  multi-GPU runs now splits one sketch batch in two and finishes.
+- Four GPUs keep only one environment on the device. The plan picks the largest
+  batches that fit first, and their working set takes each GPU's budget (59.0 GB
+  planned peak), although a GPU holds only 2.05 GB of each environment.
+- Three or four environments go to disk on one and four GPUs: the host budget, 90%
+  of the available memory (about 470 GB), falls short of the 487-496 GB that
+  keeping them all in host memory takes. On two GPUs, 14 go to disk with a planned
+  host peak of 381.7 GB against 462.6 GB on four; why the plan leaves that host
+  memory unused was not investigated.
+- The stall times the sweep logs include waiting for the GPU's own queue, so they
+  cannot tell slow transfers from a GPU-bound pass.

benches/large/bench_large.py

@@ -21,7 +21,7 @@
 import json
 import logging
 import resource
-from pathlib import Path  # noqa: TC003  (cyclopts reads the annotations at runtime)
+from pathlib import Path
 from time import perf_counter
 from typing import Annotated
 
@@ -215,6 +215,31 @@ def _distance(a: list[np.ndarray], b: list[np.ndarray]) -> float:
     return float(np.sqrt(max((aa + bb - 2 * ab).real, 0.0) / aa.real))
 
 
+def _qr_distance(a: list[np.ndarray], b: list[np.ndarray]) -> float:
+    """Relative distance ``|a - b| / |a|`` from a QR sweep over the MPO ``a - b``.
+
+    `_distance` cancels down to about ``sqrt(eps)``. Here ``a - b`` is a direct sum
+    of the two trains, and the QR sweep keeps rounding relative to ``|a|``, which
+    resolves differences down to ``eps``. Gauge-free, unlike comparing cores.
+    """
+    n = len(a)
+    carry = np.ones((1, 1), dtype=np.result_type(a[0], b[0]))
+    for j, (x, y) in enumerate(zip(a, b)):
+        x = x[None] if j == 0 else x[:, None] if j == n - 1 else x
+        y = y[None] if j == 0 else y[:, None] if j == n - 1 else y
+        if j == 0:
+            site = np.concatenate([x, -y], axis=1)
+        elif j == n - 1:
+            site = np.concatenate([x, y], axis=0)
+        else:
+            (la, ra, *phys), (lb, rb, _, _) = x.shape, y.shape
+            site = np.zeros((la + lb, ra + rb, *phys), dtype=carry.dtype)
+            site[:la, :ra], site[la:, ra:] = x, y
+        site = np.tensordot(carry, site, axes=(1, 0)).transpose(0, 2, 3, 1)
+        carry = np.linalg.qr(site.reshape(-1, site.shape[-1]), mode="r")
+    return float(np.linalg.norm(carry) / np.sqrt(_inner(a, a).real))
+
+
 @app.command
 def scaling(
     directory: Path,
@@ -257,24 +282,28 @@ def scaling(
         out, seconds = _run(directory, chi_out, resources, seed, device)
         if reference is None:
             reference, ref_seconds, ref_devices = out, seconds, count
-            distance = 0.0
+            distance = qr_distance = 0.0
         else:
             distance = _distance(reference, out)
+            qr_distance = _qr_distance(reference, out)
         efficiency = ref_seconds * ref_devices / (seconds * count)
         logger.info(
-            "GPUs %d: %.1f s, speed-up %.2f, efficiency %.0f %%, distance %.3e",
+            "GPUs %d: %.1f s, speed-up %.2f, efficiency %.0f %%, distance %.3e, "
+            "QR distance %.3e",
             count,
             seconds,
             ref_seconds / seconds,
             100 * efficiency,
             distance,
+            qr_distance,
         )
         if results is not None:
             record = {
                 "gpus": count,
                 "seconds": seconds,
                 "efficiency": efficiency,
                 "distance": distance,
+                "qr_distance": qr_distance,
                 "chi_out": chi_out,
             }
             with results.open("a") as f:

benches/primitives/README.md

@@ -8,13 +8,17 @@ controlled bond dimension, and compare to a `quimb` reference.
 ## Leonardo
 
 The Leonardo [supercomputer](https://docs.hpc.cineca.it/general/getting_started.html) at CINECA is
-a pre-exascale Tier-0 EuroHPC system. The [`leonardo/`](leonardo/) folder holds the Slurm
-script and the logs; see the [Leonardo section](../README.md#leonardo) for the environment.
+a pre-exascale Tier-0 EuroHPC system; see
+[HPC clusters](https://algorithmiq.github.io/src-method/hpc/leonardo) for how to run there.
 
 ### MPO-MPO benchmark
 
 The benchmark contracts two length-`n_sites` matrix product operators (MPOs) whose initial bond dimension is `chi_out` (unless `chi_id` is specified for the second MPO), then compresses the resulting MPO back to bond dimension `chi_out`. One MPO is fully random; the other is an identity MPO perturbed by summing it with another random MPO with small elements, making the contraction highly compressible. The table below compares `quimb` (CPU, `rsvd` compression) with `src` on the CPU and on one GPU: wall time of the contraction-compression, peak host memory of the job (MaxRSS), speedup over `quimb`, and relative distance to the uncompressed product.
 
+```bash
+uv run python benches/primitives/bench_mpo_mpo.py --n-sites 25 --chi-out 1000 --compare no --device gpu
+```
+
 One Booster node (32 cores, A100-64GB), complex128 unless noted, `chi_id=4` at `chi=1000`:
 
 | Experiment        | Library             | Time (s) | MaxRSS (GB) | Speedup | Distance to reference |
@@ -60,7 +64,7 @@ Before the refactoring, on Booster (32 CPUs) and DCGP (112 CPUs) nodes:
 
 ### MPO-MPO strong scaling (v0.3.2, DCGP)
 
-For this benchmark, we fix the problem size to `n_sites=25`, `chi=1000` and `chi_id=4`, and vary the number of CPUs in a single node. We tweak the OpenMP environment variables in the `run.sh` script to optimize performance, for instance `OMP_PLACES` and `OMP_PROC_BIND`, so the NUMA domain is taken into account. We do this both manually and automatically.
+For this benchmark, we fix the problem size to `n_sites=25`, `chi=1000` and `chi_id=4`, and vary the number of CPUs in a single node. We tweak the OpenMP environment variables in the job script to optimize performance, for instance `OMP_PLACES` and `OMP_PROC_BIND`, so the NUMA domain is taken into account. We do this both manually and automatically.
 
 | CPUs | Time auto (s)   | Time maual (s)|
 |------|-----------------|---------------|