Document multi-GPU support

.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

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) -

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}

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 = (
-        "<p>Tested against "
-        '<a href="https://github.com/abTEM/abTEM">abTEM</a>'
-        f" v{version}."
-    )
     if _is_dev():
         theme_options["announcement"] = (
-            "This is the <strong>development</strong> documentation, built from "
-            f'the latest main branch. <a href="{STABLE_URL}">Switch to the '
-            "stable version</a>."
+            "This is the <strong>development</strong> documentation, "
+            f'including unreleased work in <a href="{ABTEM_DEV_BRANCH_URL}">'
+            f'abTEM/dev</a>. <a href="{STABLE_URL}">Switch to the stable '
+            "version</a>."
+        )
+        footer = (
+            "<p>Tested against pre-release "
+            f'<a href="{ABTEM_DEV_BRANCH_URL}">abTEM v{version}</a>.'
         )
-        footer += " Development build."
     else:
-        footer += (
+        footer = (
+            "<p>Tested against "
+            '<a href="https://github.com/abTEM/abTEM">abTEM</a>'
+            f" v{version}."
             f' Also available: <a href="{DEV_URL}">development version</a>.'
         )
     footer += "</p>"

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
 

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 <walkthrough:parallelization:multigpu>` 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/). 

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

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 <walkthrough:parallelization:multigpu>`."
    ]
   },
   {
@@ -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",
     "```"
    ]
   },