# 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".

• This paper from Nvidia may be interesting for this question: arxiv.org/pdf/2205.05198.pdf Commented Jan 11, 2023 at 11:19

## 1 Answer

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:

1. Storing the model $$M_{model}$$
2. Storing the activations $$M_{activations}$$
3. 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}$$

• Your last few steps go activations/model = x/y therefore activations = x/y. I think you just mean x there, i.e. BT²? I think it is worth using 4C not C for the MLP layer, as that is the de facto standard. Commented Jul 13, 2022 at 7:42
• I've upvoted for a putting the effort in to give an answer, which others can then argue with. (I have something in my notes where I found it is a bit more complicated: yes, sequence length T is important, but for smaller values of T, all of C, D and H become more important. So I don't think you can drop those terms.) Commented Jul 13, 2022 at 7:49
• Something seems amiss here. If I have a sequence length of 10000 it should use up a lot of memory but this calculation suggest it's only using up a few hundred megabytes Commented Mar 24 at 0:29
• This can help clarify things: nadavb.com/Memory-Footprint-of-Neural-Net Commented Apr 17 at 9:45