Skip to content

Commit fb62638

Browse files
feat(sweep): ✨ split the SRC sweep among the GPUs of a node (#44)
* 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 * chore: remove implementation plans. --------- Co-authored-by: Panadestein <rpana92@gmail.com>
1 parent bcc0fb5 commit fb62638

18 files changed

Lines changed: 1655 additions & 270 deletions

File tree

‎README.md‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,9 @@ out = src(
105105
)
106106
```
107107

108-
See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory and how to read the logged plan.
108+
On a node with several GPUs, `Resources(devices=4)` splits the sweep among four of them, in the same process.
109+
110+
See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs and how to read the logged plan.
109111

110112
### Logging
111113

‎benches/large/README.md‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
# Out-of-core benchmark
22

33
`bench_large.py` compresses the stack `N . V . M . U` of Pauli-transfer-matrix MPOs
4-
(physical legs of 4) on one GPU, where `M` has a large bond and is read from disk.
4+
(physical legs of 4) on one GPU or several, where `M` has a large bond and is read
5+
from disk.
56

67
```bash
78
# Write the stack to node-local NVMe (205 GB for M at the defaults, complex128).
@@ -15,9 +16,13 @@ uv run python benches/large/bench_large.py run /local/stack --chi-out 2000 \
1516
uv run python benches/large/bench_large.py generate /local/medium --bond-m 1000
1617
uv run python benches/large/bench_large.py compare /local/medium --chi-out 500 \
1718
--small 40GB --large 80GB --scratch-dir /local/scratch
19+
20+
# Split the sweep among the GPUs of the node: 1, 2, 4 and 8 of them, in turn.
21+
uv run python benches/large/bench_large.py scaling /local/stack --chi-out 2000 \
22+
--devices 1 2 4 8 --scratch-dir /local/scratch --results scaling.jsonl
1823
```
1924

20-
`run` logs the wall time, the size of CuPy's pool at the end (its high-water mark,
25+
`run` logs the wall time, the size of CuPy's pool on each GPU at the end (its high-water mark,
2126
since the pool keeps its blocks), the host peak (`ru_maxrss`), the planned peaks,
2227
the bytes spilled and the tiers. With `--debug` before the command, `src_method` also
2328
logs its plan, the time of each pass and the stall time (`SRC stalls`).
@@ -32,6 +37,15 @@ memory and node-local NVMe:
3237
3. The stall time is below 10% of the wall time.
3338
4. `compare` at `D_M = 1000`, `chi_out = 500` reports a distance below `1e-10`.
3439

40+
`scaling` logs, for each GPU count, the wall time, the speed-up and the parallel
41+
efficiency `T_1 / (G T_G)` against the first count, and the relative distance to
42+
its output; `--results` appends the same as JSON lines. `--devices` with
43+
`run` splits a single run. Acceptance of the multi-GPU sweep, at the reference
44+
size on 8xH100 or 8xH200 (and 4xA100 on Leonardo):
45+
46+
1. Every distance is below `1e-10`.
47+
2. The parallel efficiency is at least 80% on all the GPUs of the node.
48+
3549
## Leonardo
3650

3751
The [`leonardo/`](leonardo/) folder holds the Slurm script and the logs; see the

‎benches/large/bench_large.py‎

Lines changed: 106 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,24 @@
11
"""Out-of-core SRC of ``N . V . M . U`` with a large ``M``, read from disk.
22
3-
Three commands:
3+
Four commands:
44
55
- ``generate``: write the four MPOs site by site as ``.npy`` files, so that ``M``
66
is never whole in memory.
7-
- ``run``: compress the stack with automatic (or given) budgets and report the
8-
plan, the wall time, the device pool size, the host peak and the stall time.
7+
- ``run``: compress the stack with automatic (or given) budgets, on one or several
8+
GPUs, and report the plan, the wall time, the device pool sizes, the host peak
9+
and the stall time.
910
- ``compare``: run twice with two GPU budgets and report the relative distance
1011
between the two outputs, which checks that the plan does not change the result.
12+
- ``scaling``: run on 1, 2, 4, ... GPUs of the node and report the wall time, the
13+
parallel efficiency and the distance to the single-GPU output.
1114
1215
The plan, the per-pass times and the stall time are also logged by `src_method`
1316
itself at ``DEBUG``; pass ``--debug`` to see them.
1417
"""
1518

1619
from __future__ import annotations
1720

21+
import json
1822
import logging
1923
import resource
2024
from pathlib import Path # noqa: TC003 (cyclopts reads the annotations at runtime)
@@ -94,8 +98,12 @@ def _load(directory: Path) -> list[list[np.ndarray]]:
9498

9599

96100
def _run(
97-
directory: Path, chi_out: int, resources: Resources, seed: int
98-
) -> list[np.ndarray]:
101+
directory: Path,
102+
chi_out: int,
103+
resources: Resources,
104+
seed: int,
105+
device: str = "gpu",
106+
) -> tuple[list[np.ndarray], float]:
99107
layers = _load(directory)
100108
plans = []
101109

@@ -111,29 +119,35 @@ def spy(*args: object, **kwargs: object) -> object:
111119
chi_out=chi_out,
112120
dtype=layers[0][0].dtype,
113121
seed=seed,
114-
device="gpu",
122+
device=device,
115123
resources=resources,
116124
)
117125
seconds = perf_counter() - start
118126
finally:
119127
sweep_module.make_plan = make_plan
120-
import cupy # noqa: PLC0415 (GPU-only benchmark)
121-
122128
(plan,) = plans
123129
tiers = [site.tier for site in plan.sites]
130+
pools = []
131+
if device == "gpu":
132+
import cupy # noqa: PLC0415 (optional dependency)
133+
134+
devices = resources.devices or 1
135+
for gpu in range(devices) if isinstance(devices, int) else devices:
136+
with cupy.cuda.Device(gpu):
137+
pools.append(cupy.get_default_memory_pool().total_bytes())
124138
logger.info(
125-
"Run complete: %.1f s, pool %d B, host peak %d B, planned device peak %d B, "
139+
"Run complete: %.1f s, pools %s B, host peak %d B, planned device peak %d B, "
126140
"planned host peak %d B, disk %d B, tiers %s, sketch batches %s",
127141
seconds,
128-
cupy.get_default_memory_pool().total_bytes(),
142+
pools,
129143
resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024,
130144
plan.device_peak,
131145
plan.host_peak,
132146
plan.disk_bytes,
133147
{tier: tiers.count(tier) for tier in ("device", "host", "disk")},
134148
sorted({site.sketch_batch for site in plan.sites[1:]}),
135149
)
136-
return out
150+
return out, seconds
137151

138152

139153
@app.command
@@ -144,19 +158,22 @@ def run(
144158
gpu_memory: str | None = None,
145159
host_memory: str | None = None,
146160
scratch_dir: Path | None = None,
161+
devices: int = 1,
147162
seed: int = 0,
148163
) -> None:
149164
"""Compress the stack written by ``generate`` on the GPU.
150165
151166
Args:
152167
directory: The directory given to ``generate``.
153168
chi_out: The output bond dimension.
154-
gpu_memory: GPU budget, e.g. ``36GB``; detected when omitted.
169+
gpu_memory: GPU budget of each device, e.g. ``36GB``; detected when omitted.
155170
host_memory: Host budget; detected when omitted.
156171
scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted.
172+
devices: The number of GPUs to split the sweep among.
157173
seed: Seed of the sketch.
158174
"""
159-
_run(directory, chi_out, Resources(gpu_memory, host_memory, scratch_dir), seed)
175+
resources = Resources(gpu_memory, host_memory, scratch_dir, devices=devices)
176+
_run(directory, chi_out, resources, seed)
160177

161178

162179
def _inner(a: list[np.ndarray], b: list[np.ndarray]) -> complex:
@@ -187,11 +204,82 @@ def compare(
187204
scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted.
188205
seed: Seed of the sketch, the same for both runs.
189206
"""
190-
first = _run(directory, chi_out, Resources(small, None, scratch_dir), seed)
191-
second = _run(directory, chi_out, Resources(large, None, scratch_dir), seed)
192-
aa, bb, ab = _inner(first, first), _inner(second, second), _inner(first, second)
193-
distance = np.sqrt(max((aa + bb - 2 * ab).real, 0.0) / aa.real)
194-
logger.info("Relative distance between the runs: %.3e", distance)
207+
first, _ = _run(directory, chi_out, Resources(small, None, scratch_dir), seed)
208+
second, _ = _run(directory, chi_out, Resources(large, None, scratch_dir), seed)
209+
logger.info("Relative distance between the runs: %.3e", _distance(first, second))
210+
211+
212+
def _distance(a: list[np.ndarray], b: list[np.ndarray]) -> float:
213+
"""Relative Frobenius distance ``|a - b| / |a|`` of two MPOs."""
214+
aa, bb, ab = _inner(a, a), _inner(b, b), _inner(a, b)
215+
return float(np.sqrt(max((aa + bb - 2 * ab).real, 0.0) / aa.real))
216+
217+
218+
@app.command
219+
def scaling(
220+
directory: Path,
221+
*,
222+
chi_out: int = 2000,
223+
devices: Annotated[tuple[int, ...], cyclopts.Parameter(consume_multiple=True)] = (
224+
1,
225+
2,
226+
4,
227+
8,
228+
),
229+
gpu_memory: str | None = None,
230+
host_memory: str | None = None,
231+
scratch_dir: Path | None = None,
232+
results: Path | None = None,
233+
seed: int = 0,
234+
device: str = "gpu",
235+
) -> None:
236+
"""Run on increasing numbers of GPUs and report the strong scaling.
237+
238+
The first count is the reference: the efficiency of ``G`` GPUs is
239+
``T_ref * G_ref / (T_G * G)``, and every output is compared with the
240+
reference output. Host memory holds two outputs at a time.
241+
242+
Args:
243+
directory: The directory given to ``generate``.
244+
chi_out: The output bond dimension.
245+
devices: The GPU counts, the reference first.
246+
gpu_memory: GPU budget of each device; detected when omitted.
247+
host_memory: Host budget of the node; detected when omitted.
248+
scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted.
249+
results: A JSON-lines file to append one record per run to.
250+
seed: Seed of the sketch, the same for every run.
251+
device: ``gpu``, or ``cpu`` to smoke-test with simulated devices.
252+
"""
253+
reference: list[np.ndarray] | None = None
254+
ref_seconds = ref_devices = 0.0
255+
for count in devices:
256+
resources = Resources(gpu_memory, host_memory, scratch_dir, devices=count)
257+
out, seconds = _run(directory, chi_out, resources, seed, device)
258+
if reference is None:
259+
reference, ref_seconds, ref_devices = out, seconds, count
260+
distance = 0.0
261+
else:
262+
distance = _distance(reference, out)
263+
efficiency = ref_seconds * ref_devices / (seconds * count)
264+
logger.info(
265+
"GPUs %d: %.1f s, speed-up %.2f, efficiency %.0f %%, distance %.3e",
266+
count,
267+
seconds,
268+
ref_seconds / seconds,
269+
100 * efficiency,
270+
distance,
271+
)
272+
if results is not None:
273+
record = {
274+
"gpus": count,
275+
"seconds": seconds,
276+
"efficiency": efficiency,
277+
"distance": distance,
278+
"chi_out": chi_out,
279+
}
280+
with results.open("a") as f:
281+
f.write(json.dumps(record) + "\n")
282+
del out
195283

196284

197285
@app.meta.default

‎docs/content/docs/contributing/testing.mdx‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ Always run tests through `pytest`, never as `python tests/test_<name>.py`.
2323
| `test_plan.py` | the planner: budgets, batch sizes and environment tiers |
2424
| `test_store.py` | environment storage on the device, in host memory and on disk |
2525
| `test_sites.py` | reading the cores one site at a time, with prefetching |
26+
| `test_group.py` | the device group: work splitting, collectives, failures |
2627
| `test_logging.py` | the package leaves the host's logging configuration alone |
2728
| `test_gpu_backend.py` | the CuPy path; skipped without CuPy or a GPU |
2829

@@ -88,6 +89,16 @@ pure function, so `tests/test_plan.py` checks batch sizes and tiers from shapes
8889
alone, and `tests/test_gpu_backend.py` derives GPU budgets from a plan made with
8990
`make_plan` to reach every tier.
9091

92+
## Exercising several devices
93+
94+
`Resources(devices=n)` on the CPU simulates `n` devices with threads that share
95+
the host, through the same code as on GPUs: the split of the sketch columns, the
96+
collectives and the shared reading of the cores. `tests/test_stack.py` compares
97+
such runs with a single device, also with uneven blocks (`chi_out` not a
98+
multiple of `n`, or smaller than it), a cutoff and spilling. Only the peer copies
99+
and device selection are GPU-specific; their tests in `tests/test_gpu_backend.py`
100+
skip unless at least two GPUs are visible.
101+
91102
## Logging in tests
92103

93104
`log_cli` is on, so log records show up live while the tests run. To assert on

‎docs/content/docs/features/gpu.mdx‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,7 @@ result = apply(H, psi, chi_out=256, device="gpu")
3434
The inputs can be NumPy or CuPy arrays; they are moved to the device for the
3535
sweep, and the result always comes back as NumPy arrays on the host. Without
3636
CuPy installed, `device="gpu"` raises an `ImportError`.
37+
38+
To split one sweep among several GPUs of a node, pass
39+
`resources=Resources(devices=n)`; see
40+
[Large problems](/features/large-problems#several-gpus).

‎docs/content/docs/features/large-problems.mdx‎

Lines changed: 59 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
---
22
title: Large problems
3-
description: Run stacks that fit neither on the GPU nor in host memory.
3+
description: Run stacks that fit neither on the GPU nor in host memory, on one GPU or several.
44
---
55

66
A single SRC sweep keeps, besides the input and output trains, one sketched
@@ -75,18 +75,74 @@ reference size of 50 sites, `D_M = 4000` and `chi_out = 2000` in complex128, up
7575
returns, also after an error; a process killed with `SIGKILL` leaves it behind, and
7676
its name identifies the process.
7777

78+
## Several GPUs
79+
80+
On a node with several GPUs, one sweep can be split among them:
81+
82+
```python notest
83+
from src_method import Resources, src
84+
85+
out = src(
86+
N, V, M, U,
87+
chi_out=2000,
88+
dtype=np.complex128,
89+
device="gpu",
90+
resources=Resources(devices=4, scratch_dir="/local/scratch"),
91+
)
92+
```
93+
94+
`devices` is a count, for the first `n` visible GPUs, or a list of device ids;
95+
the first one gathers the sketches and runs the QR. Everything runs in the calling
96+
process, one thread per GPU, so nothing changes for the caller: the call returns
97+
the output cores as with one GPU, errors are raised as usual, and the inputs are
98+
read once and shared.
99+
100+
The sketch index is split: each GPU owns a contiguous block of the `chi_out`
101+
sketch columns and computes only those columns of every environment and of every
102+
sketch, so the environments, the dominant memory, are split among the GPUs too.
103+
At every site of the right-to-left pass, the first GPU gathers the sketch, runs
104+
the QR and gives each GPU a block of rows of the output core; each GPU projects
105+
its rows of the next projected environment, and those rows are gathered on every
106+
GPU. The random draws are those of one GPU, so the result is that of one GPU up to
107+
rounding.
108+
109+
What is not split:
110+
111+
- every GPU holds a whole site: the cores, the projected environment and the
112+
output core of the site, so the largest site that fits is that of one GPU;
113+
- every core is read from the source once, but every GPU needs a copy. When the
114+
GPUs reach each other directly (NVLink, or PCIe peer-to-peer), each uploads a
115+
slice of the core and gathers the rest from the others; otherwise each uploads
116+
the whole core;
117+
- the QR runs on the first GPU while the others wait.
118+
119+
These costs grow more slowly with the problem than the contractions do, so large
120+
problems, the ones this is for, scale close to linearly, while small ones may run
121+
slower on several GPUs than on one.
122+
123+
The budgets apply as follows: `gpu_memory` is the budget of each GPU (detected:
124+
the smallest free memory of the GPUs), while `host_memory` and `scratch_dir` are
125+
shared by all of them. Every GPU writes its environments to its own scratch
126+
directory.
127+
128+
On the CPU, `devices` simulates the GPUs with threads that share the host; the
129+
result is the same, and the mode exists to test the multi-GPU sweep without GPUs.
130+
78131
## Reading the plan
79132

80133
With `DEBUG` logging on for the `src_method` logger (see
81134
[Logging](/features/logging)), every sweep logs its plan as `SRC plan: ...`:
82135

83136
- `prefetch`: whether the next site is read ahead;
84-
- `device peak`, `host peak`, `disk`: the planned peaks, in bytes;
137+
- `devices`: the device ids the sweep is split among;
138+
- `device peak`, `host peak`, `disk`: the planned peaks, in bytes; the device
139+
peak is that of each device, the others are totals;
85140
- `scratch`: the scratch directory, when anything spills to disk;
86141
- `tiers`: where the environment of each site lives (`device`, `host` or `disk`);
87142
- `batches`: for each site, the batch of the environment, sketch and projection
88143
steps.
89144

90145
It then logs the time of each pass and, at the end, `SRC stalls`: the seconds
91146
the sweep waited for input cores and for environments. If they are a large
92-
fraction of the run, the disk is too slow for the compute.
147+
fraction of the run, the disk is too slow for the compute. With several devices,
148+
these lines start with the device they refer to.

0 commit comments

Comments
 (0)