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/_config.yml b/docs/_config.yml
index 3cb0fc0f..c5b52435 100644
--- a/docs/_config.yml
+++ b/docs/_config.yml
@@ -39,15 +39,34 @@ 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
+ # 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..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
@@ -19,6 +42,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}
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 += "
"
diff --git a/docs/abtem/changelog.md b/docs/abtem/changelog.md
index a21eb125..a941e988 100644
--- a/docs/abtem/changelog.md
+++ b/docs/abtem/changelog.md
@@ -111,6 +111,27 @@ Planned for this release (not yet merged):
- 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))
+- Multi-GPU hardening, and the configuration fix found while chasing it
+ ([PR #346](https://github.com/abTEM/abTEM/pull/346))
+ - **The client's configuration now reaches `distributed` workers.** *ab*TEM resolves configuration
+ inside each task, and worker processes start fresh and previously saw only the YAML defaults, so
+ any distributed computation with a non-default configuration silently used the defaults instead —
+ most consequentially `precision`, which meant `float64` runs were computed in `float32`.
+ Distributed results obtained with a non-default configuration are worth repeating. Applies to any
+ distributed client, CPU clusters included
+ - `to_zarr()` on a lazy result honours `dask.multi-gpu`; it previously ignored the flag and ran the
+ whole computation on a single device
+ - `cupy.fft-cache-size` defaults to `auto` — 25 % of each device's memory, resolved per device —
+ rather than unlimited. `-1` restores unlimited, `0 MB` disables the cache, and a size such as
+ `512 MB` sets a fixed bound. A single plan larger than the bound runs uncached with a warning
+ instead of raising
+ - New config 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
+ - Automatically sized scan batches are halved on grid sizes that force cuFFT's Bluestein fallback,
+ which needs a much larger FFT workspace
+ - Warnings replace silent fallbacks: multi-GPU requested but declined (with the reason), a missing
+ `if __name__ == "__main__"` guard, and a grid size that forces the Bluestein fallback (naming the
+ next fast size)
## 1.0.10
diff --git a/docs/getting_started/install.md b/docs/getting_started/install.md
index 0c06f5d1..21b5c960 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:
@@ -100,6 +101,24 @@ 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.
+
+```{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/).
diff --git a/docs/reference/default_config.yaml b/docs/reference/default_config.yaml
index 9e515364..efaecc10 100644
--- a/docs/reference/default_config.yaml
+++ b/docs/reference/default_config.yaml
@@ -21,14 +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:
# Maximum GPU memory (bytes) used by the cuFFT plan cache.
# https://docs.cupy.dev/en/stable/user_guide/fft.html#fft-plan-cache
+ # 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).
- fft-cache-size: -1
+ # -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
diff --git a/docs/user_guide/appendix/performance_tips.ipynb b/docs/user_guide/appendix/performance_tips.ipynb
index 5a6a2724..4a62dbeb 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": "b67eea2232fd",
@@ -369,6 +418,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",
"```"
]
},
diff --git a/docs/user_guide/walkthrough/parallelization.ipynb b/docs/user_guide/walkthrough/parallelization.ipynb
index 6088380c..e53bcbff 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`."
]
},
{
@@ -2259,15 +2264,69 @@
"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",
+ "(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",
- "cluster = LocalCUDACluster()\n",
- "client = Client(cluster)\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) 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",
"```"
]
}
@@ -3899,7 +3958,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 +3980,7 @@
"user_panning_enabled": true,
"user_zooming_enabled": true,
"wheel_sensitivity": 0.1,
- "zoom": 2.0,
+ "zoom": 2,
"zooming_enabled": true
}
},
@@ -5004,7 +5063,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 +5085,7 @@
"user_panning_enabled": true,
"user_zooming_enabled": true,
"wheel_sensitivity": 0.1,
- "zoom": 2.0,
+ "zoom": 2,
"zooming_enabled": true
}
},