Scaling 3-D FWI past a single GPU’s memory

Troubleshooting’s “Running out of memory” section lists the individual memory levers (shot batching, model_gradient_sampling_interval, storage_mode, storage_compression) and Distributed (multi-GPU) execution covers splitting shots across GPUs with DataParallel/ DistributedDataParallel. This page shows what happens when you need all of those at once – a 3-D problem too large for one GPU even after applying any single lever on its own – with a measured decision table instead of arithmetic projections, and points at Flagship 3-D FWI at a size that does not fit on one GPU as a complete, working composition.

The four levers, composed

  • DDP shot-splitting. torch.nn.parallel.DistributedDataParallel divides the shots between the two GPUs of the machine measured below, halving every per-GPU term below.

  • Shot batching. Within a rank’s shots, propagate a smaller group at a time and accumulate the gradient over groups (as in Reverse-Time Migration (RTM)), scaling the per-call cost down further.

  • ``storage_mode``/``storage_compression``. Offload the sampled wavefield history kept for the model gradient to host memory (storage_mode="cpu") and/or compress it. Compression is lossy in the model gradient; see deepwave.scalar() for what it costs.

  • ``checkpoint_segments``. Every propagator takes this argument: it splits the time axis into that many segments and checkpoints all but the last, bounding backward-pass memory to roughly one segment’s history plus the segment auxiliary states instead of the whole trace’s history, at the cost of one extra forward pass per checkpointed segment. It is shorthand for a deepwave.checkpointing.run_segments() call, which spells the same thing as n_segments= plus a separate boolean checkpoint=True. The two names are not interchangeable: checkpoint_segments is a segment count on a propagator, checkpoint is a boolean on the driver. This page says checkpoint_segments throughout, and n_segments only where it means the driver’s argument.

Each lever is independently useful, but they compose: DDP and shot batching shrink the shot dimension a call sees, checkpoint_segments shrinks the time dimension, and storage_mode/storage_compression shrink the bytes per stored step. Flagship 3-D FWI at a size that does not fit on one GPU uses all four together.

Measured decision table

Scalar propagator, accuracy 4, float32, forward + backward, one homogeneous-ish model (values do not depend on the model’s content, only its shape and the call parameters). Measured on a two-GPU workstation with 31 GB of host RAM and no swap – cuda:0 (RTX 4000 Ada, 20 GB) and cuda:1 (RTX A4000, 16 GB) gave identical GPU-peak figures to the precision shown wherever both were checked (rows 4 and 7 below), consistent with Troubleshooting’s “pure allocation-size effect” finding for this propagator – so most rows below were only measured once, on cuda:0, which has more headroom for the rows that are meant to fail. “Shots/call” is the shot count passed to a single propagator call (after DDP splitting and, where used, shot batching); “Segs” is that call’s checkpoint_segments (1, or an omitted argument, means no checkpointing). The two OOM rows report the CUDA allocator’s attempted allocation size at the point of failure (the run never completed, so there is no peak to report); every other row is torch.cuda.max_memory_allocated after a real, completed forward+backward pass, plus host RSS (resource.getrusage().ru_maxrss) for storage_mode="cpu" rows.

Grid

nt

Shots/call

Segs

storage_mode

compression

GPU peak

Host RSS

Result

403

200

8 (no split)

1

device

–

4.39 GB

–

Fits with no levers at all.

1003

600

8 (no split, single GPU)

1

device

–

OOM (needed ~57.3 GB)

–

Neither card has 57 GB; this is the “hurts” baseline.

1003

600

4 (DDP split only, single GPU)

1

device

–

OOM (needed ~28.7 GB)

–

DDP alone is not enough at this size.

1003

600

4 (DDP split)

5

device

–

14.33 GB

1.48 GB

Fits cuda:0 (72%); only ~14% headroom on cuda:1 (16 GB) – checkpointing alone is not comfortable on the smaller card.

1003

600

4 (DDP split)

5

cpu

False

3.05 GB

17.86 GB

GPU now trivial, but host RSS alone is more than half of that machine’s 31 GB per rank – two DDP ranks would risk the no-swap ceiling.

1003

600

4 (DDP split)

5

cpu

True

3.08 GB

5.57 GB

Compression cuts this rank’s host RSS ~3.2x (the offloaded history itself by more, since about 1.5 GB of that RSS is the process rather than the buffer); comfortable on both axes. It is lossy in the model gradient – see deepwave.scalar().

1003

600

2 (DDP split + batch)

5

cpu

True

1.55 GB

3.53 GB

Flagship 3-D FWI at a size that does not fit on one GPU’s actual per-call cost; ample margin on either card.

1003

600

2 per rank, both ranks concurrently

5

cpu

True

~1.7 GB each card

well under 31 GB combined

The real, literal Flagship 3-D FWI at a size that does not fit on one GPU run: both GPUs at once, 2 epochs, loss decreased, finished in about a minute.

Reading the table

The baseline (row 2) needs roughly 37x the memory the fully-composed recipe (row 7) uses – that gap is what “a size that hurts” means here: not the largest grid that can be coaxed onto one card, but a size that unambiguously does not fit unmodified and comfortably does once the levers are composed (each card sits under 10% utilisation in the final row, well inside the 80-85% margin worth leaving for a real optimiser run, which touches more of the graph than this table’s fixed-velocity measurement pass).

Two findings worth designing around, not just reading past:

  • Checkpointing alone is not enough headroom on the smaller card. Row 4 (14.33 GB) is 72% of cuda:0’s 20 GB but 86% of cuda:1’s 16 GB – above the comfortable range, and Troubleshooting already reports that real runs measure higher than any static estimate, so this is not a configuration to ship without also offloading storage.

  • ``storage_mode=”cpu”`` without compression can cost more host RAM than the whole point of offloading was meant to save. Row 5’s 17.86 GB host RSS, doubled across two concurrent DDP ranks, would use more than that machine’s entire 31 GB with no swap – offload and compression are a package deal at this scale, not two independently optional levers.

estimate_memory under checkpointing: still a lower bound

Troubleshooting’s “Measured accuracy of estimate_memory” table already documents that the estimate under-predicts real peak memory everywhere it was checked (device/uncompressed: -4% to -20%; compression, high sampling intervals, or offload: up to -70%). That table does not exercise checkpointing.

estimate_memory() does have a model of it: n_segments=k divides the stored history by k and adds a auxiliary_state_bytes term for the segment auxiliary states held alive – k of them where a model gradient is being computed, one where none is. Without that argument the estimate assumes the whole trace is stored at once and over-predicts a checkpointed run enormously (29.10 GB against a measured 14.33 GB for the first row below). With it, the estimate becomes a much closer – but still low – lower bound:

Row (levers)

Estimated, n_segments=5

Measured (GPU + host)

Relative error

Error without the term

checkpoint only (device)

7.05 GB

14.33 GB

-51%

-57%

checkpoint + cpu offload

7.07 GB

20.91 GB (3.05 + 17.86)

-66%

-70%

checkpoint + cpu + compression

2.79 GB

8.65 GB (3.08 + 5.57)

-68%

-78%

full recipe (+ batching)

1.39 GB

5.08 GB (1.55 + 3.53)

-73%

-81%

The measured column is the same measurement as the decision table above. The “error without the term” column is what the estimate gave before it had one, when the only way to account for checkpointing was to hand it a per-segment nt of 600 / 5 = 120 and ignore the auxiliary states. An independent single-process re-run of the same four configurations (one GPU, no DDP) measured 14.62, 22.07, 9.21 and 5.39 GB, i.e. within 7% of the figures above, so the relative errors are not an artefact of one run.

Read the result honestly: the term helps, and it is not enough. It recovers between 4 and 10 percentage points of a gap that was -57% to -81%, and leaves one of -51% to -73%. What it does not model accounts for the rest: the fixed per-call overhead that Troubleshooting’s own table identifies (allocator fragmentation, the compiled backend’s scratch, the host-side double buffer for offload, and roughly 1.5 GB of process baseline in each host-RSS figure here) is unmodelled and grows in relative terms exactly as the levers shrink the modelled ones – which is why the row with the fewest levers has the smallest relative error (-51%, device only) and the row with the most has the largest (-73%, the full recipe).

The auxiliary-state term itself is a floor on the auxiliary states, like every other term here rather than an exception to them. It sizes each state field by field: the finite-difference halo is not carried across a boundary, and the PML memory variables are stored as the two end slabs of the axis they vary along, not as full grids, so both reductions are priced in. Where a run cannot use the slab form – a propagator with no compiled backend, an explicitly requested Python backend, a non-default boundary, M-PML – the real states are larger than the term says, which leaves it a floor there too. And the count is the number of states genuinely held at once: one per segment while a model gradient is being computed, because each is an input to a checkpoint node that autograd keeps until that segment is recomputed, but exactly one when no model requires a gradient, because then there is nothing to checkpoint and the driver threads a single state forward and drops the previous one. Scaling that case by the segment count was how this estimate used to exceed the true peak by more than fourteen times at sixty segments; it does not now.

So total_bytes_lower_bound means what it says, at every segment count. Treat estimate_memory as a first-pass sizing tool – use it to rule out configurations that are obviously too large, and pass n_segments= so it is not wrong by a factor when you are checkpointing – then confirm your actual target size with a real measurement, as this page’s table and Flagship 3-D FWI at a size that does not fit on one GPU both do, before committing to it. deepwave.plan() gives the same estimate from the real resolved call (deepwave.plan(scalar, ..., checkpoint_segments=5).memory), which removes the risk of describing your configuration to the estimator incorrectly, but does not make it any less of a lower bound.

See also