Multi-GPU and MPI

GeoSWE scales across multiple GPUs with mpi4py. The domain is split into strips, each rank owns one subgrid, and ranks exchange ghost (“halo”) rows each step. One MPI rank drives one GPU.

Running

pip install "geoswe[gpu,mpi]"              # "geoswe[gpu-rocm,mpi]" on AMD GPUs
mpirun -n 4 python -m mpi4py my_run.py     # 4 ranks -> 4 GPUs

python -m mpi4py installs mpi4py’s abort hook, so a rank that raises takes the job down instead of leaving the others waiting in the next collective until the scheduler’s wall clock kills them. geoswe.runlib.driver.main and geoswe.runlib.replay.main do that for failures inside themselves; the hook covers the rest of your script.

Under Slurm the scheduler can hand each rank its own device, on either vendor: srun -n 4 --gpus-per-task=1 --gpu-bind=closest python -m mpi4py my_run.py. The pinning line in the script below then sees one device per rank and is a no-op.

In the script, pass the communicator and a process grid to the solver:

from mpi4py import MPI
import numpy as np
import cupy as cp
from geoswe import Mesh2D, Config, Solver2D

comm = MPI.COMM_WORLD
cp.cuda.Device(comm.rank % cp.cuda.runtime.getDeviceCount()).use()   # pin a GPU

NX, NY_GLOBAL, dx = 2000, 2000, 10.0        # NY_GLOBAL must divide by the rank count
NY_LOCAL = NY_GLOBAL // comm.size
j0 = comm.rank * NY_LOCAL                   # this rank's first global row

# each rank builds ONLY its own subgrid (no global array)
mesh = Mesh2D(nx=NX, ny=NY_LOCAL, dx=dx, dy=dx, ngh=4)
cfg = Config(dtype="float32")               # ... and the rest of your configuration
jj = np.arange(j0, j0 + NY_LOCAL)           # global y index of each local row
bed_local = np.broadcast_to(0.002 * dx * jj, (NX, NY_LOCAL)).copy()
q0_local = np.zeros((3, NX, NY_LOCAL)); q0_local[0] = 1.0
s = Solver2D(mesh, cfg, q0_local, bed_local, comm=comm, dims=[1, comm.size])
s.step(dt=float(s.cfl_dt()))         # cfl_dt() does the global all-reduce

  • dims=[Px, Py] sets the process grid; [1, N] is a 1-D split in \(y\). Pass None to let MPI choose.

  • cfl_dt() performs a global all-reduce so every rank advances with the same stable dt (lockstep).

  • Halo exchange of the conserved state happens inside step(); you do not call it explicitly.

A complete, runnable weak/strong scaling benchmark is examples/ex05_scaling_bench.py. It also runs on one GPU (mpirun -n 1, or plain python). It pins rank r to GPU r modulo the number of visible GPUs, so with more ranks than GPUs the ranks share devices.

Scaling

GeoSWE’s per-step cost is dominated by the fused right-hand-side kernel; the halo exchange is a small, fixed overhead that does not grow with the per-rank domain. The result is near-ideal scaling: in the synthetic benchmark, weak scaling (fixed work per GPU) holds 99.5 % efficiency at 16 H100 GPUs and 10.24 billion cells (98.8 % at 32 Blackwell MIG slices and 20.48 billion cells), and strong scaling (fixed total problem) reaches 15.5x on 16 H100 GPUs for the compressed path. The only cost that does not shrink with the per-rank domain is the inter-rank communication (the scalar dt all-reduce plus the halo), so strong-scaling efficiency tapers once each rank holds too few cells, the classic surface-to-volume trade-off. The application runs carry 0.4 to 1.1 billion active cells per rank and sit deep in the weak-scaling regime. See benchmarks.

Tip

On hardware without GPU peer access (e.g. MIG slices), force host-staged halo exchange with SWE_HALO_CUDA_AWARE=0. With =1 the halo buffers go straight to MPI as device pointers, with no host round-trip, which is what the published multi-GPU benchmark runs use. See the configuration reference.

That variable, and every other GEOSWE_*/SWE_* setting, is read from each rank’s own environment, so a launch that spans nodes has to forward it (-x per variable with Open MPI). A rank that misses one either fails the run or quietly computes something else, depending on the variable; see performance, which also covers measuring where a run’s per-step time goes before changing anything.

On AMD GPUs the same variable selects GPU-aware MPI. With HPE Cray MPICH that also needs MPICH_GPU_SUPPORT_ENABLED=1 and an mpi4py linked against the GPU transport library; GeoSWE falls back to host staging, with a warning, when Cray MPICH is running without its GPU support. See AMD GPUs.

Large domains

For continental-scale problems that do not fit a dense grid, combine MPI with the compressed active-cell mesh, which partitions the active cells evenly across ranks and supports checkpoint/resume across job-time limits.