Will Ahmad

Linear algebra for LLMsPart 8: Gradients through matrices

Training cost

6 min

6N FLOPs per token

2N⏟forward+4N⏟backward=6N\displaystyle \underbrace{2N}_{\text{forward}} + \underbrace{4N}_{\text{backward}} = 6N

Forward: one multiply-add per weight per token (Part 2). Backward: each matmul becomes two of the same size, one for the input gradient and one for the weight gradient (two decks ago). Attention over the context adds a little more; 6N6N is the standard estimate.

Write

Training: ≈6N\approx 6N FLOPs per token.

A whole run

C≈6 N D\displaystyle C \approx 6\,N\,D
N=7×109, D=2×1012 ⁣:C=8.4×1022\displaystyle N = 7 \times 10^9,\ D = 2 \times 10^{12}\!: \quad C = 8.4 \times 10^{22}

DD is the number of training tokens. A 7B model on 2 trillion tokens takes about 8.4×10228.4 \times 10^{22} FLOPs.

Write

Training compute C≈6NDC \approx 6ND.

Into GPU-hours

8.4×10224×1014 FLOP/s=2.1×108 GPU-seconds\displaystyle \frac{8.4 \times 10^{22}}{4 \times 10^{14} \text{ FLOP/s}} = 2.1 \times 10^8 \text{ GPU-seconds}
≈58,000 GPU-hours≈57 hours on 1,024 GPUs\displaystyle \approx 58{,}000 \text{ GPU-hours} \approx 57 \text{ hours on } 1{,}024 \text{ GPUs}

Assume each GPU sustains 400 TFLOP/s of useful work, a plausible fraction of a modern accelerator’s BF16 peak. Real runs vary a lot with utilization, so treat this as an order of magnitude.

Write

Time ≈ C / (GPUs×sustained FLOP/s)C \,/\, (\text{GPUs} \times \text{sustained FLOP/s}).

Memory per parameter

What Bytes per parameter
BF16 weights 2
BF16 gradients 2
FP32 master weights 4
Adam’s two FP32 moments 8
total 16

Standard mixed-precision training with Adam keeps about 16 bytes per parameter, before activations: 7B parameters need about 112 GB just for this. That’s why training is spread across many GPUs.

Write

Mixed-precision Adam: ≈ 16 bytes/param (+ activations).

✏️ Your turn

Size a run

N=109 parameters,D=1011 tokens\displaystyle N = 10^9 \text{ parameters}, \quad D = 10^{11} \text{ tokens}

How many training FLOPs? How much memory for weights, gradients and Adam state?

Work it out on paper, then go to the next card to check.

Answer

C=6×109×1011=6×1020\displaystyle C = 6 \times 10^9 \times 10^{11} = 6 \times 10^{20}
16 bytes×109=16 GB\displaystyle 16 \text{ bytes} \times 10^9 = 16 \text{ GB}

Inference is cheaper

In a transformer

Generating a token takes only the forward pass, 2N2N FLOPs, and no gradients or optimizer state: weights (and the KV cache) are all that must sit in memory. Part 9 shows why decoding is limited by reading those weights, not by the FLOPs.

📌 Recap

Training cost in one card

  • ≈6N\approx 6N FLOPs per training token; C≈6NDC \approx 6ND for a run.
  • Time ≈C\approx C / (GPUs × sustained FLOP/s).
  • Mixed-precision Adam: about 16 bytes per parameter before activations.

Write

The three lines above.