Activation Memory Analysis

Activations are the intermediate tensors saved during forward pass for backward computation. Unlike static memory (parameters + optimizer), activation memory scales with batch size and sequence length — it’s the “variable cost” that determines your maximum batch size and sequence length.

---

1. What Gets Saved and Why

PyTorch’s autograd saves tensors needed for gradient computation. For each operation $y = f(x)$, the backward pass needs either $x$, $y$, or both:

Operation	Saved for backward	Why
Linear ($y = Wx$)	Input $x$	$\frac{\partial L}{\partial W} = \delta \cdot x^T$, $\frac{\partial L}{\partial x} = W^T \cdot \delta$
ReLU ($y = \max(0, x)$)	Boolean mask ($x > 0$)	Gradient is 0 or pass-through
SiLU/GELU	Input $x$	Non-trivial derivative
Softmax	Output $y$	$\frac{\partial y_i}{\partial z_j} = y_i(\delta_{ij} - y_j)$
LayerNorm	Input $x$, mean, variance	Needs original statistics
Dropout	Binary mask	Same mask applied in backward
Attention ($\text{softmax}(QK^T/\sqrt{d})V$)	Q, K, V, attention weights	Full chain rule through softmax + matmul

---

2. Per-Layer Activation Memory

Detailed Breakdown for One Transformer Layer

With hidden dim $h$, sequence length $s$, batch size $b$, attention heads $a$, FFN dim $d_{ff}$, all in BF16 (2 bytes per element unless noted):

Tensor	Shape	Bytes	70B example
Self-Attention
LayerNorm 1 input	$b \times s \times h$	$2bsh$	67 MB
QKV projection input	$b \times s \times h$	$2bsh$	67 MB
Q, K, V tensors	$3 \times b \times s \times h$	$6bsh$	201 MB
Attention weights (softmax output)	$b \times a \times s \times s$	$2bas^2$	2,147 MB
Attention dropout mask	$b \times a \times s \times s$	$bas^2$ (1 byte/elem)	1,074 MB
Output projection input	$b \times s \times h$	$2bsh$	67 MB
FFN
LayerNorm 2 input	$b \times s \times h$	$2bsh$	67 MB
FFN gate+up input	$b \times s \times h$	$2bsh$	67 MB
Gate output (for SwiGLU backward)	$b \times s \times d_{ff}$	$2bs \cdot d_{ff}$	235 MB
Up projection output	$b \times s \times d_{ff}$	$2bs \cdot d_{ff}$	235 MB
Down projection input	$b \times s \times d_{ff}$	$2bs \cdot d_{ff}$	235 MB
Per-layer total			~4.5 GB

70B: $b=1$, $s=4096$, $h=8192$, $a=64$, $d_{ff}=28672$

Note: The exact tensors saved depend on the autograd implementation. The table above is the standard set for a pre-norm Transformer with SwiGLU. FlashAttention changes this significantly — see below.

Per-Layer Formula

$$M_{\text{act/layer}} = \underbrace{14 \cdot bsh + 6 \cdot bs \cdot d_{ff}}_{\text{linear in } s} + \underbrace{3 \cdot bas^2}_{\text{quadratic in } s}$$

The first two terms scale linearly with sequence length. The third term — attention weights plus dropout mask — scales quadratically.

The Attention Quadratic Term

The attention scores tensor has shape $b \times a \times s \times s$. This quadratic scaling is the dominant memory cost for long sequences:

Sequence length	Attention + mask per layer (64 heads)	80 layers
2,048	0.8 GB	64 GB
4,096	3.2 GB	258 GB
8,192	12.9 GB	1,030 GB
32,768	206 GB	16,500 GB
131,072	3,300 GB	264,000 GB

This is why long-context training is impossible without FlashAttention — attention activations alone would exceed entire cluster memory.

FlashAttention eliminates the $O(s^2)$ activation storage by recomputing attention during backward. See Flash Attention.

---

3. Total Activation Memory

General Formula

$$M_{\text{activations}} = \frac{L}{\text{PP}} \times \left( 14 \cdot bsh + 6 \cdot bs \cdot d_{ff} + 3 \cdot bas^2 \right)$$

With FlashAttention, the $3bas^2$ term drops to approximately $O(bsh)$ (only Q, K, V and small row-max/sum statistics are saved instead of the full $s \times s$ matrix).

Practical Numbers (with FlashAttention, no checkpointing)

Model	$h$	$d_{ff}$	Layers	Per-layer (b=1, s=4K)	Total
7B (Llama-3.1)	4,096	14,336	32	~0.6 GB	~19 GB
13B	5,120	13,824	40	~0.8 GB	~33 GB
70B (Llama-3.1)	8,192	28,672	80	~1.2 GB	~99 GB

Still large. 99 GB for 70B with FlashAttention and batch=1 — this is why activation checkpointing is essential at scale.

Without FlashAttention (Standard Attention)

Model	Per-layer (b=1, s=4K)	Total	Ratio vs FlashAttention
7B	~1.4 GB	~46 GB	2.4×
70B	~4.5 GB	~357 GB	3.6×

---

4. Reducing Activation Memory

Strategy 1: Activation Checkpointing (Recomputation)

Instead of saving all intermediate tensors, save only the input to each layer and recompute everything during backward:

Approach	Saved per layer	Memory	Extra compute
No checkpointing	All tensors	$L \times 4.5$ GB (70B, standard)	0%
Full checkpointing	Layer input only ($2bsh$)	$L \times 67$ MB (70B)	~33%
Selective checkpointing	Cheap tensors only	~40-60% of full	~10-15%

Full checkpointing for 70B: $80 \times 67\text{ MB} = 5.4$ GB — down from 357 GB (standard) or 99 GB (FlashAttention). A 66× or 18× reduction at the cost of 33% extra compute.

from torch.utils.checkpoint import checkpoint

class TransformerLayer(nn.Module):
    def forward(self, x):
        # Checkpoint this layer — only input x is saved
        return checkpoint(self._forward, x, use_reentrant=False)
    
    def _forward(self, x):
        # Everything here is recomputed during backward
        h = self.attention(self.norm1(x)) + x
        return self.ffn(self.norm2(h)) + h

Strategy 2: FlashAttention

Eliminates the $O(s^2)$ attention activation:

	Standard attention	FlashAttention
Saved for backward	Full $s \times s$ attention matrix + dropout mask	Only Q, K, V + row-max/sum stats
Memory per layer (70B, s=4K)	3.2 GB (attention only)	~0.2 GB (Q,K,V stats)
Impact	Dominant cost at long sequences	Removes the quadratic term

For details on how FlashAttention achieves this, see FlashAttention v1 and v2.

Strategy 3: Sequence Parallelism

Split the sequence dimension across TP ranks for LayerNorm and Dropout regions:

$$M_{\text{act,SP}} = M_{\text{act}} / \text{SP degree}$$

Each TP rank handles $s / \text{TP}$ tokens for non-attention operations. Combined with TP, this reduces activation memory proportionally.

See Sequence Parallelism for details.

Strategy 4: Reduce Batch Size + Gradient Accumulation

Activation memory scales linearly with micro-batch size:

$$M_{\text{act}} \propto b_\mu$$

Use micro-batch $b_\mu = 1$ with $K$ gradient accumulation steps to achieve effective batch $b = b_\mu \times K$. Same training semantics, but $K\times$ less activation memory.

---

5. Activation Memory Budget in Practice

For 70B, TP=8, FSDP, FlashAttention enabled:

Config	Parameters	Optimizer	Activations	Total per GPU
No checkpointing, b=1	17.5 GB	13.1 GB	~12 GB	~46 GB
Full checkpointing, b=1	17.5 GB	13.1 GB	~0.7 GB	~34 GB
Full checkpointing, b=4	17.5 GB	13.1 GB	~2.7 GB	~36 GB
Selective checkpointing, b=2	17.5 GB	13.1 GB	~5 GB	~39 GB

Activation values are approximate — TP=8 reduces some activations (FFN) but not all (LayerNorm is replicated without SP).

Checkpointing buys you either more batch size or more sequence length within the same GPU memory.

---

Self-Check

• Name the tensors saved for backward in a Transformer layer and why each is needed
• Calculate per-layer activation memory for a given architecture (with and without FlashAttention)
• Explain the $O(s^2)$ attention activation problem and how FlashAttention solves it
• Calculate activation memory with full checkpointing vs none
• Determine max batch size given a fixed memory budget after accounting for activations
• Combine checkpointing + FlashAttention + sequence parallelism for a complete activation strategy

---

References

• Training Deep Nets with Sublinear Memory Cost (Chen et al., 2016)
• Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022)
• FlashAttention (Dao et al., 2022)