12
$\begingroup$

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

$\endgroup$
1

1 Answer 1

9
$\begingroup$

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}$

$\endgroup$
4
  • 1
    $\begingroup$ 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. $\endgroup$ Commented Jul 13, 2022 at 7:42
  • $\begingroup$ 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.) $\endgroup$ Commented Jul 13, 2022 at 7:49
  • 1
    $\begingroup$ 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 $\endgroup$ Commented Mar 24 at 0:29
  • $\begingroup$ This can help clarify things: nadavb.com/Memory-Footprint-of-Neural-Net $\endgroup$
    – Nathan G
    Commented Apr 17 at 9:45

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.