Activations, accumulation, and checkpointing
Checkpointing trades compute for activation memory; accumulation trades update frequency for batch memory—measure and validate each trade separately.
Training repeatedly OOMs and misses its nightly completion window.
This chapter moves from build a memory ledger to preserve update semantics. We will make the idea concrete in a small code example before placing it back inside a production system with 24 GB device and nightly completion.
Scale microbatch losses, clip once after accumulation, and advance the learning-rate schedule once per optimizer update. Compare against a real large batch on a short run.
Verify the behavior illustrated in memory_tradeoffs.py, then measure decision utility, slice reliability, and system latency.
Unsynchronized timing reports impossible speedups while retained tensors keep computation graphs alive and leak memory. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
When a model does not fit, first ask what occupies memory. Parameters are only part of the answer; saved activations often dominate training.
Backpropagation needs intermediate activations from the forward pass. Activation checkpointing saves only selected boundaries and recomputes the missing intermediates during backward, trading extra compute for lower peak memory. It is most useful for expensive stacks whose activations are large relative to their checkpointed inputs.
Gradient accumulation attacks a different limit: it splits a large effective batch into microbatches and delays the optimizer step. Dividing each microbatch loss by the accumulation count preserves the average-gradient scale. Equivalence can still break through batch normalization, random layers, gradient clipping, or schedules that advance per microbatch instead of per optimizer update.
- FP16/BF16
- loss scaling
- activation memory
- profilers
Follow the reasoning before touching the code.
- 01
Build a memory ledger
Measure parameters, gradients, optimizer state, activations, and temporary buffers at the target shape. Optimize the largest category first.
- 02
Checkpoint selectively
Place checkpoints around blocks with large saved activations and acceptable recompute cost. Verify outputs and gradients against the uncheckpointed path.
- 03
Preserve update semantics
Scale microbatch losses, clip once after accumulation, and advance the learning-rate schedule once per optimizer update. Compare against a real large batch on a short run.
import torch
from torch.utils.checkpoint import checkpoint
block = torch.nn.Sequential(torch.nn.Linear(32, 64), torch.nn.GELU(),
torch.nn.Linear(64, 32)).cuda()
head = torch.nn.Linear(32, 1).cuda()
optimizer = torch.optim.AdamW([*block.parameters(), *head.parameters()])
microbatches = [torch.randn(8, 32, device='cuda') for _ in range(4)]
optimizer.zero_grad()
for x in microbatches:
hidden = checkpoint(block, x, use_reentrant=False)
loss = head(hidden).square().mean() / len(microbatches)
loss.backward()
torch.nn.utils.clip_grad_norm_([*block.parameters(), *head.parameters()], 1.0)
optimizer.step()Combine recomputed activations with four-microbatch gradient accumulation.
What to notice
- Checkpointing reduces saved intermediate activations but performs the block’s forward computation again during backward.
- The optimizer steps once after four microbatches, and loss division keeps the accumulated gradient on an averaged-batch scale.
YOUR TURN Pause and predict before revealing the answer.
If each microbatch loss is not divided by four, what other hyperparameter effectively changes?
Need a hint?
The accumulated gradient becomes four times larger.
Check your reasoning
The effective learning-rate scale increases by roughly four for plain SGD-like updates, and clipping may trigger more often. That is not equivalent to the intended large-batch mean loss.
Lessons 1, 2, 3 in this course.
Connect the chapter to the course’s quantitative spine.
Checkpointing trades compute for activation memory; accumulation trades update frequency for batch memory—measure and validate each trade separately. The roofline bound below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
Roofline bound
| Symbol | Meaning / shape / unit |
|---|---|
P_peak | peak compute |
I | arithmetic intensity |
B_mem | memory bandwidth |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: P_peak, I, B_mem.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c12(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
GPU memory ledger
Change microbatch size and calculate a simplified training-memory budget.
Assumption: 3.2 GB weights+grads+optimizer, 0.34 GB activations per sample, 24 GB device, 10% reserve.
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| block 1 | 1.09 GB activations | resident |
| block 2 | 2.18 GB activations | resident |
| block 3 | 3.26 GB activations | resident |
| block 4 | 4.35 GB activations | resident |
| block 5 | 5.44 GB activations | resident |
Trace the complete operating path.
Carry the chapter’s central idea into the system: Checkpointing trades compute for activation memory; accumulation trades update frequency for batch memory—measure and validate each trade separately. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Build a memory ledger” into a written invariant and a failing test.
- 02
Promote memory_tradeoffs.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Preserve update semantics” as the release checkpoint: Scale microbatch losses, clip once after accumulation, and advance the learning-rate schedule once per optimizer update. Compare against a real large batch on a short run.
- 04
Serve or materialize behind explicit 24 GB device and nightly completion.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
Unsynchronized timing reports impossible speedups while retained tensors keep computation graphs alive and leak memory. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
Single-GPU optimization plan
Profile and optimize a correctness-locked trainer without sacrificing model quality.
- Profiler diagnosis
- Mixed precision parity
- Memory plan
- Throughput/cost report
Read primary material with a purpose.
Return to the opening failure.
Record hardware/software, shapes, precision, seeds, throughput method, peak memory, and cost per sample.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce a profiled, mixed-precision, memory-budgeted training loop.