Formula to compute approximate memory requirements of Transformer models I would like to roughly estimate the memory requirement of training an arbitrary Transformer model $M$, with $l$ layers, $h$ attention heads, an embedding dimension of $d$, and an input dimension of $t$ tokens.
What is the formula to compute this estimate?
If you happen to know the formula for a specific architecture (say BERT, or GPT) that would also be fine.
Note: I am not interested in precisely knowing how many bytes will be used on a specific GPU by a specific implementation with some library. I would just like a general formula to get a sense of the dimensions that is more principled than just "changing the batch size until it fits".

I was surprised that afaik there are no good answers for this (and similar) questions on the internet. I'm going to derive the following approximate formula for GPT:
$M \approx M_{activatons} \approx \frac{BT^2}{4ND^2}$
M = memory
B = batch size
T = sequence length
N = # of attention heads
D = dimension per head

Let's get started. The GPT transformer block has the following form:
Multi-head Attention -> LayerNorm -> MLP LayerNorm
To simplify the problem, let's exclude the layer norm and bias terms from our parameter count.
Assume we have $N$ heads, a hidden dimension of $D$ per each head, and data of batch size $B$ and sequence length $T$. Let's represent the total dimension as $C = N * D$ and assume the MLP has dimension $C$ also.
We want to express the memory footprint in terms of $C, B, T$.
There are three components that will contribute to the overall footprint:

*

*Storing the model $M_{model}$

*Storing the activations $M_{activations}$

*Storing the gradients $M_{gradients}$
So the total memory is $M = M_{model} + M_{activations} + M_{gradients}$. Unless you are computing higher order gradients $ M_{model} \geq  M_{gradients}$.
For transformers $M_{activations} >> M_{model}$ so the term we care about most is $M_{activations}$. I'll derive both though to show you why:
The model:
Each transformer block will have query, key, value networks and an MLP. We're ignoring layer norms and biases so the total parameters per block are $3C^2 + C^2 = 4C^2$. If the transformer has $L$ layers this means:
$M_{model} = 4LC^2 = 4 L N^2 D^2$
The activations:
Attention is the following operation $\text{Attention}(Q, K, V) = \text{softmax}(Q K^T / \sqrt{d}) V$. The $Q K^T$ operation has the following shape:
[B, N, T, D] @ [B, N, D, T] = [B, N, T, T]
Then the multiplication by $V$ and the MLP both output [B, N, T, D] activations. So the total memory per block is:
$BNT^2 + 2 BNTD = BNT(T + 2D)$
This happens at each layer so
$M_{activations} = BNLT(T+2D)$
Relative activation-to-model memory ratio is
$M_{activations} / M_{model} = BT(T+2D)/4N D^2$
Now let's assume we're modelling long sequences, then $T >> D$ and we have
$M_{activations} / M_{model} \approx \frac{BT^2}{4ND^2}$
Meaning that $M_{activations} >> M_{model}$$ so the total memory is dominated by activations:
$M \approx M_{activatons} \approx M_{model}\frac{BT^2}{4ND^2}$
