Chapter 5 · Section 5.6 · 55 min

One Transformer block, counted

You now own every part of a Transformer block. The residual stream carries each token up through the model (§5.1), norms keep what each sublayer reads at a fixed scale (§5.2), the feed-forward network transforms each token on its own (§5.3), and RoPE or a learned table tells positions apart (§5.4, §5.5). This section bolts them together and asks the questions an engineer asks of any block before building a model from it. How many parameters does it hold? How many floating-point operations does one token cost, forward and backward? Which of those numbers change when you swap GPT-2's recipe for LLaMA's?

Here is the kind of answer you should be able to give on a whiteboard. One block of GPT-2 small holds exactly 7,087,872 parameters. Pushing one token through it costs about 14.2 million FLOPs in its weight matrices plus 3.1 million in attention over a 1,024-token context (counting the full score matrix), and training on that token costs three times as much. A whole model is $L$ copies of the block plus a fixed overhead for the vocabulary, so these per-block numbers are how you budget a run, sanity-check a paper, and catch bugs. A block whose count is off by 1,536 has lost a LayerNorm. Interviewers ask these questions because they separate people who have built models from people who have read about them.

Assembling a decoder block

A GPT-2 block takes the stream $\mathbf{X} \in \R^{T\times d}$ for $T$ tokens of width $d$ and applies two pre-LN residual sublayers (§5.2),

\mathbf{H} = \mathbf{X} + \operatorname{MHA}\big(\operatorname{LN}_1(\mathbf{X})\big), \qquad \mathbf{Y} = \mathbf{H} + \operatorname{FFN}\big(\operatorname{LN}_2(\mathbf{H})\big),

where $\operatorname{MHA}$ is causal multi-head attention with $h$ heads of size $d_k = d/h$ (§4.5, §4.4) and $\operatorname{FFN}$ is the GELU MLP of §5.3 with $d_{\text{ff}} = 4d$ . Every shape along the way is worth writing out once.

step	operation	shape
input	stream $\mathbf{X}$	$T \times d$
norm	$\mathbf{U} = \operatorname{LN}_1(\mathbf{X})$	$T\times d$
projections	$[\mathbf{Q}, \mathbf{K}, \mathbf{V}] = \mathbf{U}\mathbf{W}_{QKV} + \mathbf{b}_{QKV}$ with $\mathbf{W}_{QKV} \in \R^{d\times 3d}$	$T \times 3d$
split heads	reshape each to $h \times T \times d_k$	$h\times T\times d_k$
scores	$\mathbf{Q}\mathbf{K}^\top/\sqrt{d_k}$ plus causal mask, softmax	$h\times T\times T$
mix	weights times $\mathbf{V}$ , merge heads	$T \times d$
write	$\mathbf{H} = \mathbf{X} + (\cdot)\,\mathbf{W}_O + \mathbf{b}_O$	$T\times d$
FFN	$\mathbf{Y} = \mathbf{H} + \operatorname{GELU}(\operatorname{LN}_2(\mathbf{H})\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2$	$T\times 4d \to T\times d$

A model stacks $L$ such blocks between an embedding and an unembedding. Token and position embeddings form the first stream $\mathbf{X}^{(0)}$ , the blocks run in sequence, a final norm $\operatorname{LN}_f$ normalizes the last stream, and the unembedding $\mathbf{W}_U \in \R^{d\times V}$ turns it into one row of next-token logits per position,

\mathbf{X}^{(\ell)} = \operatorname{Block}_\ell\big(\mathbf{X}^{(\ell-1)}\big), \qquad \text{logits} = \operatorname{LN}_f\big(\mathbf{X}^{(L)}\big)\,\mathbf{W}_U \in \R^{T\times V}.

Here the superscript counts blocks (unlike §5.1, where it counted sublayers), and $V$ is the vocabulary size. GPT-2 reuses its token table as the unembedding, $\mathbf{W}_U = \mathbf{E}^\top$ , which is called weight tying (Press & Wolf, 2017). Building that whole model end to end, embeddings, head, and every last parameter, is the job of §6.2. This section stays with the block, because everything else is a fixed overhead on top of $L$ copies of it.

The LLaMA recipe changes five things and keeps the skeleton. LayerNorm becomes RMSNorm (§5.2), learned positions are replaced by RoPE inside attention (§5.5), the GELU MLP becomes SwiGLU with $d_{\text{ff}} \approx 8d/3$ (§5.3), every bias is removed, and the output matrix is a separate parameter rather than tied. Larger LLaMA models also share keys and values across groups of query heads, which is called grouped-query attention (Ainslie et al., 2023, and §12.5), and which matters for counting.

Three families of Transformer

The block is the same in every family. What differs is the attention mask and how the stacks are wired.

family	attention	trained to	strengths	examples
encoder-only	bidirectional, every token sees every token	fill in masked tokens	representations for classification, retrieval	BERT, RoBERTa (§7.1)
decoder-only	causal, each token sees its prefix	predict the next token	generation, in-context learning, one stack for everything	GPT, LLaMA, Mistral
encoder-decoder	bidirectional encoder, causal decoder with cross-attention to it	map an input sequence to an output sequence	translation, summarization, explicit input/output split	original Transformer, T5, Whisper

Decoder-only models now dominate because one causal stack covers every task once tasks are phrased as text, because every token of every document provides a training signal, and because generation with a KV cache (§6.5) is simple and efficient. Section 8.1 compares the training objectives behind the three families. An encoder-decoder's decoder blocks contain a third sublayer, cross-attention (§4.6), with its own four projections and its own norm, which shows up in the count below.

Counting one block exactly

Count one GPT-2 block, piece by piece.

Attention. $\mathbf{W}_{QKV}$ is $d \times 3d$ with a $3d$ bias, and $\mathbf{W}_O$ is $d\times d$ with a $d$ bias, so $4d^2 + 4d$ .
FFN. $\mathbf{W}_1$ is $d\times 4d$ with a $4d$ bias, and $\mathbf{W}_2$ is $4d \times d$ with a $d$ bias, so $8d^2 + 5d$ .
Two LayerNorms. A gain and a shift each, so $4d$ .

Writing $N_{\text{block}}$ for the parameters of one block, the total is

N_{\text{block}} = 12d^2 + 13d .

In plain words, each block costs twelve $d\times d$ matrices' worth of weights, four in attention and eight in the FFN, plus a few vectors of size $d$ that hardly matter.

GPT-2 small, by hand. With $d = 768$ , one block has $12 \cdot 768^2 + 13 \cdot 768 = 7{,}077{,}888 + 9{,}984 = 7{,}087{,}872$ parameters, and the 12 blocks hold $85{,}054{,}464$ . The rule of thumb $12Ld^2 = 84{,}934{,}656$ is within 0.15% of that. Yet the released checkpoint has about 124 million parameters. The other 39 million live outside the blocks, and §6.2 finds every one of them.

LLaMA-2 7B, by hand. Now there are no biases, the norms are RMSNorm with a gain only ( $d$ each), and the FFN is SwiGLU with three matrices of $d\times d_{\text{ff}}$ . With $d = 4096$ and $d_{\text{ff}} = 11{,}008$ , one block has $4 \cdot 4096^2 = 67{,}108{,}864$ attention weights, $3\cdot 4096\cdot 11{,}008 = 135{,}266{,}304$ FFN weights, and $2\cdot4096 = 8{,}192$ norm weights, a total of $202{,}383{,}360$ . The 32 blocks hold $6{,}476{,}267{,}520$ , almost all of the model's 6.7 billion. Because $d_{\text{ff}}$ is $8d/3$ rounded up, the block is $4d^2 + 3d\cdot\tfrac{8d}{3} \approx 12d^2$ again (12.06 $d^2$ after the rounding). The $12d^2$ rule survives the change of recipe.

Grouped-query attention breaks it. With $h$ query heads but only $n_{kv}$ key-value heads, $\mathbf{W}_K$ and $\mathbf{W}_V$ shrink to $d \times n_{kv}d_k$ , so attention costs $2d^2 + 2d\,n_{kv}d_k$ . LLaMA-2 70B ( $d = 8192$ , $h = 64$ , $n_{kv} = 8$ , $d_k = 128$ ) has $n_{kv}d_k = d/8$ , so its attention is $2.25d^2$ instead of $4d^2$ . At the same time its FFN is wider than the 8/3 rule ( $d_{\text{ff}} = 28{,}672 = 3.5d$ ), so the FFN is $3 \cdot 3.5 = 10.5\,d^2$ . The block comes to $12.75\,d^2$ plus $2d$ of norm gains, $855{,}654{,}400$ parameters, and more than four fifths of it is FFN. LLaMA 3 8B makes the same trade at $d = 4096$ ( $n_{kv} = 8$ of 32 heads, so attention is $2.5d^2$ , and $d_{\text{ff}} = 3.5d$ ) and lands at $13d^2$ per block, plus its two norm gains.

The encoder-decoder block. A decoder block with cross-attention has a second set of four attention projections and a third norm. With $d_{\text{ff}} = 4d$ and biases it holds $16d^2 + 19d$ , so the original Transformer's base model ( $d = 512$ ) has 3,152,384 parameters per encoder block and 4,204,032 per decoder block.

The cell below prints the split for real configurations.

Python · cell 1

Interview

Count one block exactly, for any recipe

Write block_params(cfg) returning the exact number of parameters in one Transformer block, for a config dictionary with these keys.

d_model, n_head, d_ff
n_kv_head (default n_head), the number of key-value heads $n_{kv}$ for grouped-query attention, each of width $d_k$ = d_model // n_head
bias (bool), biases in every linear layer and, for LayerNorm, the shift
norm, "layernorm" (gain, plus shift when bias is true) or "rmsnorm" (gain only)
ffn, "mlp" (two matrices) or "swiglu" (three matrices)
cross_attn (default False), whether the block also has a cross-attention sublayer with its own projections and its own norm, as in an encoder-decoder's decoder

A block has one norm per sublayer. The tests check GPT-2 small and XL, LLaMA-2 7B and 70B, LLaMA 3 8B, and a decoder block of the original Transformer, all exactly.

def block_params(cfg):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

FLOPs per token and the 6N rule

The arithmetic cost of a Transformer is dominated by matrix multiplications, and those are easy to count. Multiplying $\mathbf{A} \in \R^{m\times k}$ by $\mathbf{B} \in \R^{k\times n}$ takes $mnk$ multiply-adds, or $2mnk$ floating-point operations (FLOPs), the convention of §4.7.

Forward, one block. A linear layer $\mathbf{Y} = \mathbf{X}\mathbf{W}$ with $\mathbf{X} \in \R^{T\times a}$ and $\mathbf{W}\in\R^{a\times b}$ costs $2Tab$ FLOPs, which is $2ab$ per token, two FLOPs per weight per token. A GPT-2 block's matrices hold $12d^2$ weights, so they cost $24d^2$ FLOPs per token, a third in attention's projections and two thirds in the FFN. The two attention products with no weights, the scores $\mathbf{q}\mathbf{K}^\top$ and the weighted sum of values, cost $2Td_k$ each per head for a token that attends to $T$ positions, so $4Td$ per token across all heads (§4.7 counted them product by product). Biases, norms, activations, and the softmax add only $O(d)$ or $O(T)$ per token, which we ignore. For GPT-2 small at $T = 1024$ that is $24 \cdot 768^2 \approx 14.2$ million FLOPs in the weights and $4 \cdot 1024 \cdot 768 \approx 3.1$ million in attention, per token and per block.

The whole forward pass. Summing over the blocks and adding the unembedding, a $T\times d$ by $d\times V$ product that costs $2Vd$ per token whether or not it is tied, the forward pass costs about $2N_{\text{mm}}$ FLOPs per token plus attention. Here $N_{\text{mm}}$ counts the weights that take part in matrix multiplications, every block matrix plus the unembedding. Embedding lookups cost no multiplications, so the input tables are not in $N_{\text{mm}}$ .

Backward. Given $\partial\mathcal{L}/\partial\mathbf{Y}$ , a linear layer needs two products (§0.2), one to pass the gradient down and one for its own weights,

\frac{\partial\mathcal{L}}{\partial\mathbf{X}} = \frac{\partial\mathcal{L}}{\partial\mathbf{Y}}\,\mathbf{W}^\top \quad (2Tab \text{ FLOPs}), \qquad \frac{\partial\mathcal{L}}{\partial\mathbf{W}} = \mathbf{X}^\top\frac{\partial\mathcal{L}}{\partial\mathbf{Y}} \quad (2Tab \text{ FLOPs}).

So the backward pass costs twice the forward, $4N_{\text{mm}}$ per token, and training costs

C_{\text{token}} \approx 6N, \qquad C_{\text{train}} \approx 6\,N\,D

for $D$ training tokens, where $N$ means $N_{\text{mm}}$ , which for large models is close to the total parameter count.

In plain words, every weight is used once going forward (a multiply and an add) and twice going backward (once to pass the gradient on, once to compute its own gradient), so training costs six operations per weight per token.

The attention term. The weightless products follow the same rule. Their backward pass is again two products of the same size (§4.3), so training triples the forward $4Td$ per layer, and the full per-token cost is

C_{\text{token}} \approx 6N_{\text{mm}} + 12\,L\,T\,d .

This counts the whole $T\times T$ score matrix. A causal kernel that skips the masked half does about half of that work, since on average each token sees about $T/2$ keys. Papers do not agree on this term, so check which convention a number uses. Kaplan et al. (2020) list a forward term of $2LTd$ in their Table 1, half of our full $4LTd$ , which §4.7 reads as the causal count, while PaLM's utilization accounting (Chowdhery et al., 2022) uses the full $12LTd$ for training.

How big is the attention term? With $N_{\text{mm}} \approx 12Ld^2$ , its ratio to the weight term is $12LTd/(6\cdot12Ld^2) = T/(6d)$ . For GPT-2 small at $T = 1024$ that ratio is 0.22, and for LLaMA-2 7B at $T = 4096$ it is 0.17. Counted against everything, unembedding included, attention is 13% and 14% of the training FLOPs of these two models, as the cell below prints. The term grows with context length, which is why long-context training (§12.7) costs more per token than short-context training.

Python · cell 2

A worked example worth memorizing. GPT-3 has 175 billion parameters and was trained on 300 billion tokens, so $6ND \approx 3.15\times10^{23}$ FLOPs, and the GPT-3 paper's own table reports $3.14\times10^{23}$ (its count uses 174.6 billion). On hardware that sustains $4\times10^{14}$ FLOP/s per GPU, that is $7.9\times10^{8}$ GPU-seconds, about 9,100 GPU-days. How close a real run comes to its hardware's peak, its model FLOPs utilization, is defined and measured in §6.7. The 6N rule is also the bridge to scaling laws, which ask how to split a compute budget $C = 6ND$ between parameters and data (§9.3).

Two refinements come up in practice. If a run recomputes activations during the backward pass to save memory (§10.1), the hardware executes a second forward pass, about $8N$ FLOPs per token, but the model's own cost is still $6N$ , and utilization is quoted against the $6N$ (§11.2). And for a mixture-of-experts model, $N$ in the rule is the number of parameters active per token, not the total (Chapter 13).

Inference runs only the forward pass, about $2N$ FLOPs per generated token plus $4LTd$ for attention over $T$ earlier tokens. But during one-token-at-a-time decoding, each weight is read from memory to do just two FLOPs, so decoding is limited by memory bandwidth, not arithmetic. Caching past keys and values avoids recomputing the past (§6.5), and serving many requests at once lets each weight read do work for many tokens, the central idea of batching and memory management in serving.

Core

Training FLOPs per token, with the attention term

Write train_flops_per_token(cfg, T) returning $6N_{\text{mm}} + 12LTd$ , where $N_{\text{mm}}$ counts the weights in matrix multiplications. The config uses the block keys of Exercise 5.6f plus vocab_size and n_layer (any other keys are ignored). $N_{\text{mm}}$ is the block matrices of all n_layer blocks (respecting n_kv_head and ffn) plus one $V\times d$ unembedding, whether or not it is tied, and no biases, norms, or position tables.

def train_flops_per_token(cfg, T):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Memory per token, in brief

Two memory costs grow with tokens, and each has its own section later. At inference, every layer keeps the keys and values of every earlier token, $2Ln_{kv}d_k$ numbers per token. Section 4.7 already found 128 KiB per token for Llama 3 8B in 16-bit precision, and §6.5 builds the cache and sizes it. In training, each block stores the activations its backward pass needs, about $34d + 5hT$ bytes per token in GPT-3-style accounting (Korthikanti et al., 2022), where the $5hT$ is the attention probabilities of the $h$ heads and their dropout mask and output. For a GPT-2 small block at $T = 1024$ that is 87,552 bytes per token, 70% of it attention matrices, which is why FlashAttention never stores them. Section 10.1 derives the 34 tensor by tensor and adds the 16 bytes per parameter that weights, gradients, and Adam's moments cost.

A block, from scratch in NumPy

Here is a GPT-2 block in about a dozen lines, with random weights at a toy size, stacked four deep. The prints check the two properties that matter most, that the parameter count matches $12d^2 + 13d$ and that the stack is causal (changing a later token never changes an earlier position's output).

Python · cell 3

Interview

A GPT-2 decoder block, from scratch

Write gpt2_block(X, p, H) for one token sequence X of shape (T, d), a dictionary of weights p with the keys used in the cell above (ln1_g, ln1_b, W_qkv, b_qkv, W_o, b_o, ln2_g, ln2_b, W_fc, b_fc, W_proj, b_proj), and H heads (the $h$ of the text). It is a pre-LN block with causal multi-head attention and a tanh-GELU MLP. W_qkv is laid out as in GPT-2, with the query, key, and value projections side by side in that order, and head i owns columns i*d_k to (i+1)*d_k of each. No loops over heads or positions.

This is the canonical interview implementation. The tests check it against a reference, check causality, check that a block with zeroed output projections is the identity, and catch the usual slips (the scale, the split order, the norm placement).

import numpy as np

def gpt2_block(X, p, H):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

A stack of these blocks is the body of GPT-2. Section 6.2 adds the token and position embeddings in front, the final norm and the tied head behind, initializes the whole model, and checks it against a reference implementation.

In PyTorch

The same block in PyTorch. Built on the meta device, a GPT-2 small block is counted without allocating any memory. Then a toy-size block runs under PyTorch's FLOP counter, which confirms the forward cost $24d^2T + 4T^2d$ exactly for one sequence of $T$ tokens.

🔥 In PyTorchfor a machine with PyTorch

🎯 Interview questions on this section

Derive the parameter count of a GPT-2 block.

Attention has $\mathbf{W}_{QKV}$ ( $3d^2 + 3d$ ) and $\mathbf{W}_O$ ( $d^2 + d$ ), the FFN has $d\times4d$ and $4d\times d$ matrices with biases ( $8d^2 + 5d$ ), and two LayerNorms add $4d$ , so $12d^2 + 13d$ per block. For $d = 768$ that is 7,087,872, and twelve blocks hold 85,054,464, within 0.15% of $12Ld^2$ . Attention is a third of it and the FFN two thirds. The model's remaining parameters, about 39 million in GPT-2 small, sit outside the blocks in the embeddings.

How do SwiGLU and grouped-query attention change the count of a block?

SwiGLU has three $d\times d_{\text{ff}}$ matrices, so at $d_{\text{ff}} = 8d/3$ it costs $8d^2$ , the same as a GELU MLP at $4d$ , and a LLaMA-2 7B block is $12.06d^2$ after rounding $d_{\text{ff}}$ up to 11,008. Grouped-query attention shrinks $\mathbf{W}_K$ and $\mathbf{W}_V$ to $d\times n_{kv}d_k$ , so attention is $2d^2 + 2d\,n_{kv}d_k$ , which is $2.25d^2$ at $n_{kv} = h/8$ . Recent models spend the savings on a wider FFN. LLaMA-2 70B uses $d_{\text{ff}} = 3.5d$ , so its block is $2.25d^2 + 10.5d^2 = 12.75d^2$ (855,654,400 parameters), with more than 80% in the FFN.

Where does the 6N rule come from, and what does it leave out?

Each weight does one multiply-add per token in the forward pass, 2 FLOPs. In the backward pass each linear layer computes $\partial\mathcal{L}/\partial\mathbf{X} = \partial\mathcal{L}/\partial\mathbf{Y}\,\mathbf{W}^\top$ and $\partial\mathcal{L}/\partial\mathbf{W} = \mathbf{X}^\top\partial\mathcal{L}/\partial\mathbf{Y}$ , each as costly as the forward matmul, so 4 more. Hence $6N$ FLOPs per token and $6ND$ for $D$ tokens, with $N$ the weights in matmuls (blocks plus unembedding, not the embedding lookup). It ignores the weightless attention matmuls, $12LTd$ per token for training at context $T$ , plus norms, softmax, and activations, which are lower order. Relative to $6N$ the attention term is about $T/(6d)$ , so it matters for long contexts and small widths. Activation recomputation makes the hardware do about $8N$ , but the model's FLOPs stay $6N$ .

Per token, how is a block's forward compute split, and at what context does attention dominate?

The weight matrices cost 2 FLOPs per weight, $24d^2$ for a GPT-2 block, a third in the attention projections ( $8d^2$ ) and two thirds in the FFN ( $16d^2$ ). The score and value products add $4Td$ if the full $T\times T$ matrix is computed, or about $2Td$ with a causal kernel. They equal the weight products at $T = 6d$ in the full count and $T = 12d$ causally, so 9,216 tokens for GPT-2 small with a causal kernel. Below that, the FFN is the biggest single cost; well above it, attention is.

How much compute did GPT-3 take, and how long would that be on modern hardware?

$6ND = 6\times175\times10^9\times300\times10^9 \approx 3.15\times10^{23}$ FLOPs, which matches the $3.14\times10^{23}$ in the paper. At a sustained $4\times10^{14}$ FLOP/s per GPU (roughly 40% utilization of a modern accelerator's dense bf16 peak), that is about $7.9\times10^8$ GPU-seconds, around 9,000 GPU-days, or about nine days on a thousand GPUs.

Encoder-only, decoder-only, or encoder-decoder, which would you choose for a new general-purpose model, and why?

Decoder-only, for most purposes. A single causal stack handles every task phrased as text, every token gives a training signal, inference with a KV cache is simple, and the whole ecosystem (kernels, serving, fine-tuning recipes) is built around it. Encoder-only models remain the right tool for cheap, high-quality embeddings and classification, where bidirectional context helps and generation is not needed. Encoder-decoders make sense when there is a clear input-output split with a long input, such as speech recognition (Whisper) or some translation systems.

🧭 What to remember

A pre-LN decoder block is $\mathbf{H} = \mathbf{X} + \operatorname{MHA}(\operatorname{LN}(\mathbf{X}))$ , $\mathbf{Y} = \mathbf{H} + \operatorname{FFN}(\operatorname{LN}(\mathbf{H}))$ . A model is $L$ of them between an embedding and an unembedding.
A GPT-2 block has $12d^2 + 13d$ parameters, a third in attention and two thirds in the FFN. SwiGLU at $d_{\text{ff}} = 8d/3$ keeps the $12d^2$ , while grouped-query attention plus a wider FFN moves recent blocks to about $13d^2$ , mostly FFN.
A block's forward pass costs 2 FLOPs per weight per token plus $4Td$ for attention. Training costs $6N$ FLOPs per token (2 forward, 4 backward per weight) plus $12LTd$ for attention, about $T/(6d)$ times the weight term.
$N$ in the 6N rule counts weights in matmuls, the unembedding included and the embedding lookup excluded, and only the active parameters of a mixture of experts.
Decoder-only models dominate. Encoders and encoder-decoders keep specific niches, and a cross-attention decoder block costs about $16d^2$ .
Next, §5.7 asks why stacking dozens of these blocks can still fail to train, and what keeps deep Transformers stable.

Radford et al. (2019). Language Models are Unsupervised Multitask Learners. GPT-2, the block counted in this section.
Touvron et al. (2023). Llama 2, Open Foundation and Fine-Tuned Chat Models. The modern recipe and the configs used in our counts.
Ainslie et al. (2023). GQA, Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Grouped-query attention, which changes both the count and the cache.
Kaplan et al. (2020). Scaling Laws for Neural Language Models. Table 1 gives the per-operation parameter and FLOP counts behind $C \approx 6N$ .
Brown et al. (2020). Language Models are Few-Shot Learners. Appendix D computes GPT-3's training compute.
Chowdhery et al. (2022). PaLM, Scaling Language Modeling with Pathways. Appendix B counts training FLOPs with the full attention term.
Korthikanti et al. (2022). Reducing Activation Recomputation in Large Transformer Models. The activation memory of a block, tensor by tensor.
Karpathy. nanoGPT. A compact GPT-2 implementation to compare with your counts.

Zero to Frontier is made for larger screens

One Transformer block, counted

Assembling a decoder block

Three families of Transformer

Counting one block exactly

Count one block exactly, for any recipe

FLOPs per token and the 6N rule

Training FLOPs per token, with the attention term

Memory per token, in brief

A block, from scratch in NumPy

A GPT-2 decoder block, from scratch

In PyTorch

Read more

Watch next