Parallel performance tips

Parallel performance tips#

abTEM aims for good performance out of the box, but the defaults cannot be optimal for every simulation and every machine. If your performance is unexpectedly low, or you are running out of memory, you can try optimizing using the tips below. Many of them are controlled through the configuration system; see the parallelization walkthrough for the basics of how abTEM distributes work using Dask.

The timings shown on this page were obtained on an ordinary laptop; the exact numbers will differ on your machine, but the trends should hold.

Optimize your simulation parameters#

The most effective way of speeding up a simulation is to not compute more than you need: the cost of the multislice algorithm is dominated by a pair of FFTs for every slice of the potential, for every wave function in the ensemble.

Number of `gpts`#

The number of grid points sets the size of those FFTs and, together with the antialiasing cutoff, the maximum scattering angle included in the simulation. Use the coarsest real-space sampling that still contains the largest scattering angles you detect (see the appendix on antialiasing), and only refine the sampling at the end to check convergence (see the appendix on convergence).

Number of slices#

Each slice of the potential costs one propagation step. The slice thickness is a convergence parameter: thinner slices are more accurate but proportionally slower. Typical values lie between 0.5 and 2 Å; converge it for your system rather than defaulting to the thinnest value.

Cell size#

The computational cost grows with the area of the simulation cell. Use the smallest cell consistent with your structure’s periodicity, the required scan area, and tolerable wrap-around (periodic-boundary) artifacts.

Number of scan samples#

In STEM simulations, the total cost is multiplied by the number of probe positions. If you do not specify the scan sampling, abTEM defaults to the Nyquist rate of the probe-forming aperture, which is sufficient for interpolating smooth images afterwards (using measurement.interpolate); scanning much more finely than Nyquist wastes time.

Number of ensemble samples (frozen phonons, partial coherence etc.)#

Ensemble averages multiply the cost linearly by the number of configurations. Converge the number of frozen-phonon configurations or partial-coherence samples against the quantity you actually measure — high-angle signals typically require more phonon configurations than low-angle ones.

The demonstration below shows how strongly the real-space sampling (equivalently, the number of gpts) affects the run time of a plane-wave multislice simulation:

atoms = bulk("Si", "diamond", a=5.43, cubic=True) * (4, 4, 20)

wave = abtem.PlaneWave(energy=200e3)

# warm up the FFT planner and the Dask scheduler before timing
wave.multislice(abtem.Potential(atoms, sampling=0.1, slice_thickness=2)).compute()

for sampling in (0.1, 0.05):
    potential = abtem.Potential(atoms, sampling=sampling, slice_thickness=2)

    start = time.perf_counter()
    wave.multislice(potential).compute()
    elapsed = time.perf_counter() - start

    print(f"sampling: {sampling:.2f} Å, gpts: {potential.gpts}, time: {elapsed:.2f} s")

sampling: 0.10 Å, gpts: (218, 218), time: 0.88 s

sampling: 0.05 Å, gpts: (435, 435), time: 2.60 s

Use PRISM#

In STEM, the multislice algorithm is repeated for every probe position, even though nearby probes propagate through almost the same potential. The PRISM algorithm factors this redundancy out into a scattering matrix that is computed once and reused for every probe position, at the cost of additional memory and — when using interpolation — a controlled approximation. For scans with many probe positions, PRISM can be orders of magnitude faster. See the PRISM tutorial to get started.

Running out of memory?#

With lazy evaluation (the default), peak memory use is not set by the size of the full simulation, but by the chunks in flight at any moment: roughly the chunk size times the number of Dask workers, plus intermediate results.

Reduce early#

Prefer detectors that reduce the data over storing everything: an AnnularDetector or FlexibleAnnularDetector reduces each diffraction pattern to a few numbers, while a PixelatedDetector keeps the full 4D dataset. If you only need a virtual dark-field image, do not store the 4D-STEM datacube. Any reduction (such as integrate_radial) applied before calling .compute() becomes part of the task graph, so the full dataset is never held in memory at once.

Lower the batch size#

Smaller chunks mean less memory per task; see changing the batch size below.

Lower the number of workers#

Each Dask worker holds its own chunks in memory, so peak memory scales with the number of workers:

dask.config.set(num_workers=2)

On the GPU#

Device memory is usually the tighter constraint, and a few options apply only there.

abTEM builds the potential in slice groups sized against the free VRAM it measures at computation time, rather than materializing every slice at once. If the automatic estimate is still too generous for your calculation, fix the group size yourself:

abtem.config.set({"potential.slice-chunk-size": 8})  # "auto" (the default) sizes it against free VRAM

CuPy caches the workspace of every FFT shape it has seen. The cache is bounded to 25 % of device memory by default; lower the bound, or disable caching entirely, when VRAM is tight:

abtem.config.set({"cupy.fft-cache-size": "512 MB"})  # or "0 MB" to disable

Finally, more GPUs mean more aggregate device memory. A calculation that does not fit on one device may fit comfortably across several — see Multiple GPUs.

Change the batch size#

abTEM groups the wave functions of an ensemble — for example, the probe positions of a scan — into batches, which map onto the chunks of the underlying Dask arrays. The target chunk size is set in the configuration; lowering it reduces the memory footprint of each task at the price of some scheduling overhead, and raising it does the opposite:

abtem.config.set({"dask.chunk-size": "16 MB"})

On the GPU, the corresponding option is dask.chunk-size-gpu. The effect on how a scan is partitioned:

probe = abtem.Probe(energy=100e3, semiangle_cutoff=25)
potential = abtem.Potential(atoms, sampling=0.1, slice_thickness=2)
detector = abtem.FlexibleAnnularDetector()

measurement = probe.scan(potential, detectors=detector)
print(f"chunk size: {abtem.config.get('dask.chunk-size'):>7}, scan chunks: {measurement.array.chunks[:2]}")

with abtem.config.set({"dask.chunk-size": "16 MB"}):
    measurement = probe.scan(potential, detectors=detector)
    print(f"chunk size: {'16 MB':>7}, scan chunks: {measurement.array.chunks[:2]}")

chunk size:  128 MB, scan chunks: ((18, 18, 18, 6), (18, 18, 18, 6))
chunk size:   16 MB, scan chunks: ((7, 7, 7, 7, 7, 7, 7, 7, 4), (6, 6, 6, 6, 6, 6, 6, 6, 6, 6))

Optimize your usage of FFTs#

The fast Fourier transform is the workhorse of the multislice algorithm, so the speed of the FFT library largely determines the speed of the whole simulation.

FFT library#

abTEM supports three FFT implementations: FFTW (the default), Intel’s MKL FFT, and NumPy’s (as a fallback). Their relative performance depends on your CPU and array sizes — on Intel processors, MKL is often faster — so it can be worth benchmarking your own simulation:

abtem.config.set({"fft": "mkl"})

Set internal thread parallelization#

Both FFTW and MKL can additionally thread each individual transform, via the fftw.threads and mkl.threads options. Because Dask already parallelizes across chunks, oversubscription can hurt: as a rule of thumb, keep the number of workers times the FFT threads at or below your physical core count. FFTW also exposes its planner through fftw.planning_effort; more patient planning can pay off for long-running simulations that reuse one grid size.

Use “good” numbers of `gpts`#

FFT implementations are most efficient for array sizes that factor into small primes (2, 3, 5 and 7). Since the exact number of grid points is rarely physically meaningful, round your gpts to a nearby product of small primes — powers of two are the safest choice, while a large prime number is the worst case:

factorizations = {486: "2 · 3⁵", 512: "2⁹", 521: "prime"}

for gpts, factors in factorizations.items():
    potential = abtem.Potential(atoms, gpts=gpts, slice_thickness=2)

    start = time.perf_counter()
    wave.multislice(potential).compute()
    elapsed = time.perf_counter() - start

    print(f"gpts: {gpts} ({factors}), time: {elapsed:.2f} s")

gpts: 486 (2 · 3⁵), time: 1.65 s

gpts: 512 (2⁹), time: 2.02 s

gpts: 521 (prime), time: 5.34 s

Note

This matters more on the GPU than on the CPU. cuFFT has no optimized kernel for a length with a prime factor above 7 and falls back to the Bluestein algorithm, which pads internally and needs a workspace several times the size of the transform — so an unlucky grid costs memory as well as time, and can turn a calculation that would fit on your card into one that does not. abTEM warns you when it meets such a grid, naming the offending size and the next good one:

UserWarning: FFT size 2623 x 2271 contains prime factors larger than 7; cuFFT falls back to the
Bluestein algorithm, which is several times slower and allocates a workspace of several times the
array size. Consider adjusting the transformed grid to 2625 x 2304, e.g. by setting gpts explicitly,
or -- if this grid was derived from a numeric sampling -- with
abtem.config.set({'grid.round-to-fast-fft': True}).

The same test is available directly, which is useful when choosing gpts programmatically:

from abtem.core.fft import is_fast_fft_size, next_fast_fft_size

is_fast_fft_size(2623)  # False
next_fast_fft_size(2623)  # 2625

Since version 1.1.0, you rarely need to do this by hand: Potential(..., sampling="auto") picks a fast size automatically wherever that is compatible with matching the crystal lattice (see the walkthrough on potentials for what “compatible” means and why it matters beyond speed). For a numeric gpts/sampling you choose yourself, Grid.round_to_fast_fft() applies next_fast_fft_size (shown above) to an existing grid in place, e.g. potential.grid.round_to_fast_fft().

This is also what the grid.round-to-fast-fft configuration option automates: 'auto' (the default) rounds only grids abTEM derives itself (sampling="auto"); True additionally rounds gpts derived from a numeric sampling, so a numeric gpts/sampling is left exactly as given unless you opt in — enabling this option never silently changes the result of an existing script. Rounding is always upward, so the realized sampling is never coarser than what you asked for.

Float precision#

By default, abTEM computes in single precision (float32/complex64), which is faster and uses half the memory of double precision. If you suspect that numerical precision matters for your result — or want to verify that it does not — you can switch:

abtem.config.set({"precision": "float64"})