Memory to Train a Transformer

Draft · 11 July 2026

We are going to estimate the memory needed to train a transformer. As in the layer-count estimate, we will take a simplified view of a decoder-style transformer. The goal is not to reproduce the exact number reported by a particular deep-learning framework. It is to see which parts of the architecture consume memory, and how that memory changes with the batch size and sequence length.

We will use a transformer with \(L\) layers, embedding dimension \(d_{\text{emb}}\), \(n_{\text{head}}\) attention heads, vocabulary size \(d_{\text{vocab}}\), batch size \(B\), and sequence length \(T\). Let \(N_{\text{param}}\) denote the total number of model parameters.

We will count the memory in the same order that training uses it:

1. load the weights;
2. perform the forward pass and save its activations; and
3. store the values needed for the backward pass and optimiser update.

The first and third parts are mainly controlled by the number of parameters. Forward-pass activation memory is controlled by the shapes of the tensors flowing through the model. This distinction matters: two training runs using the same model can have very different memory requirements merely because they use different values of \(B\) and \(T\).

Loading the Weights

The model starts with \(N_{\text{param}}\) weights. If the weights are stored in FP32, each one occupies 4 bytes. Merely loading the model therefore requires

In FP16 or BF16, the working weights occupy 2 bytes each, so this part would instead be \(2N_{\text{param}}\) bytes. That is not the full training requirement: the gradients and optimiser state arrive when we consider the backward pass.

Forward-Pass Activations

The forward pass produces intermediate tensors called activations. We have to retain enough of them to calculate gradients during the backward pass. Let \(s\) be the number of bytes used by each saved activation: \(s=2\) for FP16 or BF16 and \(s=4\) for FP32. We will first count the number of scalar values and multiply by \(s\) at the end.

The input embedding has shape \(B \times T \times d_{\text{emb}}\), so it contributes \(BTd_{\text{emb}}\) values. We next count the activations inside one transformer layer.

LayerNorm

We will use a pre-norm transformer block. In this arrangement, the input to the layer is normalised before self-attention. After the attention output is added back to the residual stream, the result is normalised again before entering the feed-forward layer:

Each LayerNorm output has shape \(B \times T \times d_{\text{emb}}\). There are two LayerNorm operations in each transformer block, so their main outputs contribute

For the backward pass, each LayerNorm may also save a mean and inverse standard deviation for every token. With two LayerNorms, that is four additional saved values per token in each transformer layer. We omit these smaller values from our first-order estimate.

Self-Attention

The query, key, and value tensors each have shape \(B \times T \times d_{\text{emb}}\) after all heads are considered. Together they contain

values. The attention score tensor \(\Q\K^\T\) has shape \(B \times n_{\text{head}} \times T \times T\). Applying softmax produces another tensor of the same shape. These two tensors contribute

values. Multiplying the attention probabilities by \(\V\) produces \(BTd_{\text{emb}}\) values, and the output projection produces another \(BTd_{\text{emb}}\). Under this bookkeeping, the self-attention part of one layer therefore contributes

The \(T^2\) term is the important one. Doubling the sequence length doubles most activations, but it quadruples the conventional attention matrices.

Feed-Forward Layer

We assume a two-matrix feed-forward layer with expansion factor 4. The upward projection produces \(4BTd_{\text{emb}}\) values. Its element-wise activation function produces another tensor of the same size, and the downward projection produces \(BTd_{\text{emb}}\) values. So,

Combining \(\eqref{eq:layernorm_activations}\), \(\eqref{eq:attention_activations}\), and \(\eqref{eq:ff_activations}\), one transformer layer contributes

This is where the linear coefficient becomes 16: two terms from LayerNorm, five from self-attention, and nine from the feed-forward layer.

Backward-Pass Values

The backward pass calculates one gradient for every weight. Adam also maintains a first-moment and second-moment estimate for every weight. If these values are stored in FP32, they require

Adding the FP32 weights from \(\eqref{eq:weights}\) gives \(16N_{\text{param}}\) bytes of persistent training state. A common mixed-precision arrangement arrives at the same total in a different way: 2 bytes for the low-precision working weight, 2 bytes for its gradient, 4 bytes for an FP32 master weight, and 8 bytes for the two FP32 Adam moments.

Putting Everything Together

We count the input embedding once, the expression in \(\eqref{eq:layer_activations}\) \(L\) times, and the final vocabulary logits once. The logits have shape \(B \times T \times d_{\text{vocab}}\). Our working estimate for saved activation memory is therefore

Finally, adding the loaded weights in \(\eqref{eq:weights}\), the backward-pass values in \(\eqref{eq:backward_values}\), and the forward-pass activations in \(\eqref{eq:activation_memory}\) gives

Equation \(\eqref{eq:training_memory}\) is in bytes.

Worked Example and Model Sizes

Consider the following approximate setup:

• \(N_{\text{param}} = 7 \times 10^9\),
• \(L=33\), \(d_{\text{emb}}=4096\), and \(n_{\text{head}}=32\),
• \(d_{\text{vocab}}=43{,}520\),
• \(B=1\), \(T=2048\), and
• BF16 activations, so \(s=2\).

Using Equation \(\eqref{eq:training_memory}\), the weights and backward-pass values occupy 112 GB. The saved activations occupy about 26.8 GB. The combined estimate is therefore about 138.8 GB.

If we change only the batch size from \(B=1\) to \(B=8\), the activation estimate grows eightfold to about 214 GB. The parameter-dependent terms do not change. This is why reducing the micro-batch size can rescue a run that is running out of memory even though the model itself has not changed.

There is another useful consequence. Two batches can contain the same number of tokens, \(BT\), without using the same amount of memory. The settings \(B=8,T=2048\) and \(B=16,T=1024\) have the same \(BT\), so their linear activation terms match. However, the first setting has twice the value of \(BT^2\), and therefore twice the conventional attention-matrix memory.

Comparing Model Sizes

Before calculating memory, we should check that each representative architecture actually has approximately the claimed number of parameters. The parameter calculation is derived in Estimating Layers in Transformers. Under the same standard transformer assumptions,

The first term counts the input and output embeddings. The second counts \(L\) transformer blocks. The head count does not appear because it cancels under standard multi-head attention; every shape below uses an integer head dimension \(d_{\text{head}}=d_{\text{emb}}/n_{\text{head}}\).

These are synthetic standard-transformer shapes chosen to match the nominal sizes. They are not the configurations of particular named models; Section 6 performs that separate check using published models.

Having checked the parameter counts, we can calculate each part of the memory separately. The machine we care about is an M2 with 8 GB of unified memory. The operating system and other programs claim a few GB of that, so we take roughly 4 GB as the hard training budget; the full 8 GB is the optimistic ceiling when little else is running. Every row below uses a micro-batch size of \(B=1\), ordinary FP32 Adam state, and FP32 activations (\(s=4\)). This is the conservative default because MLX uses FP32 as its default floating-point type. BF16 only applies if the model and its computation are explicitly converted to BF16.

All values in the table are decimal GB, where \(1\text{ GB}=10^9\) bytes. A GiB is a binary unit: \(1\text{ GiB}=2^{30}\) bytes, or about \(1.074\) GB. To avoid mixing the two units, we use GB throughout the calculations below.

We can check the two rows nearest the 4 GB boundary directly. For 40 M parameters at \(T=2048\), \(0.160 + 3.429 + 0.480 = 4.069\) GB. For 100 M parameters at the same context length, \(0.400 + 6.126 + 1.200 = 7.726\) GB. In each sum, the terms are weights, forward activations, and backward-pass values in that order.

This explains why a 40 M model can fail on an M2 with only 4 GB available for training. At \(T=2048\), the equation itself is already slightly above 4 GB. At \(T=512\), the equation is below 1 GB, but that still does not guarantee that the training program will fit.

The equation counts the main stored tensors. The training framework also needs working space while operations are running, and MLX may keep recently freed memory in a cache rather than immediately returning it to the system. The MLX memory documentation calls this cache memory “not currently used” memory that has not been returned to the system allocator. A micro-batch larger than \(B=1\) multiplies the activation column again.

Fitting in memory is necessary, but it is not sufficient for practical training. Training time also depends on the number of training tokens, the arithmetic required for each token, the M2's sustained compute and memory bandwidth, and how efficiently MLX executes the model. A 10 M parameter model can fit within the memory budget and still take too long for a particular experiment. The distinction between memory capacity and computation is developed further in Estimating the Compute and Memory Needed for Language Models.

A 40 M model is borderline only under the 4 GB budget. The measured run below peaks at 3.75 GB at \(T=2048\), so it trains on the full 8 GB machine without any of the usual rescue techniques. Under a hard 4 GB budget it needs a shorter context, micro-batch size 1, and careful dtype control.

MLX supports converting module parameters with set_dtype, and BF16 can reduce memory. However, this is a different training configuration: the dtypes of the weights, gradients, Adam state, and activations should all be checked rather than assumed.

The final test is to measure a real training step. MLX provides get_peak_memory for this purpose. The measurement should be taken after the Adam state has been created and after the compiled training step has warmed up.

Measuring on an 8 GB M2

We ran exactly this test on an 8 GB M2: the representative shapes from the table above, implemented as plain MLX modules, trained with ordinary Adam in FP32 at \(B=1\), with peak memory read from get_peak_memory after a warm-up step.

The estimate lands within about 10% of the measurement in each case: two rows slightly below the prediction, one slightly above, the difference being framework working space that the equation does not count. The practical reading for an 8 GB M2 is that even the 40 M model trains at the full 2,048-token context, and a 100 M model trains at \(T=512\). The 10 M figure is the answer to the 4 GB question, not the 8 GB one.

The throughput column carries its own lesson. At roughly 5,800 tokens per second, a ten-minute run shows the 10 M model about 3.5 M tokens; the 40 M model sees about 1.5 M. Memory is therefore not what limits training on an M2 — the clock is. This is the opposite of the large-cluster setting, where the accelerators have memory to spare and the game is spending compute efficiently.

The 70 B row should be read as a scale check, not a proposed single machine. Its model state alone is about 1.12 TB under ordinary Adam. Training it requires the state and activations to be distributed across many accelerators, or a different optimiser and precision setup.

Does This Hold in the Real World?

We can apply Equation \(\eqref{eq:training_memory}\) to the published configurations of several real models. For a fair comparison, every row below uses \(B=1\), \(T=2048\), and BF16 activations, so \(s=2\). We use the same conventional attention calculation for every model, even when its published maximum context is much longer.

The parameter term immediately gives the right scale. Even Qwen2.5-0.5B requires 7.84 GB for weights, gradients, and Adam state under our assumptions. Full-parameter Adam training of any model in this table therefore cannot fit within 4 GB without changing the optimiser, precision, or location of the training state.

There is also a measured result with which we can compare our scale. Qwen reports full-parameter training of its earlier 7 B model using two 80 GB A100 GPUs. Its published profiling table reports 139.2 GB at 256 tokens, 148.0 GB at 512 tokens, and 162.0 GB at 1024 tokens. The architecture and implementation are not identical to our 7 B example, but these measurements are in the same neighbourhood as our estimate.

The activation column is less exact. These real models use RMSNorm, gated feed-forward layers, and grouped-query attention, whereas our calculation uses one simplified transformer block. Modern training also uses techniques such as FlashAttention and activation checkpointing to retain fewer activations.

Distributed training changes the per-device number too. Fully sharded training divides parameters, gradients, and optimiser state across devices. LoRA and QLoRA reduce memory in a different way by freezing the original model and training only a small adapter; they are therefore not full-parameter training runs described by Equation \(\eqref{eq:training_memory}\).

What Changes the Estimate?

The expression above is intentionally simple. Several common techniques change it substantially:

1. Activation checkpointing saves only selected layer outputs and recomputes the missing activations during the backward pass. It trades additional computation for lower memory.
2. Memory-efficient attention avoids materialising the full \(T \times T\) score and probability tensors in accelerator memory. This removes the most painful part of the quadratic memory term, even though the arithmetic of standard dense attention remains quadratic in \(T\).
3. A fused cross-entropy loss can avoid retaining the full \(B \times T \times d_{\text{vocab}}\) logit tensor.
4. Sharding distributes parameters, gradients, optimiser state, or activations across devices. It reduces memory per device; ordinary data parallelism instead keeps a complete model replica on every device.
5. Different architectures change the constants. RMSNorm retains different statistics from LayerNorm; gated feed-forward layers, mixture-of-experts layers, dropout, and different attention implementations do not save exactly the set of tensors counted above.

Despite these qualifications, the decomposition is useful. Parameter memory grows with \(N_{\text{param}}\). Most activation memory grows with \(LBTd_{\text{emb}}\). Conventional attention adds a term growing with \(LBn_{\text{head}}T^2\). Once these three pieces are visible, the effect of changing the model, batch size, or context length becomes much easier to reason about.