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

[Panadestein]

🤖 AI text below 🤖

Runs #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. Targets feat/src-multi-gpu.

Results

GPU tests: tests/test_gpu_backend.py, 19/19 pass, including the five new 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 (264 tests).

Reference size (D_M = 4000, chi_out = 2000, 50 sites, complex128, 184 GB of input):

GPUs	Time (s)	Speed-up	Parallel efficiency	Right-to-left (s)	Environment stalls (s)
1	3925.3	1.00	100%	3112.1	2542.3
4	1099.9	3.57	89%	832.6-844.1	697.3-698.5

This is above the 80% acceptance target. Every run at this size failed before the fix, including #59's single-GPU run on main. Waiting for environments read back from host memory and disk is 82-84% of the right-to-left pass; with 4 GPUs each GPU reads only a quarter of each environment, so the wait drops 3.65x.

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, per-GPU numbers and analysis are in the "Several GPUs (Leonardo, 4xA100-64GB)" 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.
• 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.
• benches/large/leonardo/run_multi.sh: runs any command with the 4 GPUs visible to one process, prints the topology and peer-access matrix, and with WARM_DIR reads the stack beforehand. The logs and JSON lines are committed.
• 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 more environment on the device than one GPU. 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.
• A few environments spill to Lustre. 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.
• The first SRC run of a job is about 22 s slower; the cause is not isolated.

Still to add

• The 2-GPU point and the output distance at the reference size, from a scaling --devices 4 2 run.

[github-actions]

Test Results

    5 files      5 suites   3m 7s ⏱️
  266 tests   259 ✅  7 💤 0 ❌
1 268 runs  1 253 ✅ 15 💤 0 ❌

Results for commit fb1b213.

♻️ 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

• Preserve Path import suppression for runtime annotations · New
• Restore PLC0415 suppression for lazy optional import · New
• Retry can still OOM during concatenated result assembly · New

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.