feat(sweep): ✨ split the SRC sweep among the GPUs of a node by robertodr · Pull Request #44 · Algorithmiq/src-method

robertodr · 2026-09-27T11:35:59Z

🤖 AI text below 🤖

Splits one SRC sweep among the GPUs of a single node, so that large stacks (the out-of-core reference problem of #42) finish close to G times faster on G GPUs. The mathematics, the random draws and the result do not change. Multi-node is out of scope for now.

src(N, V, M, U, chi_out=2000, device="gpu",
    resources=Resources(devices=4, scratch_dir="/local/scratch"))

Design

One process, one thread per GPU. The earlier MPI/SPMD design (still in docs/superpowers/, to be dropped) is replaced: nothing in it paid off on a single node. The call works as before: it returns the cores, raises errors normally, needs no mpirun, and reads in-memory inputs once.

Sketch split. Each GPU owns a contiguous block of the chi_out sketch columns and builds only those columns of every environment, so the environments (the dominant memory) are split too.
- The left-to-right pass needs no communication.
- Per site of the right-to-left pass: GPU 0 gathers the sketch, runs the QR and scatters blocks of rows of the output core. Each GPU projects its rows of the next S, which is then gathered on every GPU.
Collectives (_group.py) are allgather, gather and scatter, built from peer copies (memcpyPeerAsync) and ordered with the kernels by CUDA events. If one GPU fails, the group is aborted, the other GPUs are woken and the original error is raised, so nothing hangs.
Cores are read once by one loader thread, in sweep order. With peer access, each GPU uploads 1/G of the core and gathers the rest over NVLink/P2P. Without it, each GPU uploads the whole core.
Planner: each GPU is planned for its block of the sketch columns. gpu_memory is per GPU (detected: the smallest free memory); host memory and scratch disk are shared.
Same draws, same result: the Gaussian draws are shared, so the result equals one GPU up to rounding. With one device, the output is bit-identical to main, including with cutoff and spilling.
CPU testing: on the CPU, devices=n simulates the GPUs with threads that share the host, through the same code. This lets the suite test the distributed logic without GPUs.

What is not split (sets the efficiency)

Every GPU holds a whole site: the cores, S and the output core. The largest D_M is therefore that of one GPU, about 10⁴ at chi_out = 2000 on an H200.
GPU 0 runs the QR while the others wait: 0.1–0.2 s per site at the reference size.
Every GPU needs a copy of every core: an NVLink allgather with peer access, a full upload over PCIe without.

Contractions scale as χ²D², these costs at most as D², so large problems should scale close to linearly; small ones will not.

Changes

Area	Files
Device helpers: selection, peer access, peer copies	`utils/_backend.py`
Device group and collectives	`_group.py` (new)
Per-device plans, `Resources.devices`	`_plan.py`
Cores read once in sweep order and shared	`_sites.py`
The split sweep	`_sweep.py`
`scaling` command (`--devices 1 2 4 8`), `--devices` on `run`, `--device cpu` for smoke tests	`benches/large/`
"Several GPUs" section, testing notes, README	`docs/content/docs/`, `README.md`

Testing

CPU, in the PR suite:
- collectives, uneven blocks, many rounds in a row, abort handling;
- sites shared by the devices, and a failed read reaching every device;
- simulated devices against one device for 2 and 3 devices, chi_out not a multiple of the device count or smaller than it, a custom device order, cutoff, spilling to disk with cleanup, and an error on one device;
- per-device plans, and the sharing of host memory and disk.
GPU, skipped with fewer than two GPUs:
- two GPUs against one, with split and full uploads, reaching every environment tier;
- inputs already on a GPU;
- an error on one GPU, after which the pool limits are restored;
- an out-of-range device count.

Status

Implementation, CPU tests, docs, scaling benchmark
GPU tests on a multi-GPU node: the peer copies, the cross-device events and CuPy's per-thread device switching have not run on hardware yet
Scaling runs: 1/2/4 GPUs on Leonardo (4×A100), 1/2/4/8 GPUs on an H100, H200 or GH200 node; results in benches/large/README.md
Drop docs/superpowers/

Left for later, if measurements call for them: a distributed CholeskyQR2 instead of the QR on GPU 0, splitting the bond of M for D_M above about 10⁴, and NUMA-local host buffers on GH200 nodes.

github-actions · 2026-10-08T14:51:47Z

Test Results

5 files ± 0 5 suites ±0 3m 56s ⏱️ +52s
262 tests + 56 256 ✅ + 51 6 💤 + 5 0 ❌ ±0
1 251 runs +265 1 238 ✅ +255 13 💤 +10 0 ❌ ±0

Results for commit 37f3324. ± Comparison against base commit bcc0fb5.

This pull request removes 2 and adds 58 tests. Note that renamed tests count towards both.

tests.test_sites ‑ test_sites_are_read_when_requested[0]
tests.test_sites ‑ test_sites_are_read_when_requested[1]

tests.test_backend ‑ test_copy_into_on_host
tests.test_backend ‑ test_copy_into_rejects_mismatches
tests.test_backend ‑ test_host_devices_are_simulated
tests.test_gpu_backend ‑ test_too_many_devices_raise
tests.test_gpu_backend ‑ test_two_gpus_match_one[full]
tests.test_gpu_backend ‑ test_two_gpus_match_one[split]
tests.test_gpu_backend ‑ test_two_gpus_raise_the_first_error
tests.test_gpu_backend ‑ test_two_gpus_take_device_inputs
tests.test_group ‑ test_a_failing_rank_releases_the_others
tests.test_group ‑ test_a_group_of_one_runs_inline
…

♻️ This comment has been updated with latest results.

Phase 1 of the plan (#44's own description) called these dead weight. Assisted-by: GitHub Copilot:Claude Opus 5.5

Copilot

🟡 Changes recommended

Cross-device CuPy inputs can fail on a selected single GPU, and obsolete or unsupported documentation claims remain.

4 open findings

Reuse CuPy cores only on the selected device · New
Remove abandoned one-process-per-GPU implementation plan · New
Remove obsolete MPI/SPMD design documentation · New
Add the missing scaling benchmark results · New

What changed in this PR

🤖 AI text below 🤖

Adds single-node multi-GPU SRC sweeps using one thread per device.

Changes:

Splits sketches, environments, and projections across devices.
Adds peer-copy collectives, shared site loading, and per-device planning.
Adds CPU/GPU tests, documentation, and scaling benchmarks.

File	Description
`src/src_method/_group.py`	Implements device groups and collectives.
`src/src_method/_sweep.py`	Distributes the SRC sweep.
`src/src_method/_sites.py`	Shares staged cores across devices.
`src/src_method/_plan.py`	Plans multi-device memory usage.
`src/src_method/utils/_backend.py`	Adds device and peer-copy helpers.
`src/src_method/utils/__init__.py`	Exports backend helpers.
`tests/test_group.py`	Tests collectives and failures.
`tests/test_stack.py`	Tests distributed sweep behavior.
`tests/test_sites.py`	Tests shared site loading.
`tests/test_plan.py`	Tests multi-device planning.
`tests/test_gpu_backend.py`	Adds multi-GPU tests.
`tests/test_backend.py`	Tests host backend helpers.
`benches/large/bench_large.py`	Adds scaling benchmarks.
`benches/large/README.md`	Documents benchmark usage.
`README.md`	Introduces multi-GPU usage.
`docs/content/docs/features/large-problems.mdx`	Documents multi-GPU execution.
`docs/content/docs/features/gpu.mdx`	Links GPU guidance.
`docs/content/docs/contributing/testing.mdx`	Documents multi-device testing.
`docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md`	Adds the obsolete MPI design.
`docs/superpowers/plans/2026-09-25-src-multi-gpu.md`	Adds the obsolete MPI plan.

🧠 Review effort: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

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

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

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

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

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

Assisted-by: Pi:claude-opus-5.5

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

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

Assisted-by: Pi:claude-opus-5.5

Phase 1 of the plan (#44's own description) called these dead weight. Assisted-by: GitHub Copilot:Claude Opus 5.5

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

# Conflicts: # benches/large/README.md

Copilot

🔵 Needs a closer look

The added superpowers documents contradict both the implemented architecture and the PR’s stated removal of the obsolete MPI design.

0 open findings

4 resolved since last review

🧠 Review effort: Balanced

* 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>

robertodr added this pull request to stack #43 September 27, 2026 11:36

robertodr requested a review from Panadestein September 27, 2026 11:38

robertodr mentioned this pull request Sep 29, 2026

Adopt the Python array API standard for backend dispatch and array typing #45

Open

robertodr force-pushed the feat/src-multi-gpu branch from fb7dca8 to e173b01 Compare September 30, 2026 19:55

Panadestein force-pushed the feat/src-multi-gpu branch 2 times, most recently from 44cb31e to a5ec1e2 Compare October 5, 2026 15:13

Base automatically changed from feat/src-out-of-core to main October 7, 2026 11:50

Panadestein force-pushed the feat/src-multi-gpu branch from a5ec1e2 to 03d26a6 Compare October 8, 2026 14:03

Panadestein changed the title ~~feat(sweep): ✨ SRC sweep over the GPUs of a node~~ feat(sweep): ✨ split the SRC sweep among the GPUs of a node Oct 8, 2026

Panadestein marked this pull request as ready for review October 8, 2026 14:48

Panadestein mentioned this pull request Oct 13, 2026

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

Merged

Panadestein added a commit that referenced this pull request Oct 13, 2026

chore: 🔥 drop the superpowers planning docs

9dd7593

Phase 1 of the plan (#44's own description) called these dead weight. Assisted-by: GitHub Copilot:Claude Opus 5.5

Panadestein requested a balanced review from Copilot October 13, 2026 08:33

Copilot started reviewing on behalf of Panadestein October 13, 2026 08:34 View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

Comment thread src/src_method/_sites.py

Comment thread docs/superpowers/plans/2026-09-25-src-multi-gpu.md Outdated

Comment thread docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md Outdated

Comment thread benches/large/README.md

robertodr and others added 9 commits October 13, 2026 14:29

docs(spec): 📝 fold the refinements found while planning into the spec

2f0a8b9

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

b31124d

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

ea6dcc6

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

c2d6ab0

Assisted-by: Pi:claude-opus-5.5

perf(bench): 📈 multi-GPU scaling command for the large benchmark

54caf30

Assisted-by: Pi:claude-opus-5.5

docs: 📝 running one sweep on several GPUs

3d5dde5

Assisted-by: Pi:claude-opus-5.5

Panadestein added a commit that referenced this pull request Oct 13, 2026

chore: 🔥 drop the superpowers planning docs

a7a339b

Phase 1 of the plan (#44's own description) called these dead weight. Assisted-by: GitHub Copilot:Claude Opus 5.5

Panadestein force-pushed the feat/src-multi-gpu branch from 7e7bdb9 to 3d5dde5 Compare October 13, 2026 12:29

merge: 🔀 main into feat/src-multi-gpu

0d337f4

# Conflicts: # benches/large/README.md

robertodr requested a balanced review from Copilot October 14, 2026 08:12

Copilot started reviewing on behalf of robertodr October 14, 2026 08:12 View session

Copilot AI reviewed Oct 14, 2026

View reviewed changes

robertodr commented Oct 14, 2026

View reviewed changes

Comment thread docs/superpowers/plans/2026-09-25-src-multi-gpu.md Outdated

robertodr commented Oct 14, 2026

View reviewed changes

Comment thread docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md Outdated

chore: remove implementation plans.

37f3324

Panadestein merged commit fb62638 into main Oct 14, 2026
15 checks passed

Panadestein deleted the feat/src-multi-gpu branch October 14, 2026 08:46

Navigation Menu

Repository navigation

feat(sweep): ✨ split the SRC sweep among the GPUs of a node - #44

feat(sweep): ✨ split the SRC sweep among the GPUs of a node#44
Panadestein merged 11 commits into
mainfrom
feat/src-multi-gpu

robertodr commented Sep 27, 2026 •

edited by Panadestein

Loading

Uh oh!

github-actions Bot commented Oct 8, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

robertodr commented Sep 27, 2026 • edited by Panadestein Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Design

What is not split (sets the efficiency)

Changes

Testing

Status

Uh oh!

github-actions Bot commented Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test Results

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

robertodr commented Sep 27, 2026 •

edited by Panadestein

Loading

github-actions Bot commented Oct 8, 2026 •

edited

Loading