Learn AI Create the Future — with Pranay
Main site

Inside a Transformer Block

FLOPs per token

Training a transformer costs about six multiply-adds per parameter per token, which is enough to price a training run on the back of an envelope - and enough to know when the rule stops being true.

Read these first

On this page 6
  1. Where the six comes from
  2. What it lets you work out
  3. When the rule stops working
  4. The other honest caveat
  5. Remember this
  6. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Training a transformer costs roughly six units of arithmetic for every parameter, for every token it reads.

Think about a photocopy shop. The rate per page is fixed and written on the wall. You have a book with a known number of pages. Multiply the two and you know the bill before you hand the book over.

Training a model has the same shape. The size of the model sets the rate. The amount of text sets the pages. Multiply and you have the bill.

Where the six comes from

Every parameter in the model is a number that gets multiplied by something and added to a running total. That pair of operations — one multiply and one add — is the basic unit of work.

Going forward through the model, each parameter is used once. Counting the multiply and the add separately, that is two units per parameter, per token.

Coming backward, the model works out how each parameter should change. That takes about twice as much work as going forward, for reasons worth their own lesson. Four units.

Two forward plus four backward is six. That is the whole rule, and it is remarkably reliable.

What it lets you work out

Take a model with eight billion parameters, trained on fifteen trillion tokens of text. Six, times eight billion, times fifteen trillion.

The answer is a number with twenty-three zeros after the leading digit. Divide by what a good training chip delivers in a second. The answer is roughly sixty-seven thousand chip-days.

That is the real reason large models are expensive, stated in one multiplication. Not licensing, not data, not electricity in isolation. Arithmetic, at a scale that needs thousands of machines running for months.

When the rule stops working

The rule counts the work done by the weight matrices. It leaves out the work of comparing every token against every other token.

For short inputs that leftover is small, and the rule is fine. For very long inputs it stops being small. At a hundred thousand tokens, the part the rule ignores is bigger than the part it counts.

So the rule is excellent for pricing a training run on ordinary-length text. It will mislead you badly about a long-context model. Knowing which situation you are in is the actual skill.

The other honest caveat

A chip's advertised speed is what it can do in perfect conditions. Real training reaches between a third and a half of that. Time also goes on moving data, waiting for other machines, and saving checkpoints.

So any estimate should be divided by that efficiency figure. Quoting a cost against peak speed understates it by two or three times. Published estimates go wrong this way often.

Remember this

  • Roughly six units of arithmetic per parameter per token, counting forward and backward.
  • Multiply by parameters and by tokens, and you have the cost of a training run.
  • The rule ignores the token-against-token comparison, which dominates at very long inputs.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

The rule, verified against PyTorch's own counter

Rather than trusting the arithmetic, count the operations with torch.utils.flop_counter and check the prediction closes exactly.

flops.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.flop_counter import FlopCounterMode

torch.manual_seed(0)
V, d, L, d_ff, T, nh = 256, 128, 4, 512, 64, 4

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o   = nn.Linear(d, d, bias=False)
        self.f1, self.f2 = nn.Linear(d, d_ff, bias=False), nn.Linear(d_ff, d, bias=False)
    def forward(self, x):
        B, t, _ = x.shape
        q, k, v = self.qkv(self.n1(x)).chunk(3, -1)
        q, k, v = (z.view(B, t, nh, d // nh).transpose(1, 2) for z in (q, k, v))
        a = (q @ k.transpose(-2, -1) / (d // nh) ** 0.5).softmax(-1) @ v   # written out, so
        x = x + self.o(a.transpose(1, 2).reshape(B, t, d))                # the counter sees it
        return x + self.f2(F.gelu(self.f1(self.n2(x))))

model = nn.Sequential(nn.Embedding(V, d), *[Block() for _ in range(L)],
                      nn.Linear(d, V, bias=False))
with FlopCounterMode(display=False) as ctr:
    model(torch.randint(0, V, (1, T)))
measured = ctr.get_total_flops()

n_all   = sum(p.numel() for p in model.parameters())
n_embed = V * d                                    # a lookup, not a multiply
n_norm  = sum(p.numel() for n, p in model.named_parameters() if ".n1." in n or ".n2." in n)
n_mat   = n_all - n_embed - n_norm                 # only these get multiplied by every token

print(f"d={d}  layers={L}  d_ff={d_ff}  heads={nh}  T={T}")
print(f"  all parameters             {n_all:>12,}")
print(f"  embedding table (a lookup) {n_embed:>12,}")
print(f"  norm gains and biases      {n_norm:>12,}")
print(f"  matrices that do the work  {n_mat:>12,}")

weight_flops = 2 * n_mat * T
attn_flops = 4 * L * T * T * (d // nh) * nh
print(f"\nforward pass over {T} tokens")
print(f"  2 x matrices x tokens      {weight_flops:>12,}")
print(f"  attention, 4 x L x T^2 x d {attn_flops:>12,}")
print(f"  predicted total            {weight_flops + attn_flops:>12,}")
print(f"  measured by torch          {measured:>12,}")
print(f"  difference                 {measured - weight_flops - attn_flops:>12,}")

q = torch.randn(1, nh, T, d // nh)
with FlopCounterMode(display=False) as c:
    F.scaled_dot_product_attention(q, q, q, is_causal=True)
print(f"\ncareful: torch's counter reports {c.get_total_flops()} FLOPs for "
      f"scaled_dot_product_attention.")
print("the fused kernel is not instrumented, so a counted total silently omits attention.")

print("\nthe attention term the 6ND rule leaves out, for Llama 3 8B (d=4096, L=32)")
N_llama, d_, L_ = 8_030_261_248, 4096, 32
print(f"{'T':>9} {'6ND per token':>18} {'attention per token':>21} {'attention share':>17}")
for T_ in (512, 2_048, 8_192, 32_768, 131_072):
    body, attn = 6 * N_llama, 12 * L_ * T_ * d_
    print(f"{T_:>9,} {body:>18,} {attn:>21,} {attn/(body+attn):>16.1%}")
print("  -> under a few thousand tokens the rule is fine. at 128k it misses most of the bill.")

print("\nwhat 6ND costs in machine-time")
for name, N_, D_ in (("GPT-2 small", 124_439_808, 10_000_000_000),
                     ("Llama 3 8B", 8_030_261_248, 15_000_000_000_000)):
    flops = 6 * N_ * D_
    days = flops / (312e12 * 0.4) / 86400        # A100 bf16 peak 312 TFLOP/s, 40% achieved
    print(f"  {name:12s} {D_:>18,} tokens -> {flops:.3e} FLOPs ~ {days:,.0f} A100-days")
Output
d=128  layers=4  d_ff=512  heads=4  T=64
  all parameters                  854,016
  embedding table (a lookup)       32,768
  norm gains and biases             2,048
  matrices that do the work       819,200

forward pass over 64 tokens
  2 x matrices x tokens       104,857,600
  attention, 4 x L x T^2 x d    8,388,608
  predicted total             113,246,208
  measured by torch           113,246,208
  difference                            0

careful: torch's counter reports 0 FLOPs for scaled_dot_product_attention.
the fused kernel is not instrumented, so a counted total silently omits attention.

the attention term the 6ND rule leaves out, for Llama 3 8B (d=4096, L=32)
        T      6ND per token   attention per token   attention share
      512     48,181,567,488           805,306,368             1.6%
    2,048     48,181,567,488         3,221,225,472             6.3%
    8,192     48,181,567,488        12,884,901,888            21.1%
   32,768     48,181,567,488        51,539,607,552            51.7%
  131,072     48,181,567,488       206,158,430,208            81.1%
  -> under a few thousand tokens the rule is fine. at 128k it misses most of the bill.

what 6ND costs in machine-time
  GPT-2 small      10,000,000,000 tokens -> 7.466e+18 FLOPs ~ 1 A100-days
  Llama 3 8B   15,000,000,000,000 tokens -> 7.227e+23 FLOPs ~ 67,026 A100-days

The zero is the whole point

difference: 0. The prediction and PyTorch's own instrumentation agree to the last floating-point operation.

Getting that exact required being precise about three things, and each one is a trap people fall into:

The embedding table is a lookup, not a multiplication. 32,768 of the 854,016 parameters are never multiplied by anything during the forward pass. Including them in the "2N" term overstates the cost.

Normalisation parameters are not matrices either. 2,048 gains and biases. They are multiplied elementwise, not contracted, so they contribute a negligible amount and are excluded here.

Attention has to be counted separately. 4 x L x T^2 x d covers the score matrix and the weighted sum. The rule-of-thumb term does not include it.

So the precise statement is this. Forward cost is twice the parameters that appear in a matrix multiply, times the tokens, plus the attention term.

The 0 FLOPs line is a warning worth reading twice

F.scaled_dot_product_attention reported zero FLOPs to the counter. It is a fused kernel and, on this build, is not instrumented.

So a model written the sensible way, using the fast attention path, produces a counted total. That total silently omits every attention operation. At short sequences you would barely notice. At long sequences your measurement would be off by half and would look plausible.

This is why the code above writes the attention matmuls out longhand. Not because that is how you should ship it. It is the only way to make the counter tell the truth on this version. Written against PyTorch 2.5.1; check on your own build before trusting a counted total.

Where the 6 comes from, precisely

For one weight matrix in a forward pass, computing $y = Wx$ costs two operations per weight. That counts the multiply and the add.

The backward pass computes two things. One is the gradient with respect to the input, so the signal continues backwards. The other is the gradient with respect to the weight. Each is another matrix multiply of the same size. That is four more operations per weight.

Two plus four is six. Kaplan et al. (2020) state it as $C \approx 6ND$, for $N$ non-embedding parameters and $D$ tokens. It has held up as the standard estimate ever since.

Using it

QuestionCalculation
Cost of a training runsix, times parameters, times tokens
Time on a clusterthat number, divided by chips times per-chip rate times efficiency
Cost of one inference tokentwo, times parameters (forward only)
Cost of processing a prompttwo, times parameters, times prompt length

The inference row explains something people find surprising. Reading a 1,000-token prompt costs about the same as generating 1,000 tokens in raw arithmetic. Generation still feels far slower. Generation is limited by memory bandwidth, not arithmetic. Each new token requires reading the entire weight set from memory. See latency and throughput.

Common mistakes

Using total parameters where the rule wants non-embedding parameters. For GPT-2 small that is a 32 percent error. For an 8B model it is under 15 percent. Small models are where this bites.

Quoting peak chip throughput as achieved throughput. Large training runs typically report 35 to 55 percent of peak. The 40 percent used above is a reasonable middle. Anything above 60 percent should be treated as a claim, not an assumption.

Applying 6ND to a mixture-of-experts model. Use activated parameters, not total. A model with 100B total and 6B activated costs like a 6B model to train. It needs the memory of a 100B one.

Ignoring the attention term for a long-context model. At 128k tokens it is four fifths of the bill. The table above exists for exactly this.

Try it yourself

Change T to 256 in the script and re-run. The weight term scales by four; the attention term scales by sixteen. Watch the two get closer, then find the T where the attention term overtakes. Predict it first: the crossover is where sequence length reaches six times the model width.

What to learn next

Researcher — Mathematics and papers.

The estimate, stated carefully

For a dense transformer with $N$ non-embedding parameters trained on $D$ tokens:

$$ C_{\text{train}} \approx 6 N D \ \text{FLOPs}, \qquad C_{\text{fwd}} \approx 2 N D $$

The factor decomposes as $2$ forward $+\ 4$ backward. For $y = Wx$ with $W \in \mathbb{R}^{m \times n}$, the forward costs $2mn$ per token. The backward computes $\partial \mathcal{L}/\partial x = W^\top \bar{y}$ and $\partial \mathcal{L}/\partial W = \bar{y} x^\top$, each $2mn$, giving $4mn$.

The full per-token count, including the term the rule drops:

$$ C_{\text{fwd}} = 2N + 4 L T d, \qquad C_{\text{train}} = 6N + 12 L T d $$

With $N \approx 12 L d^2$, the attention term is a fraction

$$ \frac{12 L T d}{6 \cdot 12 L d^2 + 12 L T d} = \frac{T}{6d + T} $$

reaching one half at $T = 6d$. For $d = 4096$ that is 24,576 tokens — consistent with the measured table above, and a useful number to carry.

Kaplan et al. (2020), arXiv:2001.08361, Table 1, list the QK product and the attention-over-values product separately. Each contributes $2 n_{\text{layer}} n_{\text{ctx}} d_{\text{attn}}$ per token. That is the same $4LTd$, split across two lines.

What the estimate omits, and by how much

  • Elementwise operations. Softmax, GELU, normalisation and residual adds cost $O(Td)$ or $O(T^2)$ per layer. Negligible in FLOPs, and frequently dominant in wall-clock, because they are memory-bandwidth bound.
  • Activation recomputation. Gradient checkpointing repeats the forward pass during the backward, raising the effective factor from 6 toward 8. Selective recomputation of only the cheap-to-recompute parts lands between the two.
  • Embedding and output head. The head is a genuine matrix multiply of cost $2Vd$ per token. For a small model with a large vocabulary that is substantial. It is excluded from $N$ by convention and then partly reintroduced through the head. Be explicit about your convention.
  • Optimiser arithmetic. Adam's per-parameter update is $O(N)$ per step, not per token. It is negligible against $6ND$ at any realistic batch size.

Model FLOPs utilisation

The reporting standard from Chowdhery et al. (2022), PaLM, arXiv:2204.02311:

$$ \text{MFU} = \frac{6 N D / t}{P_{\text{peak}}} $$

where $t$ is wall-clock time and $P_{\text{peak}}$ the aggregate hardware peak. PaLM reported 46.2 percent MFU on 6144 TPU v4 chips. Published large-scale runs generally fall between 35 and 55 percent. Their companion metric, hardware FLOPs utilisation, includes recomputation. It is therefore higher for the same run, and the two are frequently confused.

The scaling-law consequence

Given a compute budget $C = 6ND$, the allocation between $N$ and $D$ is a free choice. Kaplan et al. (2020) concluded that most of a marginal budget should go to $N$. Hoffmann et al. (2022), Chinchilla, arXiv:2203.15556, corrected this with a much larger sweep. They found $N$ and $D$ should scale roughly in proportion. That gives approximately 20 tokens per parameter at the compute-optimal point.

Two caveats that are routinely dropped when the Chinchilla ratio is quoted:

  • It optimises training compute alone. A model that will serve many requests should be smaller and trained longer than the compute-optimal point. Inference cost scales with $N$, not with $D$. Llama 3 8B's roughly 15 trillion tokens is about 1,875 tokens per parameter. That is nearly a hundred times past compute-optimal. It is a deliberate inference-cost decision, not a mistake.
  • The fit depends on the learning-rate schedule. Hoffmann et al. note that Kaplan's estimates used a schedule that did not decay to the end of training. That biased the result toward larger models.

Mixture of experts

For a sparse model, $6ND$ takes $N$ as activated parameters per token, not total. The gap between activated and total is the entire point of the design. Training and inference compute follow the activated count. Capacity follows the total. Memory, bandwidth and routing overhead all follow the total. The arithmetic saving is real; the systems saving is not.

Papers

What to learn next