fix(sweep): 🐛 split a batch that runs out of GPU memory by Panadestein · Pull Request #60 · Algorithmiq/src-method

Panadestein · 2026-10-13T07:31:17Z

🤖 AI text below 🤖

Runs the multi-GPU sweep of #44 on real GPUs for the first time (one Leonardo Booster node: 4x A100-SXM-64GB, all pairs linked by NVLink) and fixes the out-of-memory error that stopped the reference size.

Results

GPU tests: tests/test_gpu_backend.py, 20/20 pass, including the multi-GPU tests. Peer access is on for every pair of GPUs, so the split-upload path is the one exercised. The fast suite passes on a GPU node too.

Reference size (D_M = 4000, chi_out = 2000, 50 sites, complex128, 184 GB of input), against the single-GPU run:

GPUs	Time (s)	Speed-up	Parallel efficiency	Right-to-left (s)	Environment stalls (s)
1	3925.3	1.00	100%	3112.1	2542.3
2	2348.8	1.67	84%	1584.0-1606.3	1856.4-1895.6
4	1099.9	3.57	89%	832.6-844.1	697.3-698.5

Accuracy: the 2- and 4-GPU outputs match to a QR distance of 7.5e-13.
Efficiency: 84% on two GPUs and 89% on four, above the 80% acceptance target.
Before the fix, every run at this size failed, including chore(bench): 📈 refresh Leonardo benchmarks and add a GPU quench experiment #59's single-GPU run on main.
Environment stalls cover both passes and include waiting for the GPU's own queue, so they don't separate slow transfers from a GPU-bound pass. They shrink as GPUs are added.

D_M = 1000, chi_out = 500:

Accuracy: the 2- and 4-GPU outputs match one GPU to 2.8e-13 and 4.4e-13.
Efficiency on 4 GPUs: 84%.
2 GPUs: superlinear (126%), not explained yet; details in the README.

Full tables and analysis are in the Leonardo section of benches/large/README.md.

Changes

fix(sweep): split a batch that runs out of GPU memory.
- Cause: at the reference size, live memory peaked at 52 GB against 56.7 GB planned, so the planner's accounting is right. CuPy's pool, however, held up to 66.7 GB in split, cached blocks, so a 12.6-14.7 GB allocation went over the 67.6 GB cap, at a different point in each run.
- Fix: the kernels now catch the out-of-memory error, release the pool's cached blocks, and compute that batch in halves joined along the batch axis. That kernel keeps the finer split for the rest of the sweep.
- Unchanged: the plan and the batches the store reads and writes, so the GPUs never need to agree on a split. Results are unchanged.
- Tests: two new tests simulate a fragmented pool on the CPU, with one and two simulated devices.
fix(sites): stage cores held on another GPU. A single-GPU sweep on Resources(devices=[d]) reused a core held on another GPU as is. It now stages it through the host, as a group does. A new GPU test covers it. (Copilot review on feat(sweep): ✨ split the SRC sweep among the GPUs of a node #44.)
Benchmark: scaling also reports a QR-based distance over the difference a - b. The inner-product distance cancels down to about sqrt(eps) (~1.5e-8), so it cannot show 1e-10; the QR distance can.
Docs: "HPC clusters" section with a Leonardo page: the CUDA 13 forward-compat setup for the nodes' older driver, an example Slurm job script for one or four GPUs, queues and performance notes.
Benchmarks: the benches/*/leonardo/ folders (Slurm scripts and run logs) are removed; they remain in the git history. bench_mpo_mpo.py moves up to benches/primitives/, and the READMEs keep only the result summaries.
Docs: large-problems.mdx describes the out-of-memory split.
Dropped the docs/superpowers/ planning docs, which feat(sweep): ✨ split the SRC sweep among the GPUs of a node #44's own description already called dead weight.

Findings for follow-up

Environments mostly stay off the GPUs. Four GPUs keep only one environment on the device. The planner picks the largest batches first, and their working set takes each GPU's memory, although each GPU holds only 2.05 GB of each environment.
Some environments spill to Lustre. On one and four GPUs, 3-4 environments go to disk because the host budget (90% of the available memory, about 470 GB) is short of the 487-496 GB needed to keep them all in RAM. 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 memory unused was not investigated.
The first SRC run of a job is about 22 s slower; the cause is not isolated.

github-actions · 2026-10-13T07:34:17Z

Test Results

5 files ± 0 5 suites ±0 3m 21s ⏱️ -1s
266 tests + 4 259 ✅ + 3 7 💤 +1 0 ❌ ±0
1 268 runs +17 1 253 ✅ +15 15 💤 +2 0 ❌ ±0

Results for commit e02ceb8. ± Comparison against base commit fb62638.

♻️ This comment has been updated with latest results.

Copilot

🟡 Changes recommended

Result concatenation can reintroduce OOM, and removed Ruff suppressions will fail linting.

3 open findings

What changed in this PR

Improves sweep resilience to fragmented GPU memory and records Leonardo multi-GPU results.

Changes:

Retries failed kernels using smaller batches.
Adds QR-based accuracy measurement and Leonardo benchmark evidence.
Documents memory handling and removes obsolete plans.

File	Description
`src/src_method/_sweep.py`	Adds adaptive batch splitting.
`src/src_method/utils/_backend.py`	Adds pool-cache release helper.
`src/src_method/utils/__init__.py`	Exports the helper.
`tests/test_stack.py`	Tests split retries and failures.
`benches/large/bench_large.py`	Adds QR distance reporting.
`benches/large/README.md`	Documents Leonardo results.
`benches/large/leonardo/run_multi.sh`	Adds four-GPU job runner.
`benches/large/leonardo/logs/scaling_m1000.jsonl`	Records warm scaling results.
`benches/large/leonardo/logs/scaling_m1000_cold.jsonl`	Records cold scaling results.
`benches/large/leonardo/logs/retry_tests_1gpu_60137742.out`	Records retry test run.
`benches/large/leonardo/logs/retry_4gpu_60137752.out`	Records reference four-GPU run.
`benches/large/leonardo/logs/multi_E_budget_59783330.out`	Records budget comparisons.
`benches/large/leonardo/logs/multi_D_m4000_1gpu_60143214.out`	Records reference single-GPU run.
`benches/large/leonardo/logs/multi_B_m1000_59782618.out`	Records warm scaling run.
`benches/large/leonardo/logs/multi_B_m1000_59776459.out`	Records cold scaling run.
`benches/large/leonardo/logs/multi_A_tests_59776456.out`	Records GPU test results.
`docs/content/docs/features/large-problems.mdx`	Documents OOM splitting.
`docs/superpowers/specs/2026-09-25-src-multi-gpu-design.md`	Removes obsolete design.
`docs/superpowers/plans/2026-09-25-src-multi-gpu.md`	Removes obsolete plan.
`.gitignore`	Preserves benchmark logs.

🧠 Review effort: Balanced

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

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

- 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

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

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

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

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

# Conflicts: # benches/large/README.md # benches/large/bench_large.py # src/src_method/_sites.py # src/src_method/_sweep.py # tests/test_gpu_backend.py # tests/test_stack.py

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

robertodr · 2026-10-14T09:09:44Z

+    """The site kernels, computed in smaller pieces after running out of memory.
+
+    A batch that fails is split along its batch axis and its pieces joined, so
+    the batches read from and written to the store stay those of the plan. The
+    split of a kernel lasts for the rest of the sweep. One instance per device.


can you rephrase the prose a bit? More human readable and less the usual, convoluted Claudio-style...

robertodr

this is really exciting stuff! Are the leonardo logs to be kept in the repo?

Assisted-by: GitHub Copilot:Claude Opus 5.5

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

Panadestein requested a review from robertodr October 13, 2026 07:53

robertodr added this pull request to stack #43 October 13, 2026 07:59

Panadestein mentioned this pull request Oct 13, 2026

chore(bench): 📈 refresh Leonardo benchmarks and add a GPU quench experiment #59

Merged

2 tasks

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 benches/large/bench_large.py

Comment thread benches/large/bench_large.py Outdated

Comment thread src/src_method/_sweep.py

Panadestein mentioned this pull request Oct 13, 2026

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

Merged

4 tasks done

robertodr and others added 15 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

docs(bench): 📝 correct what the site stalls measure

3c3fe95

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

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 chore/multi-gpu-leonardo branch from 0392e7e to fb1b213 Compare October 13, 2026 12:29

Panadestein added 2 commits October 13, 2026 16:20

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

0d337f4

# Conflicts: # benches/large/README.md

merge: 🔀 feat/src-multi-gpu into chore/multi-gpu-leonardo

087ca38

# Conflicts: # benches/large/README.md

Base automatically changed from feat/src-multi-gpu to main October 14, 2026 08:46

Panadestein changed the title ~~fix(sweep): 🐛 multi-GPU on Leonardo: survive pool fragmentation, 89% efficiency at the reference size~~ fix(sweep): 🐛 split a batch that runs out of GPU memory Oct 14, 2026

Merge remote-tracking branch 'origin/main' into chore/multi-gpu-leonardo

b0741e1

# Conflicts: # benches/large/README.md # benches/large/bench_large.py # src/src_method/_sites.py # src/src_method/_sweep.py # tests/test_gpu_backend.py # tests/test_stack.py

robertodr reviewed Oct 14, 2026

View reviewed changes

Comment thread benches/large/leonardo/run_multi.sh Outdated

robertodr reviewed Oct 14, 2026

View reviewed changes

robertodr previously approved these changes Oct 14, 2026

View reviewed changes

docs(bench): 📝 keep the best 4-GPU point at the reference size

1da7b3e

Assisted-by: GitHub Copilot:Claude Opus 5.5

Panadestein dismissed robertodr’s stale review via 1da7b3e October 14, 2026 09:16

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

Panadestein deleted the chore/multi-gpu-leonardo branch October 14, 2026 09:33

Navigation Menu

Repository navigation

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

fix(sweep): 🐛 split a batch that runs out of GPU memory#60
Panadestein merged 21 commits into
mainfrom
chore/multi-gpu-leonardo

Panadestein commented Oct 13, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 13, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

robertodr Oct 14, 2026

Uh oh!

robertodr left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Panadestein commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Results

Changes

Findings for follow-up

Uh oh!

github-actions Bot commented Oct 13, 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!

robertodr Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

robertodr left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Panadestein commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading