---
title: "Gradient Checkpointing for Long DDE Simulations"
subtitle: "Trading Recompute for Memory When Differentiating Through Brain Network Models"
format:
html:
code-fold: false
toc: true
echo: false
embed-resources: true
fig-width: 8
out-width: "100%"
jupyter: python3
execute:
cache: true
---
Try this notebook interactively:
[Download .ipynb](https://github.com/virtual-twin/tvboptim/blob/main/docs/advanced/gradient_checkpointing.ipynb){.btn .btn-primary download="gradient_checkpointing.ipynb"}
[Download .qmd](gradient_checkpointing.qmd){.btn .btn-secondary download="gradient_checkpointing.qmd"}
[Open in Colab](https://colab.research.google.com/github/virtual-twin/tvboptim/blob/main/docs/advanced/gradient_checkpointing.ipynb){.btn .btn-warning target="_blank"}
## Introduction
Long simulations of delay-coupled brain network models are cheap to run forward
but expensive to differentiate through. Every step of `jax.lax.scan` saves its
carry state for the backward pass, and for DDEs that carry includes the per-coupling
history buffer of shape `[history_length, n_states, n_nodes]`. Backward memory
therefore grows as
$$ \text{memory} \;\propto\; n_\text{steps} \times \text{history length} \times n_\text{states} \times n_\text{nodes} $$
For a BOLD/FC fitting workload at `dt = 1 ms`, `T = 60 s`, ~80 regions and ~20 ms
maximum delay, this works out to **hundreds of megabytes of activations just
for the history buffers** — enough to push a gradient computation over the edge
of RAM on a workstation, even though the forward pass fits comfortably.
```{python}
#| output: false
#| echo: false
try:
import google.colab
print("Running in Google Colab - installing dependencies...")
!pip install -q tvboptim
print("✓ Dependencies installed!")
except ImportError:
pass
```
The standard remedy is **gradient checkpointing**: instead of saving every
step's activations for the backward pass, save only a sparse subset and
recompute the missing ones on demand. TVB-Optim implements this for the native
solver path as a single optional knob on `NativeSolver`:
```python
solver = Heun(checkpoint_every=256)
```
When `checkpoint_every` is `None` (the default) the integration runs as a
single `jax.lax.scan` exactly as before — there is no overhead and no behaviour
change. When set to an integer `K`, the scan is split into an outer scan over
blocks of `K` steps wrapped in `jax.checkpoint`, with an inner scan running
the `K` steps inside each block. Backward memory then scales as
`O(n_steps/K + K)` instead of `O(n_steps)`, at the cost of a **modest gradient
overhead — typically 1.3–1.7× depending on workload**, since backward is
usually already several times more expensive than forward and one extra
forward pass adds only a fraction to that total. Forward time is unchanged. The optimum for memory minimisation lies near `K ≈ √n_steps`.
::: {.callout-note}
## Scope and limitations
- **Native solvers only.** `DiffraxSolver` is not affected; Diffrax exposes
its own `RecursiveCheckpointAdjoint` for adaptive ODE solves, but it does
not support delays.
- **No effect when `checkpoint_every is None`.** The call site falls through to
the original `jax.lax.scan(op, state0, scan_inputs)` line. The default
behaviour is bit-exact with prior versions.
- **Forward is unaffected by memory savings.** Forward simulations do not retain
step activations regardless; checkpointing only matters when you take a
gradient.
:::
```{python}
#| output: false
#| code-fold: true
#| code-summary: "Environment Setup and Imports"
#| echo: true
import time
import gc
import os
import threading
import numpy as np
import matplotlib.pyplot as plt
import jax
import jax.numpy as jnp
import equinox as eqx
try:
import psutil
_HAS_PSUTIL = True
except ImportError:
_HAS_PSUTIL = False
# Enable float64 for numerically stable comparisons.
jax.config.update("jax_enable_x64", True)
from tvboptim.experimental.network_dynamics import Network, prepare
from tvboptim.experimental.network_dynamics.dynamics.tvb import ReducedWongWang
from tvboptim.experimental.network_dynamics.coupling import DelayedLinearCoupling
from tvboptim.experimental.network_dynamics.graph import DenseDelayGraph
from tvboptim.experimental.network_dynamics.noise import AdditiveNoise
from tvboptim.experimental.network_dynamics.solvers import Heun
from tvboptim.observations.tvb_monitors.bold import HRFBold
from tvboptim.observations.observation import compute_fc, rmse
from tvboptim.data import load_structural_connectivity, load_functional_connectivity
from tvboptim.utils import set_cache_path, cache
set_cache_path("./gradient_checkpointing_benchmark")
```
## Workload: RWW + Delays + BOLD FC Fitting
We use the same Reduced Wong-Wang / BOLD / FC workflow as
[`RWW.qmd`](../workflows/RWW.qmd), but with **`DelayedLinearCoupling`** in place of
`FastLinearCoupling`. This is the configuration where gradient memory
typically becomes the bottleneck for empirical fits. The structural
connectivity is the `dk_average` parcellation (68 regions); tract lengths are
converted to delays assuming a conduction speed of 4 mm/ms.
```{python}
#| echo: true
#| output: false
DT = 1.0 # Integration step (ms)
T1 = 60_000.0 # Total simulation length (ms) — 60 s
N_STEPS = int(T1 / DT) # 60_000 integration steps
CONDUCTION_SPEED = 4.0 # mm/ms
# Load empirical structural and functional connectivity.
weights, lengths, region_labels = load_structural_connectivity(name="dk_average")
weights = weights / np.max(weights)
delays = jnp.asarray(lengths / CONDUCTION_SPEED)
n_nodes = weights.shape[0]
fc_target = load_functional_connectivity(name="dk_average")
# Build the network: RWW dynamics + delayed linear coupling + additive noise.
graph = DenseDelayGraph(
weights=jnp.asarray(weights),
delays=delays,
region_labels=region_labels,
)
dynamics = ReducedWongWang(w=0.5, I_o=0.32, INITIAL_STATE=(0.3,))
coupling = DelayedLinearCoupling(
incoming_states="S",
G=0.5,
buffer_strategy="roll",
)
noise = AdditiveNoise(sigma=0.00283, apply_to="S", key=jax.random.key(0))
network = Network(
dynamics=dynamics,
coupling={"delayed": coupling},
graph=graph,
noise=noise,
)
# BOLD monitor — TR = 1 s, intermediate downsample matches dt.
bold_monitor = HRFBold(period=1000.0, downsample_period=DT, voi=0)
max_delay = float(delays.max())
history_length = int(np.ceil(max_delay / DT)) + 1
print(f"n_nodes={n_nodes} n_steps={N_STEPS} history_length={history_length}")
print(f"max delay = {max_delay:.2f} ms")
```
The history buffer for a single coupling is therefore roughly
`history_length × n_states × n_nodes × 8 bytes ≈ history_length × 544 bytes` per
step. With ~60 000 steps the total *forward-saved* activation footprint for
the coupling state alone runs into the hundreds of megabytes — and that is on
top of the dynamics state, the noise tensor, and the auxiliary tape.
## Benchmark
We benchmark forward time, gradient time, and (where the backend supports it)
peak device memory across a sweep of `checkpoint_every` values. The sweep
covers:
- `None` — the default, single `jax.lax.scan`. Reference for performance.
- Small `K` — frequent checkpoints, maximal recompute, minimal saved memory.
- `K ≈ √n_steps` — theoretical memory minimum.
- Large `K` — sparse checkpoints, close to no-checkpointing in cost.
- A non-divisor `K` — exercises the main-scan + tail-scan path.
```{python}
#| echo: true
#| output: false
#| code-fold: true
#| code-summary: "Benchmark Setup"
# K = None is the baseline. The dense middle (128, 256, 512, 1024, 2048)
# brackets sqrt(n_steps) so the U-shape near the minimum is well-resolved,
# while the wings (32, 8192, 30000) cover the asymptotic regimes. K = 30000
# is a clean divisor of n_steps (no tail). Most other values do not divide
# n_steps exactly and therefore exercise the main-scan + tail-scan path,
# which matters for the memory story — see "Reading the memory curve".
CHECKPOINT_VALUES = [None, 32, 128, 256, 512, 1024, 2048, 8192, 30000, N_STEPS]
N_FORWARD_RUNS = 3
N_GRADIENT_RUNS = 3
G_INIT = jnp.asarray(0.5)
class RSSPeakMonitor:
"""Context manager that records peak process RSS during the with-block.
Background thread polls ``psutil.Process.memory_info().rss`` at
``sample_interval`` seconds and tracks the maximum observed. On exit
``peak_delta_bytes`` holds the peak minus the baseline RSS taken just
before entry — i.e. the transient memory added by the block.
This is a *pragmatic CPU proxy*, not an accelerator profile:
- Linux RSS is process-resident memory and includes Python objects,
JIT artifacts, XLA scratch, and pooled CPU allocations. JAX on CPU
uses the system allocator, so transient activations show up here.
- ~50 ms sampling can miss sub-50 ms peaks; gradient passes through
tens of thousands of steps run for many seconds, so the sampler
catches the activation peak comfortably.
- **Pool effects matter.** XLA's CPU allocator pools pages and does
not always release them between configs. The reported delta is the
*additional* RSS the process had to allocate during the call —
configs whose peak fits inside memory already pooled by a previous
config will report a small or zero delta even though their
absolute requirement is non-trivial. To get clean per-config peaks
anyway, the sweep below is ordered with the most memory-hungry
configs *first*, so subsequent smaller-K configs are measured
against the already-grown pool and their deltas represent only
the marginal storage they add (which is zero or small if they fit
— i.e. exactly the success case for checkpointing).
- On GPU/TPU the activation tape lives in device memory, not host
RSS — use ``jax.devices()[0].memory_stats()['peak_bytes_in_use']``
there instead. This monitor is the CPU fallback.
"""
def __init__(self, sample_interval: float = 0.05):
self.sample_interval = sample_interval
self.peak_delta_bytes = None
def __enter__(self):
if not _HAS_PSUTIL:
return self
self._process = psutil.Process()
self._baseline = self._process.memory_info().rss
self._peak = self._baseline
self._stop = threading.Event()
self._thread = threading.Thread(target=self._sample, daemon=True)
self._thread.start()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if not _HAS_PSUTIL:
return False
self._stop.set()
self._thread.join()
self.peak_delta_bytes = max(0, self._peak - self._baseline)
return False
def _sample(self):
while not self._stop.is_set():
try:
rss = self._process.memory_info().rss
if rss > self._peak:
self._peak = rss
except Exception:
break
self._stop.wait(self.sample_interval)
def benchmark_one(checkpoint_every, fc_target):
"""Time forward + gradient, capture peak RSS during gradient, and return
the gradient value for cross-check."""
solver = Heun(checkpoint_every=checkpoint_every)
solve_fn, state = prepare(network, solver, t0=0.0, t1=T1, dt=DT)
solve_fn = jax.jit(solve_fn)
def loss(G):
cfg = eqx.tree_at(lambda c: c.coupling.delayed.G, state, G)
result = solve_fn(cfg)
bold = bold_monitor(result)
fc = compute_fc(bold, skip_t=20)
return rmse(fc, jnp.asarray(fc_target))
grad_fn = jax.jit(jax.value_and_grad(loss))
# Warm up (JIT compile both paths) so allocations from compilation do
# not contaminate the peak-RSS measurement below.
jax.block_until_ready(solve_fn(state).ys)
v0, g0 = grad_fn(G_INIT)
jax.block_until_ready(g0)
del g0
# Capture peak RSS delta during one fresh gradient call. The activation
# tape for the backward pass is the headline memory cost, so we measure
# exactly that. gc.collect() drops any temporaries from the warmup so
# the baseline is as flat as possible.
gc.collect()
monitor = RSSPeakMonitor(sample_interval=0.05)
with monitor:
v_mem, g_mem = grad_fn(G_INIT)
jax.block_until_ready(g_mem)
peak_delta = monitor.peak_delta_bytes
g_value_for_check = float(g_mem)
del v_mem, g_mem
gc.collect()
fwd_times = []
for _ in range(N_FORWARD_RUNS):
t = time.perf_counter()
r = solve_fn(state)
jax.block_until_ready(r.ys)
fwd_times.append(time.perf_counter() - t)
grad_times = []
for _ in range(N_GRADIENT_RUNS):
t = time.perf_counter()
v, g = grad_fn(G_INIT)
jax.block_until_ready(g)
grad_times.append(time.perf_counter() - t)
return {
"fwd_mean": float(np.mean(fwd_times)),
"fwd_std": float(np.std(fwd_times)),
"grad_mean": float(np.mean(grad_times)),
"grad_std": float(np.std(grad_times)),
"loss": float(v0),
"grad_value": g_value_for_check,
"peak_bytes_delta": peak_delta,
}
@cache("checkpoint_sweep", redo=False)
def run_sweep():
results = {}
for k in CHECKPOINT_VALUES:
label = "None" if k is None else str(k)
print(f"checkpoint_every = {label} ...", flush=True)
results[label] = benchmark_one(k, fc_target)
gc.collect()
return results
sweep_results = run_sweep()
```
## Results
```{python}
#| label: fig-checkpoint-benchmark
#| fig-cap: "**Gradient checkpointing benchmark.** Top row — *time*. Top-left: Forward and gradient wall time as a function of `checkpoint_every`; dotted horizontals mark the `None` baseline, dashed vertical marks `√n_steps`. Top-right: Per-call gradient-to-forward ratio. Bottom row — *memory*. Bottom-left: Peak RSS delta during a gradient call vs `checkpoint_every`, showing the `O(n_steps/K + K)` minimum near `√n_steps`. Bottom-right: Memory–time Pareto, with the `None` star at the low-time / high-memory extreme and checkpointed points tracing the front."
#| echo: true
#| code-fold: true
#| code-summary: "Plotting code"
baseline = sweep_results["None"]
sqrt_n = np.sqrt(N_STEPS)
# K-axis panels drop "None" — it has no x-coordinate on a checkpoint_every
# axis, only a horizontal-reference role. The Pareto panel keeps it as a
# distinct star marker because its axes are (time, memory) and there is no
# overlap risk.
ck_labels = [l for l in sweep_results if l != "None"]
xs_raw = np.array([float(l) for l in ck_labels])
order = np.argsort(xs_raw)
ck_labels = [ck_labels[i] for i in order]
xs = xs_raw[order]
fwd = np.array([sweep_results[l]["fwd_mean"] for l in ck_labels])
fwd_err = np.array([sweep_results[l]["fwd_std"] for l in ck_labels])
grad = np.array([sweep_results[l]["grad_mean"] for l in ck_labels])
grad_err = np.array([sweep_results[l]["grad_std"] for l in ck_labels])
peaks_all = [sweep_results[l]["peak_bytes_delta"] for l in sweep_results]
has_memory = all(p is not None for p in peaks_all)
if has_memory:
mem_ck_mb = np.array(
[sweep_results[l]["peak_bytes_delta"] for l in ck_labels], dtype=float
) / 1e6
def _mark_sqrt_n(ax):
"""Vertical reference line + label at √n_steps, anchored near the top."""
ax.axvline(sqrt_n, color="gray", linestyle="--", alpha=0.5, zorder=0)
ymin, ymax = ax.get_ylim()
y = ymax / ((ymax / ymin) ** 0.05) if ax.get_yscale() == "log" else ymax - 0.05 * (ymax - ymin)
ax.text(sqrt_n, y, r"$\sqrt{n_\mathrm{steps}}$",
color="gray", fontsize=12, ha="center", va="top",
bbox=dict(facecolor="white", edgecolor="none", alpha=0.8, pad=2))
def _pareto_front(times, mems):
"""Return boolean mask of Pareto-optimal points (minimise time AND memory).
A point is dominated if some other point has time<= and memory<= with at
least one strict inequality. The remaining points form the Pareto front.
"""
n = len(times)
keep = np.ones(n, dtype=bool)
for i in range(n):
for j in range(n):
if i == j:
continue
if (times[j] <= times[i] and mems[j] <= mems[i]
and (times[j] < times[i] or mems[j] < mems[i])):
keep[i] = False
break
return keep
# Bump default font sizes for the whole figure via a context manager so other
# notebook plots are not affected.
with plt.rc_context({
"font.size": 12,
"axes.titlesize": 14,
"axes.labelsize": 13,
"xtick.labelsize": 11,
"ytick.labelsize": 11,
"legend.fontsize": 11,
}):
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
# === Top row: time ===
# --- Top-left: time vs checkpoint_every (twin y-axes) ---
# Forward and gradient times are an order of magnitude apart, so a
# shared y-axis would collapse the forward curve to a flat line near
# the bottom. Twin axes let each curve fill its own range; we use
# linear scaling on both because each axis spans well under a decade.
# Axis spines and tick colours are tinted to indicate which curve goes
# with which side.
ax = axes[0, 0]
ax_g = ax.twinx()
fwd_color = "steelblue"
grad_color = "firebrick"
ax.errorbar(xs, fwd, yerr=fwd_err, marker="o", color=fwd_color,
label="forward", lw=1.8, markersize=7, capsize=3)
ax.axhline(baseline["fwd_mean"], color=fwd_color, linestyle="dashed",
alpha=0.7, label="forward (None)")
ax_g.errorbar(xs, grad, yerr=grad_err, marker="s", color=grad_color,
label="gradient", lw=1.8, markersize=7, capsize=3)
ax_g.axhline(baseline["grad_mean"], color=grad_color, linestyle="dashed",
alpha=0.7, label="gradient (None)")
ax.set_xscale("log")
ax.set_xlabel("checkpoint_every")
ax.set_ylabel("forward wall time (s)", color=fwd_color)
ax_g.set_ylabel("gradient wall time (s)", color=grad_color)
ax.set_title("Time vs checkpoint_every")
# Tint the axis ticks and spines to match the data they describe.
ax.tick_params(axis="y", colors=fwd_color)
ax.spines["left"].set_color(fwd_color)
ax_g.tick_params(axis="y", colors=grad_color)
ax_g.spines["right"].set_color(grad_color)
ax_g.spines["left"].set_visible(False)
# Combined legend from both axes.
h1, l1 = ax.get_legend_handles_labels()
h2, l2 = ax_g.get_legend_handles_labels()
ax.legend(h1 + h2, l1 + l2, loc="best", framealpha=0.9)
ax.grid(alpha=0.3, which="both")
_mark_sqrt_n(ax)
# --- Top-right: grad/forward ratio ---
ax = axes[0, 1]
ratio = grad / fwd
baseline_ratio = baseline["grad_mean"] / baseline["fwd_mean"]
ax.plot(xs, ratio, marker="^", color="darkgreen", lw=1.8, markersize=8,
label="grad / forward")
ax.axhline(baseline_ratio, color="darkgreen", linestyle="dashed", alpha=0.7,
label=f"None baseline ({baseline_ratio:.2f}×)")
ax.set_xscale("log")
ax.set_xlabel("checkpoint_every")
ax.set_ylabel("grad / forward")
ax.set_title("Gradient overhead")
ax.grid(alpha=0.3, which="both")
ax.legend(loc="best", framealpha=0.9)
_mark_sqrt_n(ax)
# === Bottom row: memory ===
# --- Bottom-left: memory vs checkpoint_every ---
ax = axes[1, 0]
if has_memory:
ax.plot(xs, mem_ck_mb, marker="D", color="purple", lw=1.8,
markersize=8, label="peak RSS delta during grad")
ax.axhline(baseline["peak_bytes_delta"] / 1e6, color="purple",
linestyle="dashed", alpha=0.7, label="None baseline")
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel("checkpoint_every")
ax.set_ylabel("peak RSS delta during grad (MB)")
ax.set_title("Memory vs checkpoint_every")
ax.grid(alpha=0.3, which="both")
ax.legend(loc="best", framealpha=0.9)
_mark_sqrt_n(ax)
else:
ax.text(0.5, 0.5,
"Peak memory unavailable\n(psutil not installed)",
transform=ax.transAxes, ha="center", va="center",
fontsize=12,
bbox=dict(boxstyle="round,pad=0.5", facecolor="lightyellow"))
ax.set_xticks([])
ax.set_yticks([])
ax.set_title("Memory vs checkpoint_every (unavailable)")
# --- Bottom-right: memory–time Pareto ---
# Two cleanup ideas vs the old "connect-by-time" line, which crossed
# itself wherever memory did not move monotonically with time:
# 1. Drop the connecting line; the scatter alone carries the points.
# 2. Compute the actual Pareto front (non-dominated points) and
# connect *only* those with a clean monotone curve.
# We do both — the front is a thin solid line, dominated points are
# plotted as scatter but not connected, and None is highlighted as a
# red star because it lies on the front but represents the no-checkpoint
# baseline.
ax = axes[1, 1]
if has_memory:
grad_all = np.array([sweep_results[l]["grad_mean"]
for l in sweep_results])
mem_all = np.array([sweep_results[l]["peak_bytes_delta"]
for l in sweep_results], dtype=float) / 1e6
label_all = list(sweep_results.keys())
pareto_mask = _pareto_front(grad_all, mem_all)
# Pareto-front line: sort the kept points by time so the line is
# monotone (memory decreases as time increases along a true front).
kept = np.where(pareto_mask)[0]
kept = kept[np.argsort(grad_all[kept])]
ax.plot(grad_all[kept], mem_all[kept], color="gray", lw=2.0,
alpha=0.6, zorder=1, label="Pareto front")
# Scatter all points, distinguishing None and Pareto vs dominated.
for i, l in enumerate(label_all):
on_front = pareto_mask[i]
x, y = grad_all[i], mem_all[i]
if l == "None":
ax.scatter([x], [y], s=240, marker="*", color="crimson",
edgecolor="black", linewidth=0.8, zorder=4,
label="None (baseline)")
elif on_front:
ax.scatter([x], [y], s=80, color="purple",
edgecolor="black", linewidth=0.5, zorder=3)
else:
ax.scatter([x], [y], s=55, facecolor="white",
edgecolor="purple", linewidth=1.3, zorder=2)
ax.annotate(l, (x, y), textcoords="offset points",
xytext=(8, 6), fontsize=11)
ax.set_xlabel("gradient time (s)")
ax.set_ylabel("peak RSS delta during grad (MB)")
ax.set_yscale("log")
ax.set_title("Memory–time Pareto")
ax.grid(alpha=0.3, which="both")
# Custom legend: front line + filled marker (on front) + hollow
# marker (dominated) + None star.
from matplotlib.lines import Line2D
legend_elems = [
Line2D([0], [0], color="gray", lw=2.0, alpha=0.6,
label="Pareto front"),
Line2D([0], [0], marker="o", color="w",
markerfacecolor="purple", markeredgecolor="black",
markersize=9, label="on front"),
Line2D([0], [0], marker="o", color="w",
markerfacecolor="white", markeredgecolor="purple",
markersize=8, markeredgewidth=1.3,
label="dominated"),
Line2D([0], [0], marker="*", color="w",
markerfacecolor="crimson", markeredgecolor="black",
markersize=14, label="None"),
]
ax.legend(handles=legend_elems, loc="best", framealpha=0.9)
else:
ax.text(0.5, 0.5,
"Peak memory unavailable\n(psutil not installed)",
transform=ax.transAxes, ha="center", va="center",
fontsize=12,
bbox=dict(boxstyle="round,pad=0.5", facecolor="lightyellow"))
ax.set_xticks([])
ax.set_yticks([])
ax.set_title("Memory–time Pareto (unavailable)")
plt.tight_layout()
plt.show()
```
## Reading the Memory Curve
The memory-vs-`checkpoint_every` panel (bottom-left) shows the classical analysis. Peak gradient
memory scales as
$$ \mathrm{peak\,memory} \;\approx\; \underbrace{\frac{n_\text{steps}}{K} \cdot c_\text{outer}}_{\text{block-boundary tape}} \;+\; \underbrace{K \cdot c_\text{inner}}_{\text{per-block inner tape during backward}} $$
with a minimum near $K \approx \sqrt{n_\text{steps} \cdot c_\text{outer} / c_\text{inner}}$.
For our workload this puts the optimum within a small factor of
$\sqrt{n_\text{steps}}$, which matches what the plot shows. Three real-world
deviations from the textbook curve all show up in the data:
**1. `jax.checkpoint` boundaries inflate the per-step inner tape.** Inside a
`jax.checkpoint` region, XLA cannot fuse intermediates across the boundary
and must preserve the per-step VJP tape in a form that the rematerialise
step can later consume. The result is that $c_\text{inner}$ is larger than
the per-step tape size XLA picks when it sees the whole scan
uncheckpointed. The inflation is biggest for short inner scans (where XLA
cannot amortise its setup) and shrinks toward the uncheckpointed value as
$K$ grows. The `None` baseline is therefore an unfair lower bound — even
the memory-optimal checkpointed `K` carries more per-step state than a
hypothetical "no checkpoint at this `K`" reference would. This is
intrinsic to how rematerialisation works and is the price you pay for
asking JAX to discard activations.
**2. Non-divisor `K` leaves a tail that holds its own tape.** When
`n_steps % K != 0`, the implementation runs the remainder as a plain
(uncheckpointed) `jax.lax.scan` after the checkpointed main scan. That
tail's activation tape is laid down on the forward pass and **stays live
in memory while the outer backward descends into the checkpointed main
block**. Peak memory then becomes roughly $K \cdot c_\text{inner} +
\mathrm{remainder} \cdot c_\text{unchecked}$ instead of just
$K \cdot c_\text{inner}$. For small remainders this is invisible; for
large remainders it can spike the peak above the `None` baseline. In this
sweep, $K = 60000$ divides $n_\text{steps}$ exactly (`remainder = 0`) and
shows clean behaviour, whereas a value like $K = 32768$ would produce a
27 232-step tail and inflate peak memory accordingly. **Practical
implication: prefer `K` values that divide `n_steps`** (or come close).
For $n_\text{steps} = 60\,000$ the divisor-friendly grid is 30, 120, 250,
500, 1000, 2000, 7500, 30 000.
**3. `K = n_steps` gives no memory benefit.** The single-block degenerate
case still wraps one `jax.lax.scan` in `jax.checkpoint`. Forward saves
only the entry carry; backward materialises the full activation tape for
the one block, runs backward, then frees it. Peak memory equals the full
uncheckpointed activation tape — about the same as `None` — except now
you also pay one extra forward recompute. `K = n_steps` is **single-pass
rematerialisation**: useful when you cannot pay even one set of saved
activations but otherwise dominated by smaller `K`.
The memory curve traces a U-shape:
- a steep left wall driven by the outer-scan boundary cost
($n_\text{steps}/K$) and the per-step inflation that hurts short inner
scans most;
- a shallow minimum bracketing $\sqrt{n_\text{steps}}$;
- a steady rise on the right driven by the linear $K \cdot c_\text{inner}$
term and, for non-divisor `K`, an extra bump from tail tape;
- the `None` reference sitting at the asymptote that the curve's right
arm approaches.
The sweet spot lands within a factor of 2 of theory's prediction (`K ≈ 244`);
the optimum cuts gradient memory roughly $10\times$ versus `None`.
## Correctness Check
A checkpointed gradient must agree with the uncheckpointed gradient to floating
point precision — the forward path is bit-exact (the scan body is identical;
only the loop nesting changes) and the backward path differs only in
floating-point rounding because activations are recomputed rather than read
from the saved tape.
```{python}
#| echo: true
baseline = sweep_results["None"]
print(f"{'checkpoint_every':<20} {'loss':<22} {'grad':<22} {'|Δgrad / grad|':<18}")
print("-" * 82)
for label, res in sweep_results.items():
rel = abs((res["grad_value"] - baseline["grad_value"]) / baseline["grad_value"])
print(f"{label:<20} {res['loss']:<22.16f} {res['grad_value']:<22.16e} {rel:<18.3e}")
```
The relative gradient error is at the level of double-precision rounding
(~1e-15 to 1e-13) for every block size — confirming that checkpointing does not change the answer.
## Summary Table
A single self-contained table of all measured quantities. Copy-pasteable into
an issue, a discussion, or back to an LLM for analysis. `fwd_ratio` and
`grad_ratio` are normalised against the `None` baseline; `peak_MB` is the
peak process-RSS delta during a single gradient call (CPU proxy via psutil),
or the device-memory delta if `jax.devices()[0].memory_stats()` is available
(GPU / TPU). `NA` when neither is available.
```{python}
#| echo: true
#| code-fold: true
#| code-summary: "Table code"
baseline = sweep_results["None"]
header = (
f"{'ckpt_every':<12} "
f"{'fwd_s':<14} "
f"{'grad_s':<14} "
f"{'grad/fwd':<10} "
f"{'fwd_ratio':<11} "
f"{'grad_ratio':<11} "
f"{'peak_MB':<10} "
f"{'loss':<22} "
f"{'grad':<14} "
f"{'|Δgrad/grad|':<14}"
)
print(header)
print("-" * len(header))
for label, r in sweep_results.items():
fwd = f"{r['fwd_mean']:.4f}±{r['fwd_std']:.4f}"
grd = f"{r['grad_mean']:.4f}±{r['grad_std']:.4f}"
ratio = r["grad_mean"] / r["fwd_mean"]
fwd_ratio = r["fwd_mean"] / baseline["fwd_mean"]
grad_ratio = r["grad_mean"] / baseline["grad_mean"]
peak = (
f"{r['peak_bytes_delta'] / 1e6:.1f}"
if r["peak_bytes_delta"] is not None
else "NA"
)
rel = abs((r["grad_value"] - baseline["grad_value"]) / baseline["grad_value"])
print(
f"{label:<12} "
f"{fwd:<14} "
f"{grd:<14} "
f"{ratio:<10.2f} "
f"{fwd_ratio:<11.2f} "
f"{grad_ratio:<11.2f} "
f"{peak:<10} "
f"{r['loss']:<22.16f} "
f"{r['grad_value']:<14.6e} "
f"{rel:<14.3e}"
)
# Compact context block (helpful when sharing the table).
print()
print(
f"# workload: n_nodes={n_nodes}, n_steps={N_STEPS}, dt={DT}, T={T1/1000:.0f}s, "
f"max_delay={max_delay:.1f}ms, history_length={history_length}"
)
print(f"# sqrt(n_steps) ≈ {int(np.sqrt(N_STEPS))} (memory-optimal block size)")
print(f"# device: {jax.devices()[0].platform} jax {jax.__version__}")
```
## No-Regression Check
Because `checkpoint_every=None` selects the original `jax.lax.scan` call site
verbatim (a literal if-branch), forward and gradient times for the default
must be **within timing noise** of the previous non-checkpointed implementation.
The benchmark above implicitly verifies this: the `None` row should be
statistically indistinguishable from any prior measurement of the unchecked
path. `K = n_steps` is **not** equivalent to `None` — it still
wraps the (single) inner scan in `jax.checkpoint`, so the backward pass
recomputes the entire forward once, costing roughly 1.3× the `None`
gradient time. Only `checkpoint_every=None` skips checkpointing entirely.
## Practical Guidance
```python
import math
from tvboptim.experimental.network_dynamics.solvers import Heun
# Default: no checkpointing. Fastest gradient when memory is not the issue.
solver = Heun()
# Memory-optimal default when gradients no longer fit in memory.
solver = Heun(checkpoint_every=int(math.sqrt(n_steps)))
# Aggressive: minimal memory, maximal recompute. Use only if the sqrt
# default still OOMs.
solver = Heun(checkpoint_every=64)
```
The same field works on `Euler`, `Heun`, `RungeKutta4`, and any
`BoundedSolver` wrapping one of those — the setting is delegated through the
wrapper to the base solver.
