A training calculation ran successfully, then failed when activation checkpointing repeated it during backward. The inputs had not changed. The attention-backend settings had.

We reduced the failure encountered while building Vesper to a small CPU-only PyTorch example. Keeping the same attention policy during the original calculation and its recomputation repaired the tested failure. This is activation checkpointing, which trades computation for memory, not saving model weights for interrupted-run recovery.

The context ended too soon

PyTorch's scaled dot-product attention can dispatch to different implementations. The sdpa_kernel context manager restricts which implementations are enabled, then restores the previous settings when its scope ends.

with sdpa_kernel(SDPBackend.MATH):
    output = checkpoint(attention, q, k, v, use_reentrant=False)

output.square().sum().backward()

Here the original attention calculation runs under a MATH-only policy. Backward happens outside that scope. When checkpointing calls attention again, the original restriction is no longer active.

In the tested environment, backward raised CheckpointError: recomputed tensors had different metadata from the saved tensors. One reported shape changed from [2, 4, 8] to [1, 2, 8, 4]. A separate CPU operator trace recorded both the MATH implementation and a CPU flash-attention operator. That operator name does not establish CUDA Flash Attention support.

Three repairs, one principle

The simplest repair keeps forward and backward inside the same MATH scope. If the training loop calls backward elsewhere, context_fn supplies separate, fresh contexts for the original calculation and recomputation:

def contexts():
    return sdpa_kernel(SDPBackend.MATH), sdpa_kernel(SDPBackend.MATH)

output = checkpoint(
    attention, q, k, v,
    use_reentrant=False,
    context_fn=contexts,
)
output.square().sum().backward()

The third repair enters a fresh MATH context inside the attention function every time it is called. All three enforce the same requirement: checkpoint recomputation must use a consistent execution policy. MATH was chosen to isolate the issue, not as a recommendation to replace faster production backends.

What the saved experiment shows

The experiment used PyTorch 2.11.0+cu128, Python 3.12.10, Windows, CPU FP32 and one compute thread. Q, K and V each had shape [1, 2, 8, 4], with seed 1729, causal attention and no dropout. No model, dataset, optimizer, GPU, compilation or distributed execution was involved.

  • MATH around checkpoint forward only: failed during backward.
  • MATH around both forward and backward: passed.
  • Fresh contexts through context_fn: passed.
  • A fresh context inside the checkpointed function: passed.
  • Default policy throughout: passed against its own uncheckpointed default-policy reference.

The three MATH repairs were compared with an uncheckpointed MATH reference for output and all Q/K/V gradients. Absolute tolerance was 1e-6 and relative tolerance was 1e-5. Observed maximum absolute differences were zero in these small tests. Three saved fresh-process result records, including a copied-source replay, agree exactly. These are same-machine checks, not independent reproduction.

A complete minimal example

Save the following as checkpoint_attention.py and run it in an existing compatible PyTorch environment. No installation step is embedded. The broken command is expected to exit with an error in the tested environment; the fixed command completes backward with finite gradients. This smaller script demonstrates failure and repair, not the full reference-comparison suite described above.

python checkpoint_attention.py broken
python checkpoint_attention.py fixed
"""Run with 'broken' or 'fixed'. Authored CPU tensors only; no optimizer."""
import sys

import torch
from torch.nn import functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.utils.checkpoint import checkpoint

torch.set_num_threads(1)
generator = torch.Generator(device="cpu").manual_seed(1729)
q, k, v = [
    torch.randn(1, 2, 8, 4, generator=generator, device="cpu",
                dtype=torch.float32, requires_grad=True)
    for _ in range(3)
]


def attention(q, k, v):
    return F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True)


def contexts():
    # Return two fresh contexts, one for forward and one for recomputation.
    return sdpa_kernel(SDPBackend.MATH), sdpa_kernel(SDPBackend.MATH)


mode = sys.argv[1]
if mode == "broken":
    with sdpa_kernel(SDPBackend.MATH):
        output = checkpoint(attention, q, k, v, use_reentrant=False)
    # The MATH-only policy has ended before recomputation begins.
elif mode == "fixed":
    output = checkpoint(attention, q, k, v, use_reentrant=False, context_fn=contexts)
else:
    raise SystemExit("Expected broken or fixed")

output.square().sum().backward()
assert all(x.grad is not None and torch.isfinite(x.grad).all() for x in (q, k, v))
print("Backward completed with finite gradients.")

Scope and sources

PyTorch's checkpoint documentation warns that different function behavior between forward and backward can produce errors or incorrect gradients, and documents context_fn. Its attention-context documentation describes restoring backend settings on exit. The original investigation checked the corresponding warning and API in the installed runtime as well as those public main-branch documents.

This is a concrete example of a documented integration pitfall, not an established new framework bug. We observed an explicit exception, not silent gradient corruption. Other versions, platforms and backend choices may dispatch differently and need not reproduce the failure. Determinism checks stayed enabled.

Implementation, investigation and writing were AI-assisted. Publication review compared saved records and source hashes; it did not rerun the experiment. The evidence summary includes the tested cases and environment. The complete private investigation package is not released. No separate code license is granted by this note.