fix(sweep): 🐛 multi-GPU on Leonardo: survive pool fragmentation, 89% efficiency at the reference size

benches/large/README.md

@@ -46,6 +46,11 @@ 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
@@ -74,3 +79,112 @@ 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 (Leonardo, 4xA100-64GB)
+
+One Booster node: 4 A100-SXM-64GB, 32 cores, 512 GB of memory, CUDA 13 through the
+forward-compatible driver (see [`../README.md`](../README.md#leonardo)). The stacks
+live on Lustre scratch. `leonardo/run_multi.sh` runs any command with the 4 GPUs
+visible to one process, prints the topology and the peer-access matrix, and with
+`WARM_DIR=<stack>` reads the stack once beforehand, so that the first run of a job
+does not read it cold. Logs and JSON lines are in [`leonardo/logs/`](leonardo/logs/).
+
+### GPU tests
+
+`pytest tests/test_gpu_backend.py` (job 59776456): all 19 tests pass, among them
+`test_two_gpus_match_one[split]`, `test_two_gpus_match_one[full]`,
+`test_two_gpus_take_device_inputs`, `test_two_gpus_raise_the_first_error` and
+`test_too_many_devices_raise`. The four GPUs are fully connected (`NV4`, four
+bonded NVLinks per pair) and peer access is on for every pair:
+
+```text
+[[0, 1, 1, 1], [1, 0, 1, 1], [1, 1, 0, 1], [1, 1, 1, 0]]
+```
+
+so each GPU uploads a slice of every core and gathers the rest over NVLink.
+
+### D_M = 1000, chi_out = 500
+
+`scaling --devices 1 2 4` (job 59782618), complex128, every environment on the
+device in all three runs:
+
+| GPUs | Time (s) | Speed-up | Efficiency | Speed-up vs 109.9 s | Efficiency vs 109.9 s | QR distance |
+|---|---|---|---|---|---|---|
+| 1 | 131.4 | 1.00 | 100% | | | |
+| 2 | 43.6 | 3.01 | 151% | 2.52 | 126% | 2.8e-13 |
+| 4 | 32.7 | 4.02 | 101% | 3.36 | 84% | 4.4e-13 |
+
+The single-GPU time depends on whether it is the first SRC run of the job. The
+same single-GPU run as the second or third run of a job (job 59783330, separate
+processes) takes 109.7 s with a 60 GB budget and 109.9 s with the detected one, as
+in the single-GPU `compare` of #59 (109.9 s); the plan
+is identical (device peak 58.8 GB, sketch batch 480, every environment on the
+device). The extra 22 s of the first run fall in both passes (left-to-right 23.8 s
+against 13.7 s, right-to-left 102.5 s against 95.4 s) with the same site stalls
+(13.3 s against 13.5 s); its cause is not isolated. The second pair of columns uses
+109.9 s.
+
+| GPUs | Left-to-right (s) | Right-to-left (s) | Site stalls (s) | Sketch batch | Planned device peak (GB) |
+|---|---|---|---|---|---|
+| 1 (second run of a job) | 13.7 | 95.4 | 13.5 | 480 | 58.8 |
+| 2, per GPU | 7.02 | 35.9 / 36.4 | 6.7 / 6.8 | 224 | 39.0 |
+| 4, per GPU | 3.93 | 28.0 (GPU 0: 28.5) | 3.6 | 96 | 32.3 |
+
+No GPU stalls on environments: they all stay on the device. The pool sizes are
+not listed: `scaling` runs every count in one process and CuPy keeps its blocks,
+so the pools report the high-water mark of all the runs so far.
+
+Where the time goes:
+
+- The left-to-right pass needs no communication and scales by 1.95 and 3.49 on 2
+  and 4 GPUs. The site stalls (the time a GPU thread waits for the loader to hand
+  over the next site's cores) shrink in proportion (13.5, 6.8, 3.6 s) and are close
+  to the left-to-right time. The timer stops before the upload, and the loader also
+  waits for the GPU to release a staging buffer, so a stall cannot tell a slow read
+  from a GPU-bound pass; the second is likely here, but it was not verified.
+- The right-to-left pass scales by 2.62 on 2 GPUs, which is superlinear, and by 3.35
+  on 4. The superlinear part is not explained by the plan or the tiers, which are
+  the same on every count; the sketch batch drops from 480 to 224 and 96 columns
+  per GPU, and smaller batches may run the contractions more efficiently, but this
+  was not isolated. On 4 GPUs the pass takes 28.5 s against 23.9 s for a quarter
+  of the single-GPU pass; GPU 0 ends 0.46 s after the others. The log does not
+  time the QR or the collectives separately, so the 4.6 s cannot be split between
+  GPU 0's QR (the others wait in the scatter) and the gathers.
+- Every QR distance is below `1e-10`. Against the first-run single-GPU time, the
+  efficiency on 4 GPUs is 101%; against the 109.9 s run it is 84%, above the 80%
+  target. This problem fits on one GPU, so it says little about the reference size.
+
+### D_M = 4000, chi_out = 2000 (reference size)
+
+50 sites, complex128, 184 GB of input on Lustre: 49 environments of 8.19 GB, 401 GB
+in all. The single-GPU `run` (job 60143214) read the stack beforehand; the
+`run --devices 4` (job 60137752) did not, so its site stalls include cold reads.
+
+| GPUs | Time (s) | Speed-up | Efficiency | Left-to-right (s) | Right-to-left (s) | Environment stalls (s) | Site stalls (s) | Environments device / host / disk |
+|---|---|---|---|---|---|---|---|---|
+| 1 | 3925.3 | 1.00 | 100% | 794.8 | 3112.1 | 2542.3 | 0.0 | 0 / 45 / 4 |
+| 4 | 1099.9 | 3.57 | 89% | 228.0-240.4 | 832.6-844.1 | 697.3-698.5 | 42.1-43.1 | 1 / 45 / 3 |
+
+Before 5fd4ecf, every run at this size failed with a CuPy out-of-memory error at a
+different point (job 59626785 of #59 on `main`, and jobs 59777837 and 59777847
+here). The live memory stayed within the plan, 52 GB against 56.7 GB planned, but
+the 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 of the four GPUs now split one sketch batch (192
+columns into two pieces) and finished; the single-GPU run did not need to.
+
+Where the time goes:
+
+- The right-to-left pass waits for environments read back from host memory and
+  disk: 82% of the pass on one GPU, 84% on four. Each GPU reads only its quarter
+  of every environment, so the wait shrinks 3.65 times. Site stalls are about 4%
+  of the run on four GPUs, and none on one, whose stack was read beforehand.
+- Four GPUs keep only one more environment on the device than one GPU. 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 Lustre: 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.
+- The efficiency on 4 GPUs is 89%, above the 80% target. The distance between the
+  outputs at this size, and the 2-GPU point, come from `scaling --devices 4 2`
+  (job 60143212), still queued.

benches/large/bench_large.py

@@ -21,13 +21,12 @@
 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
 
 import cyclopts
 import numpy as np
-
 import src_method._sweep as sweep_module
 from src_method import Resources, src
 from src_method._plan import make_plan
@@ -129,7 +128,7 @@ def spy(*args: object, **kwargs: object) -> object:
     tiers = [site.tier for site in plan.sites]
     pools = []
     if device == "gpu":
-        import cupy  # noqa: PLC0415  (optional dependency)
+        import cupy
 
         devices = resources.devices or 1
         for gpu in range(devices) if isinstance(devices, int) else devices:
@@ -215,6 +214,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 +281,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/large/leonardo/logs/multi_A_tests_59776456.out

@@ -0,0 +1,70 @@
+NODE: lrdn0172, CPUS: 32, COMMIT: 7e7bdb9
+COMMAND: python -m pytest tests/test_gpu_backend.py -v -o log_cli=false -p no:randomly
+GPU 0: NVIDIA A100-SXM-64GB (UUID: GPU-49a582c5-5db4-49ce-ed6c-bff7283dd5bd)
+GPU 1: NVIDIA A100-SXM-64GB (UUID: GPU-eb681d88-e88f-a83e-fa2e-ec7a23f37a35)
+GPU 2: NVIDIA A100-SXM-64GB (UUID: GPU-01563636-9a99-60dd-a632-b0ffc18f4638)
+GPU 3: NVIDIA A100-SXM-64GB (UUID: GPU-71a4b225-e397-6550-84a9-77c5254110fb)
+	GPU0	GPU1	GPU2	GPU3	NIC0	NIC1	NIC2	NIC3	CPU Affinity	NUMA Affinity	GPU NUMA ID
+GPU0	 X 	NV4	NV4	NV4	PXB	NODE	NODE	NODE	0-15	0		N/A
+GPU1	NV4	 X 	NV4	NV4	NODE	PXB	NODE	NODE	0-15	0		N/A
+GPU2	NV4	NV4	 X 	NV4	NODE	NODE	PXB	NODE	0-15	0		N/A
+GPU3	NV4	NV4	NV4	 X 	NODE	NODE	NODE	PXB	0-15	0		N/A
+NIC0	PXB	NODE	NODE	NODE	 X 	NODE	NODE	NODE				
+NIC1	NODE	PXB	NODE	NODE	NODE	 X 	NODE	NODE				
+NIC2	NODE	NODE	PXB	NODE	NODE	NODE	 X 	NODE				
+NIC3	NODE	NODE	NODE	PXB	NODE	NODE	NODE	 X 				
+
+Legend:
+
+  X    = Self
+  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
+  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
+  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
+  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
+  PIX  = Connection traversing at most a single PCIe bridge
+  NV#  = Connection traversing a bonded set of # NVLinks
+
+NIC Legend:
+
+  NIC0: mlx5_0
+  NIC1: mlx5_1
+  NIC2: mlx5_2
+  NIC3: mlx5_3
+
+4 [[0, 1, 1, 1], [1, 0, 1, 1], [1, 1, 0, 1], [1, 1, 1, 0]]
+============================= test session starts ==============================
+platform linux -- Python 3.14.2, pytest-9.1.0, pluggy-1.6.0 -- /leonardo_work/ALQ_prod2526/rpanades/src-multi-gpu/.venv/bin/python
+cachedir: .pytest_cache
+benchmark: 5.3.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
+rootdir: /leonardo_work/ALQ_prod2526/rpanades/src-multi-gpu
+configfile: pyproject.toml
+plugins: markdown-docs-0.9.2, cov-7.1.0, benchmark-5.3.0
+collecting ... collected 19 items
+
+tests/test_gpu_backend.py::test_apply_mpo_mps_gpu_matches_cpu[cpu] PASSED [  5%]
+tests/test_gpu_backend.py::test_apply_mpo_mps_gpu_matches_cpu[gpu] PASSED [ 10%]
+tests/test_gpu_backend.py::test_apply_mpo_mpo_gpu_matches_cpu[cpu] PASSED [ 15%]
+tests/test_gpu_backend.py::test_apply_mpo_mpo_gpu_matches_cpu[gpu] PASSED [ 21%]
+tests/test_gpu_backend.py::test_compress_mpo_gpu_matches_cpu[cpu] PASSED [ 26%]
+tests/test_gpu_backend.py::test_compress_mpo_gpu_matches_cpu[gpu] PASSED [ 31%]
+tests/test_gpu_backend.py::test_apply_mpo_mpo_cpu_gpu_equivalent PASSED  [ 36%]
+tests/test_gpu_backend.py::test_compress_mpo_cpu_gpu_equivalent PASSED   [ 42%]
+tests/test_gpu_backend.py::test_invalid_device_raises PASSED             [ 47%]
+tests/test_gpu_backend.py::test_src_stack_gpu_matches_reference[cpu] PASSED [ 52%]
+tests/test_gpu_backend.py::test_src_stack_gpu_matches_reference[gpu] PASSED [ 57%]
+tests/test_gpu_backend.py::test_gpu_tiers_and_batches_match_cpu PASSED   [ 63%]
+tests/test_gpu_backend.py::test_gpu_pool_limit_is_restored PASSED        [ 68%]
+tests/test_gpu_backend.py::test_gpu_pool_cap_leaves_room_for_fragmentation PASSED [ 73%]
+tests/test_gpu_backend.py::test_two_gpus_match_one[split] PASSED         [ 78%]
+tests/test_gpu_backend.py::test_two_gpus_match_one[full] PASSED          [ 84%]
+tests/test_gpu_backend.py::test_two_gpus_take_device_inputs PASSED       [ 89%]
+tests/test_gpu_backend.py::test_two_gpus_raise_the_first_error PASSED    [ 94%]
+tests/test_gpu_backend.py::test_too_many_devices_raise PASSED            [100%]
+
+======================== 19 passed in 67.77s (0:01:07) =========================
+JobID           JobName    Elapsed     MaxRSS  MaxDiskRead MaxDiskWrite ExitCode 
+------------ ---------- ---------- ---------- ------------ ------------ -------- 
+59776456     multi_A_t+   00:01:49                                           0:0 
+59776456.ba+      batch   00:01:49                                           0:0 
+59776456.ex+     extern   00:01:49                                           0:0 
+59776456.0       python   00:01:17      0.90G        0.04G        0.02G      0:0