Activation checkpointing saves memory by recomputing intermediates during backward. The repeated calculation needs the same execution settings as the original. When an attention-backend policy ends before recomputation begins, the result is a documented integration failure, not a new framework bug.
The context ended before the repeated calculation
PyTorch's scaled dot-product attention can use different implementations, called backends. The sdpa_kernel context manager restricts which backends are enabled, then restores the previous settings when its scope exits.
The failing pattern restricts attention to the MATH backend during the forward pass, then triggers backward after leaving that scope:
with sdpa_kernel(SDPBackend.MATH):
output = checkpoint(attention, q, k, v, use_reentrant=False)
output.square().sum().backward()The backward pass triggers a second call to attention, but the MATH-only restriction is no longer active. In the tested environment (PyTorch 2.11.0+cu128, Windows, CPU FP32, one thread) this raised CheckpointError reporting that recomputed tensors had different metadata from those saved during forward. One reported shape changed from [2, 4, 8] to [1, 2, 8, 4].
A separate CPU operator trace recorded both the MATH attention implementation and PyTorch's CPU flash-attention operator. The word "flash" in that CPU operator's name does not establish CUDA Flash Attention availability. No GPU was used.
Three ways to keep the settings consistent
The simplest repair keeps both forward and backward inside the same MATH scope. When backward happens elsewhere in a training loop, PyTorch provides context_fn. It returns two context managers: one for the original calculation and one for 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()A third option is to enter a fresh MATH context inside the checkpointed function each time it runs. That applies the same repair principle used in Vesper: the backend settings must be applied again when the function is called again. These examples intentionally select MATH to isolate the issue; they are not a recommendation to replace faster attention implementations with MATH in production.
What we tested
Inputs: three tensors of shape [1, 2, 8, 4], fixed seed 1729, causal attention, dropout 0.0. No dataset, model training, optimizer, GPU, compilation, or distributed execution. Each case comparison used absolute tolerance 1e-6 and relative tolerance 1e-5.
| Case | Result |
|---|---|
| MATH scope around checkpoint forward only | CheckpointError during backward |
| MATH scope around forward and backward | Passed output and gradient comparison |
Fresh MATH contexts through context_fn | Passed output and gradient comparison |
| MATH context inside checkpointed function | Passed output and gradient comparison |
| Default backend policy throughout | Passed against uncheckpointed default-policy reference |
For the repaired cases, output and gradients of all three input tensors were compared against their appropriate uncheckpointed reference. Observed maximum absolute differences were zero. Default-policy checkpointing was compared against its own uncheckpointed default-policy reference, not assumed numerically identical to MATH.
Three fresh-process runs, including a copied-source replay, produced identical result records. A separate minimal broken/fixed script confirmed broken mode raised the expected error and repaired mode completed backward with finite gradients.
A documented requirement, demonstrated concretely
PyTorch's checkpoint documentation warns that changing how a function behaves between forward and backward can produce errors or incorrect gradients, and documents context_fn. The installed runtime's source describes the same forward/recompute consistency requirement. The attention-context documentation explains that leaving the context restores previous backend flags.
This example demonstrates one way to violate that documented requirement. It does not establish a new framework bug. A bounded related-issue search did not establish an exact duplicate or novelty; no issue was filed. We observed an explicit exception, not silent gradient corruption.
Publication disclosure
Implementation, investigation and writing were AI-assisted. This note is based on a September 12 CPU experiment reproduced without Vesper imports using only small synthetic tensors. The complete private comparison script and source receipts remain outside this repository. The embedded minimal script (checkpoint_attention_reproducer.py) is byte-identical to the tested minimal script, with SHA-256 matching. Before release, word choice, runnable code, and code-distribution terms were reviewed by the user. No separate code license is granted by this note. No new dataset, model training, GPU work, or installation command was executed.
The reproduced failure confirms a documented integration pitfall that can arise when a context manager exits before checkpoint recomputation begins. For anyone debugging this error, the useful question is direct: when checkpointing calls this function again, are the same execution settings still in effect?