Activation Memory: Why the Forward Pass Costs More Than the Weights

Tell an activation apart from a weight, using code you already write

Start somewhere completely familiar. Here is a function you could have written on any ordinary Tuesday.

def f(x):
    a = x * 3
    b = a + 2
    return b * b

Call it with x = 5. Inside, a becomes 15, b becomes 17, and the function returns 289. a and b are intermediate values. They exist only while the function is running. When it returns they are gone, and nobody misses them.

A neural network is that same idea, much longer. Data goes in at one end, a few hundred operations run in order, and an answer comes out the other end. Running the model in that direction, from input to answer, is called the forward pass. Every intermediate value the forward pass computes along the way is an activation.

That is the whole definition. An activation is an intermediate. It is not a setting you configured, and it is not something saved to disk.

The other kind of number in the building

In the function above, the 3 and the 2 are constants baked into the code. A neural network has the same category of number, except there are rather a lot of them and they were learned rather than typed. Those are the weights, also called parameters. Qwen2.5-1.5B, the model used throughout this series, has about 1.54 billion.

Weights and activations behave in opposite ways, and holding that difference clearly is most of the battle.

• Weights. Their count is fixed the moment you choose a model. They are loaded once and stay alive for the whole run.
• Activations. Their size is set by how much data you push through in one go. They appear and vanish inside a single step, and two of the numbers that set their size are training flags you type yourself.

Nearly every memory surprise in fine-tuning comes from applying weight intuition to activations. You can run out of memory on the exact same model by changing one number in a config file.

One more word: tensor

The intermediates in a neural network are not single numbers like a and b. They are tensors. A tensor is a block of numbers with a shape, and that is the entire concept.

A single number is a tensor with no shape. A list of 1,536 numbers is a tensor of shape 1536. A tensor of shape 4 x 1024 x 1536 is a three-level nested array: 4 blocks, each holding 1,024 rows, each row holding 1,536 numbers. If you have ever allocated a multi-dimensional array, you already know what a tensor is.

The one extra fact worth carrying is that a tensor is stored as one contiguous run of bytes, so its memory cost is exactly the number of elements multiplied by the bytes per element. No overhead, no surprises. We will use that constantly.

See why the backward pass cannot throw an activation away

If activations are only intermediates, why not free each one the moment the next operation has consumed it? That is exactly what happens when you run a model to get an answer. Training is different, and the reason is worth building carefully, because everything else here follows from it.

Training does three things in a loop. It shows the model an example. It measures how wrong the answer was. Then it nudges every weight slightly in whichever direction would have made the answer less wrong. Three pieces there, each with a name.

1. The loss is a single number saying how wrong the answer was. Big is bad. All of training is an effort to make it smaller.
2. The gradient of a weight is one number that answers one question: if I nudge this weight up by a hair, does the loss go up or down, and how fast? One gradient per weight, so about 1.54 billion of them for our model. If you want that idea developed slowly with pictures, the walkthrough of how a neural network learns builds it from nothing.
3. The backward pass, also called backpropagation, is the procedure that computes all of those gradients. It starts at the loss and works back towards the input, which is where the name comes from.

Nothing there mentions memory yet. Here is where memory arrives.

The whole argument, in two numbers

Take the smallest imaginable step of a network. One number goes in, it is multiplied by one weight, one number comes out.

• The input is x = 4.
• The weight is w = 0.5.
• So the output is y = w * x = 2.

The backward pass has already worked its way back through everything after this step, and it arrives here carrying exactly one number: how much the loss changes for each unit of change in y. Say that number is -1. Read it as “push y up by 1 and the loss comes down by 1″. So we would like y to be larger.

The question this step has to answer is: how should w change? Work it out by hand. Raise w from 0.5 to 0.6, a change of 0.1. Then y goes from 2 to 0.6 * 4 = 2.4, a change of 0.4. Raise it again from 0.6 to 0.7 and y becomes 2.8, another 0.4. Every single unit of change in w moves y by 4.

Four. Which is x.

The sensitivity of the output to the weight is the input. Combine that with the number arriving from above and you have this weight’s gradient: the loss changes by -1 * 4 = -4 for each unit of change in w. That is the instruction the optimizer acts on, the optimizer being the component that decides how far each weight actually moves once it knows which way to go.

Now the point of the whole exercise. That calculation needed x. Not the shape of x, not a summary of x, the actual value 4. Watch what happens if only the input differs. Keep w = 0.5, keep the same message from above, and set x = 100. The sensitivity is now 100, and the gradient is -100 instead of -4. Same weight, same message, an instruction twenty-five times larger, purely because the input was different.

There is no way to recover that number after the fact. Either the forward pass kept x alive in memory, or this weight’s gradient simply cannot be computed. That is the reason activations are stored. They are an operand of a multiplication that happens later.

The same idea with grids instead of single numbers

A real layer multiplies a tensor of inputs by a grid of weights, doing millions of multiply-and-add operations at once. The recipe does not change. The gradient for the weight grid is the message arriving from above, combined with the stored input tensor.

You will see it written compactly as dL/dW = dL/dy * x^T. Read it straight back in the words you just derived. dL/dW is the gradient for the weights. dL/dy is the gradient that arrived at this step’s output, the message from above. x is the stored input. The ^T means transpose, which is flipping a grid on its diagonal so its rows become its columns, bookkeeping to make the two grids line up for multiplication. It adds no new idea and no extra memory.

If you want that derivation with every intermediate step spelled out, the full calculus behind backpropagation is there for you. It is not required reading, because the two-number version above is the same fact.

What actually gets saved, and why the byte count is approximate

The framework decides which intermediates to keep, and it is more selective than you might expect. As you run the forward pass, PyTorch, the library most training code is written in, quietly records each operation and the values that went into it. That recording is the autograd graph, and it behaves like an undo history. When you ask for gradients it replays the list in reverse, and each recorded operation asks for what it needs.

Values that some backward step will need are saved, meaning they stay resident in GPU memory instead of being freed. A layer that multiplies by a weight grid saves its input, for the reason you just worked through. Values nothing will need are dropped at once, and some cheap operations prefer to recompute their result later rather than store it.

So the exact byte count depends on the framework version, the attention implementation and your flags. What does not depend on any of that is the shape of the answer: roughly a dozen saved tensors per layer, each about the size of the data flowing through. Treat every activation figure in this part as order of magnitude. The arithmetic below is right to within maybe 30 percent, and 30 percent is enough to tell you whether a run fits.

Predict the exact moment in a step when memory peaks

Knowing that activations are kept is half the story. Knowing how long each one is kept is the other half, and it is what gives the memory curve its shape.

Take a network with three layers and follow the bookkeeping. During the forward pass, layer 1 runs and saves its input, call it x1. Layer 2 runs and saves x2. Layer 3 runs and saves x3. Nothing has been freed, because nothing has been used yet.

Then the backward pass starts, and it runs in reverse. It reaches layer 3 first, uses x3, and frees it. Then layer 2, uses x2, frees it. Then layer 1, uses x1, frees it.

Look at x1. It was created first and it is consumed last. In a 28-layer model, the first layer’s saved input has to survive 27 more forward layers and 27 backward layers before anything needs it.

Now draw the memory. It climbs in steps through the forward pass, one step per layer, and nothing comes back down. It reaches its highest point the instant the forward pass finishes, because that is the one moment when every layer’s saved values are alive at once. Then it drains as the backward pass consumes them.

Two practical consequences fall straight out of that curve.

First, you cannot free memory between the forward pass and the backward pass. They are not two phases with a gap between them. They share one live set of values, and the handover is the peak.

Second, this gives you a free diagnostic. If a run loads the model fine and then dies a second or two into the first step, that is an activation problem almost every time. Weights run out of memory at load. Activations run out at the forward and backward boundary, so the timing of the crash tells you which tenant to go after.

Turn a tensor’s shape into a number of bytes

Time to put real numbers on all of this. Sizing an activation needs four quantities, and three of them are probably new words.

Token. Text is chopped into tokens before the model sees it. A token is roughly a word or a piece of one. “Fine-tuning” might be two or three tokens.

Sequence length. How many tokens are in one training example. Our running example uses 1,024.

Batch size. How many examples the model processes at the same time. Our running example uses 4. You type this one into a config, which is the first hint that activation memory is under your control in a way weight memory is not.

Hidden size. Inside the model, each token is carried as a list of numbers rather than as a word. The length of that list is the hidden size. For Qwen2.5-1.5B it is 1,536. So a token is not “cat” inside the model, it is 1,536 numbers that encode everything the model currently thinks about that position.

Put those together and a tensor flowing between two operations has shape 4 x 1024 x 1536: four examples, each 1,024 tokens long, each token carrying 1,536 numbers.

Count the elements. 4 * 1024 = 4,096. Then 4,096 * 1,536 = 6,291,456 numbers.

Now bytes per number. This run stores its numbers in bf16, a compact 16-bit floating-point format that uses 2 bytes each instead of the 4 bytes a standard float uses. The next part covers why bf16 is the format everyone reaches for. For now, 2 bytes is all you need.

So 6,291,456 * 2 = 12,582,912 bytes, or about 12.6 MB for one single intermediate tensor. Hold on to that number. We are about to multiply it a great many times.

Notice what never appeared in that calculation: 1.54 billion. An activation’s size is set by how much data you push through the model, never by how large its weight file is.

Count what one transformer layer leaves behind

Twelve point six megabytes is one tensor. To reach a whole model we need to know how many tensors one layer produces, which means a quick tour of what a layer does.

Qwen2.5-1.5B is 28 identical layers stacked on top of each other, with a lookup table at the bottom that turns each token into its first list of 1,536 numbers, and a small output stage at the top. Each of the 28 layers is a transformer block, and a transformer block is a two-stage workshop. The layer count, the hidden size and the head counts below are all from the Qwen team’s technical report.

Stage one: attention

Every token needs to look at the other tokens, because the meaning of a word depends on its neighbours. Attention is the machinery for that, and it works by giving each token three roles.

The layer converts each token’s 1,536 numbers into three shorter lists. A query says what this token is looking for. A key says what it has to offer. A value is what it hands over if it gets picked. Three tensors, three activations. The query is about the size of the tensor that flowed in. The key and the value are smaller in this particular model, for a reason that comes back at the end of this part.

Then every token’s query is compared against every token’s key, producing one score for each pair of tokens. With 1,024 tokens that is 1,024 x 1,024 scores. The comparison is run 12 times in parallel over different slices of the numbers, so the model can attend to several kinds of relationship at once. Each parallel run is called a head, and this model has 12 query heads.

So the score tensor has shape 4 x 12 x 1024 x 1024. That is batch, then heads, then a full square of token-against-token. Take a good look at that shape, because it is the villain of the second half of this article.

Stage two: the feed-forward network

After attention has mixed information between tokens, each token is processed on its own. The feed-forward network widens each token’s 1,536 numbers out to 8,960, passes them through a simple fixed rule applied to each number in turn, then squeezes them back down to 1,536. That widening is where a transformer does most of its per-token work, and the widened tensor is by far the largest ordinary activation in the layer.

The small parts that still cost bytes

Two more things happen around those stages, and they each leave an intermediate behind. Layer normalisation rescales a tensor so its numbers stay in a workable range instead of drifting to extremes. A residual connection adds a stage’s input back onto its output, so the original signal is never lost no matter how deep the stack gets. Each appears twice per layer, and each emits its own tensor.

Add it all up and a single transformer block emits roughly a dozen tensors of the 4 x 1024 x 1536 size, one fat 4 x 1024 x 8960 tensor, and, depending on how attention is implemented, one 4 x 12 x 1024 x 1024 score tensor.

For the operation-by-operation detail of what each stage computes, the full walkthrough of one transformer block covers it. Here we only care what it leaves behind.

Work out the running example’s 8 GB by hand

Everything is now in place. Three multiplications and we have the answer.

The feed-forward intermediate. Shape 4 x 1024 x 8960. Elements: 4,096 * 8,960 = 36,700,160. Bytes at 2 each: 73,400,320, or about 73 MB.

The hidden-width tensors. We said roughly a dozen at 12.6 MB each. Call it ten to stay conservative: about 126 MB.

The attention scores. Shape 4 x 12 x 1024 x 1024. Elements: 4 * 12 = 48, and 1,024 * 1,024 = 1,048,576, so 48 * 1,048,576 = 50,331,648. Bytes at 2 each: 100,663,296, or about 100 MB.

Roughly 300 MB per layer, times 28 layers, gives about 8,400 MB. Call it 8 GB of activations, order of magnitude, for a batch of 4 sequences of 1,024 tokens.

Put that next to the weights, carefully

Qwen2.5-1.5B’s weights in bf16 are 1.54 billion * 2 bytes, so about 3 GB. The activations at this batch and sequence are close to three times the size of the weights themselves. For training that is completely normal.

Be careful how you quote that 3 GB, because this is where memory estimates go wrong in public. Three gigabytes is the weights alone. A full fine-tune also holds one gradient per weight, plus two optimizer statistics and a high-precision master copy of every weight. That is the standard mixed-precision AdamW accounting of 16 bytes per parameter, AdamW being the optimizer almost every language model is trained with. It comes to about 24.6 GB of static state for this model, meaning memory that stays occupied for the whole run. Part 2 derives all sixteen of those bytes.

Activations are not inside that 24.6 GB. Add the roughly 8 GB from the table and a full fine-tune of this model needs about 33 GB. A 32 GB card gives you about 29.8 GiB usable, and every real run wants 2 to 3 GB of headroom on top for memory fragmentation and temporary buffers. So it does not fit, and no amount of optimism closes that gap.

Whenever you meet one of these figures, ask which kind it is. A static-state number and a total are different claims, and the distance between them is the subject of this article.

Name the four dials that set the bill, and spot the quadratic one

Weights depend on one thing: how many parameters the model has. Activations depend on four things, and none of them is the parameter count.

1. Batch size. Yours to set.
2. Sequence length. Yours to set.
3. Hidden size. Fixed by the model you chose.
4. Layer count. Fixed by the model you chose.

Two of the four are training flags. That is the good news in this whole article: when activations are what is killing you, you can usually fix it tonight without changing models.

Some of those words are new. LoRA means freezing most of the weights and training a small patch instead. Quantisation means storing each number in fewer bits. Checkpointing and flash attention each get a section of their own below.

The one activation that does not play fair

Look again at the attention score tensor: 4 x 12 x 1024 x 1024. Sequence length appears in it twice, because it is every token compared against every token. Every other activation has sequence in it once.

That difference is enormous. Double the sequence length from 1,024 to 2,048 and every ordinary activation doubles, which is annoying but survivable. The score tensor goes up four times, from about 100 MB per layer to about 400 MB. Double again to 4,096 and it is about 1.6 GB per layer. Across 28 layers that is roughly 45 GB of attention scores alone, on a card that has 32 GB in total, before a single weight is loaded.

This is the long-context wall, and it is why “just train on longer examples” is never a free change.

The order to pull the levers in

When a run dies on activations, work down this list. It is ordered by how much relief you get per unit of pain.

1. Cap the sequence length to what your data actually needs. If 95 percent of your examples are under 600 tokens, training at 2,048 just buys padding, the filler tokens added to make every example the same length. The quadratic term rewards this lever more than any other.
2. Cut the batch size, then recover the effective batch through gradient accumulation. That means running several small batches one after another, adding their gradients together, and only updating the weights once at the end. Gradients are one number per weight regardless of how many examples produced them, so accumulating costs no extra memory. Activations only ever hold one small batch at a time. Hold the effective batch fixed while you do this, because Goyal et al. showed the learning rate has to scale with that number and not with the small batch in front of it. Part 9 sets it up in a real run.
3. Turn on gradient checkpointing. The structural lever, and the next section.
4. Only then reach for a smaller model or a quantised one. Changing the model changes your results. The first three do not.

Cut the peak with gradient checkpointing, and measure the cut

Imagine a long calculation on paper, where you keep every line of working because you will need it later. You run out of paper. The fix is obvious: keep every tenth line, throw the rest away, and when you need line 47 again, start from line 40 and redo seven lines. A little more time, a fraction of the paper.

That is gradient checkpointing, also called activation checkpointing, in full. During the forward pass you keep activations only at a sparse set of positions, called checkpoints, and let the rest go. When the backward pass needs a discarded value, you re-run the forward pass for that short stretch, use the value, and drop it again.

The trade is compute for memory, and the exchange rate is good. You pay roughly one extra forward pass per step, which in practice costs 20 to 30 percent more wall-clock time. In return, Chen and colleagues showed that with checkpoints spaced well, the peak stops being proportional to the number of layers and becomes proportional to the square root of it. For 28 layers, the square root is a little over 5.

Where the checkpoints go

In practice frameworks put the checkpoints at transformer-block boundaries. It needs no tuning and it captures most of the saving. During the forward pass, only the tensor entering each block is kept. During the backward pass, when the work reaches a block, that block’s forward pass is re-run from its saved input to rebuild its dozen interior tensors, which are used and then released before the next block is touched.

So instead of 28 blocks’ worth of interiors alive at once, you hold 28 small boundary tensors plus one block’s interiors. That is where the order-of-magnitude drop comes from.

One thing checkpointing does not touch: weights, gradients and optimizer state. It is an activations-only lever, so if your problem is the 24.6 GB of static state it will not save you.

And one habit worth forming. Checkpointing costs throughput every step, so leaving it on for a run that already fits is paying 25 percent for nothing. Turn it on when activations are the binding constraint, meaning long sequences, a large batch, or a model near the edge of the card. Turn it off when you have headroom and want speed. It is a lever, not a virtue.

Measuring the saving instead of trusting it

You do not have to take any of this on faith. PyTorch keeps its own high-water mark of GPU memory, and you can read it. This runs one training step with checkpointing off, then one with it on, and prints both peaks.

# torch 2.x
import torch

def peak_gb(model, batch, checkpointing):
    torch.cuda.reset_peak_memory_stats()
    if checkpointing:
        model.gradient_checkpointing_enable()
    model(**batch).loss.backward()
    return torch.cuda.max_memory_allocated() / 1e9

print(f"without checkpointing: {peak_gb(model, batch, False):.1f} GB")
print(f"with checkpointing:    {peak_gb(model, batch, True):.1f} GB")

Line by line, assuming you have never written PyTorch before.

• import torch pulls in PyTorch, the library that owns the tensors and the GPU allocations.
• torch.cuda.reset_peak_memory_stats() zeroes a counter. PyTorch tracks the largest amount of GPU memory it has had allocated at once, and this resets that high-water mark so the next measurement starts clean.
• model.gradient_checkpointing_enable() switches on the recompute behaviour described above. Hugging Face model objects provide this method, so it works on any transformer loaded through that library without your writing the segmenting logic.
• model(**batch) runs the forward pass. batch is a dictionary of named inputs, and the ** is ordinary Python that spreads a dictionary into keyword arguments. It has nothing to do with PyTorch.
• .loss pulls the single how-wrong-were-we number out of what the model returned.
• .backward() runs the backward pass. This is the line where all those stored activations are finally consumed, and it is also the line where an activation-driven OOM almost always fires.
• torch.cuda.max_memory_allocated() reads the high-water mark back, in bytes. Dividing by 1e9 converts it to gigabytes.
• The two print lines use Python f-strings. {...:.1f} formats the result to one decimal place.

Expect the second number to be clearly lower than the first, and the run itself to take noticeably longer. Notice also that the two calls run in the order False first and True second. That ordering is not cosmetic, and the exercises at the end of this part ask you to work out why.

Pin your PyTorch and Transformers versions before you build anything on top of this. Memory-reporting helpers and checkpointing entry points have both moved between releases.

Remove the quadratic term entirely with flash attention

Checkpointing shrinks activations by agreeing to recompute them. A second, completely different move is available against the worst offender, and it does better than shrinking the score matrix. It never builds it.

Some hardware background first. A GPU has two very different kinds of memory. There is the large pool, tens of gigabytes of it, which is what people mean by VRAM, and it is comparatively slow to reach. Then there is a tiny scratchpad on the chip itself, measured in tens of kilobytes, which is extremely fast. Standard attention builds the whole 1024 x 1024 score grid in the large pool, which is why it costs about 100 MB per layer here.

Flash attention, which Dao et al. introduced, chops the work into tiles small enough to fit in the on-chip scratchpad. It computes one tile’s contribution to the final answer right there, accumulates it, and discards the tile. Then the next tile. The output is mathematically identical to standard attention. The full score grid simply never exists anywhere.

The memory effect is the structural one: the sequence-squared term disappears and attention memory becomes linear in sequence length. There is a speed benefit too, because moving less data to and from slow memory is most of what made attention slow.

The two levers do different jobs and are meant to be used together. Flash attention removes the quadratic score activation. Checkpointing shrinks the ordinary per-layer activations across the whole stack. Serious recipes run both: a fused attention backend by default, and checkpointing switched on when the remaining linear activations still will not fit.

One caution about defaults. Modern stacks usually select a fused attention backend, either flash attention itself or PyTorch’s scaled dot product attention, rather than a hand-written score matrix. Confirm which one your run picked rather than assuming, because the attention implementation is a setting on the model and a fallback path can quietly cost you 100 MB per layer.

Explain why inference is cheap, and what the KV cache really is

Everything so far has been about training. Serving a trained model is a different world, and the difference explains why a model that will not fine-tune on a card serves happily on it.

Inference has no backward pass. Nothing computed in the forward pass will ever be needed again, so each layer’s output is freed the moment the next layer has read it. Only one layer’s worth of activations is live at any moment, and that 8 GB peak collapses to a sliver.

There is exactly one deliberate exception. A model generates text one token at a time, and each new token attends to every token before it. That means using each earlier token’s key and value, the two lists from the attention stage, which were computed when that token was first processed. Recomputing them at every step would mean redoing the entire prompt for every single token.

So they are kept. That retained slice is the KV cache. It is not a new species of object. It is an activation that inference chooses to remember because remembering it pays for itself.

Its size is straightforward: 2 (one for keys, one for values) times the layer count, times the number of key and value heads, times the width each head works on, times the context length, times bytes per number. Run it for our model. Each head works on a slice of 128 numbers, called the head dimension, and the model has 28 layers and 2 key and value heads. So 2 * 28 * 2 * 128 * 2 bytes = 28,672 bytes, about 28 KB per token. Fill a 256K-token context and you hold roughly 7 GB of cache.

Two engineering responses fall straight out of that formula. Grouped-query attention lets several query heads share one key and value head, which is why this model has 12 query heads but only 2 key and value heads. Each of those serves six query heads, so the cache is six times smaller than it would otherwise be. Separately, paged key and value storage keeps the cache in non-contiguous blocks, so a serving system need not reserve the maximum possible length for every request up front. The part on continuous batching and paged attention covers both.

Place activations among the other tenants, and see why LoRA does not help

Step back and put all five memory tenants in one view.

Now the row people misread, which decides real hardware purchases. Freezing a weight means the optimizer will never update it. That removes the weight’s gradient and its optimizer state, which is where LoRA’s saving comes from: static state for this model drops from about 24.6 GB to about 4 GB at a typical adapter size. Genuine, large, and worth having.

It does nothing at all to activations, for two reasons. First, the forward pass still runs through every frozen layer, because the frozen layers are what computes the answer. Freezing does not skip a layer, it only stops that layer from changing.

Second, the backward pass still runs back through every frozen layer too. LoRA places a small trainable adapter inside every layer, including the earliest ones, so the gradient message has to travel the whole way down the stack to reach the adapter in layer 1. Travelling through layer 20 means performing layer 20’s backward step, and that step needs layer 20’s stored input. Frozen means “do not update”, and it never means “do not participate”.

So a LoRA run pays roughly the same 8 GB activation bill as a full fine-tune. Its total is about 4 GB of static state plus about 8 GB of activations, so around 12 GB, which fits a 32 GB card with room to spare. The saving is real and it is entirely on the static side. The part on LoRA returns to this where it bites.