Activation Memory, With and Without Recompute

Megatron-LM's activation-recomputation paper gives the activation memory one transformer layer must keep for its backward pass, in bytes:

$s \cdot b \cdot h \left(34 + 5\frac{a \cdot s}{h}\right)$

with $s$ the sequence length, $b$ the microbatch size, $h$ the hidden size, and $a$ the number of attention heads.

For $s = 2048$ , $b = 1$ , $h = 4096$ , $a = 32$ in a 32-layer model, how much activation memory does one microbatch hold, and what does full recomputation change?

About 30.6 GB, and full recomputation is free, because the recomputed forward pass runs inside the same kernels as the backward pass.

About 61 GB, since the backward pass needs its own copy of every stored activation.

About 30.6 GB, roughly 0.96 GB per layer across 32 layers. Full recomputation stores only each layer's input tensor, about 0.54 GB in total, and pays roughly 30-40% extra execution time.

About 0.96 GB: the formula already accounts for the whole model, not one layer.

About 15.3 GB, since the constant 34 counts fp32 bytes and this run is in bf16.