GPU memory maths for LLMs and SLMs

Learn in your browser · Chapter 01 of 06

Every fit check starts with the weights, the number you can compute exactly.

A model is a long list of numbers called parameters, each stored in a set number of bytes chosen by its precision. The rule never changes: weights in bytes = parameters x bytes per parameter. For a mixture-of-experts model (one split into many expert blocks, of which only a few run for each token), count every expert.

Here 1 GB is 1,000,000,000 bytes. Some tools report GiB (1,073,741,824 bytes, about 7.4% more), so check the unit. Card sizes here are examples only: read the vendor's specification page for yours.

Worked examples use hypothetical models invented for teaching. Model A has 7 billion parameters, 36 layers, 32 query heads, 4 key-value heads and head dimension 128.

Scroll sideways to see every column.

At 16-bit: 7 billion x 2 = 14 billion bytes = 14.00 GB. Real 4-bit files are a little larger than the simple sum, because each block of weights also stores a scale, so trust the file size on the download page.

Learn in your browser · Chapter 02 of 06

While it generates, a model stores two vectors per layer for every earlier token. That store can rival the weights.

These are keys and values, so the store is the key-value cache (KV cache).

kv_per_token = 2 x layers x kv_heads x head_dim x bytes_per_value
kv_cache     = kv_per_token x tokens_in_context x simultaneous_chats

The 2 is for keys and values. Use key-value heads, not attention heads: grouped-query attention shares each key-value head across several query heads, and multi-query attention shares one across all. The paged-attention paper's example, 2 x 5,120 (hidden size) x 40 (layers) x 2 (bytes) per token, is this sum with full multi-head attention.

Worked example: Model A, 16-bit values, 16,384 tokens.

per token = 2 x 36 x 4 x 128 x 2 = 73,728 bytes
one chat  = 73,728 x 16,384 = 1,207,959,552 bytes = 1.21 GB

With full multi-head attention (32 key-value heads), each token costs 2 x 36 x 32 x 128 x 2 = 589,824 bytes, and 589,824 x 16,384 = 9.66 GB: eight times the cache, from one line of a model card.

Learn in your browser · Chapter 03 of 06

Weights and cache are exact. The rest is not, so budget an allowance, then measure.

Activations are the temporary tensors each layer produces. In inference they need not be kept for a training step, so they are usually small, but a long prompt needs more scratch space than one generation step. Framework overhead adds the runtime's own memory and memory the allocator cannot reuse. The ZeRO paper reports out-of-memory failures with over 30% of memory still free, in some extreme cases, when training very large models.

None of this has a published constant, so this workbook assumes an allowance of the larger of 1 GB or 10% of everything else counted. It is our planning assumption, not a published rule: measure and replace it. It also assumes a separate safety margin, keeping 10% of capacity free for other programs and for errors in your inputs.

fit rule:  weights + cache + allowance  <=  0.9 x capacity

Worked example: Model A, 16-bit weights, one 16,384-token chat, on a card with 16 GB (an example size). Allowance = max(1.00, 10% x (14.00 + 1.21)) = 1.52 GB. Need = 14.00 + 1.21 + 1.52 = 16.73 GB. Usable = 0.9 x 16.00 = 14.40 GB. It is short by 2.33 GB.

Scroll sideways to see every column.

Cutting the context to 4,096 tokens leaves the cache at 0.30 GB but the 16-bit total at 14.00 + 0.30 + 1.43 = 15.73 GB. The weights dominate: change the biggest term first.

Learn in your browser · Chapter 04 of 06

Inference only reads the weights. Training changes them, and the model states alone need eight times the size of a 16-bit model file.

In a training step, the forward pass keeps activations, the backward pass uses them to compute a gradient for every trainable parameter, and the optimiser updates the weights using the gradients and its own state. An Adam-style optimiser keeps a running average of the gradients (momentum) and of their squares (variance). Mixed precision runs both passes in 16-bit but keeps a 32-bit master copy of the weights.

Scroll sideways to see every column.

The ZeRO paper writes this as 2P + 2P + KP = 16P bytes, with P the parameter count and K = 12 for mixed-precision Adam. Its example: 1.5 billion x 16 = 24 billion bytes, which it calls at least 24 GB.

Activations are kept for the backward pass and scale with batch size, sequence length and depth. One published estimate for a plain 16-bit transformer layer is s x b x h x (34 + 5 x a x s / h) bytes, where s, b, h and a are sequence length, micro-batch, hidden size and attention heads. The ZeRO paper's rule is roughly 12 x h x b x s x layers values, or 24 x h x b x s x layers bytes at 16-bit. Here they differ mostly because the first counts attention scores, which some implementations avoid storing.

Worked example: hypothetical Model B has 1 billion parameters, 24 layers, hidden size 2,048, 16 attention heads. Train it at sequence length 1,024, micro-batch 2.

states    = 1 billion x 16 = 16.00 GB
s x b x h = 1,024 x 2 x 2,048 = 4,194,304
per layer = 4,194,304 x (34 + 5 x 16 x 1,024 / 2,048)
          = 4,194,304 x 74 = 310,378,496 bytes = 0.31 GB
24 layers = 7,449,083,904 bytes = 7.45 GB

The second rule gives 2.42 GB here: plan for a range, then measure.

Activation checkpointing keeps only some activations and recomputes the rest in the backward pass. The original paper's method needs memory growing with the square root of the layer count, for one extra forward pass. A backward pass costs about twice a forward pass, so that adds about a third more compute, matching the ZeRO paper's 33%. One simple scheme stores each layer's input, 2 x s x b x h bytes, plus one layer's full activations.

inputs      = 2 x 4,194,304 x 24 = 0.20 GB
activations = 0.20 + one layer 0.31 = 0.51 GB   (was 7.45 GB)

With the allowance: 16.00 + 7.45 + 2.34 = 25.79 GB without checkpointing, and 16.00 + 0.51 + 1.65 = 18.16 GB with it. The states are untouched, so neither fits the 16 GB card.

Learn in your browser · Chapter 05 of 06

Full fine-tuning pays 16 bytes for every parameter. Adapters train a tiny fraction and skip most of that bill.

LoRA freezes a pretrained matrix W and adds a small trainable correction, B x A, where B is d x r, A is r x k and the rank r is far smaller than d and k. Frozen weights need no gradients or optimiser state, which the paper credits for its saving. The adapters pay the full 16 bytes each, but there are only r x (d + k) per matrix.

Worked example: Model B, rank 16 adapters on the query and value matrices of every layer, both assumed square at 2,048 x 2,048. Use true shapes for real models.

per matrix     = 16 x (2,048 + 2,048) = 65,536 parameters
matrices       = 2 x 24 = 48
adapters       = 48 x 65,536 = 3,145,728 parameters (0.31% of 1 billion)
adapter states = 3,145,728 x 16 = 0.05 GB
frozen base    = 1 billion x 2 = 2.00 GB

Scroll sideways to see every column.

Adapters cut the states from 16.00 GB to 2.05 GB but not most of the activations, because gradients still flow back through the frozen layers, so the table keeps the full estimate as a cautious plan. Storing the frozen base at 4-bit too is the QLoRA paper's approach. The LoRA paper reports results on par with full fine-tuning on its tasks, and no added inference delay.

Learn in your browser · Chapter 06 of 06

Turn the method into a repeatable procedure, then be straight about what one GPU can do.

1. Note the device's free memory and unit, and the model card's parameters, layers, key-value heads, head dimension and context.
2. Choose a precision: weights, then cache (for training, states and activations), then the allowance. Compare with 0.9 x capacity, or with the free memory if that is lower.
3. Load it, read the peak memory your framework reports, and replace the allowance with that.

Save the calculator below as memcalc.py, edit the numbers at the bottom, and run python memcalc.py. It should print the two lines that follow it.

# memcalc.py: rough GPU fit check
GB = 1e9  # decimal GB

def weights(params, bits):
    return params * bits / 8

def kv_cache(layers, kv_heads, head_dim, tokens, chats=1, bits=16):
    return 2 * layers * kv_heads * head_dim * tokens * chats * bits / 8

def activations(s, b, h, a, layers, checkpoint=False):
    layer = s * b * h * (34 + 5 * a * s / h)  # plain attention, 16-bit
    return (2 * s * b * h * layers + layer) if checkpoint else layer * layers

def need(parts):
    base = sum(parts)
    return base + max(1 * GB, 0.10 * base)  # base plus the allowance

if __name__ == "__main__":
    w = weights(7e9, 16)
    kv = kv_cache(36, 4, 128, 16384)
    usable = 0.9 * 16 * GB
    print(f"weights {w/GB:.2f} kv {kv/GB:.2f} need {need([w, kv])/GB:.2f} usable {usable/GB:.2f}")
    act = activations(1024, 2, 2048, 16, 24, checkpoint=True)
    print(f"states {1e9 * 16 / GB:.2f} activations {act/GB:.2f}")

weights 14.00 kv 1.21 need 16.73 usable 14.40
states 16.00 activations 0.51

Scroll sideways to see every column.

The honest scope: a laptop or one GPU can run models that fit its memory, fine-tune small ones (especially with adapters) and train toy ones from scratch. It cannot pretrain a production frontier-scale model, and memory is only one reason. 7 billion parameters need 112.00 GB of states before activations, and splitting them across devices divides that by about the device count, at a cost in communication.

Compute time. The scaling-laws paper estimates about 6 operations per parameter per training token. The compute-optimal paper found that parameters and tokens should grow together: its model used 1.4 trillion tokens for 70 billion parameters, 20 per parameter, a rough guide. The table assumes a made-up sustained 100 trillion operations per second. Replace it with a rate you have measured.

Scroll sideways to see every column.

For 7 billion (7e9 is shorthand for 7 followed by nine zeros): tokens = 20 x 7e9 = 1.4e11. Operations = 6 x 7e9 x 1.4e11 = 5.88e21. Seconds = 5.88e21 / 1e14 = 5.88e7, and 5.88e7 / 86,400 = 680.56 days.

Even at ten times that rate, the 7-billion run takes 68.1 days and the 70-billion run 18.6 years. Then come data to collect, clean and license, and engineering: keeping many devices running for weeks and surviving failures.