From f1cfc9ca40fa41ddaecf38bd62f4f66dd5335a2b Mon Sep 17 00:00:00 2001 From: Shaowen Wang Date: Mon, 26 Oct 2026 11:17:39 +0300 Subject: [PATCH] test: faster CPU suite, FD checks that check Grids a test steps once run uncompiled (a cold torch.compile cost more than the run), runs and references several tests share are computed once, and the halo checks of one world size share a spawn: test/*.py 208 -> 117 s on the CI's CPU, the same pass/skip counts. test_biot's FD misfit weighted p 1e13 over vx/vz, which left the solid-velocity adjoint unchecked: each receiver type is now in its own unit. The curvilinear and visco FD steps are cut to where truncation is well under the tolerance. --- test/test_acoustic_tti.py | 6 +- test/test_biot.py | 45 ++- test/test_compiled_stencil_slices.py | 9 +- test/test_eager_compile_fallback.py | 15 +- test/test_eager_snapshot_buffer.py | 6 +- test/test_elastic_curvilinear.py | 39 ++- test/test_kwarg_validation.py | 5 +- test/test_model_parallel_halo.py | 420 +++++++++++++++------------ test/test_visco_elastic.py | 13 +- 9 files changed, 329 insertions(+), 229 deletions(-) diff --git a/test/test_acoustic_tti.py b/test/test_acoustic_tti.py index fa60634f..779ec4d5 100644 --- a/test/test_acoustic_tti.py +++ b/test/test_acoustic_tti.py @@ -10,6 +10,8 @@ radian angle, not the ~1° the bug produced. """ +from functools import lru_cache + import numpy as np import pytest import torch @@ -43,9 +45,11 @@ def _src(nz, nx): return torch.tensor(np.array([[nx // 2, nz // 2]], dtype=np.int64)[None]) +@lru_cache(maxsize=None) def _wavefront_tilt_deg(theta_rad, nz=NZ, nx=NX): """Propagate AcousticTTI with a constant ``theta`` (radians) and return the - principal-axis tilt of the final |h1| snapshot in degrees.""" + principal-axis tilt of the final |h1| snapshot in degrees. Cached: both + tests start from theta = 0, and the run is deterministic.""" shape = (nz, nx) models = [ torch.full(shape, VP_BG, dtype=torch.float32), diff --git a/test/test_biot.py b/test/test_biot.py index 74464dfb..8345f26d 100644 --- a/test/test_biot.py +++ b/test/test_biot.py @@ -9,6 +9,7 @@ from sweep.equations import Biot, Elastic from sweep.equations.biot import biot_from_biot_constants, biot_from_rock, biot_wave_velocities +from sweep.propagator.options import EagerOptions from sweep.propagator.torch import PropTorch from sweep.signal import ricker @@ -23,9 +24,13 @@ def _wavelet(nt, dt, f0): return torch.from_numpy(ricker(t, f=f0)) -def _prop(eq, shape, dh, dt, src, rec, abcn=20, free_surface=False): +def _prop(eq, shape, dh, dt, src, rec, abcn=20, free_surface=False, compiled=True): + # A cold torch.compile of the step costs seconds per grid, source field and gradient mode, + # as much as a few thousand uncompiled steps of these grids: only the runs with many steps + # per compiled graph (wave speeds, wavenumber integration) compile. return PropTorch(eq, shape=shape, dev=DEV, dh=dh, dt=dt, source_type=src, receiver_type=rec, - abcn=abcn, free_surface=free_surface, backend="torch", impl="eager") + abcn=abcn, free_surface=free_surface, backend="torch", impl="eager", + eager_options=EagerOptions(use_compile=compiled)) def _fields(coef, shape): @@ -63,8 +68,10 @@ def test_no_relative_flow_limit_is_undrained_elastic(): src = np.array([[35, 20]]) rec = np.stack([np.arange(5, 65, 3), np.full(20, 40)], 1)[None] w = _wavelet(nt, dt, 12.0) - db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"])(w, src, rec, models=coef) - de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"])(w, src, rec, models=[vp, vs, rho]) + db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], compiled=False)( + w, src, rec, models=coef) + de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], compiled=False)( + w, src, rec, models=[vp, vs, rho]) err = float((db - de).abs().max() / de.abs().max()) assert err < 1e-4, err @@ -103,11 +110,16 @@ def test_gradient_matches_finite_difference_inside(free_surface): src = np.array([[24, 10]]) rec = np.stack([np.arange(8, 40, 4), np.full(8, 30)], 1)[None] w = _wavelet(nt, dt, 10.0) - prop = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz", "p"], free_surface=free_surface) + prop = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz", "p"], free_surface=free_surface, + compiled=False) obs = prop(w, src, rec, models=coef).detach() + # each receiver type in its own unit: as recorded, p (~1e-2 Pa) outweighs vx and vz + # (~1e-8 m/s) by 1e13 in the misfit, which left the solid-velocity adjoint unchecked + weight = 1.0 / obs.abs().amax(dim=(0, 1, 2)) def J(mu): - return 0.5 * (prop(w, src, rec, models=[coef[0], mu, *coef[2:]]) - 1.05 * obs).double().square().sum() + r = (prop(w, src, rec, models=[coef[0], mu, *coef[2:]]) - 1.05 * obs) * weight + return 0.5 * r.double().square().sum() mu0 = coef[1].clone().requires_grad_() (g,) = torch.autograd.grad(J(mu0), mu0) @@ -135,10 +147,10 @@ def test_free_surface_no_relative_flow_limit_is_elastic(): rec = np.concatenate([np.stack([np.arange(5, 65, 3), np.zeros(20, int)], 1), np.stack([np.arange(5, 65, 3), np.full(20, 12)], 1)])[None] w = _wavelet(nt, dt, 12.0) - db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True)( + db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True, compiled=False)( w, src, rec, models=coef) - de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True)( - w, src, rec, models=[vp, vs, rho]) + de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True, + compiled=False)(w, src, rec, models=[vp, vs, rho]) err = float((db - de).abs().max() / de.abs().max()) assert err < 1e-4, err @@ -160,7 +172,7 @@ def s_time(tt): return np.exp(-(math.pi * f0 * (tt - t0)) ** 2) w = torch.from_numpy(dt * np.gradient(s_time(t), t) / dx**2).float() - prop = _prop(Biot(device=DEV), shape, dx, dt, ["p"], ["p"], abcn=30, free_surface=True) + prop = _prop(Biot(device=DEV), shape, dx, dt, ["p"], ["p"], abcn=30, free_surface=True, compiled=False) with torch.no_grad(): num = prop(-w, np.array([[xs, zs]]), rec_xz[None], models=coef)[0, :, :, 0].double().numpy() @@ -245,7 +257,7 @@ def s_time(tt): src, rec = np.array([[pad, pad]]), np.array([[[pad + off, pad + off]]]) num = 0.0 for fields, scale in ((["sxx", "szz"], 1.0), (["p"], -qpart / phi)): - prop = _prop(Biot(device=DEV), shape, dx, dt, fields, ["vz"], abcn=30) + prop = _prop(Biot(device=DEV), shape, dx, dt, fields, ["vz"], abcn=30, compiled=False) with torch.no_grad(): num = num + prop(w * scale, src, rec, models=coef)[0, :, 0, 0].double().numpy() ana = _analytic_line_source_vz(models, t, off * dx, off * dx + 0.5 * dx, s_time, qpart, phi) @@ -331,10 +343,13 @@ def s_time(tt): prop = _prop(Biot(device=DEV), shape, dx, dt, fields, ["vx", "vz"], abcn=30, free_surface=True) with torch.no_grad(): num = num + prop(w * scale, np.array([[xs, zs]]), np.array([[[xr, 0]]]), models=coef)[0, :, 0].double().numpy() - # vx sits at (x + dx/2, 0), vz at (x, dx/2) on the staggered grid + # vx sits at (x + dx/2, 0), vz at (x, dx/2) on the staggered grid. One integration for both + # receivers, which share its per-frequency solves (most of its time), and only up to + # kx = 1 rad/m: past it the band's propagating modes are evanescent over the 20 m source + # depth (exp(-17) at 100 Hz), and the reference moves by < 1e-5 relative from kx_max = 2 off = (xr - xs) * dx - ref_vx, _ = _dwn_halfspace(models, (1.0, qpart / phi), zs * dx, [(off + 0.5 * dx, 0.0)], t, s_time) - _, ref_vz = _dwn_halfspace(models, (1.0, qpart / phi), zs * dx, [(off, 0.5 * dx)], t, s_time) - for k, ref in enumerate((ref_vx[:, 0], ref_vz[:, 0])): + ref_vx, ref_vz = _dwn_halfspace(models, (1.0, qpart / phi), zs * dx, [(off + 0.5 * dx, 0.0), (off, 0.5 * dx)], + t, s_time, kx_max=1.0) + for k, ref in enumerate((ref_vx[:, 0], ref_vz[:, 1])): err = np.linalg.norm(num[:, k] - ref) / np.linalg.norm(ref) assert err < tol, (("vx", "vz")[k], err) diff --git a/test/test_compiled_stencil_slices.py b/test/test_compiled_stencil_slices.py index da73b2b4..6619dd98 100644 --- a/test/test_compiled_stencil_slices.py +++ b/test/test_compiled_stencil_slices.py @@ -30,11 +30,13 @@ DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu") -@pytest.fixture(autouse=True) +@pytest.fixture def _release_compiled_state(): """Dynamo's cache keeps the compiled steps' propagators -- and their GPU buffers -- alive; release them so a later test that reads the absolute - peak memory (test_eager_boundary_saving) does not count them.""" + peak memory (test_eager_boundary_saving) does not count them. Taken by + the tests that compile only: the reset and a full gc.collect() cost + 0.1-0.25 s a test, and the kernel-level tests build no compiled state.""" yield if hasattr(torch, "_dynamo"): torch._dynamo.reset() @@ -177,6 +179,7 @@ def _assert_matches(run, ref): @requires_compile +@pytest.mark.usefixtures("_release_compiled_state") def test_the_conv_counter_sees_a_conv(monkeypatch): """With the slice form off the counter must find the convs, or the zero the tests below assert proves nothing.""" @@ -188,6 +191,7 @@ def test_the_conv_counter_sees_a_conv(monkeypatch): @requires_compile +@pytest.mark.usefixtures("_release_compiled_state") @pytest.mark.parametrize("cls,shape,surface", EQUATIONS) def test_compiled_rollout_is_one_conv_free_graph_matching_uncompiled(cls, shape, surface): """A rollout compiles one graph (fullgraph=True raises on any break) with @@ -206,6 +210,7 @@ def test_compiled_rollout_is_one_conv_free_graph_matching_uncompiled(cls, shape, @requires_compile +@pytest.mark.usefixtures("_release_compiled_state") def test_plain_warmup_prefix_changes_nothing(monkeypatch): """Before torch 2.5 a checkpointed rollout runs its first, compiling steps on the plain tape, because compiling inside the checkpoint breaks its diff --git a/test/test_eager_compile_fallback.py b/test/test_eager_compile_fallback.py index 858bbe63..ed83b879 100644 --- a/test/test_eager_compile_fallback.py +++ b/test/test_eager_compile_fallback.py @@ -7,6 +7,7 @@ uncompiled, warns once, and returns exactly what ``use_compile=False`` does. """ import warnings +from functools import lru_cache import numpy as np import pytest @@ -46,6 +47,14 @@ def _run(use_compile, memory=None): return solver, out.detach(), vp.grad +@lru_cache(maxsize=None) +def _reference(boundary_saving): + """``(record, gradient)`` of the plain uncompiled run every fallback must + reproduce, per memory mode: deterministic, so one run serves every test.""" + _, out, grad = _run(False, BoundarySaving() if boundary_saving else None) + return out, grad + + @pytest.fixture def compiles_on_cpu(monkeypatch): """These tests stand in their own torch.compile; on torch 2.3 the CPU step @@ -57,7 +66,7 @@ def compiles_on_cpu(monkeypatch): @pytest.mark.usefixtures("compiles_on_cpu") @pytest.mark.parametrize("memory", [None, BoundarySaving()], ids=["full", "boundary_saving"]) def test_falls_back_to_the_uncompiled_step(monkeypatch, memory): - _, ref_out, ref_grad = _run(False, memory) + ref_out, ref_grad = _reference(memory is not None) calls = [] monkeypatch.setattr(torch, "compile", _compile_that_cannot_build(calls)) @@ -90,7 +99,7 @@ def test_a_torch_without_compile_runs_uncompiled_and_says_so_once(monkeypatch, m """torch < 2.0 has no torch.compile at all. The propagator used to skip compiling there in silence, leaving users on the slow path unawares.""" import sweep.propagator._torch_eager as TE - _, ref_out, ref_grad = _run(False, memory) + ref_out, ref_grad = _reference(memory is not None) monkeypatch.delattr(torch, "compile") monkeypatch.setattr(TE, "_NO_COMPILE_WARNED", False) @@ -119,7 +128,7 @@ def test_torch_23_runs_uncompiled_on_cpu_and_says_so_once(monkeypatch, memory): while compiling the backward, where no fallback catches it), and nothing turns the faulty fusion off.""" import sweep.propagator._torch_eager as TE - _, ref_out, ref_grad = _run(False, memory) + ref_out, ref_grad = _reference(memory is not None) calls = [] monkeypatch.setattr(TE, "INDUCTOR_23", True) diff --git a/test/test_eager_snapshot_buffer.py b/test/test_eager_snapshot_buffer.py index 9d678ead..91a43322 100644 --- a/test/test_eager_snapshot_buffer.py +++ b/test/test_eager_snapshot_buffer.py @@ -11,6 +11,7 @@ import torch from sweep.equations import Acoustic +from sweep.propagator.options import EagerOptions from sweep.propagator.torch import PropTorch NZ, NX, NT = 24, 28, 12 @@ -24,7 +25,10 @@ def prop(): receiver_type=["h1"], # snapshots need the full tape; chunk checkpointing, the # eager default, refuses return_wavefield outright - use_ckpt=False) + use_ckpt=False, + # the buffer is the rollout's, not the step's: a cold + # torch.compile of a 12-step run is nearly all its time + eager_options=EagerOptions(use_compile=False)) def _shoot(prop): diff --git a/test/test_elastic_curvilinear.py b/test/test_elastic_curvilinear.py index fbb7b30b..cee92ca9 100644 --- a/test/test_elastic_curvilinear.py +++ b/test/test_elastic_curvilinear.py @@ -10,6 +10,8 @@ what the staggered chain-rule predecessor could not do. """ +from functools import lru_cache + import numpy as np import pytest import torch @@ -73,13 +75,20 @@ def _hill(height_cells=6.0, sigma_cells=12.0): return np.round(z.max() - z).astype(np.int64) +@lru_cache(maxsize=None) +def _flat_record(): + """The flat-grid run of the default geometry; two tests read it (a run is + deterministic, so one serves both).""" + sources, receivers = _geometry() + _, prop = _make_prop(topography=None) + return prop(_wavelet(), sources, receivers, models=_models()) + + # -------------------------------------------------------------------------- def test_flat_topography_runs_and_is_stable(): - sources, receivers = _geometry() - _, prop = _make_prop(topography=None) - syn = prop(_wavelet(), sources, receivers, models=_models()) + syn = _flat_record() assert torch.isfinite(syn).all(), "flat curvilinear run produced NaN/Inf" assert syn.abs().max().item() > 0, "record is identically zero" @@ -87,8 +96,7 @@ def test_flat_topography_runs_and_is_stable(): def test_zero_topography_matches_no_topography(): """A flat ``topography`` array must be the same problem as no array.""" sources, receivers = _geometry() - _, prop_none = _make_prop(topography=None) - a = prop_none(_wavelet(), sources, receivers, models=_models()) + a = _flat_record() _, prop_zero = _make_prop(topography=np.zeros(NX, dtype=np.int64)) b = prop_zero(_wavelet(), sources, receivers, models=_models()) assert torch.equal(a, b) @@ -317,7 +325,10 @@ def misfit(vp): g = vp.grad.detach() iz, ix = np.unravel_index(int(g.abs().argmax()), g.shape) - h = 10.0 + # that cell is next to the source, where the misfit is far from linear in vp: h = 10 m/s + # left a 4.8% truncation error against the 5% tolerance (1.1% at 5 m/s, 0.27% at 2.5; + # float64 keeps the round-off orders below) + h = 2.5 js = [] for sgn in (+1.0, -1.0): vpp = vp0.clone() @@ -417,14 +428,22 @@ def _snap_run(prop, topo, **kw): return snaps.detach().cpu() +@lru_cache(maxsize=None) +def _smooth_hill_physical(): + """A smooth-hill propagator and its default (physical-grid) snapshots; two + tests read them (a run is deterministic, so one serves both).""" + topo = _smooth_hill() + _, prop = _make_prop(topography=topo) + return prop, _snap_run(prop, topo) + + def test_snapshots_are_the_grid_rows_interpolated_per_column(): """The default snapshot is the computational one, each Cartesian row interpolated from the grid's rows around its depth -- checked against a plain per-column np.interp, and against a second run that asks for it.""" topo = _smooth_hill() - _, prop = _make_prop(topography=topo) + prop, phys = _smooth_hill_physical() comp = _snap_run(prop, topo, snapshot_grid="computational") - phys = _snap_run(prop, topo) mapped = prop._curvilinear_snapshots(comp.clone()) scale = float(comp.abs().max()) assert float((phys - mapped).abs().max()) <= 1e-6 * scale @@ -444,9 +463,7 @@ def test_snapshots_are_the_grid_rows_interpolated_per_column(): def test_air_cells_are_zero_in_every_field(): - topo = _smooth_hill() - _, prop = _make_prop(topography=topo) - phys = _snap_run(prop, topo) + prop, phys = _smooth_hill_physical() surf = prop.curvilinear_grid.h_surface.numpy() / DH air = np.arange(NZ)[:, None] < surf[None, :] assert air.any() diff --git a/test/test_kwarg_validation.py b/test/test_kwarg_validation.py index cf652d05..83d9e002 100644 --- a/test/test_kwarg_validation.py +++ b/test/test_kwarg_validation.py @@ -16,6 +16,7 @@ import torch from sweep.equations.acoustic import Acoustic +from sweep.propagator.options import EagerOptions from sweep.propagator.torch import PropTorch DEV = "cuda" if torch.cuda.is_available() else "cpu" @@ -60,7 +61,9 @@ def test_call_time_typo_raises_and_names_the_key(): def test_call_time_recognised_extras_still_work(): - p = _build() + # uncompiled: the extras reach init_abc, not the step, and a cold torch.compile + # of the step would cost ~30x this 50-step run + p = _build(eager_options=EagerOptions(use_compile=False)) wavelet = torch.zeros(50, device=DEV) src = np.array([[[24, 28]]], dtype=np.int64) rec = np.array([[[10, 28]]], dtype=np.int64) diff --git a/test/test_model_parallel_halo.py b/test/test_model_parallel_halo.py index 8c729787..4fa10931 100644 --- a/test/test_model_parallel_halo.py +++ b/test/test_model_parallel_halo.py @@ -2,7 +2,8 @@ Single-process tests cover the degenerate "no neighbour" case where HaloExchange must be a true no-op. Multi-process tests spawn 2 or 4 gloo -ranks on CPU and verify the actual halo-content exchange and its adjoint. +ranks on CPU and verify the actual halo-content exchange and its adjoint; +the checks of one world size share one spawn (see ``_worker_checks``). We avoid ``torch.autograd.gradcheck`` here because the natural finite- difference setup requires perturbing the input on ONE rank while the other @@ -17,6 +18,7 @@ import os import socket import sys +import traceback import pytest import torch @@ -158,217 +160,255 @@ def _cleanup_pg(): dist.destroy_process_group() -# --------------------------------------------------------------------------- -# 2-rank forward / backward (2-D, py=1, px=2) -# --------------------------------------------------------------------------- - -def _worker_forward_2d_1x2(rank, world_size, port): - try: - _init_pg(rank, world_size, port) - device = _worker_device(rank) - mesh = ModelParallelMesh(grid=(1, 2), shot_groups=1) - halo = 2 - Nz, Nx_loc = 4, 6 - total_x = Nx_loc + 2 * halo - - w = torch.full((1, 1, Nz, total_x), -1.0, device=device) - positions = torch.arange(Nx_loc, device=device).float().view(1, 1, 1, Nx_loc) - w[..., halo:-halo] = (rank + 1) * 100.0 + positions - - out = HaloExchange.apply(w, mesh, halo, ('x',)) - - if rank == 0: - assert (out[..., -halo:] == - torch.tensor([200.0, 201.0], device=device)).all(), \ - out[..., -halo:] - assert (out[..., :halo] == -1.0).all() - else: - assert (out[..., :halo] == - torch.tensor([104.0, 105.0], device=device)).all(), \ - out[..., :halo] - assert (out[..., -halo:] == -1.0).all() - - expected_interior = (rank + 1) * 100.0 + positions - assert (out[..., halo:-halo] == expected_interior).all() - finally: - _cleanup_pg() - -def test_halo_forward_2d_1x2(): - _run_spawn(_worker_forward_2d_1x2, 2) +def _worker_checks(rank, world_size, port, names, out_dir): + """Run the checks ``names`` in ONE process group, writing each one's + outcome (``ok`` or its traceback) to ``out_dir/.``. -def _worker_backward_2d_1x2(rank, world_size, port): + Starting the ranks -- a fresh interpreter importing torch, then the + rendezvous -- is nearly all of a check's time, so the checks of one world + size share a spawn. Each check builds its own mesh (``new_group`` on every + rank, in the same order) and finishes its exchange before it asserts, so + one failing on a rank leaves the collectives of every rank in step; the + barrier after it keeps a check that raised mid-exchange from pairing its + messages with the next check's (that hangs into the spawn timeout + instead, as it did with a spawn per check).""" try: _init_pg(rank, world_size, port) - device = _worker_device(rank) - mesh = ModelParallelMesh(grid=(1, 2), shot_groups=1) - halo = 2 - Nz, Nx_loc = 4, 6 - total_x = Nx_loc + 2 * halo - - # NCCL does not support fp64 reductions on all GPUs; use float32 there. - dtype = torch.float32 if _test_backend() == "nccl" else torch.float64 - w = torch.randn(1, 1, Nz, total_x, dtype=dtype, device=device, - requires_grad=True) - out = HaloExchange.apply(w, mesh, halo, ('x',)) - # loss = out.sum() — every cell of `out` (interior + halo) gets unit grad. - out.sum().backward() - - # Expected gradient w.r.t. w: - # interior cells get 1 from own loss; cells adjacent to a neighbour also - # receive that neighbour's halo gradient (= 1) accumulated in. Halo - # cells facing a neighbour are zeroed (forward overwrote them). - expected = torch.ones_like(w.detach()) - if mesh.neighbour_rank('x', +1) is not None: - expected[..., -2 * halo:-halo] += 1.0 # interior_high gets +1 from rank+1 - expected[..., -halo:] = 0.0 # halo_high is overwritten, grad=0 - if mesh.neighbour_rank('x', -1) is not None: - expected[..., halo:2 * halo] += 1.0 # interior_low gets +1 from rank-1 - expected[..., :halo] = 0.0 # halo_low overwritten - - assert torch.equal(w.grad, expected), ( - f"rank {rank} grad mismatch:\n got\n{w.grad}\n want\n{expected}" - ) + for name in names: + try: + globals()[name](rank, world_size) + outcome = "ok" + except Exception: + outcome = traceback.format_exc() + with open(os.path.join(out_dir, f"{name}.{rank}"), "w") as f: + f.write(outcome) + dist.barrier() finally: _cleanup_pg() -def test_halo_backward_2d_1x2(): - _run_spawn(_worker_backward_2d_1x2, 2) +def _spawn_checks(world_size, names, out_dir): + """``{name: {rank: outcome}}`` of ``names`` run on ``world_size`` ranks.""" + _run_spawn(_worker_checks, world_size, list(names), str(out_dir)) + outcomes = {} + for name in names: + outcomes[name] = {} + for rank in range(world_size): + path = out_dir / f"{name}.{rank}" + outcomes[name][rank] = path.read_text() if path.exists() else "no outcome written" + return outcomes -# --------------------------------------------------------------------------- -# 4-rank 3-D (py=2, px=2): forward and backward exercise both x and y axes -# --------------------------------------------------------------------------- +def _assert_passed(outcomes, name): + failed = {rank: out for rank, out in outcomes[name].items() if out != "ok"} + assert not failed, "\n".join(f"rank {rank}:\n{out}" for rank, out in failed.items()) -def _worker_forward_3d_2x2(rank, world_size, port): - try: - _init_pg(rank, world_size, port) - device = _worker_device(rank) - mesh = ModelParallelMesh(grid=(2, 2), shot_groups=1) - halo = 1 - Nz, Ny_loc, Nx_loc = 2, 3, 3 - total_y = Ny_loc + 2 * halo - total_x = Nx_loc + 2 * halo - - w = torch.full((1, 1, Nz, total_y, total_x), -1.0, device=device) - w[..., halo:-halo, halo:-halo] = (rank + 1) * 100.0 - out = HaloExchange.apply(w, mesh, halo, ('x', 'y')) - - # Expected neighbour tags for this rank - tag_self = (rank + 1) * 100.0 - n_xm = mesh.neighbour_rank('x', -1) - n_xp = mesh.neighbour_rank('x', +1) - n_ym = mesh.neighbour_rank('y', -1) - n_yp = mesh.neighbour_rank('y', +1) - - # Interior preserved - assert torch.all(out[..., halo:-halo, halo:-halo] == tag_self) - - # x halos (along last axis), inspected in the interior y-strip so that - # we test x exchange independently of y exchange. - if n_xm is not None: - assert torch.all(out[..., halo:-halo, :halo] == (n_xm + 1) * 100.0) - else: - assert torch.all(out[..., halo:-halo, :halo] == -1.0) - if n_xp is not None: - assert torch.all(out[..., halo:-halo, -halo:] == (n_xp + 1) * 100.0) - else: - assert torch.all(out[..., halo:-halo, -halo:] == -1.0) - - # y halos (along axis -2), inspected in the interior x-strip - if n_ym is not None: - assert torch.all(out[..., :halo, halo:-halo] == (n_ym + 1) * 100.0) - else: - assert torch.all(out[..., :halo, halo:-halo] == -1.0) - if n_yp is not None: - assert torch.all(out[..., -halo:, halo:-halo] == (n_yp + 1) * 100.0) - else: - assert torch.all(out[..., -halo:, halo:-halo] == -1.0) - finally: - _cleanup_pg() +@pytest.fixture(scope="module") +def two_ranks(tmp_path_factory): + return _spawn_checks(2, ["_check_forward_2d_1x2", "_check_backward_2d_1x2", + "_check_multi_field_2rank"], tmp_path_factory.mktemp("halo_2ranks")) -def test_halo_forward_3d_2x2(): - _run_spawn(_worker_forward_3d_2x2, 4) +@pytest.fixture(scope="module") +def four_ranks(tmp_path_factory): + return _spawn_checks(4, ["_check_forward_3d_2x2", "_check_backward_3d_2x2"], + tmp_path_factory.mktemp("halo_4ranks")) -def _worker_backward_3d_2x2(rank, world_size, port): - try: - _init_pg(rank, world_size, port) - device = _worker_device(rank) - mesh = ModelParallelMesh(grid=(2, 2), shot_groups=1) - halo = 1 - Nz, Ny_loc, Nx_loc = 2, 3, 3 - total_y = Ny_loc + 2 * halo - total_x = Nx_loc + 2 * halo - - dtype = torch.float32 if _test_backend() == "nccl" else torch.float64 - w = torch.randn(1, 1, Nz, total_y, total_x, - dtype=dtype, device=device, requires_grad=True) - out = HaloExchange.apply(w, mesh, halo, ('x', 'y')) - out.sum().backward() - - expected = torch.ones_like(w.detach()) - # x axis adjustments (last dim) - if mesh.neighbour_rank('x', +1) is not None: - expected[..., -2 * halo:-halo] += 1.0 - expected[..., -halo:] = 0.0 - if mesh.neighbour_rank('x', -1) is not None: - expected[..., halo:2 * halo] += 1.0 - expected[..., :halo] = 0.0 - # y axis adjustments (second-to-last dim) - if mesh.neighbour_rank('y', +1) is not None: - expected[..., -2 * halo:-halo, :] += 1.0 - expected[..., -halo:, :] = 0.0 - if mesh.neighbour_rank('y', -1) is not None: - expected[..., halo:2 * halo, :] += 1.0 - expected[..., :halo, :] = 0.0 - - assert torch.equal(w.grad, expected), ( - f"rank {rank}: grad mismatch\n{w.grad}\nvs\n{expected}" - ) - finally: - _cleanup_pg() +# --------------------------------------------------------------------------- +# 2-rank forward / backward (2-D, py=1, px=2) +# --------------------------------------------------------------------------- -def test_halo_backward_3d_2x2(): - _run_spawn(_worker_backward_3d_2x2, 4) +def _check_forward_2d_1x2(rank, world_size): + device = _worker_device(rank) + mesh = ModelParallelMesh(grid=(1, 2), shot_groups=1) + halo = 2 + Nz, Nx_loc = 4, 6 + total_x = Nx_loc + 2 * halo + + w = torch.full((1, 1, Nz, total_x), -1.0, device=device) + positions = torch.arange(Nx_loc, device=device).float().view(1, 1, 1, Nx_loc) + w[..., halo:-halo] = (rank + 1) * 100.0 + positions + + out = HaloExchange.apply(w, mesh, halo, ('x',)) + + if rank == 0: + assert (out[..., -halo:] == + torch.tensor([200.0, 201.0], device=device)).all(), \ + out[..., -halo:] + assert (out[..., :halo] == -1.0).all() + else: + assert (out[..., :halo] == + torch.tensor([104.0, 105.0], device=device)).all(), \ + out[..., :halo] + assert (out[..., -halo:] == -1.0).all() + + expected_interior = (rank + 1) * 100.0 + positions + assert (out[..., halo:-halo] == expected_interior).all() + + +def test_halo_forward_2d_1x2(two_ranks): + _assert_passed(two_ranks, "_check_forward_2d_1x2") + + +def _check_backward_2d_1x2(rank, world_size): + device = _worker_device(rank) + mesh = ModelParallelMesh(grid=(1, 2), shot_groups=1) + halo = 2 + Nz, Nx_loc = 4, 6 + total_x = Nx_loc + 2 * halo + + # NCCL does not support fp64 reductions on all GPUs; use float32 there. + dtype = torch.float32 if _test_backend() == "nccl" else torch.float64 + w = torch.randn(1, 1, Nz, total_x, dtype=dtype, device=device, + requires_grad=True) + out = HaloExchange.apply(w, mesh, halo, ('x',)) + # loss = out.sum() — every cell of `out` (interior + halo) gets unit grad. + out.sum().backward() + + # Expected gradient w.r.t. w: + # interior cells get 1 from own loss; cells adjacent to a neighbour also + # receive that neighbour's halo gradient (= 1) accumulated in. Halo + # cells facing a neighbour are zeroed (forward overwrote them). + expected = torch.ones_like(w.detach()) + if mesh.neighbour_rank('x', +1) is not None: + expected[..., -2 * halo:-halo] += 1.0 # interior_high gets +1 from rank+1 + expected[..., -halo:] = 0.0 # halo_high is overwritten, grad=0 + if mesh.neighbour_rank('x', -1) is not None: + expected[..., halo:2 * halo] += 1.0 # interior_low gets +1 from rank-1 + expected[..., :halo] = 0.0 # halo_low overwritten + + assert torch.equal(w.grad, expected), ( + f"rank {rank} grad mismatch:\n got\n{w.grad}\n want\n{expected}" + ) + + +def test_halo_backward_2d_1x2(two_ranks): + _assert_passed(two_ranks, "_check_backward_2d_1x2") # --------------------------------------------------------------------------- -# exchange_halos: multi-field convenience +# 4-rank 3-D (py=2, px=2): forward and backward exercise both x and y axes # --------------------------------------------------------------------------- -def _worker_multi_field_2rank(rank, world_size, port): - try: - _init_pg(rank, world_size, port) - device = _worker_device(rank) - mesh = ModelParallelMesh(grid=(1, 2), shot_groups=1) - halo = 1 - Nx_loc = 4 - total_x = Nx_loc + 2 * halo - - f0 = torch.full((1, 1, 2, total_x), -1.0, device=device) - f0[..., halo:-halo] = (rank + 1) # 1 or 2 - f1 = torch.full((1, 1, 2, total_x), -1.0, device=device) - f1[..., halo:-halo] = (rank + 1) * 10 # 10 or 20 - - outs = exchange_halos([f0, f1], mesh, halo, ('x',)) - - # Each output's halo should reflect the OTHER rank's interior of that SAME field. - if rank == 0: - assert torch.all(outs[0][..., -halo:] == 2) # rank 1's f0 = 2 - assert torch.all(outs[1][..., -halo:] == 20) # rank 1's f1 = 20 - assert torch.all(outs[0][..., :halo] == -1) - assert torch.all(outs[1][..., :halo] == -1) - else: - assert torch.all(outs[0][..., :halo] == 1) - assert torch.all(outs[1][..., :halo] == 10) - finally: - _cleanup_pg() +def _check_forward_3d_2x2(rank, world_size): + device = _worker_device(rank) + mesh = ModelParallelMesh(grid=(2, 2), shot_groups=1) + halo = 1 + Nz, Ny_loc, Nx_loc = 2, 3, 3 + total_y = Ny_loc + 2 * halo + total_x = Nx_loc + 2 * halo + + w = torch.full((1, 1, Nz, total_y, total_x), -1.0, device=device) + w[..., halo:-halo, halo:-halo] = (rank + 1) * 100.0 + out = HaloExchange.apply(w, mesh, halo, ('x', 'y')) + + # Expected neighbour tags for this rank + tag_self = (rank + 1) * 100.0 + n_xm = mesh.neighbour_rank('x', -1) + n_xp = mesh.neighbour_rank('x', +1) + n_ym = mesh.neighbour_rank('y', -1) + n_yp = mesh.neighbour_rank('y', +1) + + # Interior preserved + assert torch.all(out[..., halo:-halo, halo:-halo] == tag_self) + + # x halos (along last axis), inspected in the interior y-strip so that + # we test x exchange independently of y exchange. + if n_xm is not None: + assert torch.all(out[..., halo:-halo, :halo] == (n_xm + 1) * 100.0) + else: + assert torch.all(out[..., halo:-halo, :halo] == -1.0) + if n_xp is not None: + assert torch.all(out[..., halo:-halo, -halo:] == (n_xp + 1) * 100.0) + else: + assert torch.all(out[..., halo:-halo, -halo:] == -1.0) + + # y halos (along axis -2), inspected in the interior x-strip + if n_ym is not None: + assert torch.all(out[..., :halo, halo:-halo] == (n_ym + 1) * 100.0) + else: + assert torch.all(out[..., :halo, halo:-halo] == -1.0) + if n_yp is not None: + assert torch.all(out[..., -halo:, halo:-halo] == (n_yp + 1) * 100.0) + else: + assert torch.all(out[..., -halo:, halo:-halo] == -1.0) + + +def test_halo_forward_3d_2x2(four_ranks): + _assert_passed(four_ranks, "_check_forward_3d_2x2") + + +def _check_backward_3d_2x2(rank, world_size): + device = _worker_device(rank) + mesh = ModelParallelMesh(grid=(2, 2), shot_groups=1) + halo = 1 + Nz, Ny_loc, Nx_loc = 2, 3, 3 + total_y = Ny_loc + 2 * halo + total_x = Nx_loc + 2 * halo + + dtype = torch.float32 if _test_backend() == "nccl" else torch.float64 + w = torch.randn(1, 1, Nz, total_y, total_x, + dtype=dtype, device=device, requires_grad=True) + out = HaloExchange.apply(w, mesh, halo, ('x', 'y')) + out.sum().backward() + + expected = torch.ones_like(w.detach()) + # x axis adjustments (last dim) + if mesh.neighbour_rank('x', +1) is not None: + expected[..., -2 * halo:-halo] += 1.0 + expected[..., -halo:] = 0.0 + if mesh.neighbour_rank('x', -1) is not None: + expected[..., halo:2 * halo] += 1.0 + expected[..., :halo] = 0.0 + # y axis adjustments (second-to-last dim) + if mesh.neighbour_rank('y', +1) is not None: + expected[..., -2 * halo:-halo, :] += 1.0 + expected[..., -halo:, :] = 0.0 + if mesh.neighbour_rank('y', -1) is not None: + expected[..., halo:2 * halo, :] += 1.0 + expected[..., :halo, :] = 0.0 + + assert torch.equal(w.grad, expected), ( + f"rank {rank}: grad mismatch\n{w.grad}\nvs\n{expected}" + ) + + +def test_halo_backward_3d_2x2(four_ranks): + _assert_passed(four_ranks, "_check_backward_3d_2x2") -def test_exchange_halos_multi_field(): - _run_spawn(_worker_multi_field_2rank, 2) +# --------------------------------------------------------------------------- +# exchange_halos: multi-field convenience +# --------------------------------------------------------------------------- + +def _check_multi_field_2rank(rank, world_size): + device = _worker_device(rank) + mesh = ModelParallelMesh(grid=(1, 2), shot_groups=1) + halo = 1 + Nx_loc = 4 + total_x = Nx_loc + 2 * halo + + f0 = torch.full((1, 1, 2, total_x), -1.0, device=device) + f0[..., halo:-halo] = (rank + 1) # 1 or 2 + f1 = torch.full((1, 1, 2, total_x), -1.0, device=device) + f1[..., halo:-halo] = (rank + 1) * 10 # 10 or 20 + + outs = exchange_halos([f0, f1], mesh, halo, ('x',)) + + # Each output's halo should reflect the OTHER rank's interior of that SAME field. + if rank == 0: + assert torch.all(outs[0][..., -halo:] == 2) # rank 1's f0 = 2 + assert torch.all(outs[1][..., -halo:] == 20) # rank 1's f1 = 20 + assert torch.all(outs[0][..., :halo] == -1) + assert torch.all(outs[1][..., :halo] == -1) + else: + assert torch.all(outs[0][..., :halo] == 1) + assert torch.all(outs[1][..., :halo] == 10) + + +def test_exchange_halos_multi_field(two_ranks): + _assert_passed(two_ranks, "_check_multi_field_2rank") diff --git a/test/test_visco_elastic.py b/test/test_visco_elastic.py index 13aae1e2..9e78f34e 100644 --- a/test/test_visco_elastic.py +++ b/test/test_visco_elastic.py @@ -201,15 +201,18 @@ def test_gradient_matches_finite_difference(free_surface): the interior (the PML profile depends on the model non-differentiably).""" NZ, NX, wav, src, rec, m0, prop = _grad_setup(free_surface) g = torch.Generator().manual_seed(0) - d0 = prop()(wav, src.copy(), rec.copy(), models=m0) - wr = torch.randn(d0.shape, generator=g).to(DEVICE) - J = lambda ms: (prop()(wav, src.copy(), rec.copy(), models=ms) * wr).sum() ms = [m.clone().requires_grad_(True) for m in m0] - J(ms).backward() + # the differentiated run sizes the weights: J(ms) without a forward of its own + d = prop()(wav, src.copy(), rec.copy(), models=ms) + wr = torch.randn(d.shape, generator=g).to(DEVICE) + J = lambda ms: (prop()(wav, src.copy(), rec.copy(), models=ms) * wr).sum() + (d * wr).sum().backward() mask = torch.zeros(NZ, NX) mask[8:-8, 8:-8] = 1 k = torch.ones(1, 1, 7, 7) / 49 - for i, scale in enumerate([50.0, 30.0, 40.0, 3.0, 2.0]): + # Q steps of 5-7%: at 10-13% the central difference's truncation alone was 1.6e-2 (Qs, + # no free surface) of the 2e-2 tolerance + for i, scale in enumerate([50.0, 30.0, 40.0, 1.5, 1.0]): dm = torch.nn.functional.conv2d(torch.randn(1, 1, NZ, NX, generator=g), k, padding=3)[0, 0] * mask dm = (scale * dm / dm.abs().max()).to(DEVICE) eps = 1.0