Document multi-GPU support

[pzeiger]

Documents opt-in multi-GPU execution (abTEM/abTEM#269) together with the GPU memory and FFT-sizing changes that came out of hardening it (abTEM/abTEM#346), and corrects two statements in the published documentation that never matched the shipped code.

Merged current main first — this branch predated autodoc2 by 18 commits. Nothing is re-executed, so the notebook diffs are prose only.

Corrections to published documentation

• The FFT plan cache. The walkthrough states that abTEM "sets the FFT plan cache size of cupy to zero, as we find that in most cases the increased memory consumption of the plans are not worth the small speedup they provide". The shipped default was -1 (CuPy's own unlimited) and is now auto — 25 % of the device's total memory, resolved per device. The recommendation is inverted too: a flat bound measured about 7 % slower on fast-radix scans, which is why the default is device-relative rather than a fixed cap. Replaced with what the cache does, all four settings, and the oversized-plan fallback.
• reference/default_config.yaml is now a verbatim copy of abtem/core/abtem.yaml. It carried the same 0 MB error and was missing dask.multi-gpu-rmm-pool, dask.multi-gpu-devices (both #346) and potential.slice-chunk-size (Major GPU update: memory-bounded potential chunking, GPU kernel improvements, and opt-in multi-GPU abTEM#269). The page introduces it as "the full default configuration file", so it is now exactly that, comments included.

Multi-GPU walkthrough

The "Multiple GPUs" section covers the dask.multi-gpu flag, the bring-your-own-cluster path, and the single-threaded-client requirement, and now also:

• to_zarr() distributes like compute(). Previously it silently ignored the flag and ran the whole scan on one device — the failure that motivated Multi gpu hardening abTEM#346.
• The if __name__ == "__main__" guard that dask-cuda requires in a script rather than a notebook.
• dask.multi-gpu-rmm-pool and dask.multi-gpu-devices as configuration, covering the two cases the section previously sent readers to a hand-built cluster for.
• The client configuration reaching the workers, so precision and device mean the same in a distributed run as in a local one.
• A note that a declined multi-GPU request is now reported with its reason rather than silently ignored.

The scaling tip no longer claims super-linear scaling. Its argument is unchanged — more aggregate device memory permits larger batches and fewer rebuilds — but the earlier "~2.0× on two GPUs, ~4.3× on four" figures predate abTEM/abTEM#346's discovery that the client configuration never reached the workers, so a float64 client dispatched float32 work to them while its single-GPU baseline ran in float64. Better-than-linear speed-ups were that bug's signature. The tip now quotes measured single-node, equal-precision numbers instead: 51 minutes on one 40 GB A100 against 13 minutes on four, 3.86×.

Performance tips appendix

• "Running out of memory?" listed only host-memory remedies. It now also covers the potential slice-chunk size, the cuFFT plan-cache bound, and the aggregate device memory more GPUs provide.
• "Use 'good' numbers of gpts" gains a GPU note: cuFFT's Bluestein fallback costs a workspace several times the transform size, not just time, so an unlucky grid can turn a calculation that fits on a card into one that does not. abTEM now warns and names the next good size; the note quotes that warning and shows is_fast_fft_size / next_fast_fft_size.
• Float precision notes that the setting is resolved inside each task and now reaches distributed workers — and that before 1.1 it did not, so double-precision results from a distributed cluster on an earlier version were actually computed in float32.

Install and changelog

• getting_started/install notes the optional dask-cuda dependency, and that it pins dask/distributed to its own RAPIDS release — --no-deps is usually the practical route. Adds a (getting_started:install)= target so the walkthrough can link to it.
• abtem/changelog had no multi-GPU entry. It now has the feature entries and, kept separate, the behaviour changes: the configuration reaching distributed workers (which changes distributed results for anyone relying on a non-default configuration), the bounded plan cache, and the halved scan batches on Bluestein grids.

Checks

scripts/check_notebook_widgets.py passes, both notebooks validate against nbformat, and every MyST cross-reference target in the repo resolves.

🤖 Generated with Claude Code

[pzeiger]

Hi @TomaSusi,

(This comment — like the commits it describes — was written by Claude; Paul reviewed and posted it.)

Updated for abTEM/abTEM#346, as discussed on that PR. The branch was 18 commits behind main (it predates autodoc2), so it starts with a merge of main, then four commits. The description above is rewritten to match.

Two things need your judgement rather than just review.

The super-linear scaling tip is removed. It quoted "~2.0× on two GPUs and about 4.3× on four — beyond the 4× that ideal linear scaling would predict", from a benchmark run in July. That predates #346's finding that abTEM's configuration never reached the dask-cuda workers: a client configured for float64 dispatched float32 work to them, while its single-GPU baseline ran in float64. Better-than-linear speed-ups are precisely that bug's signature — it is what made us re-examine the "12.9× faster than serial" figure in #346, which collapsed to a clean 3.86× once the configuration reached the workers. The tip keeps its argument, which is sound and is why the production case runs at all — more aggregate device memory permits larger batches and fewer rebuilds — but the numbers are now the measured single-node, equal-precision ones: 51 minutes on one 40 GB A100 against 13 minutes on four, 3.86×. If you would like a genuinely super-linear example in the docs, it is worth re-measuring one on the fixed code; I did not want to publish an unverified one next to a PR that explains why such figures were suspect.

Should this retarget from main to v1.1.0? Multi-GPU is a 1.1 feature and you are actively updating the public pages, so that may be where it belongs. v1.1.0 already contains main, so the merge is compatible either way — I left the base alone rather than moving it under you.

The rest, by commit:

Sync the configuration reference with the shipped abtem.yaml. default_config.yaml had drifted well past the one key this PR originally added. cupy.fft-cache-size was documented as 0 MB when the shipped default was -1, and is now auto; dask.multi-gpu-rmm-pool and dask.multi-gpu-devices from #346 were absent; and so was potential.slice-chunk-size from abTEM/abTEM#269. The file is now a verbatim copy of abtem/core/abtem.yaml, which also brings across the per-key comments. Flagging the #269 key explicitly since it is not this PR's — syncing only part of a file the page introduces as "the full default configuration file" seemed worse than syncing all of it, but say the word if you would rather it came out.

Correct the FFT plan-cache documentation and expand the multi-GPU section. The "Using GPUs" section states that abTEM "sets the FFT plan cache size of cupy to zero, as we find that in most cases the increased memory consumption of the plans are not worth the small speedup they provide". This is already published and has never matched the code — the default was -1, CuPy's own unlimited, and is now auto. #346's A/B measurement also inverts the recommendation: a flat bound cost about 7 % on fast-radix scans, which is why the default ended up device-relative rather than a fixed cap. The paragraph now explains what the cache does, gives all four settings, and covers the oversized-plan fallback.

"Multiple GPUs" gains what changed in #346: to_zarr() distributing like compute() — it was not mentioned at all, and it is the failure that started the whole PR; the if __name__ == "__main__" guard that dask-cuda needs in a script rather than a notebook; dask.multi-gpu-rmm-pool and dask.multi-gpu-devices, which cover the two cases the section previously sent readers to a hand-built cluster for; the client configuration now reaching the workers; and a note that a declined multi-GPU request is reported rather than silently ignored.

Add GPU memory and FFT-size guidance to the performance tips. "Running out of memory?" listed only host-memory remedies, so it now also covers the potential slice-chunk size, the plan-cache bound, and the aggregate device memory more GPUs provide. The "good numbers of gpts" discussion gains a GPU note: cuFFT's Bluestein fallback costs a workspace several times the transform size, not merely time, so an unlucky grid can turn a calculation that fits on a card into one that does not — and abTEM now warns and names the next good size. The float-precision section notes that the setting is resolved inside each task and now reaches distributed workers, and that before 1.1 it did not, so float64 on a distributed cluster silently computed in float32.

Document the dask-cuda dependency conflict and add changelog entries. dask-cuda pins dask and distributed to its own RAPIDS release, which regularly conflicts with what abTEM was installed with; --no-deps is the practical route and the install page now says so (plus a cross-reference target so the walkthrough can link to it). The changelog had no multi-GPU entry at all — it now has the feature entries and, kept separate, the behaviour changes: the configuration reaching workers, the bounded plan cache, and the halved scan batches on Bluestein grids. The first of those changes distributed results, not just performance, so it seemed worth calling out rather than filing under bug fixes.

Nothing was re-executed, so the notebook diffs are prose only. scripts/check_notebook_widgets.py passes, both notebooks validate against nbformat, and every MyST reference target in the repo resolves (I added (getting_started:install)=, which the walkthrough now links to).

One small question: the Bluestein note shows is_fast_fft_size and next_fast_fft_size, which currently live in abtem.core.fft. Would you like them promoted to the public API? They are genuinely useful for choosing gpts programmatically, and abTEM/abTEM#347 builds on the same helpers.

Best,
Paul

🤖 Written by Claude Code on Paul's behalf

[TomaSusi]

Thank you!

One small question: the Bluestein note shows is_fast_fft_size and next_fast_fft_size, which currently live in abtem.core.fft. Would you like them promoted to the public API? They are genuinely useful for choosing gpts programmatically, and abTEM/abTEM#347 builds on the same helpers.

In the new documentation (which now includes a /dev version that is roughly kept up to date with abTEM/dev: https://abtem.github.io/doc/dev/intro.html) the API reference is generated with autodoc2, which includes all non-private methods. So I think they will automatically be part of the public API.

[TomaSusi]

I can address the merge conflict easily enough, but one minor issue cropped up during review:

parallelization.ipynb:2329 — The multi-GPU benchmark tip claims a 3.86× speed-up from 51 minutes (1 GPU) vs. 13 minutes (4 GPU), but 51/13 = 3.92×, not 3.86×. Either the displayed minute values are rounded from more precise underlying numbers that do produce 3.86×, or one of the three figures is wrong — worth double-checking against the raw benchmark data before merge.

[TomaSusi]

Reviewed — looks good to merge. Confirmed the merge conflict against v1.1.0 is resolved (the fuller multi-GPU config in default_config.yaml was correctly kept over the version that had already landed on v1.1.0), and the earlier benchmark arithmetic mismatch (51/13 min vs. the quoted 3.86× speed-up) was fixed by quoting precise raw seconds.