LM8-T3 — Memory and Compute for Training

Learning objectives

After this lesson, you should be able to:

• List the four terms of training memory and compute each in bytes.
• Explain how precision, LoRA and quantisation change each term.
• Estimate compute time from 6·P·D and the achievable utilisation.
• Estimate inference memory including the KV cache.

The idea: four buckets of bytes

During training the GPU holds:

• Weights: P parameters, b bytes each.
• Gradients: one number per trainable parameter, produced by the backward pass.
• Optimizer state: for Adam, two numbers (first and second moment) per trainable parameter, usually in float32.
• Activations: the intermediate results of the forward pass that the backward pass needs. They scale with batch size and sequence length, not with the number of parameters.

Mixed precision

Training in 16-bit halves weights, gradients and activations and speeds up matrix products. But small updates are lost when added to 16-bit weights, so mixed-precision training keeps a float32 master copy of the weights in the optimizer. The usual accounting is 16 bytes per parameter with Adam: 2 for the 16-bit weight, 2 for the 16-bit gradient, 4 for the float32 master weight, 4 + 4 for the two Adam moments. A 7-billion-parameter model then needs 7 · 10⁹ × 16 = 112 GB before activations. That is why full fine-tuning of such a model needs several GPUs, with the state split across them.

What PEFT and quantisation change

With LoRA, gradients and optimizer state exist only for the adapter: a fraction of a percent. The weights stay (stored, but frozen), and can be quantised to 8 or 4 bits. Activations remain, because the backward pass must still flow through the frozen layers to reach the adapters. This is why QLoRA lets a 7B model train on a 12 GB GPU, and a 65B model on a single 48 GB GPU in the QLoRA paper.

Activations and checkpointing

Activations are the term that grows with batch size and sequence length. Gradient checkpointing stores only some of them (for example one per block) and recomputes the rest during the backward pass, trading about a third more compute for a large memory saving. It is the standard way to fit longer sequences.

Inference

At inference there are no gradients or optimizer state, but there is the KV cache: the keys and values of all past tokens at all layers, kept so that each new token costs one position of work (Module 10). Its size is proportional to the sequence length and the number of concurrent requests.

The mathematics: bytes and FLOPs

With P parameters, T trainable, bytes per weight b_w:

b_w

bytes per stored weight: 4 (float32), 2 (bf16/fp16), 1 (int8), 0.5 (4-bit).

b_g

bytes per gradient: usually 2 or 4.

b_o

bytes per parameter of optimizer state: 8 for Adam moments in float32 (plus 4 for a float32 master copy in mixed precision).

M_act

activations; grows with batch × sequence length × width × layers.

Examples (decimal gigabytes, activations taken as 1.5 GB). GPT-2 small (P = 124.4 M) in float32, full fine-tuning: 0.50 + 0.50 + 1.00 + 1.5 = 3.5 GB. 7B model (P = 6.7 · 10⁹), 16-bit, full, Adam: 13.4 + 13.4 + 53.6 + 1.5 = 81.9 GB. Same model, LoRA with rank 8 on four matrices per block (T = 32 · 16 · 4,096 · 8 = 16.8 M), adapter in float32 (16 bytes per parameter), base in 4 bits: 3.35 + 0.27 + 1.5 = 5.1 GB; with a 16-bit base: 13.4 + 0.27 + 1.5 = 15.2 GB.

Compute. Training FLOPs ≈ 6 · P · D for D tokens (Lesson LM5-T4). With LoRA the weight-gradient part for frozen layers is skipped, so the backward cost falls from about 4P to about 2P per token plus the adapter terms: a training step costs about 4 · P per token instead of 6 · P, a third less. Time = FLOPs ÷ (peak FLOP/s × utilisation); realistic utilisation is 30–50% on GPUs for well-tuned code and much lower for small batches.

KV cache. For L layers, width d, sequence length n, batch B and b bytes per element (keys and values both stored):

GPT-2 small in 16-bit: 2 · 12 · 768 · 2 = 36,864 bytes per token, 37.7 MB for 1,024 tokens. A 7B-class model (L = 32, d = 4,096) in 16-bit: 524,288 bytes per token, 2.1 GB for a single 4,096-token sequence. The cache, not the weights, limits the number of concurrent long requests.

A worked example: calculators you can reuse

The code implements the training-memory formula and the KV cache formula and reproduces the numbers above.

GB = 1e9

def train_memory(P, T, bw, bg=None, opt_bytes=8, act_gb=1.5, full=False):
    """Estimated training memory in GB. For full fine-tuning T = P and gradients have the size of the weights."""
    if full:
        T, bg = P, bw
    elif bg is None:
        bg = 4
    weights = P * bw
    grads = T * bg
    opt = T * opt_bytes
    adapter = 0 if full else T * 4            # the adapter's own weights, kept in float32
    return (weights + adapter + grads + opt) / GB + act_gb

print(round(train_memory(124.4e6, 124.4e6, 4, full=True), 2))          # 3.49   GPT-2 small, float32, full
print(round(train_memory(6.7e9, 6.7e9, 2, full=True), 1))              # 81.9   7B, 16-bit, full, Adam
T7 = 32 * 16 * 4096 * 8
print(T7, round(train_memory(6.7e9, T7, 0.5), 1), round(train_memory(6.7e9, T7, 2), 1))   # 16777216 5.1 15.2

def kv_cache_bytes(layers, d, n, batch=1, bytes_per=2):
    return 2 * layers * d * n * batch * bytes_per

print(kv_cache_bytes(12, 768, 1), round(kv_cache_bytes(12, 768, 1024) / 1e6, 1))          # 36864 37.7
print(kv_cache_bytes(32, 4096, 1), round(kv_cache_bytes(32, 4096, 4096) / GB, 2))         # 524288 2.15

# training time from 6 * P * D
P, D = 6.7e9, 20e6                    # fine-tune 7B on 20 million tokens
flops = 6 * P * D
peak, util = 1e14, 0.4                # a GPU with 100 TFLOP/s at 40 % utilisation
print(f"{flops:.1e} FLOPs, {flops / (peak * util) / 3600:.1f} hours")   # 8.0e+17 FLOPs, 5.6 hours

The last line shows the other budget: even when memory fits, a 7B model fine-tuned on 20 million tokens needs 5.6 hours on a 100-TFLOP/s GPU at 40% utilisation; with LoRA, a third less arithmetic, about 3.7 hours.

Reading an out-of-memory error

Engineering implications and limitations

Estimate before you launch. A two-minute calculation avoids an hour of queueing and crashing. Keep a small script like the one above.

Gradient accumulation. Splits a batch into micro-batches processed one after another; activations scale with the micro-batch only, so you can reach any effective batch size at the memory of a small one (Lesson LM5-T3).

Gradient checkpointing. Enables much longer sequences or bigger batches for roughly a third more compute; switch it on if activations are the bottleneck.

8-bit and paged optimizers. Storing Adam’s moments in 8 bits, or paging optimizer state to CPU memory when the GPU is full, are practical ways to cut the optimizer term further.

Measure, then trust the formula. Compare the estimate with the allocator’s actual peak (torch.cuda.max_memory_allocated() on a GPU); differences are fragmentation, workspaces and temporary buffers. Keep 10–20% headroom.

Sequence length is quadratic in attention. Without fused attention kernels, the score tensor grows with n² (Lesson LM3-T6); with them, memory is linear in n.

Budget time as well. Throughput in tokens per second on your hardware, measured on a short run, converts a token budget into hours.

Energy and cost. A GPU-hour has a price and an energy cost; estimating FLOPs first keeps experiments proportionate.

Common misconceptions

Engineering takeaways

1. Training memory = weights + gradients + optimizer state + activations; with Adam in mixed precision that is about 16 bytes per trained parameter plus the stored weights.
2. 7B model: about 82 GB (16-bit, full, Adam) against about 5 GB (LoRA r=8 on a 4-bit base) with 1.5 GB for activations.
3. LoRA saves gradients and optimizer state; quantisation saves weight storage; checkpointing and accumulation save activations.
4. Compute ≈ 6·P·D FLOPs; time = FLOPs / (peak × utilisation); frozen layers cut the backward cost by about a third.
5. KV cache = 2·L·d·n·B·bytes: 36,864 bytes per token for GPT-2 small, 524,288 for a 7B-class model in 16 bits.

References and further reading

See Chapter 4 of the book (computing memory requirements) and Appendix E (memory use with LoRA). The activation allowance of 1.5 GB used in the labs is a simplification, and real activation memory depends on batch, sequence length and implementation.

• Rajbhandari, S. et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054 (the 16 bytes per parameter accounting for mixed-precision Adam).
• Dettmers, T. et al. (2023). QLoRA: Efficient Finetuning of Quantized LLMs. arXiv:2305.14314.
• Build a Large Language Model (From Scratch), Appendix E: Parameter-efficient Finetuning with LoRA (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: quantisation, which shrinks the weight term itself.

Related content

Nothing below is required to finish this lesson.

• Before
  LM8-T2: Low-Rank Adaptation (LoRA)

  Derive the low-rank update W + (alpha/r) B A and show why it needs few parameters.

• Lab
  LM8-L1: Estimate Trainable Parameters and Memory

  Write a calculator for parameters, optimizer state and memory under full fine-tuning, LoRA and quantization.