Skip to content

Commit c91f31f

Browse files
fix(sweep): πŸ› split a batch that runs out of GPU memory (#60)
* docs(spec): πŸ“ design for an SRC sweep over the GPUs of a node 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 * docs(spec): πŸ“ fold the refinements found while planning into the spec 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 * docs(plan): πŸ“ implementation plan for the multi-GPU SRC sweep 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 * docs(plan): πŸ“ align the multi-GPU spec and plan with main Phase 1 is merged, logging is the standard library, the docs moved to the Fumadocs site, and the CUDA 13 CuPy needs the cu13 NCCL wheel. Assisted-by: Pi:claude-opus-5.5 * feat(utils): ✨ device selection and peer copies between GPUs Assisted-by: Pi:claude-opus-5.5 * feat(group): ✨ one thread per device and the collectives between them Allgather, gather and scatter over peer copies, ordered by events, in a single process. A failing rank aborts the group so that no rank waits forever; on the CPU the devices are simulated for testing. Assisted-by: Pi:claude-opus-5.5 * feat(sweep): ✨ split the sweep among the GPUs of a node Resources(devices=...) gives each device a contiguous block of the sketch columns. The left-to-right pass needs no communication; per site of the right-to-left pass the sketch is gathered on the first device for the QR, the output core scattered by rows, and the projected environment gathered on every device. Cores are read once and, with peer access, uploaded in slices and gathered over the peer links. One device keeps the previous sweep, bit for bit. Assisted-by: Pi:claude-opus-5.5 * perf(bench): πŸ“ˆ multi-GPU scaling command for the large benchmark Assisted-by: Pi:claude-opus-5.5 * docs: πŸ“ running one sweep on several GPUs Assisted-by: Pi:claude-opus-5.5 * perf(bench): πŸ“ˆ multi-GPU tests and scaling on Leonardo (4xA100-64GB) - Add `leonardo/run_multi.sh`: any command with 4 GPUs in one process, printing the topology and peer-access matrix; `WARM_DIR` reads the stack beforehand. - `scaling` also reports a QR distance over `a - b`, which resolves differences below the sqrt(eps) floor of the inner-product distance. - GPU tests pass (19/19); D_M = 1000 scaling with QR distances below 1e-12. The reference size is queued; its results follow. - Keep `benches/*/leonardo/logs/*.out` in git. Assisted-by: GitHub Copilot:Claude Opus 5.5 * docs(bench): πŸ“ correct what the site stalls measure The stall timer ends before the upload to the device, so it does not include the split upload or the NVLink gather; it also counts the loader waiting for the GPU. Assisted-by: GitHub Copilot:Claude Opus 5.5 * fix(sweep): πŸ› split a batch that runs out of GPU memory At D_M = 4000, chi_out = 2000 on an A100-64GB, CuPy's pool fragments beyond the margin the plan leaves for it: live memory peaked at 52 GB against the planned 56.7 GB, but the pool held up to 66.7 GB and a 12.6-14.7 GB allocation failed, at a different point in each run. On out of memory, a kernel now releases the pool's cached blocks and computes the batch in halves, joined along the batch axis, keeping the split for the rest of the sweep. The plan and the batches the store reads and writes are unchanged, so the devices need not agree on it. Assisted-by: GitHub Copilot:Claude Opus 5.5 * perf(bench): πŸ“ˆ reference size on 1 and 4 GPUs of Leonardo At D_M = 4000, chi_out = 2000: 3925.3 s on one A100-64GB, 1099.9 s on four, a speed-up of 3.57 and an efficiency of 89%. Reading environments back from host memory and disk takes 82-84% of the right-to-left pass. The docs describe the out-of-memory split. The 2-GPU point and the distance at this size are queued. Assisted-by: GitHub Copilot:Claude Opus 5.5 * chore: πŸ”₯ drop the superpowers planning docs Phase 1 of the plan (#44's own description) called these dead weight. Assisted-by: GitHub Copilot:Claude Opus 5.5 * fix(sites): πŸ› stage cores held on another GPU A single-GPU sweep reused any CuPy core in place, so with Resources(devices=[d]) a core on another GPU met environments and sketches on GPU d. Reuse a core only when it sits on the sweep's GPU; stage the others through the host buffer as a group does. Addresses the Copilot review on #44. Assisted-by: GitHub Copilot:Claude Opus 5.5 * perf(bench): πŸ“ˆ 2-GPU point at the reference size Add the 2-GPU run and a second 4-GPU run (scaling --devices 4 2) at D_M = 4000, chi_out = 2000: 2348.8 s on two GPUs (84% against one GPU), 1290.4 s on four, QR distance 7.5e-13 between them. Correct what the environment stalls measure: they cover both passes and include waiting for the GPU. Assisted-by: GitHub Copilot:Claude Opus 5.5 * docs(bench): πŸ“ keep the best 4-GPU point at the reference size Assisted-by: GitHub Copilot:Claude Opus 5.5 * docs(hpc): πŸ“ HPC clusters section; drop the Leonardo logs and scripts Remove the benches/*/leonardo folders: the Slurm scripts and the run logs, which stay in the git history. bench_mpo_mpo.py moves up to benches/primitives/. The bench READMEs keep the result summaries, without job IDs. Add an "HPC clusters" section to the docs with a Leonardo page: the CUDA 13 forward-compat setup, an example Slurm job script for one or four GPUs, queues and performance notes. Assisted-by: GitHub Copilot:Claude Opus 5.5 --------- Co-authored-by: Roberto Di Remigio EikΓ₯s <roberto@totaltrash.xyz>
1 parent fb62638 commit c91f31f

52 files changed

Lines changed: 427 additions & 1323 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

β€Ž.gitignoreβ€Ž

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ build*/
3131
# Executables
3232
*.exe
3333
*.out
34-
!benches/*/leonardo/logs/*.out
3534
*.app
3635

3736
### Python ###

β€Žbenches/README.mdβ€Ž

Lines changed: 7 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -12,38 +12,18 @@ different SRC variants. They are grouped by the kind of workload exercised:
1212
wall time on small chains, and quench dynamics of 50 and 100 qubits on the GPU.
1313
See [`stack/README.md`](stack/README.md).
1414
- [`large/`](large/) β€” Out-of-core SRC of `N . V . M . U` with a large `M` read
15-
from disk, on one GPU: plan, wall time, memory peaks and stall time. See
16-
[`large/README.md`](large/README.md).
15+
from disk, on one GPU or several: plan, wall time, memory peaks, stall time and
16+
parallel efficiency. See [`large/README.md`](large/README.md).
1717

1818
The scripts need the `bench` dependency group (included in `dev`):
1919

2020
```bash
2121
uv sync --group bench
2222
```
2323

24-
## Leonardo
24+
## HPC clusters
2525

26-
Each folder has a `leonardo/` subfolder with a Slurm script for the Booster
27-
partition of [Leonardo](https://docs.hpc.cineca.it/general/getting_started.html)
28-
at CINECA (one node: 32 Xeon 8358 cores, 4 A100-64GB, 512 GB of memory), and the
29-
logs of the recorded runs. Set up the environment once, from the repository
30-
root on a login node:
31-
32-
```bash
33-
# Keep the uv cache and Python installs out of $HOME.
34-
export UV_CACHE_DIR=$PWD/.uv-cache UV_PYTHON_INSTALL_DIR=$PWD/.uv-python
35-
uv sync --all-groups --extra gpu-nvidia
36-
37-
# The Booster driver (535) predates CUDA 13: unpack NVIDIA's forward-compat libcuda.
38-
mkdir -p .cuda-compat && cd .cuda-compat
39-
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
40-
rpm2cpio cuda-compat-13-4-*.rpm | cpio -idm && cd ..
41-
```
42-
43-
The scripts prepend `.cuda-compat/usr/local/cuda-13.4/compat` to
44-
`LD_LIBRARY_PATH` (override with `CUDA_COMPAT`) and keep CuPy's kernel cache in
45-
`.cupy-cache/`. Submit from the `leonardo/` folder, passing the benchmark's own
46-
arguments, e.g. `sbatch run.sh --n-sites 25 --chi-out 1000 --device gpu`. The
47-
default QoS allows long runs but can queue for hours; for jobs under 30 minutes
48-
add `--qos=boost_qos_dbg --time=00:30:00`, which starts quickly but takes at
49-
most two jobs per user.
26+
The recorded results come from the Booster partition of Leonardo at CINECA. The
27+
[HPC clusters](https://algorithmiq.github.io/src-method/hpc/leonardo) section of
28+
the documentation covers the setup and has an example job script; pass it the
29+
benchmark command, e.g. `python benches/large/bench_large.py run ...`.

β€Žbenches/large/README.mdβ€Ž

Lines changed: 57 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -46,25 +46,21 @@ size on 8xH100 or 8xH200 (and 4xA100 on Leonardo):
4646
1. Every distance is below `1e-10`.
4747
2. The parallel efficiency is at least 80% on all the GPUs of the node.
4848

49+
`scaling` reports two distances. `distance` comes from inner products and cancels
50+
down to about `sqrt(eps)`, near `1.5e-8` in complex128, so it cannot show `1e-10`.
51+
`QR distance` sweeps a QR over the MPO `a - b` (the direct sum of the two outputs)
52+
and resolves differences down to rounding; the acceptance applies to it.
53+
4954
## Leonardo
5055

51-
The [`leonardo/`](leonardo/) folder holds the Slurm script and the logs; see the
52-
[Leonardo section](../README.md#leonardo) for the environment. Booster nodes have
53-
no local disk and a 10 GB `/tmp`, so the script points `TMPDIR` at
54-
`$CINECA_SCRATCH`, and the stacks live there too (Lustre, not NVMe):
56+
Recorded on Booster nodes of Leonardo at CINECA; see
57+
[HPC clusters](https://algorithmiq.github.io/src-method/hpc/leonardo) for the setup
58+
and a job script. The nodes have no local disk, so the stacks and the spilled
59+
environments live on Lustre scratch rather than NVMe.
5560

56-
```bash
57-
D=$CINECA_SCRATCH/src-large
58-
sbatch run.sh generate $D/m4000 --bond-m 4000
59-
sbatch run.sh --debug run $D/m4000 --chi-out 2000 --scratch-dir $D/scratch
60-
sbatch run.sh generate $D/m1000 --bond-m 1000
61-
sbatch run.sh --debug compare $D/m1000 --chi-out 500 --small 4GB --large 60GB \
62-
--scratch-dir $D/scratch
63-
```
61+
### One GPU: D_M = 1000, chi_out = 500
6462

65-
`generate` writes the `D_M = 1000` stack (12 GB) in 51 s.
66-
67-
`compare` at `D_M = 1000`, `chi_out = 500`, complex128, on one A100-64GB:
63+
`compare --small 4GB --large 60GB`, complex128, on one A100-64GB:
6864

6965
| GPU budget | Wall time (s) | Pool (GB) | Planned device peak (GB) | Host peak (GB) | Tiers (device/host/disk) |
7066
|---|---|---|---|---|---|
@@ -74,3 +70,49 @@ sbatch run.sh --debug compare $D/m1000 --chi-out 500 --small 4GB --large 60GB \
7470
The two plans keep opposite tiers yet give bit-identical outputs (relative distance
7571
`0.000e+00`). Spilling 49 of 50 environments to the host costs 18 % of the wall
7672
time. Both pools stay within their caps (budget plus 10 % of the card, 6.9 GB).
73+
74+
### Several GPUs: D_M = 1000, chi_out = 500
75+
76+
One node, 4 A100-SXM-64GB fully connected by NVLink, with peer access on every
77+
pair, so each GPU uploads a slice of every core and gathers the rest from its
78+
peers. `scaling --devices 1 2 4`, complex128, every environment on the device:
79+
80+
| GPUs | Time (s) | Speed-up | Efficiency | QR distance |
81+
|---|---|---|---|---|
82+
| 1 | 109.9 | 1.00 | 100% | |
83+
| 2 | 43.6 | 2.52 | 126% | 2.8e-13 |
84+
| 4 | 32.7 | 3.36 | 84% | 4.4e-13 |
85+
86+
- The single-GPU time is that of a second SRC run in its job; the first SRC run of
87+
a job takes about 22 s longer, for a reason not isolated.
88+
- The 2-GPU point is superlinear. The plan and tiers are the same on every count,
89+
and the sketch batch drops from 480 to 224 columns per GPU, which may run the
90+
contractions more efficiently; this was not isolated.
91+
- This problem fits on one GPU, so it says little about the reference size.
92+
93+
### Several GPUs: D_M = 4000, chi_out = 2000 (reference size)
94+
95+
50 sites, complex128, 184 GB of input: 49 environments of 8.19 GB, 401 GB in all.
96+
Speed-up and efficiency are against the single-GPU run:
97+
98+
| GPUs | Time (s) | Speed-up | Efficiency | Left-to-right (s) | Right-to-left (s) | Environments device / host / disk | QR distance |
99+
|---|---|---|---|---|---|---|---|
100+
| 1 | 3925.3 | 1.00 | 100% | 794.8 | 3112.1 | 0 / 45 / 4 | |
101+
| 2 | 2348.8 | 1.67 | 84% | 730.6 | 1606.3 | 1 / 35 / 14 | 7.5e-13 (vs 4 GPUs) |
102+
| 4 | 1099.9 | 3.57 | 89% | 240.4 | 844.1 | 1 / 45 / 3 | |
103+
104+
- Both efficiencies are above the 80% target.
105+
- Before the out-of-memory split, every run at this size failed. The live memory
106+
stayed within the plan, but CuPy's pool held up to 66.7 GB in split blocks and a
107+
12.6-14.7 GB allocation did not fit under its 67.6 GB cap. Each GPU of the
108+
multi-GPU runs now splits one sketch batch in two and finishes.
109+
- Four GPUs keep only one environment on the device. The plan picks the largest
110+
batches that fit first, and their working set takes each GPU's budget (59.0 GB
111+
planned peak), although a GPU holds only 2.05 GB of each environment.
112+
- Three or four environments go to disk on one and four GPUs: the host budget, 90%
113+
of the available memory (about 470 GB), falls short of the 487-496 GB that
114+
keeping them all in host memory takes. On two GPUs, 14 go to disk with a planned
115+
host peak of 381.7 GB against 462.6 GB on four; why the plan leaves that host
116+
memory unused was not investigated.
117+
- The stall times the sweep logs include waiting for the GPU's own queue, so they
118+
cannot tell slow transfers from a GPU-bound pass.

β€Žbenches/large/bench_large.pyβ€Ž

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
import json
2222
import logging
2323
import resource
24-
from pathlib import Path # noqa: TC003 (cyclopts reads the annotations at runtime)
24+
from pathlib import Path
2525
from time import perf_counter
2626
from typing import Annotated
2727

@@ -215,6 +215,31 @@ def _distance(a: list[np.ndarray], b: list[np.ndarray]) -> float:
215215
return float(np.sqrt(max((aa + bb - 2 * ab).real, 0.0) / aa.real))
216216

217217

218+
def _qr_distance(a: list[np.ndarray], b: list[np.ndarray]) -> float:
219+
"""Relative distance ``|a - b| / |a|`` from a QR sweep over the MPO ``a - b``.
220+
221+
`_distance` cancels down to about ``sqrt(eps)``. Here ``a - b`` is a direct sum
222+
of the two trains, and the QR sweep keeps rounding relative to ``|a|``, which
223+
resolves differences down to ``eps``. Gauge-free, unlike comparing cores.
224+
"""
225+
n = len(a)
226+
carry = np.ones((1, 1), dtype=np.result_type(a[0], b[0]))
227+
for j, (x, y) in enumerate(zip(a, b)):
228+
x = x[None] if j == 0 else x[:, None] if j == n - 1 else x
229+
y = y[None] if j == 0 else y[:, None] if j == n - 1 else y
230+
if j == 0:
231+
site = np.concatenate([x, -y], axis=1)
232+
elif j == n - 1:
233+
site = np.concatenate([x, y], axis=0)
234+
else:
235+
(la, ra, *phys), (lb, rb, _, _) = x.shape, y.shape
236+
site = np.zeros((la + lb, ra + rb, *phys), dtype=carry.dtype)
237+
site[:la, :ra], site[la:, ra:] = x, y
238+
site = np.tensordot(carry, site, axes=(1, 0)).transpose(0, 2, 3, 1)
239+
carry = np.linalg.qr(site.reshape(-1, site.shape[-1]), mode="r")
240+
return float(np.linalg.norm(carry) / np.sqrt(_inner(a, a).real))
241+
242+
218243
@app.command
219244
def scaling(
220245
directory: Path,
@@ -257,24 +282,28 @@ def scaling(
257282
out, seconds = _run(directory, chi_out, resources, seed, device)
258283
if reference is None:
259284
reference, ref_seconds, ref_devices = out, seconds, count
260-
distance = 0.0
285+
distance = qr_distance = 0.0
261286
else:
262287
distance = _distance(reference, out)
288+
qr_distance = _qr_distance(reference, out)
263289
efficiency = ref_seconds * ref_devices / (seconds * count)
264290
logger.info(
265-
"GPUs %d: %.1f s, speed-up %.2f, efficiency %.0f %%, distance %.3e",
291+
"GPUs %d: %.1f s, speed-up %.2f, efficiency %.0f %%, distance %.3e, "
292+
"QR distance %.3e",
266293
count,
267294
seconds,
268295
ref_seconds / seconds,
269296
100 * efficiency,
270297
distance,
298+
qr_distance,
271299
)
272300
if results is not None:
273301
record = {
274302
"gpus": count,
275303
"seconds": seconds,
276304
"efficiency": efficiency,
277305
"distance": distance,
306+
"qr_distance": qr_distance,
278307
"chi_out": chi_out,
279308
}
280309
with results.open("a") as f:

β€Žbenches/large/leonardo/logs/large_59625448.outβ€Ž

Lines changed: 0 additions & 15 deletions
This file was deleted.

β€Žbenches/large/leonardo/logs/large_compare_59625755.outβ€Ž

Lines changed: 0 additions & 28 deletions
This file was deleted.

β€Žbenches/large/leonardo/logs/large_generate_59625937.outβ€Ž

Lines changed: 0 additions & 15 deletions
This file was deleted.

β€Žbenches/large/leonardo/run.shβ€Ž

Lines changed: 0 additions & 41 deletions
This file was deleted.

β€Žbenches/primitives/README.mdβ€Ž

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,17 @@ controlled bond dimension, and compare to a `quimb` reference.
88
## Leonardo
99

1010
The Leonardo [supercomputer](https://docs.hpc.cineca.it/general/getting_started.html) at CINECA is
11-
a pre-exascale Tier-0 EuroHPC system. The [`leonardo/`](leonardo/) folder holds the Slurm
12-
script and the logs; see the [Leonardo section](../README.md#leonardo) for the environment.
11+
a pre-exascale Tier-0 EuroHPC system; see
12+
[HPC clusters](https://algorithmiq.github.io/src-method/hpc/leonardo) for how to run there.
1313

1414
### MPO-MPO benchmark
1515

1616
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.
1717

18+
```bash
19+
uv run python benches/primitives/bench_mpo_mpo.py --n-sites 25 --chi-out 1000 --compare no --device gpu
20+
```
21+
1822
One Booster node (32 cores, A100-64GB), complex128 unless noted, `chi_id=4` at `chi=1000`:
1923

2024
| 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:
6064

6165
### MPO-MPO strong scaling (v0.3.2, DCGP)
6266

63-
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.
67+
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.
6468

6569
| CPUs | Time auto (s) | Time maual (s)|
6670
|------|-----------------|---------------|
File renamed without changes.

0 commit comments

Comments
Β (0)