Skip to content

v0.2.0 — Domain decomposition

Choose a tag to compare

@GeophyAI GeophyAI released this 24 Aug 14:21
· 916 commits to dev since this release
34b6e0e

Domain decomposition, and the boundary-memory work that came out of running it
at scale.

pip install --upgrade sweepx     # then: import sweep

Domain 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 autograd

Forward 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/zeta and 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 legacy use_ckpt /
    boundary_saving_config knobs 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_surface used 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.
  • AcousticVRZ3D low-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.