25 · Domain Decomposition (DD): split one model across GPUs¶

ModelParallel runs one forward/backward over a model that is sliced into tiles — one tile per GPU — exchanging a thin halo every time step. Reach for it when a single model is too large for one GPU, or to speed up a single shot across cards.

DD is not shot-parallel. Notebook 12_multi_gpu is shot-parallel (data parallel): each GPU runs a different shot on a full copy of the model. DD is model parallel: all GPUs cooperate on one shot of one sliced model. They are different axes — you can even combine them.

In [2]:

import sweep, torch
from sweep.parallel import ModelParallel, MeshTopology   # DD entry points
print("torch      :", torch.__version__, "| GPU:", torch.cuda.get_device_name(0))

import sweep, torch from sweep.parallel import ModelParallel, MeshTopology # DD entry points print("torch :", torch.__version__, "| GPU:", torch.cuda.get_device_name(0))

torch      : 2.9.1+cu128 | GPU: Tesla V100-SXM2-32GB

How ModelParallel works¶

ModelParallel(prop, mesh) wraps a global single-domain PropTorch (the spec carrier — global shape, dh, dt, equation, source/receiver types, …) plus a mesh. You then hand it the global problem (global source/receiver coordinates, global model). Internally each rank:

1. slices its tile out of the global model and pads it (PML on the outer edges, only a stencil halo on interior cuts — the cut-aware compact pad);
2. remaps the source/receiver coordinates that fall inside its tile;
3. runs the normal stepped forward/backward, exchanging the M-cell halo with neighbour tiles every step (NCCL);
4. on a cut face the PML/interior predicates collapse so cut-adjacent cells use the same kernel expression as single-domain interior cells → the assembled result is bit-exact versus the single-domain run.

It is autograd-transparent: pass a model tensor with requires_grad=True and the returned record carries a grad_fn, so loss.backward() fills the model's .grad — exactly like single-domain PropTorch (no manual adjoint).

The mesh is MeshTopology(py, px, shot_groups, world_size, rank): px tiles along x, py along y (3-D); px * py * shot_groups == world_size (with the default shot_groups=1, px * py == world_size).

Single-GPU walk-through (world=1)¶

world=1 is the degenerate case — one tile = the whole model, no NCCL, single process — ideal for learning the API and checking correctness right in a notebook. The API is identical to the multi-GPU case; only MeshTopology and the launch differ. We run the same problem through ModelParallel and a plain PropTorch and confirm the forward record and gradient are bit-exact.

In [3]:

import numpy as np

dev = torch.device("cuda:0")
DT, nt, so, abcn = 0.0015, 60, 4, 10
nz, nx = 48, 56
shape  = (nz, nx)

g  = np.linspace(0, 1, nz * nx, dtype=np.float32).reshape(shape)
vp = (1800.0 + 600.0 * g).copy()                      # global model (physical grid)

def ricker(nt, dt, fm=10.0, delay=0.06):
    t = np.arange(nt, dtype=np.float32) * dt - delay; a = np.pi * fm * t
    return ((1 - 2 * a**2) * np.exp(-a**2)).astype(np.float32)

wav = ricker(nt, DT)
src = np.array([[[nx // 2, nz // 4]]], dtype=np.int32)               # GLOBAL coords
rec = np.array([[[ix, 2] for ix in range(2, nx - 2, 5)]], dtype=np.int32)

import numpy as np dev = torch.device("cuda:0") DT, nt, so, abcn = 0.0015, 60, 4, 10 nz, nx = 48, 56 shape = (nz, nx) g = np.linspace(0, 1, nz * nx, dtype=np.float32).reshape(shape) vp = (1800.0 + 600.0 * g).copy() # global model (physical grid) def ricker(nt, dt, fm=10.0, delay=0.06): t = np.arange(nt, dtype=np.float32) * dt - delay; a = np.pi * fm * t return ((1 - 2 * a**2) * np.exp(-a**2)).astype(np.float32) wav = ricker(nt, DT) src = np.array([[[nx // 2, nz // 4]]], dtype=np.int32) # GLOBAL coords rec = np.array([[[ix, 2] for ix in range(2, nx - 2, 5)]], dtype=np.int32)

In [4]:

from sweep.equations import Acoustic
from sweep.propagator.torch import PropTorch
from sweep.propagator.options import BoundarySaving

# Single-domain reference (the run DD must reproduce).
prop = PropTorch(Acoustic(spatial_order=so, device=dev, backend="torch"),
                 backend="torch", impl="c", shape=shape, dev=dev, dh=10.0, dt=DT,
                 source_type=["h1"], receiver_type=["h1"], abcn=abcn,
                 free_surface=False, nt=nt, B=1,
                 memory=BoundarySaving(storage="gpu"))
m    = [torch.tensor(vp, device=dev, requires_grad=True)]
recR = prop(wav, src, rec, models=m)
recR = recR[0] if isinstance(recR, (tuple, list)) else recR
(0.5 * recR.pow(2).sum()).backward()        # adjoint source = the record itself
gR   = m[0].grad.detach()
print("PropTorch  record", tuple(recR.shape), " grad", tuple(gR.shape))

from sweep.equations import Acoustic from sweep.propagator.torch import PropTorch from sweep.propagator.options import BoundarySaving # Single-domain reference (the run DD must reproduce). prop = PropTorch(Acoustic(spatial_order=so, device=dev, backend="torch"), backend="torch", impl="c", shape=shape, dev=dev, dh=10.0, dt=DT, source_type=["h1"], receiver_type=["h1"], abcn=abcn, free_surface=False, nt=nt, B=1, memory=BoundarySaving(storage="gpu")) m = [torch.tensor(vp, device=dev, requires_grad=True)] recR = prop(wav, src, rec, models=m) recR = recR[0] if isinstance(recR, (tuple, list)) else recR (0.5 * recR.pow(2).sum()).backward() # adjoint source = the record itself gR = m[0].grad.detach() print("PropTorch record", tuple(recR.shape), " grad", tuple(gR.shape))

PropTorch  record (1, 60, 11, 1)  grad (48, 56)

In [5]:

# ModelParallel wraps a GLOBAL single-domain PropTorch (the spec carrier) + a
# mesh; it builds the per-tile cut-aware solver internally. world=1 -> one tile.
# (Acoustic / PropTorch already imported in the reference cell above.)
topo   = MeshTopology(py=1, px=1, shot_groups=1, world_size=1, rank=0)
prop_g = PropTorch(Acoustic(spatial_order=so, device=dev, backend="torch"),
                   backend="torch", impl="c", shape=shape, dev=dev, dh=10.0, dt=DT,
                   source_type=["h1"], receiver_type=["h1"], abcn=abcn,
                   free_surface=False, nt=nt, B=1)
ddp = ModelParallel(prop_g, topo)

vp_leaf = torch.tensor(vp, device=dev, requires_grad=True)   # GLOBAL model leaf
recD = ddp.forward(wav, src, rec, models=[vp_leaf])          # GLOBAL inputs -> tiled
(0.5 * recD.pow(2).sum()).backward()        # autograd: same loss as the reference
gD   = vp_leaf.grad.detach()
print("ModelParallel record", tuple(recD.shape), " grad", tuple(gD.shape))

# ModelParallel wraps a GLOBAL single-domain PropTorch (the spec carrier) + a # mesh; it builds the per-tile cut-aware solver internally. world=1 -> one tile. # (Acoustic / PropTorch already imported in the reference cell above.) topo = MeshTopology(py=1, px=1, shot_groups=1, world_size=1, rank=0) prop_g = PropTorch(Acoustic(spatial_order=so, device=dev, backend="torch"), backend="torch", impl="c", shape=shape, dev=dev, dh=10.0, dt=DT, source_type=["h1"], receiver_type=["h1"], abcn=abcn, free_surface=False, nt=nt, B=1) ddp = ModelParallel(prop_g, topo) vp_leaf = torch.tensor(vp, device=dev, requires_grad=True) # GLOBAL model leaf recD = ddp.forward(wav, src, rec, models=[vp_leaf]) # GLOBAL inputs -> tiled (0.5 * recD.pow(2).sum()).backward() # autograd: same loss as the reference gD = vp_leaf.grad.detach() print("ModelParallel record", tuple(recD.shape), " grad", tuple(gD.shape))

ModelParallel record (1, 60, 11, 1)  grad (48, 56)

In [6]:

# ModelParallel hands back the SAME record layout as a single-card PropTorch,
# (B, nt, nrec, nfield), so the two are directly comparable -- no re-alignment.
rec_bit  = torch.equal(recD.detach().cpu(), recR.detach().cpu())
grad_bit = torch.equal(gD.cpu(), gR.cpu())
print("forward record bit-exact :", rec_bit)
print("vp gradient    bit-exact :", grad_bit)
assert rec_bit and grad_bit, "world=1 DD must match single-domain bit-for-bit"
print("OK - ModelParallel(world=1) == single-domain PropTorch, bit-for-bit")

# ModelParallel hands back the SAME record layout as a single-card PropTorch, # (B, nt, nrec, nfield), so the two are directly comparable -- no re-alignment. rec_bit = torch.equal(recD.detach().cpu(), recR.detach().cpu()) grad_bit = torch.equal(gD.cpu(), gR.cpu()) print("forward record bit-exact :", rec_bit) print("vp gradient bit-exact :", grad_bit) assert rec_bit and grad_bit, "world=1 DD must match single-domain bit-for-bit" print("OK - ModelParallel(world=1) == single-domain PropTorch, bit-for-bit")

forward record bit-exact : True
vp gradient    bit-exact : True
OK - ModelParallel(world=1) == single-domain PropTorch, bit-for-bit

Multi-GPU — the real use case¶

A notebook is a single process, so the actual multi-tile run goes through torchrun (one process per GPU). The only changes versus the cell above: set px (and/or py) to the GPU count, and launch with torchrun. The ModelParallel calls are byte-for-byte the same.

# 4 GPUs, x-split into 4 tiles — PASS means bit-exact vs single domain
torchrun --standalone --nproc-per-node=4 test/dd_api_check.py --family acoustic --ndim 2

# weak-scaling benchmark (per-card tile fixed at 4096^2)
torchrun --standalone --nproc-per-node=4 test/dd_nccl_bench.py --nz 4096 --nxp 4096 --nt 400

# 3-D 2x2 mesh (px=2, py=2) on 4 GPUs, elastic
torchrun --standalone --nproc-per-node=4 test/dd_api_check.py --family elastic --ndim 3 --px 2 --py 2

Inside those scripts the mesh is MeshTopology(py=py, px=px, shot_groups=1, world_size=world, rank=rank) with rank/world from torch.distributed. A custom driver skeleton (autograd path):

import torch.distributed as dist
from sweep.parallel import ModelParallel, MeshTopology
from sweep.propagator.torch import PropTorch
dist.init_process_group("nccl")
rank, world = dist.get_rank(), dist.get_world_size()
torch.cuda.set_device(rank % torch.cuda.device_count())
topo = MeshTopology(py=1, px=world, shot_groups=1, world_size=world, rank=rank)
prop = PropTorch(equation, shape=global_shape, dh=..., dt=..., nt=..., abcn=...,
                 source_type=..., receiver_type=..., dev=torch.device(f"cuda:{rank}"))
ddp  = ModelParallel(prop, topo)
vp   = torch.tensor(model_global, device=f"cuda:{rank}", requires_grad=True)
rec  = ddp.forward(wavelet, sources_global, receivers_global, models=[vp])
(0.5 * (rec - obs).pow(2).sum()).backward()   # autograd -> vp.grad (the DD adjoint)
full = ddp.gather_record(rec)                  # rank 0 assembles the global record

Run multi-tile DD straight from this notebook¶

NCCL needs one process per GPU, so the notebook's own (single) process can't hold multiple ranks. The trick: launch dd_api_check.py through a subprocess (torch.distributed.run, i.e. what torchrun runs). On a single-GPU box this runs world=1 (still a bit-exact PASS); on an N-GPU box set nproc-per-node=N for a real N-tile DD run — no code change.

Needs a git clone. test/dd_api_check.py lives in the repository, not in the pip wheel, so this cell only runs from a source checkout; on a pip install it prints a skip message instead.

In [7]:

import os, sys, subprocess
from pathlib import Path
import torch, sweep

REPO = Path(sweep.__file__).resolve().parents[2]          # repo root (robust to cwd)
SCRIPT = REPO / "test" / "dd_api_check.py"                # in a git clone only, not in the wheel
ngpu = min(torch.cuda.device_count(), 4)                  # up to 4 GPUs; 1 GPU -> world=1
if not SCRIPT.exists():
    print(f"skipped: {SCRIPT} not found -- this cell needs a git clone of sweep "
          "(test/ is not shipped in the pip wheel)")
else:
    cmd = [sys.executable, "-m", "torch.distributed.run", "--standalone",
           f"--nproc-per-node={ngpu}", "test/dd_api_check.py",
           "--family", "acoustic", "--ndim", "2"]
    print(f"visible GPUs: {torch.cuda.device_count()}  ->  {ngpu}-way DD")
    env = dict(os.environ, PYTHONPATH=str(REPO / "src"))
    res = subprocess.run(cmd, cwd=str(REPO), env=env, capture_output=True, text=True)
    for ln in (res.stdout + res.stderr).splitlines():
        if any(k in ln for k in ("DD_API_CHECK", "bit=", "PASS", "FAIL")):
            print(ln)
    print("exit code:", res.returncode)

import os, sys, subprocess from pathlib import Path import torch, sweep REPO = Path(sweep.__file__).resolve().parents[2] # repo root (robust to cwd) SCRIPT = REPO / "test" / "dd_api_check.py" # in a git clone only, not in the wheel ngpu = min(torch.cuda.device_count(), 4) # up to 4 GPUs; 1 GPU -> world=1 if not SCRIPT.exists(): print(f"skipped: {SCRIPT} not found -- this cell needs a git clone of sweep " "(test/ is not shipped in the pip wheel)") else: cmd = [sys.executable, "-m", "torch.distributed.run", "--standalone", f"--nproc-per-node={ngpu}", "test/dd_api_check.py", "--family", "acoustic", "--ndim", "2"] print(f"visible GPUs: {torch.cuda.device_count()} -> {ngpu}-way DD") env = dict(os.environ, PYTHONPATH=str(REPO / "src")) res = subprocess.run(cmd, cwd=str(REPO), env=env, capture_output=True, text=True) for ln in (res.stdout + res.stderr).splitlines(): if any(k in ln for k in ("DD_API_CHECK", "bit=", "PASS", "FAIL")): print(ln) print("exit code:", res.returncode)

visible GPUs: 4  ->  4-way DD

[rank0] record: bit=True max|d|=0.000e+00 rel=0.000e+00
[rank0] tile0 grad[0]: bit=True max|d|=0.000e+00 rel=0.000e+00
[rank0] tile1 grad[0]: bit=True max|d|=0.000e+00 rel=0.000e+00
[rank0] tile2 grad[0]: bit=True max|d|=0.000e+00 rel=0.000e+00
[rank0] tile3 grad[0]: bit=True max|d|=0.000e+00 rel=0.000e+00
DD_API_CHECK: PASS
exit code: 0

Measured scaling (ibex V100, padopt build)¶

Weak scaling — per-card load fixed, efficiency = T1 / TN (1→8 GPU):

equation	per-card tile	1→8 GPU efficiency
acoustic 2D	4096^2	84.6%
acoustic 3D	320^3	97.7%
elastic 2D	4096^2	96.1%
elastic 3D	256^3	~100% (comm-light)

Strong scaling — global fixed 4096x16384, to 8 GPU: acoustic 2D 5.1x; per-card peak memory drops ~1/N (3.30→0.41 GB). That memory split is the core DD payoff: a model too big for one card fits across many. (A balanced 2-D cut via sweep.parallel.balanced_grid lifts strong scaling further on cubic 3-D models.)

Elastic (compute-heavy, small communication fraction) scales best; acoustic (communication-bound) is lower but still ~85-90%. The DD path is bit-exact versus the single-domain run for both forward and gradient (acoustic + elastic, 2-D/3-D, including free surface and 2-D cuts), so correctness is independent of the GPU count.