From 7a9f813a40992aee1b4434f350b4c7beb2167a7a Mon Sep 17 00:00:00 2001 From: Paul Zeiger Date: Mon, 13 Jul 2026 07:21:14 +0000 Subject: [PATCH 01/11] Document multi-GPU support - walkthrough/parallelization: expand the "Multiple GPUs" section with the new `dask.multi-gpu` config flag (automatic dask-cuda cluster over all GPUs), the bring-your-own-cluster path, the single-threaded-client requirement, and when multi-GPU actually helps; add a label for cross-referencing. - reference/default_config.yaml: add the `dask.multi-gpu` option (was out of sync with abtem.yaml). - getting_started/install: note the optional `dask-cuda` dependency and link to the walkthrough. Co-Authored-By: Claude Opus 4.8 --- docs/getting_started/install.md | 5 ++++ docs/reference/default_config.yaml | 3 +++ .../walkthrough/parallelization.ipynb | 23 +++++-------------- 3 files changed, 14 insertions(+), 17 deletions(-) diff --git a/docs/getting_started/install.md b/docs/getting_started/install.md index 53f32bf6..85648355 100644 --- a/docs/getting_started/install.md +++ b/docs/getting_started/install.md @@ -76,6 +76,11 @@ where * should be substituted for the CUDA Toolkit version. ```` ````` +To distribute a calculation across **several GPUs** on a node, also install +[`dask-cuda`](https://docs.rapids.ai/api/dask-cuda/stable/install/) (NVIDIA GPUs, +Linux only), matched to your CUDA/RAPIDS version. See +{ref}`Multiple GPUs ` in the walkthrough. + ### Metal on Apple silicon (experimental) A subset of features in *ab*TEM can be accelerated on Apple silicon processors using their [Metal API](https://developer.apple.com/metal/). diff --git a/docs/reference/default_config.yaml b/docs/reference/default_config.yaml index 33baffb9..1963552e 100644 --- a/docs/reference/default_config.yaml +++ b/docs/reference/default_config.yaml @@ -18,6 +18,9 @@ dask: chunk-size: 128 MB # The target chunk size to use for dask arrays on the gpu chunk-size-gpu: 512 MB + # Automatically distribute gpu computations over all visible gpus by starting a + # dask-cuda cluster. Requires the optional dask-cuda package. + multi-gpu: false cupy: # The size of the fft cache in MB used by cupy # https://docs.cupy.dev/en/stable/user_guide/fft.html#fft-plan-cache diff --git a/docs/user_guide/walkthrough/parallelization.ipynb b/docs/user_guide/walkthrough/parallelization.ipynb index f0d6bab7..0b272708 100644 --- a/docs/user_guide/walkthrough/parallelization.ipynb +++ b/docs/user_guide/walkthrough/parallelization.ipynb @@ -2258,18 +2258,7 @@ }, "tags": [] }, - "source": [ - "### Multiple GPUs \n", - "The above is enough for running *ab*TEM on a single GPU; if you are using an NVidia GPU, you may want to install `dask_cuda` (this is currently only supported on Linux). However, `dask_cuda` is currently necessary for multi-GPU calculations with *ab*TEM.\n", - "\n", - "```python\n", - "from dask_cuda import LocalCUDACluster\n", - "from dask.distributed import Client\n", - "\n", - "cluster = LocalCUDACluster()\n", - "client = Client(cluster)\n", - "```" - ] + "source": "(walkthrough:parallelization:multigpu)=\n### Multiple GPUs\n\nBy default *ab*TEM uses a **single GPU** — setting `device=\"gpu\"` runs on the current CUDA device. To distribute a calculation across **all GPUs on a node**, set the `dask.multi-gpu` configuration flag. *ab*TEM then starts a [`dask_cuda`](https://docs.rapids.ai/api/dask-cuda/stable/) cluster with one worker process per GPU and spreads the (chunked) computation across them:\n\n```python\nabtem.config.set({\"device\": \"gpu\", \"dask.multi-gpu\": True})\n\nhaadf_images.compute() # now runs across every visible GPU\n```\n\nThis requires the optional `dask_cuda` package (NVIDIA GPUs, Linux only), installed separately and matched to your CUDA/RAPIDS version. Multi-GPU is opt-in: with `dask.multi-gpu` disabled (the default), GPU computations run on a single device.\n\nFor more control — an [RMM](https://docs.rapids.ai/api/rmm/stable/) memory pool, NVLink/UCX transport, or a specific subset of GPUs — you can start the cluster yourself. *ab*TEM detects an active `dask_cuda` client and uses it:\n\n```python\nfrom dask_cuda import LocalCUDACluster\nfrom dask.distributed import Client\n\nclient = Client(LocalCUDACluster())\nabtem.config.set({\"device\": \"gpu\"})\n\nhaadf_images.compute() # uses the active cluster\n```\n\n```{note}\nOnly single-threaded, GPU-pinned clients (as produced by `LocalCUDACluster`) are used for GPU work. A plain multi-threaded `Client`/`LocalCluster` is not suitable for CuPy, so *ab*TEM falls back to single-GPU execution in that case.\n```\n\nAs with any Dask parallelism, the speed-up requires the computation to split into **at least as many independent chunks as there are GPUs** — for example frozen-phonon configurations or batches of scan positions (see the *Chunks* section above). A single plane-wave multislice has no ensemble axis to distribute and will not benefit. Results are gathered back to the client GPU." } ], "metadata": { @@ -3899,7 +3888,7 @@ "headless": false, "hide_edges_on_viewport": false, "layout": "IPY_MODEL_347632dfad1e4340b4d811303c39869e", - "max_zoom": 4.0, + "max_zoom": 4, "min_zoom": 0.2, "motion_blur": false, "motion_blur_opacity": 0.2, @@ -3921,7 +3910,7 @@ "user_panning_enabled": true, "user_zooming_enabled": true, "wheel_sensitivity": 0.1, - "zoom": 2.0, + "zoom": 2, "zooming_enabled": true } }, @@ -5004,7 +4993,7 @@ "headless": false, "hide_edges_on_viewport": false, "layout": "IPY_MODEL_e6f4e200ec394d87a2d6916697e782b1", - "max_zoom": 4.0, + "max_zoom": 4, "min_zoom": 0.2, "motion_blur": false, "motion_blur_opacity": 0.2, @@ -5026,7 +5015,7 @@ "user_panning_enabled": true, "user_zooming_enabled": true, "wheel_sensitivity": 0.1, - "zoom": 2.0, + "zoom": 2, "zooming_enabled": true } }, @@ -9205,4 +9194,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file From b15125094ca0e37277abc702d9bcb49744667206 Mon Sep 17 00:00:00 2001 From: Paul Zeiger Date: Thu, 16 Jul 2026 08:14:49 +0000 Subject: [PATCH 02/11] Document super-linear multi-GPU scaling for memory-bound calculations Add a tip to the "Multiple GPUs" section noting that more GPUs provide more aggregate device memory, so memory-bound simulations can scale super-linearly. Quote a large benchmark: ~2.0x on two GPUs and ~4.3x on four (beyond the 4x linear ideal), with the caveat that the exact speed-up depends on the calculation parameters and available hardware. Co-Authored-By: Claude Opus 4.8 --- docs/user_guide/walkthrough/parallelization.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user_guide/walkthrough/parallelization.ipynb b/docs/user_guide/walkthrough/parallelization.ipynb index 0b272708..dc635567 100644 --- a/docs/user_guide/walkthrough/parallelization.ipynb +++ b/docs/user_guide/walkthrough/parallelization.ipynb @@ -2258,7 +2258,7 @@ }, "tags": [] }, - "source": "(walkthrough:parallelization:multigpu)=\n### Multiple GPUs\n\nBy default *ab*TEM uses a **single GPU** — setting `device=\"gpu\"` runs on the current CUDA device. To distribute a calculation across **all GPUs on a node**, set the `dask.multi-gpu` configuration flag. *ab*TEM then starts a [`dask_cuda`](https://docs.rapids.ai/api/dask-cuda/stable/) cluster with one worker process per GPU and spreads the (chunked) computation across them:\n\n```python\nabtem.config.set({\"device\": \"gpu\", \"dask.multi-gpu\": True})\n\nhaadf_images.compute() # now runs across every visible GPU\n```\n\nThis requires the optional `dask_cuda` package (NVIDIA GPUs, Linux only), installed separately and matched to your CUDA/RAPIDS version. Multi-GPU is opt-in: with `dask.multi-gpu` disabled (the default), GPU computations run on a single device.\n\nFor more control — an [RMM](https://docs.rapids.ai/api/rmm/stable/) memory pool, NVLink/UCX transport, or a specific subset of GPUs — you can start the cluster yourself. *ab*TEM detects an active `dask_cuda` client and uses it:\n\n```python\nfrom dask_cuda import LocalCUDACluster\nfrom dask.distributed import Client\n\nclient = Client(LocalCUDACluster())\nabtem.config.set({\"device\": \"gpu\"})\n\nhaadf_images.compute() # uses the active cluster\n```\n\n```{note}\nOnly single-threaded, GPU-pinned clients (as produced by `LocalCUDACluster`) are used for GPU work. A plain multi-threaded `Client`/`LocalCluster` is not suitable for CuPy, so *ab*TEM falls back to single-GPU execution in that case.\n```\n\nAs with any Dask parallelism, the speed-up requires the computation to split into **at least as many independent chunks as there are GPUs** — for example frozen-phonon configurations or batches of scan positions (see the *Chunks* section above). A single plane-wave multislice has no ensemble axis to distribute and will not benefit. Results are gathered back to the client GPU." + "source": "(walkthrough:parallelization:multigpu)=\n### Multiple GPUs\n\nBy default *ab*TEM uses a **single GPU** — setting `device=\"gpu\"` runs on the current CUDA device. To distribute a calculation across **all GPUs on a node**, set the `dask.multi-gpu` configuration flag. *ab*TEM then starts a [`dask_cuda`](https://docs.rapids.ai/api/dask-cuda/stable/) cluster with one worker process per GPU and spreads the (chunked) computation across them:\n\n```python\nabtem.config.set({\"device\": \"gpu\", \"dask.multi-gpu\": True})\n\nhaadf_images.compute() # now runs across every visible GPU\n```\n\nThis requires the optional `dask_cuda` package (NVIDIA GPUs, Linux only), installed separately and matched to your CUDA/RAPIDS version. Multi-GPU is opt-in: with `dask.multi-gpu` disabled (the default), GPU computations run on a single device.\n\nFor more control — an [RMM](https://docs.rapids.ai/api/rmm/stable/) memory pool, NVLink/UCX transport, or a specific subset of GPUs — you can start the cluster yourself. *ab*TEM detects an active `dask_cuda` client and uses it:\n\n```python\nfrom dask_cuda import LocalCUDACluster\nfrom dask.distributed import Client\n\nclient = Client(LocalCUDACluster())\nabtem.config.set({\"device\": \"gpu\"})\n\nhaadf_images.compute() # uses the active cluster\n```\n\n```{note}\nOnly single-threaded, GPU-pinned clients (as produced by `LocalCUDACluster`) are used for GPU work. A plain multi-threaded `Client`/`LocalCluster` is not suitable for CuPy, so *ab*TEM falls back to single-GPU execution in that case.\n```\n\nAs with any Dask parallelism, the speed-up requires the computation to split into **at least as many independent chunks as there are GPUs** — for example frozen-phonon configurations or batches of scan positions (see the *Chunks* section above). A single plane-wave multislice has no ensemble axis to distribute and will not benefit. Results are gathered back to the client GPU.\n\n```{tip}\nMore GPUs also mean more *aggregate* device memory, which can make multi-GPU worthwhile even when raw compute is not the bottleneck. In **memory-bound** simulations — where a single GPU is forced into small batches (or has to recompute intermediate results such as potential slices) to stay within its memory — spreading the work over more devices lets each one run larger, more efficient batches, so scaling holds up well and can even become **super-linear** at higher GPU counts. In one large benchmark calculation (about $4.2$ hours on a single GPU) we observed a near-ideal $\\sim\\!2.0\\times$ speed-up on two GPUs and about $4.3\\times$ on four — beyond the $4\\times$ that ideal linear scaling would predict. The exact speed-up will depend on the specific parameters of your calculation and the available hardware, but for memory-hungry calculations, moving to more GPUs can pay off in wall-clock time beyond what the extra compute alone would suggest.\n```" } ], "metadata": { From 147d60ef7cbce5c99ae71db115c5d420471e8552 Mon Sep 17 00:00:00 2001 From: Paul Zeiger Date: Thu, 27 Aug 2026 14:41:27 +0000 Subject: [PATCH 03/11] Sync the configuration reference with the shipped abtem.yaml The reference page states that it shows "the full default configuration file", but the copy had drifted: cupy.fft-cache-size was documented as "0 MB" when the shipped default was -1 (and is now "auto"), and the potential.slice-chunk-size, dask.multi-gpu-rmm-pool and dask.multi-gpu-devices keys were missing entirely. The file is now a verbatim copy of abtem/core/abtem.yaml, so it also picks up the per-key comments explaining the fft-cache-size settings. Co-Authored-By: Claude Opus 5 --- docs/reference/default_config.yaml | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/docs/reference/default_config.yaml b/docs/reference/default_config.yaml index 1963552e..d95f57c3 100644 --- a/docs/reference/default_config.yaml +++ b/docs/reference/default_config.yaml @@ -8,7 +8,7 @@ fft: fftw precision: float32 diagnostics: # Show the progress bar. Options are 'true', 'false' or 'tqdm' - progress_bar: tqdm + progress_bar: "tqdm" # Show the progress of each task. Options are 'true' or 'false' task_progress: false dask: @@ -21,10 +21,27 @@ dask: # Automatically distribute gpu computations over all visible gpus by starting a # dask-cuda cluster. Requires the optional dask-cuda package. multi-gpu: false + # Optional RMM memory-pool size per multi-gpu worker (e.g. "20 GB"). + # null disables the RMM pool. + multi-gpu-rmm-pool: null + # Optional subset of GPUs for the multi-gpu cluster, as a list of device + # indices (e.g. [0, 1]) or a comma-separated string (e.g. "0,1"). + # null spans all visible GPUs. + multi-gpu-devices: null cupy: - # The size of the fft cache in MB used by cupy + # Maximum GPU memory (bytes) used by the cuFFT plan cache. # https://docs.cupy.dev/en/stable/user_guide/fft.html#fft-plan-cache - fft-cache-size: 0 MB + # auto — 25% of the device's total memory, resolved per device: scales + # with the card like the auto-sized batches whose plans it holds, + # keeping live plans hot while capping stale-plan retention. A + # ceiling, not a reservation — unused headroom costs no memory. + # 0 MB — disable caching entirely (workspace freed after every FFT call; + # use when VRAM is tight and large-grid OOMs occur). + # > 0 — bound cached workspace to this many bytes (e.g. "512 MB"). + # -1 — unlimited (CuPy default; fastest but accumulates workspace, + # which on Bluestein-fallback grid sizes can reach tens of GB). + # null — same as -1 (no bound). + fft-cache-size: auto mkl: # The number of threads to use for mkl threads: 2 @@ -44,6 +61,9 @@ warnings: dask-blockwise-performance: false # Show a warning when the grid is overspecified overspecified-grid: true +potential: + # Number of slices to build at once during multislice. "auto" = memory-budget-aware. + slice-chunk-size: "auto" antialias: # The antialias cutoff in reciprocal space cutoff: 0.6666666 @@ -63,4 +83,4 @@ visualize: # Scale the values of interactive plots automatically autoscale: false # Use tex rendering in plots - use_tex: true \ No newline at end of file + use_tex: true From 569c0a28db66aa4d17545919eb0a1411b82e3b49 Mon Sep 17 00:00:00 2001 From: Paul Zeiger Date: Thu, 27 Aug 2026 14:41:27 +0000 Subject: [PATCH 04/11] Correct the FFT plan-cache documentation and expand the multi-GPU section The "Using GPUs" section claimed that abTEM sets the CuPy FFT plan cache size to zero because the plans are not worth their memory. That was never the shipped default -- it was -1 (unlimited, CuPy's own default) and is now "auto", 25 % of the device's total memory resolved per device. The recommendation was inverted too: a flat bound measurably costs performance on fast-radix grids, which is why the default is device-relative. The paragraph now describes what the cache does, all four settings, and the uncached-oversized-plan fallback. The "Multiple GPUs" section gains what changed with the multi-GPU hardening work: - to_zarr() distributes like compute(); previously it silently ran the whole scan serially on one device, which is the failure that motivated the work. - The multiprocessing entry-point guard that dask-cuda requires 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 reaching the workers, so that 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 rather than silently ignored. The scaling tip keeps its argument -- aggregate device memory permits larger batches -- but replaces the earlier super-linear figures with the measured single-node numbers at equal precision: 51 minutes on one 40 GB A100 against 13 minutes on four, 3.86x. The earlier numbers predate the discovery that the client configuration never reached the workers, so a float64 client dispatched float32 work to them; better-than-linear speed-ups were that bug's signature and should not be published. Co-Authored-By: Claude Opus 5 --- .../walkthrough/parallelization.ipynb | 80 +++++++++++++++++-- 1 file changed, 75 insertions(+), 5 deletions(-) diff --git a/docs/user_guide/walkthrough/parallelization.ipynb b/docs/user_guide/walkthrough/parallelization.ipynb index dc635567..af9e23e6 100644 --- a/docs/user_guide/walkthrough/parallelization.ipynb +++ b/docs/user_guide/walkthrough/parallelization.ipynb @@ -2241,11 +2241,16 @@ "The batch size only determines the maximum number of plane waves in a batch, so you need to leave room in the memory for any intermediate overhead. \n", "```\n", "\n", - "Finally, *ab*TEM by default sets the [FFT plan cache](https://docs.cupy.dev/en/stable/user_guide/fft.html#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. You can however also change this through the *ab*TEM config.\n", + "Finally, *ab*TEM bounds the [cuFFT plan cache](https://docs.cupy.dev/en/stable/user_guide/fft.html#fft-plan-cache) of `cupy`. CuPy caches the workspace of every transform shape it has seen, which makes repeated FFTs on the same grid fast but accumulates memory across the several batch shapes a scan goes through — on grid sizes that force cuFFT's Bluestein fallback, a single plan can reach several GB. The default is `auto`, meaning 25 % of the device's total memory, resolved per device: a ceiling rather than a reservation, so unused headroom costs nothing, and it scales with the card just like the auto-sized batches whose plans it holds.\n", "\n", "```python\n", - "abtem.config.set({\"cupy.fft-cache-size\" : \"1024 MB\"})\n", - "```" + "abtem.config.set({\"cupy.fft-cache-size\": \"auto\"}) # the default: 25 % of device memory\n", + "abtem.config.set({\"cupy.fft-cache-size\": \"512 MB\"}) # a fixed bound\n", + "abtem.config.set({\"cupy.fft-cache-size\": \"0 MB\"}) # disable caching entirely\n", + "abtem.config.set({\"cupy.fft-cache-size\": -1}) # unlimited (CuPy's own default)\n", + "```\n", + "\n", + "Lower the bound (or disable the cache) if you are running out of VRAM on large grids; raise it if you see the warning below. Note that CuPy does not degrade gracefully when a *single* plan exceeds the bound — it raises rather than falling back — so *ab*TEM runs such plans uncached instead, at replanning cost, and warns you with the shape, the current bound, and both remedies: raise `cupy.fft-cache-size`, or reduce `max_batch`." ] }, { @@ -2258,7 +2263,72 @@ }, "tags": [] }, - "source": "(walkthrough:parallelization:multigpu)=\n### Multiple GPUs\n\nBy default *ab*TEM uses a **single GPU** — setting `device=\"gpu\"` runs on the current CUDA device. To distribute a calculation across **all GPUs on a node**, set the `dask.multi-gpu` configuration flag. *ab*TEM then starts a [`dask_cuda`](https://docs.rapids.ai/api/dask-cuda/stable/) cluster with one worker process per GPU and spreads the (chunked) computation across them:\n\n```python\nabtem.config.set({\"device\": \"gpu\", \"dask.multi-gpu\": True})\n\nhaadf_images.compute() # now runs across every visible GPU\n```\n\nThis requires the optional `dask_cuda` package (NVIDIA GPUs, Linux only), installed separately and matched to your CUDA/RAPIDS version. Multi-GPU is opt-in: with `dask.multi-gpu` disabled (the default), GPU computations run on a single device.\n\nFor more control — an [RMM](https://docs.rapids.ai/api/rmm/stable/) memory pool, NVLink/UCX transport, or a specific subset of GPUs — you can start the cluster yourself. *ab*TEM detects an active `dask_cuda` client and uses it:\n\n```python\nfrom dask_cuda import LocalCUDACluster\nfrom dask.distributed import Client\n\nclient = Client(LocalCUDACluster())\nabtem.config.set({\"device\": \"gpu\"})\n\nhaadf_images.compute() # uses the active cluster\n```\n\n```{note}\nOnly single-threaded, GPU-pinned clients (as produced by `LocalCUDACluster`) are used for GPU work. A plain multi-threaded `Client`/`LocalCluster` is not suitable for CuPy, so *ab*TEM falls back to single-GPU execution in that case.\n```\n\nAs with any Dask parallelism, the speed-up requires the computation to split into **at least as many independent chunks as there are GPUs** — for example frozen-phonon configurations or batches of scan positions (see the *Chunks* section above). A single plane-wave multislice has no ensemble axis to distribute and will not benefit. Results are gathered back to the client GPU.\n\n```{tip}\nMore GPUs also mean more *aggregate* device memory, which can make multi-GPU worthwhile even when raw compute is not the bottleneck. In **memory-bound** simulations — where a single GPU is forced into small batches (or has to recompute intermediate results such as potential slices) to stay within its memory — spreading the work over more devices lets each one run larger, more efficient batches, so scaling holds up well and can even become **super-linear** at higher GPU counts. In one large benchmark calculation (about $4.2$ hours on a single GPU) we observed a near-ideal $\\sim\\!2.0\\times$ speed-up on two GPUs and about $4.3\\times$ on four — beyond the $4\\times$ that ideal linear scaling would predict. The exact speed-up will depend on the specific parameters of your calculation and the available hardware, but for memory-hungry calculations, moving to more GPUs can pay off in wall-clock time beyond what the extra compute alone would suggest.\n```" + "source": [ + "(walkthrough:parallelization:multigpu)=\n", + "### Multiple GPUs\n", + "\n", + "By default *ab*TEM uses a **single GPU** — setting `device=\"gpu\"` runs on the current CUDA device. To distribute a calculation across **all GPUs on a node**, set the `dask.multi-gpu` configuration flag. *ab*TEM then starts a [`dask_cuda`](https://docs.rapids.ai/api/dask-cuda/stable/) cluster with one worker process per GPU and spreads the (chunked) computation across them:\n", + "\n", + "```python\n", + "abtem.config.set({\"device\": \"gpu\", \"dask.multi-gpu\": True})\n", + "\n", + "haadf_images.compute() # now runs across every visible GPU\n", + "```\n", + "\n", + "Saving a lazy result distributes in the same way, so a large scan never has to be brought back to the client at all:\n", + "\n", + "```python\n", + "haadf_images.to_zarr(\"scan.zarr\") # also runs across every visible GPU\n", + "```\n", + "\n", + "````{warning}\n", + "`dask_cuda` starts its workers as separate processes. In a **script** (as opposed to a notebook) this requires the usual multiprocessing entry-point guard, or the script will re-import itself in every worker:\n", + "\n", + "```python\n", + "if __name__ == \"__main__\":\n", + " abtem.config.set({\"device\": \"gpu\", \"dask.multi-gpu\": True})\n", + " ...\n", + "```\n", + "\n", + "*ab*TEM detects the resulting spawn failure and translates it into a message naming this cause, but the guard has to be added by you.\n", + "````\n", + "\n", + "This requires the optional `dask_cuda` package (NVIDIA GPUs, Linux only), installed separately and matched to your CUDA/RAPIDS version — see the {ref}`installation instructions `. Multi-GPU is opt-in: with `dask.multi-gpu` disabled (the default), GPU computations run on a single device. Whenever *ab*TEM cannot honour the flag — `dask_cuda` missing, no GPUs visible, an unsuitable client already active — it says so with a warning naming the reason, rather than silently falling back.\n", + "\n", + "Two common refinements are available as configuration rather than requiring a hand-built cluster: an [RMM](https://docs.rapids.ai/api/rmm/stable/) memory pool per worker, and restricting the cluster to a subset of the visible GPUs.\n", + "\n", + "```python\n", + "abtem.config.set({\n", + " \"dask.multi-gpu\": True,\n", + " \"dask.multi-gpu-rmm-pool\": \"20 GB\", # null (the default) disables the pool\n", + " \"dask.multi-gpu-devices\": [0, 1], # null (the default) spans all visible GPUs\n", + "})\n", + "```\n", + "\n", + "For anything beyond that — NVLink/UCX transport, a multi-node cluster, custom worker options — start the cluster yourself. *ab*TEM detects an active `dask_cuda` client and uses it:\n", + "\n", + "```python\n", + "from dask_cuda import LocalCUDACluster\n", + "from dask.distributed import Client\n", + "\n", + "client = Client(LocalCUDACluster(protocol=\"ucx\", enable_nvlink=True))\n", + "abtem.config.set({\"device\": \"gpu\"})\n", + "\n", + "haadf_images.compute() # uses the active cluster\n", + "```\n", + "\n", + "```{note}\n", + "Only single-threaded, GPU-pinned clients (as produced by `LocalCUDACluster`) are used for GPU work. A plain multi-threaded `Client`/`LocalCluster` is not suitable for CuPy, so *ab*TEM falls back to single-GPU execution in that case.\n", + "```\n", + "\n", + "Because *ab*TEM resolves configuration *inside* each task, your configuration has to reach the worker processes, which start fresh and would otherwise see only the defaults from your YAML files. *ab*TEM therefore pushes the client's merged configuration to the workers on every dispatch, to any active distributed client — the multi-GPU cluster or one you started yourself, CPU or GPU. Settings such as `precision` and `device` consequently mean the same thing in a distributed run as in a local one.\n", + "\n", + "As with any Dask parallelism, the speed-up requires the computation to split into **at least as many independent chunks as there are GPUs** — for example frozen-phonon configurations or batches of scan positions (see the *Chunks* section above). A single plane-wave multislice has no ensemble axis to distribute and will not benefit. Results are gathered back to the client GPU.\n", + "\n", + "```{tip}\n", + "More GPUs also mean more *aggregate* device memory, which can make multi-GPU worthwhile even when raw compute is not the bottleneck. In **memory-bound** simulations a single GPU is forced into small batches — or has to rebuild intermediate results such as potential slices — to stay within its memory, and in the worst case cannot run the calculation at all. Spreading the work over more devices lets each one use larger, more efficient batches. In a large STEM benchmark (101,871 scan positions on a $2623\\times2271$ grid in double precision) one 40 GB A100 completed the scan in 51 minutes and four completed it in 13 — a $3.86\\times$ speed-up on four devices. The exact speed-up depends on the parameters of your calculation and on your hardware.\n", + "```" + ] } ], "metadata": { @@ -9194,4 +9264,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} From ac4f5c27bc20ae195c1fd13074d79ff8c617ccec Mon Sep 17 00:00:00 2001 From: Paul Zeiger Date: Thu, 27 Aug 2026 14:41:27 +0000 Subject: [PATCH 05/11] Add GPU memory and FFT-size guidance to the performance tips "Running out of memory?" listed only reduce-early, smaller batches and fewer workers, all of which are host-memory answers. Device memory is usually the tighter constraint, so the section now also covers the potential slice-chunk size, the cuFFT plan-cache bound, and the aggregate device memory that more GPUs provide. The "good numbers of gpts" discussion gains a note that the point is sharper on GPU: cuFFT has no kernel for a length with a prime factor above 7 and falls back to Bluestein, which needs a workspace several times the transform size, so an unlucky grid costs memory as well as time. abTEM now warns when it meets such a grid and names the next good size; the note quotes that warning and shows the helpers behind it. 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 double-precision results from a distributed run on an earlier version were actually computed in float32. Co-Authored-By: Claude Opus 5 --- .../appendix/performance_tips.ipynb | 59 ++++++++++++++++++- 1 file changed, 58 insertions(+), 1 deletion(-) diff --git a/docs/user_guide/appendix/performance_tips.ipynb b/docs/user_guide/appendix/performance_tips.ipynb index 7fbf40c2..abf39714 100644 --- a/docs/user_guide/appendix/performance_tips.ipynb +++ b/docs/user_guide/appendix/performance_tips.ipynb @@ -172,7 +172,29 @@ "\n", "```python\n", "dask.config.set(num_workers=2)\n", - "```" + "```\n", + "\n", + "### On the GPU\n", + "\n", + "Device memory is usually the tighter constraint, and a few options apply only there.\n", + "\n", + "*ab*TEM builds the potential in slice groups sized against the free VRAM it measures at computation time, rather than\n", + "materializing every slice at once. If the automatic estimate is still too generous for your calculation, fix the group\n", + "size yourself:\n", + "\n", + "```python\n", + "abtem.config.set({\"potential.slice-chunk-size\": 8}) # \"auto\" (the default) sizes it against free VRAM\n", + "```\n", + "\n", + "CuPy caches the workspace of every FFT shape it has seen. The cache is bounded to 25 % of device memory by default;\n", + "lower the bound, or disable caching entirely, when VRAM is tight:\n", + "\n", + "```python\n", + "abtem.config.set({\"cupy.fft-cache-size\": \"512 MB\"}) # or \"0 MB\" to disable\n", + "```\n", + "\n", + "Finally, more GPUs mean more *aggregate* device memory. A calculation that does not fit on one device may fit\n", + "comfortably across several — see {ref}`Multiple GPUs `." ] }, { @@ -311,6 +333,33 @@ " print(f\"gpts: {gpts} ({factors}), time: {elapsed:.2f} s\")" ] }, + { + "cell_type": "markdown", + "id": "gpu-fft-size-note", + "metadata": {}, + "source": [ + "````{note}\n", + "This matters more on the GPU than on the CPU. cuFFT has no optimized kernel for a length with a prime factor above 7\n", + "and falls back to the Bluestein algorithm, which pads internally and needs a workspace several times the size of the\n", + "transform — so an unlucky grid costs memory as well as time, and can turn a calculation that would fit on your card\n", + "into one that does not. *ab*TEM warns you when it meets such a grid, naming the offending size and the next good one:\n", + "\n", + " UserWarning: FFT size 2623 x 2271 contains prime factors larger than 7; cuFFT falls back to the\n", + " Bluestein algorithm, which is several times slower and allocates a workspace of several times the\n", + " array size. Consider adjusting the grid to 2625 x 2304, e.g. by setting gpts explicitly instead of\n", + " the sampling.\n", + "\n", + "The same test is available directly, which is useful when choosing `gpts` programmatically:\n", + "\n", + "```python\n", + "from abtem.core.fft import is_fast_fft_size, next_fast_fft_size\n", + "\n", + "is_fast_fft_size(2623) # False\n", + "next_fast_fft_size(2623) # 2625\n", + "```\n", + "````" + ] + }, { "cell_type": "markdown", "id": "eace7d67", @@ -324,6 +373,14 @@ "\n", "```python\n", "abtem.config.set({\"precision\": \"float64\"})\n", + "```\n", + "\n", + "```{note}\n", + "*ab*TEM resolves the precision inside each task rather than at the point you set it, so in a distributed computation\n", + "the setting has to reach the worker processes. It does: the client's configuration is pushed to the workers on every\n", + "dispatch. Before *ab*TEM 1.1 it was not, and a distributed run silently computed in the default `float32` no matter\n", + "what the client had configured — so double-precision results obtained with a `distributed` cluster on an earlier\n", + "version are worth repeating.\n", "```" ] }, From add6921c4c82f8058c149a048ef685e4f1bfcb80 Mon Sep 17 00:00:00 2001 From: Paul Zeiger Date: Thu, 27 Aug 2026 14:41:27 +0000 Subject: [PATCH 06/11] Document the dask-cuda dependency conflict and add changelog entries Installing dask-cuda pulls in the dask and distributed pins of its own RAPIDS release, which frequently conflict with the versions abTEM was installed with; --no-deps is the practical route, and the install page now says so. The page also gains a cross-reference target so the walkthrough can link to it. The changelog gains the multi-GPU feature entries and, separately, the behaviour changes: the client configuration now reaching distributed workers (which changes distributed results for anyone who relied on a non-default configuration), the bounded cuFFT plan cache, and the halved scan batches on Bluestein grids. Co-Authored-By: Claude Opus 5 --- docs/abtem/changelog.md | 8 ++++++++ docs/getting_started/install.md | 14 ++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/docs/abtem/changelog.md b/docs/abtem/changelog.md index 61f9016a..bab439f5 100644 --- a/docs/abtem/changelog.md +++ b/docs/abtem/changelog.md @@ -12,6 +12,14 @@ pseudo-Voigtian (L + G) source-size distributions and filters ([PR #270](https:/ - Radially variable detector sensitivity ([PR #283](https://github.com/abTEM/abTEM/pull/283)) - Plasmons: fast `PhaseScramblePlasmons` for multislice and PRISM, and `MonteCarloPlasmons` for Bloch wave - CBED patterns for Bloch waves ([PR #254](https://github.com/abTEM/abTEM/pull/254)) +- Opt-in multi-GPU execution: setting `dask.multi-gpu` starts a `dask-cuda` cluster with one worker per GPU and distributes the computation across all visible devices, on both the `compute()` and the `to_zarr()` path ([PR #269](https://github.com/abTEM/abTEM/pull/269), [PR #346](https://github.com/abTEM/abTEM/pull/346)) +- New configuration keys `dask.multi-gpu-rmm-pool` and `dask.multi-gpu-devices` for an RMM memory pool per worker and for restricting the cluster to a subset of GPUs ([PR #346](https://github.com/abTEM/abTEM/pull/346)) +- *ab*TEM now warns instead of silently falling back: multi-GPU requested but declined (with the reason), a missing `if __name__ == "__main__"` guard, a grid size that forces cuFFT's Bluestein fallback (naming the next fast size), and a cuFFT plan too large for the cache bound ([PR #346](https://github.com/abTEM/abTEM/pull/346)) + +Behavior changes: +- **The client's configuration now reaches `distributed` workers.** *ab*TEM resolves configuration inside each task, and worker processes previously saw only the defaults from the YAML files — so any distributed computation with non-default configuration silently used the defaults. Most consequentially, `precision: float64` was ignored and distributed runs computed in `float32`. Distributed results from earlier versions that relied on non-default configuration should be repeated ([PR #346](https://github.com/abTEM/abTEM/pull/346)) +- `cupy.fft-cache-size` now defaults to `auto`, bounding the cuFFT plan cache to 25 % of each device's memory (it was previously unlimited). Set `-1` to restore unlimited caching, `0 MB` to disable it, or a size such as `512 MB` for a fixed bound ([PR #346](https://github.com/abTEM/abTEM/pull/346)) +- On GPU grids that force cuFFT's Bluestein fallback, automatically sized scan batches are halved to reflect the larger FFT workspace such grids require ([PR #346](https://github.com/abTEM/abTEM/pull/346)) ## 1.0.10 diff --git a/docs/getting_started/install.md b/docs/getting_started/install.md index 85648355..323dc963 100644 --- a/docs/getting_started/install.md +++ b/docs/getting_started/install.md @@ -1,3 +1,4 @@ +(getting_started:install)= # Installation There are many ways to install the *ab*TEM package, for example conda or pip: @@ -81,6 +82,19 @@ To distribute a calculation across **several GPUs** on a node, also install Linux only), matched to your CUDA/RAPIDS version. See {ref}`Multiple GPUs ` in the walkthrough. +```{note} +`dask-cuda` pins `dask` and `distributed` to the versions of its own RAPIDS +release, which may not be the versions *ab*TEM was installed with. If pip +insists on downgrading `dask` (or refuses to resolve the environment at all), +install `dask-cuda` without its dependencies and keep the versions you already +have: + + pip install --no-deps dask-cuda + +Both packages track `dask` closely, so it is worth checking that a simple +multi-GPU computation runs after installing this way. +``` + ### Metal on Apple silicon (experimental) A subset of features in *ab*TEM can be accelerated on Apple silicon processors using their [Metal API](https://developer.apple.com/metal/). From ae0fb10759c4703d4c4348a1087ebc3d7ab699d4 Mon Sep 17 00:00:00 2001 From: TomaSusi Date: Fri, 28 Aug 2026 15:20:34 +0200 Subject: [PATCH 07/11] Fix API reference Parameters sections rendering as a single paragraph autodoc2 has no napoleon integration (it never fires the autodoc-process-docstring event napoleon hooks into), so abTEM NumPy-style Parameters/Returns/etc docstring sections were passed through as plain text, and MyST collapsed each section into one unbroken paragraph. Adds a builder-inited hook that runs every collected docstring through sphinx.ext.napoleon.NumpyDocstring (converting to RST field lists) and forces those to be parsed as RST rather than in the ambient MyST context, restoring the old sphinx.ext.napoleon per-parameter rendering. Also works around a crash this exposed in autodoc2 DocstringRenderer: it can raise TypeError (unsupported operand for +: NoneType and int) when docutils emits a system_message without a line number while parsing under an explicit parser option, aborting the whole build. Patched to fail soft (render that one docstring empty, with a warning) instead, matching how docutils own default reporter already handles this case. --- docs/_config.yml | 5 +++ docs/_ext/abtem_autodoc2.py | 68 +++++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/docs/_config.yml b/docs/_config.yml index 3cb0fc0f..f2e3f0b4 100644 --- a/docs/_config.yml +++ b/docs/_config.yml @@ -48,6 +48,11 @@ sphinx: autodoc2_render_plugin: "myst" autodoc2_output_dir: "reference/api/apidocs" autodoc2_sort_names: true + # Docstrings are NumPy-style, converted to RST field lists by the + # abtem_autodoc2 extension (autodoc2 has no napoleon integration) - + # parse them as RST rather than in the ambient MyST context, which + # doesn't understand RST field-list syntax. + autodoc2_docstring_parser_regexes: [[".*", "rst"]] autodoc2_hidden_objects: ["private", "dunder", "inherited"] autodoc2_skip_module_regexes: [".*\\._.*"] # Skip the generated index.rst wrapper page (just a toctree + footnote) - diff --git a/docs/_ext/abtem_autodoc2.py b/docs/_ext/abtem_autodoc2.py index 927e4721..45614682 100644 --- a/docs/_ext/abtem_autodoc2.py +++ b/docs/_ext/abtem_autodoc2.py @@ -19,6 +19,74 @@ def _set_autodoc2_packages(app, config): ] +def _convert_numpydoc_to_rst(app): + """Run every collected docstring through napoleon to get RST field lists. + + autodoc2 has no napoleon integration (it never fires the + autodoc-process-docstring event napoleon hooks into), so abTEM's + NumPy-style "Parameters" sections are otherwise passed straight through + as plain text, with docutils/MyST collapsing each section into one + unbroken paragraph. Rewriting the docstrings in-place, after autodoc2 + has finished analysing the package, restores the per-parameter + rendering the old sphinx.ext.napoleon setup produced. + """ + from autodoc2.sphinx.utils import get_database + from sphinx.ext.napoleon import Config as NapoleonConfig, NumpyDocstring + + napoleon_config = NapoleonConfig() + db = get_database(app.env) + for item in db._items.values(): + if item.get("doc"): + item["doc"] = str(NumpyDocstring(item["doc"], napoleon_config)) + + +def _patch_autodoc2_source_line_bug(): + """Work around a crash in autodoc2's DocstringRenderer. + + When autodoc2 is given an explicit `:parser:` (which we set via + autodoc2_docstring_parser_regexes, to parse the napoleon-converted RST), + it points the parsed document's reporter at a `lambda li: (source_path, + li + source_offset)` for source/line lookups. docutils sometimes emits + a system_message (e.g. the harmless "Enumerated list start value not + ordinal-1" INFO notice) without a line number, passing `li=None`, which + makes that lambda raise `TypeError: unsupported operand type(s) for +: + 'NoneType' and 'int'` - an unhandled exception that aborts the whole + build. docutils' own system_message() already tolerates this pattern + for its default reporter (it catches AttributeError and falls back to + an unknown source/line), so we only need autodoc2's directive to fail + the same soft way instead of raising. Reported upstream; remove once + fixed: https://github.com/sphinx-extensions2/sphinx-autodoc2/issues + """ + from autodoc2.sphinx.docstring import DocstringRenderer + + if getattr(DocstringRenderer.run, "_abtem_patched", False): + return + + original_run = DocstringRenderer.run + + def patched_run(self): + try: + return original_run(self) + except TypeError as exc: + if "NoneType" not in str(exc): + raise + from sphinx.util import logging + + logging.getLogger(__name__).warning( + f"[[abtem_autodoc2]] Could not render docstring for " + f"{self.arguments[0]!r} (hit a known autodoc2/docutils " + f"line-number bug); rendering it empty instead." + ) + return [] + + patched_run._abtem_patched = True + DocstringRenderer.run = patched_run + + def setup(app): app.connect("config-inited", _set_autodoc2_packages) + _patch_autodoc2_source_line_bug() + # Must run after autodoc2's own builder-inited handler (default + # priority 500), which populates the database this reads from. + app.connect("builder-inited", _convert_numpydoc_to_rst, priority=900) return {"parallel_read_safe": True, "parallel_write_safe": True} From 84ec8bb70e1f99e9e2e5b1c1c9ddda85bcc80ddf Mon Sep 17 00:00:00 2001 From: TomaSusi Date: Fri, 28 Aug 2026 16:11:03 +0200 Subject: [PATCH 08/11] Make dev docs actually build against abTEM's unreleased dev branch The dev-docs workflow installed the same PyPI release as the stable build, so the development banner and version footer were misleading - there was no unreleased code being tested against. Install abTEM from its dev branch for this workflow, and update the banner/footer wording to match. Co-Authored-By: Claude Sonnet 5 --- .github/workflows/deploy-dev.yml | 11 ++++++++--- docs/_ext/abtem_version_footer.py | 23 +++++++++++++---------- 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/.github/workflows/deploy-dev.yml b/.github/workflows/deploy-dev.yml index 9a292c14..8c94758b 100644 --- a/.github/workflows/deploy-dev.yml +++ b/.github/workflows/deploy-dev.yml @@ -1,6 +1,7 @@ -# Publishes the development documentation (built from main) under /dev/ on -# gh-pages. The stable docs at the site root are published by deploy-book.yml -# on each release and are the default readers see. +# Publishes the development documentation (this repo's main branch, built +# against abTEM's unreleased `dev` branch) under /dev/ on gh-pages. The +# stable docs at the site root are published by deploy-book.yml on each +# release and are the default readers see. name: deploy-dev-docs @@ -37,6 +38,10 @@ jobs: run: | pip install -r docs/requirements.txt + - name: Install abTEM from the dev branch (unreleased) + run: | + pip install --no-deps --force-reinstall git+https://github.com/abTEM/abTEM.git@dev + - name: Strip orphaned ipywidgets state run: | python scripts/check_notebook_widgets.py --fix diff --git a/docs/_ext/abtem_version_footer.py b/docs/_ext/abtem_version_footer.py index dd598497..cbeaa556 100644 --- a/docs/_ext/abtem_version_footer.py +++ b/docs/_ext/abtem_version_footer.py @@ -13,6 +13,7 @@ STABLE_URL = "https://abtem.github.io/doc/" DEV_URL = STABLE_URL + "dev/" +ABTEM_DEV_BRANCH_URL = "https://github.com/abTEM/abTEM/tree/dev" def _is_dev(): @@ -28,20 +29,22 @@ def _set_abtem_version_footer(app, config): version = "unknown" theme_options = dict(config.html_theme_options or {}) - footer = ( - "

Tested against " - 'abTEM' - f" v{version}." - ) if _is_dev(): theme_options["announcement"] = ( - "This is the development documentation, built from " - f'the latest main branch. Switch to the ' - "stable version." + "This is the development documentation, " + f'including unreleased work in ' + f'abTEM/dev. Switch to the stable ' + "version." + ) + footer = ( + "

Tested against pre-release " + f'abTEM v{version}.' ) - footer += " Development build." else: - footer += ( + footer = ( + "

Tested against " + 'abTEM' + f" v{version}." f' Also available: development version.' ) footer += "

" From 013241e8164a2971ea725ea77fde12befe473364 Mon Sep 17 00:00:00 2001 From: TomaSusi Date: Fri, 28 Aug 2026 16:11:10 +0200 Subject: [PATCH 09/11] Shorten API reference signatures and link return types Function/method/class headings in the generated API reference repeated every argument with its fully-qualified type and default, duplicating the Parameters list rendered just below. Add a custom autodoc2 MyST renderer that collapses argument lists to (...) in headings, and enable python_use_unqualified_type_names to shorten the remaining return-type annotation. Also add intersphinx mappings for numpy and python so well-formed external type references (e.g. numpy.ndarray) link out instead of staying plain text. Co-Authored-By: Claude Sonnet 5 --- docs/_config.yml | 16 +++++++++++++++- docs/_ext/abtem_autodoc2.py | 23 +++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/docs/_config.yml b/docs/_config.yml index f2e3f0b4..c5b52435 100644 --- a/docs/_config.yml +++ b/docs/_config.yml @@ -39,13 +39,27 @@ sphinx: extra_extensions: - 'sphinx.ext.viewcode' - 'sphinx.ext.inheritance_diagram' + - 'sphinx.ext.intersphinx' - 'matplotlib.sphinxext.roles' - 'autodoc2' local_extensions: abtem_version_footer: _ext abtem_autodoc2: _ext config: - autodoc2_render_plugin: "myst" + # Lets return/parameter types like np.ndarray in docstrings resolve to a + # link instead of staying plain text. + intersphinx_mapping: + python: ["https://docs.python.org/3", null] + numpy: ["https://numpy.org/doc/stable/", null] + # Collapses the full argument list in autodoc2 headings down to + # `name(...)` - see AbtemMystRenderer in _ext/abtem_autodoc2.py. The + # full, per-argument detail is already rendered in the Parameters list + # just below each heading. + autodoc2_render_plugin: "abtem_autodoc2.AbtemMystRenderer" + # Strips leading module qualifiers (abtem.waves.Waves -> Waves) from + # the remaining return-type arrow in headings, since the args list + # itself is already collapsed by AbtemMystRenderer above. + python_use_unqualified_type_names: true autodoc2_output_dir: "reference/api/apidocs" autodoc2_sort_names: true # Docstrings are NumPy-style, converted to RST field lists by the diff --git a/docs/_ext/abtem_autodoc2.py b/docs/_ext/abtem_autodoc2.py index 45614682..50ea2b5e 100644 --- a/docs/_ext/abtem_autodoc2.py +++ b/docs/_ext/abtem_autodoc2.py @@ -7,6 +7,29 @@ import os +from autodoc2.render.myst_ import MystRenderer + + +class AbtemMystRenderer(MystRenderer): + """MyST renderer that collapses argument lists in signature headings. + + autodoc2 puts the full, fully-qualified argument list (and defaults) in + the heading of every function/method/class - e.g. + ``transition_potential_multislice_and_detect(waves: abtem.waves.Waves, + potential: abtem.potentials.iam.BasePotential, ...)``. That's redundant + with the Parameters field list rendered from the docstring immediately + below, and makes headings unreadably long for abTEM's many-argument + functions. Collapsing the heading to ``name(...)`` (or ``name()`` for + no-arg callables) keeps the heading scannable; full argument details + still live in the Parameters section. + """ + + def format_args(self, args_info, include_annotations=True, ignore_self=None): + args = list(args_info) + if args and ignore_self is not None and args[0][1] == ignore_self: + args = args[1:] + return "..." if args else "" + def _set_autodoc2_packages(app, config): import abtem From af948d0fde677ecad2e6d749ecc789a9d738afb4 Mon Sep 17 00:00:00 2001 From: Paul Zeiger Date: Fri, 28 Aug 2026 19:06:46 +0000 Subject: [PATCH 10/11] Quote the benchmark's raw seconds so the speed-up is checkable Review (TomaSusi): the tip gave 51 minutes against 13 minutes and called it 3.86x, but 51/13 = 3.92. The minute figures were rounded from the raw measurement, and rounding both of them down inflated the implied ratio. The underlying numbers are 3088.8 s and 799.4 s, whose ratio is 3.8639 -- so the 3.86x was right and the minutes were the lossy part. Quoting the seconds makes all three figures consistent and lets a reader verify the ratio, which is the point of putting a benchmark in the documentation at all. Co-Authored-By: Claude Opus 5 --- docs/user_guide/walkthrough/parallelization.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user_guide/walkthrough/parallelization.ipynb b/docs/user_guide/walkthrough/parallelization.ipynb index af9e23e6..2912f8da 100644 --- a/docs/user_guide/walkthrough/parallelization.ipynb +++ b/docs/user_guide/walkthrough/parallelization.ipynb @@ -2326,7 +2326,7 @@ "As with any Dask parallelism, the speed-up requires the computation to split into **at least as many independent chunks as there are GPUs** — for example frozen-phonon configurations or batches of scan positions (see the *Chunks* section above). A single plane-wave multislice has no ensemble axis to distribute and will not benefit. Results are gathered back to the client GPU.\n", "\n", "```{tip}\n", - "More GPUs also mean more *aggregate* device memory, which can make multi-GPU worthwhile even when raw compute is not the bottleneck. In **memory-bound** simulations a single GPU is forced into small batches — or has to rebuild intermediate results such as potential slices — to stay within its memory, and in the worst case cannot run the calculation at all. Spreading the work over more devices lets each one use larger, more efficient batches. In a large STEM benchmark (101,871 scan positions on a $2623\\times2271$ grid in double precision) one 40 GB A100 completed the scan in 51 minutes and four completed it in 13 — a $3.86\\times$ speed-up on four devices. The exact speed-up depends on the parameters of your calculation and on your hardware.\n", + "More GPUs also mean more *aggregate* device memory, which can make multi-GPU worthwhile even when raw compute is not the bottleneck. In **memory-bound** simulations a single GPU is forced into small batches — or has to rebuild intermediate results such as potential slices — to stay within its memory, and in the worst case cannot run the calculation at all. Spreading the work over more devices lets each one use larger, more efficient batches. In a large STEM benchmark (101,871 scan positions on a $2623\\times2271$ grid in double precision) the scan took $3088.8\\ \\mathrm{s}$ on one 40 GB A100 and $799.4\\ \\mathrm{s}$ on four — a $3.86\\times$ speed-up on four devices. The exact speed-up depends on the parameters of your calculation and on your hardware.\n", "```" ] } From a790eeb64a77dd713d3d2e9563646ea69dad4824 Mon Sep 17 00:00:00 2001 From: TomaSusi Date: Sat, 29 Aug 2026 12:27:08 +0200 Subject: [PATCH 11/11] Fix JSON broken by merge conflict resolution in performance_tips.ipynb The merge of v1.1.0 (which added the sampling="auto" rounding cells) dropped the closing bracket/brace/cell-type lines between the GPU Bluestein-note cell and the next inserted cell, corrupting the notebook's JSON and failing the check-notebook-widgets CI check. Co-Authored-By: Claude Sonnet 5 --- docs/user_guide/appendix/performance_tips.ipynb | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/user_guide/appendix/performance_tips.ipynb b/docs/user_guide/appendix/performance_tips.ipynb index 695f9024..4a62dbeb 100644 --- a/docs/user_guide/appendix/performance_tips.ipynb +++ b/docs/user_guide/appendix/performance_tips.ipynb @@ -358,6 +358,10 @@ "next_fast_fft_size(2623) # 2625\n", "```\n", "````" + ] + }, + { + "cell_type": "markdown", "id": "b67eea2232fd", "metadata": {}, "source": [