Design for the memory tail, not the average
Variable-length batches turn activation memory into a distribution. A configuration that is comfortable at the median can fail when a long-sequence bucket, an unusually dense packed batch, or a rare shape reaches the allocator. Reporting one peak from one friendly batch hides the operational question: which shapes determine admission, and how often do they occur?
Start with a trace keyed by shape. For each optimizer step, capture maximum sequence length, summed valid tokens, padding fraction, microbatch size, peak allocated memory, peak reserved memory, forward time, backward time, and whether an out-of-memory retry occurred. Reset the peak statistic at a deliberate boundary. PyTorch exposes maximum allocated memory since program start or the last reset, so an unreset counter cannot identify the offending bucket.
Checkpointing trades stored activations for replay
Activation checkpointing omits selected intermediate tensors during the original forward pass and reruns the chosen region during backward. The benefit is lower saved-activation memory. The cost is additional computation, possible random-number-state handling, and more complicated debugging. It does not reduce parameter, gradient, optimizer-state, communication-buffer, or every workspace allocation, so it cannot be assumed to solve any memory failure.
The original sublinear-memory work by Chen and colleagues formalized the compute-memory trade. In a framework implementation, details are concrete: PyTorch warns that a checkpointed function must behave the same during recomputation and the original forward, and it describes random-number-state preservation for operations such as dropout. A global variable or shape-dependent branch that changes between passes can produce errors or incorrect gradients.
Build a memory envelope by bucket
Use the production collator and sampler to replay representative length buckets, including the p50, p90, p99, and configured maximum. Warm the kernels before measuring, then run enough repeated steps to expose stable allocator behavior. Measure allocated and reserved memory separately. Reserved memory includes allocator-held segments and can stay high after a transient large batch, while allocated memory tracks live tensor allocations more closely.
Do not extrapolate a straight line from one short sequence. Attention implementations, padding strategy, compilation guards, kernel workspaces, and packing can create different scaling regimes. For every bucket, record the entire software and precision configuration. A memory curve collected under one attention kernel is not a promise for another. Also retain the raw shape trace so a later regression can be replayed. Include several batches per bucket because equal maximum lengths can still have different packing density and valid-token totals.
A worked selective policy limits the tax
Consider an illustrative trace of 10,000 microbatches. Suppose 7,000 have lengths up to 2,048 and peak at 45 GiB, 2,500 fall between 2,049 and 4,096 and peak at 61 GiB, and 500 exceed 4,096 and peak at 78 GiB on a device with a 75 GiB operational ceiling. The percentages are 70, 25, and 5. Only the tail violates the policy; blanket checkpointing would charge recomputation to all 10,000 microbatches.
Suppose a policy known before the forward pass checkpoints selected attention blocks only for the 5 percent tail. If that policy reduces the illustrative tail peak to 69 GiB and increases tail step time from 2.0 to 2.5 seconds, its weighted added time is 0.05 x 0.5 = 0.025 seconds per microbatch across this fixed trace. If blanket checkpointing adds 0.30 seconds to every microbatch, its average added time is 0.30 seconds. This comparison assumes the other buckets and the selection mechanism add no time. These figures are arithmetic examples, not measured product claims. Actual effects depend on model, kernels, shapes, and checkpoint boundaries.
| Illustrative length bucket | Share | Peak before | Policy |
|---|---|---|---|
| Up to 2,048 | 70% | 45 GiB | No checkpoint |
| 2,049 to 4,096 | 25% | 61 GiB | Profile, normally none |
| Above 4,096 | 5% | 78 GiB | Selective checkpoint |
Choose boundaries by bytes saved and replay cost
Blanket layer-by-layer recomputation is simple, but it can replay cheap and expensive operations indiscriminately. Profile candidate regions for saved bytes and incremental backward time at the tail shapes. A useful ranking is bytes of peak reduction divided by added milliseconds, provided correctness is unchanged. Regions whose outputs are large and whose forward computation is relatively cheap are natural candidates; the exact answer is model-specific.
Checkpoint boundaries interact with pipeline stages, distributed wrappers, compilation, and unused-parameter behavior. PyTorch's DistributedDataParallel documentation recommends the non-reentrant checkpoint implementation and describes limitations of the reentrant variant. Test the actual distributed topology, not only a single-process surrogate. Keep one uncheckpointed small-shape path available for gradient comparison and debugging.
Diagnose before adding more recompute
When a tail batch fails, first reproduce its exact sample IDs and shapes. Second compare allocated versus reserved memory around forward and backward. Third identify the phase and operator associated with the high-water mark. Fourth remove unrelated retention, such as logging tensors that still reference graphs. Fifth test a candidate checkpoint region and measure both peak reduction and added time. Only then broaden the region or reduce microbatch size.
This order matters because checkpointing cannot repair every cause. A leak across steps, a persistent compilation cache, an oversized communication workspace, or optimizer state growth may remain after activations are replayed. If the failure comes from one malformed sample beyond the supported length, input validation is a better control than recomputing the whole model.
Set a tail-aware decision rule
Accept a selective policy if every supported replay bucket stays under a declared ceiling, such as 92 percent of usable device memory, with zero failures across the validation trace, while p50 step time regresses by no more than 2 percent and p99 by no more than the serviceable training budget. These thresholds are illustrative. Check process-visible free memory as well as allocated and reserved framework memory, because allocator statistics do not account for every library, driver, or non-model allocation.
The honest limitation is that a finite trace cannot prove safety for an unbounded length distribution or future software revisions. Enforce a maximum supported shape, test after framework or kernel upgrades, and keep out-of-memory handling explicit. ClusterBid capacity matching is most useful when the request carries the measured tail envelope and recomputation policy, not only parameter count or an average memory figure.