Repository navigation
v0.2.0 — Domain decomposition
Domain decomposition, and the boundary-memory work that came out of running it
at scale.
pip install --upgrade sweepx # then: import sweepDomain decomposition — sweep.parallel
ModelParallel splits one model into tiles, one GPU per tile, and exchanges
a halo every time step, so a single shot is solved cooperatively instead of
replicated. It is the model-parallel counterpart to shot-parallel DDP, and the
two compose: MeshTopology(py, px, shot_groups=...) describes a
shot_groups × (py × px) rank grid.
from sweep.parallel import MeshTopology, ModelParallel, pad_to_mesh
topo = MeshTopology(py=2, px=2, shot_groups=1, world_size=world, rank=rank)
ddp = ModelParallel(prop, topo) # prop is the GLOBAL problem spec
rec = ddp(wavelet, src, rec_global, models=[vp_tile])
rec.pow(2).sum().backward() # plain autogradForward and backward are ordinary autograd, and on fp32 gpu-direct boundaries
the DD gradient is bit-identical to the single-domain gradient. Supported:
Acoustic (2-D), Acoustic3D, AcousticVRZ3D, Elastic (2-D), Elastic3D;
anything without stepped kernels is refused at construction.
Memory
- CPML aux strip allocation.
psi/zetaand the elastic memory variables
live in per-axis slabs — the absorbing band plus stencil reach — instead of
full grids. Gradients are bit-for-bit unchanged; only the allocation shrinks
(3.46 → 2.50 GiB forward at 176³/abcn=30). BoundaryOptions.tail_steps. Keep only the last K steps of the boundary
ring and stop the reverse loop there. For steady-state objectives whose
adjoint source is zero outside a probe window, the truncated gradient is the
same gradient. Composes with DD.- CPU-staged boundary storage under DD.
storage="cpu"for the Acoustic
2-D/3-D DD backward, so a tile whose ring does not fit has a fallback.
API
- One three-way gradient-memory mode —
'full'/'boundary'/'ckpt'—
resolved identically for the eager and CUDA backends via
memory=MemoryOptions(strategy=...). The legacyuse_ckpt/
boundary_saving_configknobs still work and resolve into the same choice;
contradictory combinations now raise instead of one path silently winning. sweep.__version__finally reports something. It reads the installed
distribution's metadata, so it tells you the version you actually have.- The boundary spec is immutable after construction. Assigning
prop.free_surfaceused to land on the wrapper and shadow the backend's
value — reads showed the new setting while every kernel kept the old one. It
now raises and points at the constructor.
Fixes
- DAS Mu 2-D/3-D were non-deterministic on
impl='c'— same input,
different answer each run, on a plain single-GPU forward. They borrow the
elastic kernels, which address the CPML memory variables through the solver's
aux slabs, but their drivers never installed them, so the row stride
collapsed and every row aliased the first. - CPML aux writes raced on a DD cut tile, where the per-axis compute band
reaches columns that carry no slab storage. AcousticVRZ3Dlow-precision boundary staging overran its staging copy.
Scope
DD cpu staging is acoustic only and storage="disk" under DD is still refused.
Per-edge free surface on impl='c' remains Acoustic/Elastic 2-D. Full notes in
CHANGELOG.md;
the DD guide is at
docs/user-guide/parallel.md.