Training memory and communication
Build and validate a time-ordered ledger for parameters, gradients, optimizer state, saved activations, workspaces, and communication buffers
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Transformer workload graph
- Backpropagation concept
- Bytes per datatype
The recurring training workload uses a decoder-only model with about 1 billion parameters on one server containing eight GPUs connected by high-speed links. Each GPU processes a microbatch of two sequences, each 2,048 tokens long, using BF16 forward and backward arithmetic. Fully sharded data parallelism divides parameter-related state across the eight workers. The training program still reports an out-of-memory failure even though one BF16 copy of 1 billion weights requires only about 2 GB.
A PyTorch memory snapshot separates live tensor allocations from memory reserved by the caching allocator and records when allocations appear during the training step. The snapshot shows BF16 compute weights, BF16 gradients, FP32 master weights, two FP32 Adam moment tensors, saved forward activations, matrix-multiplication workspaces, and buffers used to exchange parameter shards between GPUs. Different allocations begin and end at different operations.
A useful capacity model places every allocation on the same forward–backward–update timeline. Peak memory is the largest sum of allocations alive at one instant, not the sum of category maxima observed at unrelated times. The worked investigation calculates the 16 bytes of parameter-related state per learned parameter, predicts the ideal eight-way shard, adds transient all-gather and activation allocations, and compares the predicted high-water mark with the framework snapshot.
Separate persistent parameter-related state
Begin with the parameter count P and the stored format of each copy. A common mixed-precision Adam configuration can include 2-byte compute weights, 2-byte gradients, a 4-byte FP32 master weight, and two 4-byte FP32 moment tensors. This example totals 16 bytes per parameter before temporary buffers.
Do not assume that every framework uses this exact arrangement. Some optimizers update the main weights directly, store states in lower precision, or fuse conversion and update work. Gradient storage can be absent, aliased, bucketed, or freed progressively. Inspect the implementation or memory snapshot and make the calculation match it.
Calculate saved activation memory
Backward computation needs values from the forward pass. Automatic differentiation saves selected inputs, outputs, masks, normalization statistics, and other intermediates. Their size often grows with microbatch size, sequence length, hidden width, layer count, and attention implementation.
An operator graph tells you which values backward actually needs. A fused implementation can save a compact statistic instead of a large intermediate. An IO-aware attention implementation avoids retaining a complete score matrix. Count saved tensors from the executed graph rather than applying one universal bytes-per-token coefficient.
Understand activation checkpointing
Activation checkpointing stores selected boundary tensors during forward and discards interior values. During backward, the framework reruns the forward operations inside a checkpointed region to reconstruct the needed intermediates. Activation checkpointing exchanges additional computation for lower retained activation memory.
The time cost depends on which operations are recomputed and whether they were on the critical path. Recomputing inexpensive elementwise work can cost little. Repeating an expensive attention operation or communication can cost much more. Selective checkpointing preserves expensive outputs and recomputes cheaper values when the framework and graph permit it.
Add workspaces and communication buffers
Libraries can allocate temporary workspace for matrix multiplication, attention, sorting, or other algorithms. Distributed training can allocate gradient buckets, all-gather outputs, reduce-scatter buffers, and prefetched parameter shards. Temporary and communication allocations can be larger than their logical payload because of alignment or implementation requirements.
Overlap affects both time and capacity. A parameter all-gather can run while the current layer computes, hiding communication time. Its gathered parameters and buffers remain live during the overlap. A schedule with better time overlap can therefore have a higher memory peak.
Model peak memory as a timeline
Draw forward, backward, optimizer, and communication phases on a time axis. For each allocation, mark creation and last use. Add the sizes of all intervals crossing each point. The maximum sum is the modeled high-water mark.
This avoids two opposite errors. Adding every category maximum can overestimate when those maxima occur at different times. Counting only logical model state can underestimate when a workspace, communication buffer, prefetched shard, and saved activations coexist.
Include allocator behavior and safety margin
A framework allocator can reserve more device memory than currently allocated to tensors. Variable shapes can leave free blocks that cannot satisfy a later large request without additional reservation. Compilation and communication libraries can also retain internal state.
Report both allocated tensor memory and reserved device memory when the framework exposes them. Capacity planning should leave margin for shape variation, asynchronous overlap, library changes, and monitoring. A configuration that fits only when every allocation reaches an exact predicted lifetime is not a safe operating point.
Sharding exchanges storage for communication
Data-parallel training traditionally replicates parameters, gradients, and optimizer state on every worker. ZeRO-style and fully sharded methods partition selected state across workers. Ideal per-rank persistent storage falls with the shard count.
Computation still needs the relevant parameters and gradients at particular times. Workers all-gather parameters before computation and reduce-scatter gradients after backward segments. Prefetching and overlap can hide some time but create transient live sets. Calculate persistent savings, communication volume, and peak gathered state together.
GiB/rank = parameters × bytes/parameter ÷ shardsThis is a floor. Add transient lifetimes before deciding whether the job fits.
Worked example: Calculating Adam parameter-state memory
A decoder has 1 billion parameters and trains on one eight-GPU server. Each GPU receives a microbatch of two 2,048-token sequences. The training path uses BF16 compute weights and gradients, FP32 master weights, two FP32 Adam moments, and fully sharded data parallelism across all eight workers.
- Add bytes per parameter: 2 for compute weights + 2 for gradients + 4 for the FP32 master copy + 4 for the first moment + 4 for the second moment = 16 bytes per parameter.
- Multiply by 1×10^9 parameters. The unsharded result is 16×10^9 bytes, or 16 GB in decimal units. The 2-GB BF16 weight copy accounts for only one eighth of this parameter-related total.
- Under ideal eight-way sharding of all five categories, persistent parameter-related state is approximately 16 GB / 8 = 2 GB per GPU. Inspect the selected sharding policy because a policy that replicates gradients or compute weights retains more than the ideal estimate.
- Add the largest gathered parameter region and prefetch. If the current layer's full 400 MB parameter region overlaps a 400 MB prefetch for the next layer plus a 100 MB communication workspace, the transient addition is 900 MB, not only one local parameter shard.
- Obtain saved-activation sizes for microbatch 2 and sequence length 2,048 from the executed graph or PyTorch memory snapshot. Repeat the snapshot with microbatch 1; activation-related allocations should approximately follow the changed token count while persistent parameter state should remain nearly unchanged.
- Place optimizer workspaces, gradient buckets, all-gather buffers, and activations on a timeline. Compare the predicted peak with measured allocated and reserved memory. Investigate the largest difference instead of adding an arbitrary correction factor.
- Leave explicit operating margin. Then evaluate whether the memory savings allow a larger microbatch or sequence and whether the resulting throughput offsets communication and recomputation costs.
Conclusion: The calculation explains why the 2-GB BF16 weight size does not describe training capacity, why ideal sharding reduces 16 GB of parameter-related state to about 2 GB per GPU, and how transient gathered parameters and saved activations can still determine the measured peak.
Common errors
- Multiplying parameter count by the compute-weight width and calling the result training memory. This counts one representation and omits all other categories.
- Adding all category maxima without considering lifetime. Maxima that occur at different times do not necessarily coexist.
- Ignoring communication and prefetch buffers after sharding. Persistent state decreases, but gathered parameters can create a significant transient peak.
- Treating gradient accumulation as a reduction in all memory. Gradient accumulation can reduce activations per microbatch, but model and optimizer state remain.
Reference summary
Training memory contains persistent state and transient state. Persistent state can include compute weights, master weights, gradients, and optimizer moments. Transient state includes saved activations, recomputed values, operator workspaces, communication buffers, and prefetched parameter shards. Allocator fragmentation and framework state consume additional capacity.
peak device memory = maximum over time of all simultaneously live allocationsActivation checkpointing saves fewer forward intermediates and recomputes selected operations during the backward pass. Sharding divides parameter-related state across workers but introduces all-gather or reduce-scatter communication. Both methods exchange one resource for another. The relevant comparison includes peak memory, step time, communication volume, and any change in batch or sequence capacity.
Knowledge check
Which forward tensors does activation checkpointing stop retaining, which forward operations must run again during backward, and which profiler timeline signature would show that the recomputation overlaps or remains off the critical path?
Show the answer guide
- Activation checkpointing stops retaining selected internal forward activations for the checkpointed region. Boundary inputs, outputs required by later regions, parameters, and any state required by the checkpoint implementation still remain available.
- During backward, the framework runs the checkpointed region's forward operations again to reconstruct the missing activations before applying the corresponding gradient operations.
- A profiler timeline that combines central processing unit (CPU) and GPU events should therefore show a second occurrence of the region's forward matrix, normalization, and activation operations during backward, followed by or interleaved with their gradient work.
- Dependency arrows or stream ordering determine whether recomputation lies on the training-step critical path. Temporal overlap with unrelated communication is not automatically removable if the gradient waits for the reconstructed activation.
- Compare peak allocated memory, total forward-operation count, step-time distributions, and identical loss or gradients. Lower retained memory is the intended mechanism; repeated forward work is its compute cost.
Trace one checkpointed block
Run a four-layer MLP training step with and without checkpointing the middle two layers, using the same batch, seed, datatype, and optimizer. Record at least 20 warmed step times and a CPU/GPU trace or framework operator profile for one step per configuration.
Evidence to keep: Keep the peak-memory comparison, raw step times, loss and gradient comparison, annotated trace showing the repeated forward operators during backward, and a causal conclusion relating saved activations to recomputation cost.
Training memory follows forward and backward lifetimes. Autoregressive inference retains a different state and alternates between prompt processing and one-token decode iterations.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.