Accelerator Architecture intermediate 8 min read 7 flashcards

The FLOPs of a Transformer Forward Pass

A systematic derivation of how many floating-point operations a single transformer forward pass costs, and why that number dictates hardware choice, batch strategy, and scaling decisions.

GPT-3 costs roughly 350 petaFLOPs to train. That number is not plucked from thin air: it follows directly from counting the multiplications and additions inside a transformer forward pass, then multiplying by the training steps. If you cannot derive that count yourself, roofline analysis, hardware selection, and cost estimation all rest on foundations you cannot see. This concept gives you the derivation from first principles.

What counts as a FLOP

One FLOP is one floating-point multiply-add. Most hardware vendors count a fused multiply-add (FMA) as two FLOPs (one multiply, one add), and most FLOPs budgets in the literature follow that convention. A matrix multiply of shape [M, K] × [K, N] costs 2MKN FLOPs: for each of the MN output elements you do K multiplications and K additions.

That single formula drives almost every number in transformer analysis.

The major contributors in one transformer layer

A standard decoder-only transformer layer (pre-norm, multi-head self-attention followed by a two-layer MLP) has the following components. Let:

  • B = batch size (tokens processed in parallel)
  • T = sequence length
  • d = model hidden dimension (sometimes d_model)
  • h = number of attention heads
  • d_h = d / h = per-head dimension
  • d_ff = 4d = MLP intermediate dimension (the conventional 4x expansion)

Self-attention projections. Each token is projected to queries, keys, and values via weight matrices W_Q, W_K, W_V ∈ R^{d×d}, and the output is projected back via W_O ∈ R^{d×d}. That is four matrix multiplies, each of shape [BT, d] × [d, d], costing 2 * BT * d^2 each. Total for the four projections:

FLOPs_QKV_proj = 4 × 2BTd² = 8BTd²

Attention score computation. For each head, computing QK^T is [BT, d_h] × [d_h, BT], costing 2BT²d_h per head, so 2BT²d total across all heads. Multiplying the softmax weights by V adds another 2BT²d. Total attention scores:

FLOPs_attn_scores = 4BT²d

Note the here: this is where long sequences become expensive. At T = 2048 and d = 4096, this term is smaller than the projection term. At T = 32768, it dominates.

MLP block. Two linear layers: [BT, d] × [d, 4d] and [BT, 4d] × [4d, d], each costing 2BT × d × 4d = 8BTd². Total:

FLOPs_MLP = 16BTd²

Per-layer total. Summing and ignoring small terms (layer norm, activation, bias, softmax, which are O(BTd) and negligible at scale):

FLOPs_layer ≈ 24BTd²  +  4BT²d

For most production models T << d (e.g. T = 2048, d = 8192), so the MLP and projection terms dominate and the quadratic attention term is a modest fraction. The rule of thumb that floats around in the literature is:

FLOPs per layer ≈ 24BTd² (sequence-length-limited regimes drop this approximation)

Scaling to the full model

A transformer with L layers, ignoring embedding and unembedding layers (which are often tied and contribute ~2BTd × V FLOPs; at large vocabulary V = 32000 and large d this is non-trivial):

FLOPs_forward ≈ L × (24BTd²  +  4BT²d)

The parameter count for the non-embedding weights is approximately 12Ld² (four d×d matrices in attention + two matrices with a 4d intermediate in the MLP). So you can write:

FLOPs_forward ≈ 2 × N_params × BT     (when T << d)

where N_params = 12Ld². This is the celebrated "2N tokens" rule: a single forward pass through a model with N non-embedding parameters costs roughly 2N FLOPs per token. Training adds a backward pass, which costs about twice the forward pass, giving the further approximation that training a model on D tokens costs ~6ND FLOPs. That is the basis for the Chinchilla compute-optimal analysis.

Let us sanity-check with GPT-3: N ≈ 175B non-embedding parameters, D ≈ 300B tokens. Predicted training FLOPs: 6 × 175e9 × 300e9 ≈ 3.15 × 10^23. The OpenAI paper cites approximately 3.14 × 10^23 FLOPs, consistent to three significant figures.

A concrete worked example

GPT-2 medium: L = 24, d = 1024, d_ff = 4096. A single forward pass with B = 1, T = 512:

Component Formula FLOPs
QKV + output projections 8 × 1 × 512 × 1024² per layer × 24 ~103 GFLOPs
Attention scores 4 × 1 × 512² × 1024 per layer × 24 ~26 GFLOPs
MLP 16 × 1 × 512 × 1024² per layer × 24 ~206 GFLOPs
Total ~335 GFLOPs

An A100 delivers 312 TFLOPS (BF16 tensor core). Ignoring memory bandwidth limits, that single forward pass would take about 1 ms in pure compute time. Real latency is higher because memory reads dominate for small batches, as the roofline model predicts.

When it falls down

The 2N approximation breaks at long context. When T grows to tens of thousands of tokens, the 4BT²d term is no longer negligible. A d = 4096 model with T = 32768 has an attention score term of roughly 18 GFLOPs per layer per sequence, comparable to the MLP block. FLOPs estimates based purely on parameter count silently under-count.

Sparse and mixture-of-experts models. The derivation assumes all parameters are active for every token. In MoE models, a fraction k/E of experts fires per token (where E is the expert count). The FLOPs count scales with active parameters, not total parameters. Confusing the two is a common source of misleading capability comparisons.

Vocabulary and embedding layers. At large vocabulary (V = 128k, as in LLaMA 3) the embedding and unembedding matrices each have Vd parameters. If weights are not tied, the unembedding projection alone costs 2BTVd FLOPs, which at large V is comparable to several transformer layers. Many FLOPs estimates omit this.

The backward pass is not 2x for all ops. The ~6ND training estimate assumes the backward pass costs ~2x the forward pass. This holds for simple linear layers but is not exact for attention (softmax backward has different structure), normalisation layers, or activation checkpointing, which adds a partial second forward pass. Treat 6ND as an order-of-magnitude estimate.

Counting FLOPs is not counting wall-clock time. FLOPs measure arithmetic work, not the time taken. Memory bandwidth, kernel fusion, and tensor core utilisation all determine how quickly those FLOPs execute. A model that is twice as cheap in FLOPs can be slower in practice if its operations are memory-bound rather than compute-bound.

Further reading

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track