Project: measure activation recomputation on receipt training

Last updated: 6 Oct 20265 min read

project

Checkpoint a receipt encoder block only after proving gradient parity, then report the measured memory saving and update-time cost on the intended training device.

Build a fixed training reference

Use a receipt-quality model, fixed physical-receipt split, preprocessing and labeled batch. Pin the optimizer and seed, then run one full-precision reference update with no recomputation. Save per-parameter gradients, output logits and next-step weights. Warm the optimizer before memory measurement so moment allocation is not confused with activations. The earlier memory project covers accumulation and precision; this project isolates recomputation.

Choose a legal segment

Wrap an encoder block whose forward is safe to replay during backward. Inspect dropout and normalization behavior, hooks and counters; side effects inside a replayed function can change state twice. Use an explicit checkpoint variant and RNG policy. Compare plain and recomputed output, input gradient, parameter gradients and one optimizer update within a declared tolerance. If parity fails, stop before reporting any memory result. The recomputation lesson gives a minimal test.

Measure like for like

On the intended accelerator, run several warm updates with identical batch size, image shape, dtype and optimizer state. Reset peak counters per step and synchronize timing. Report median update time, peak allocated bytes and peak reserved bytes for both paths; also report device-level process memory when available. The code below gates supplied measurements and does not pretend to collect GPU data in this environment. The profiler lesson explains each counter.

Test a second use of saved memory

If recomputation frees enough memory for a larger microbatch, treat that as a separate run. Preserve effective batch and accepted-example weighting if the aim is an optimization comparison; otherwise state that the training objective changed. Measure examples per second and rare-defect validation behavior. A smaller allocated peak can still yield lower throughput because backward reruns part of the forward. Do not claim a win based only on the largest batch that fits.

Release a reproducible decision

Predeclare a minimum memory saving, maximum update-time increase and gradient tolerance. Keep the original training path when the technique misses these gates. Package model and block revisions, input shape, optimizer state, device, precision, RNG policy and profiler trace. Re-run the fixed-batch parity test after library upgrades. The deliverable is a measured tradeoff with a rollback path, not a statement that activation checkpointing is always beneficial.

Implementation

from dataclasses import dataclass

@dataclass(frozen=True)
class MemoryRun:
    peak_allocated_mb: float
    median_update_ms: float
    largest_gradient_gap: float

def release_failures(reference: MemoryRun, candidate: MemoryRun) -> list[str]:
    failures = []
    if candidate.largest_gradient_gap > 0.0001:
        failures.append("gradient parity")
    if candidate.peak_allocated_mb > reference.peak_allocated_mb * 0.82:
        failures.append("memory saving")
    if candidate.median_update_ms > reference.median_update_ms * 1.35:
        failures.append("update latency")
    return failures

plain_run = MemoryRun(2140, 38.2, 0.0)
recomputed_run = MemoryRun(1600, 55.4, 0.00003)
assert release_failures(plain_run, recomputed_run) == ["update latency"]

Performance and operating cost

A checkpointed block trades stored activations for extra recomputation. The release gate itself is O(1), but each candidate needs multiple measured full updates and a gradient-parity comparison across model parameters. Peak allocated bytes are only one memory view; record reserved and device-level usage too. The example shows a candidate that saves memory yet fails its update-time gate. Its numbers are illustrative, not benchmarks from this machine. A larger-batch follow-up adds another training and validation cost.

Common Mistakes

Do not use illustrative values as measured GPU results.
Do not report memory savings before matching gradients and update boundaries.
Do not assume a larger fitting batch preserves the original effective-batch objective.