AI Engineer Exhaustive
HardPlusGPT-2

Full Forward Pass

GPT-2

Hard

Complete GPT-2 forward pass

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

GPT-2 is a decoder-only Transformer language model. Section 2.3 of Language Models are Unsupervised Multitask Learners says the architecture largely follows the original GPT, with three changes that matter for the full stack: layer normalization moves to the input of each sub-block (pre-norm), a final layer norm is applied after the last block, and residual projections are initialized with a depth-dependent scale. The released graph in src/model.py is the inference form of that stack: token ids go in, a sequence of hidden states is produced, and a vocab-sized logit tensor comes out.

The training objective is ordinary next-token language modeling. Equation (1) in the paper factorizes a sequence as

p(x)=\prod_{i=1}^{n}p(s_i\mid s_1,\ldots,s_{i-1}).

A single forward pass computes all of those conditionals in parallel by applying a causal mask inside attention. Generation then consumes the last-position logits one token at a time. Table 2 lists four widths; the official defaults in default_hparams() match the smallest setting (n_layer=12, n_embd=768, n_head=12, n_ctx=1024). The vocabulary size in the paper is 50,257; the checkpoint hparams.json files in the repo set n_vocab when weights are loaded.

How it works

Let (X\in\mathbb{Z}^{B\times T}) be token ids. The official model() builds the residual stream as token embeddings plus learned absolute positions:

H^{(0)}=\mathrm{gather}(W_e,X)+\mathrm{gather}(W_p,\;\mathrm{pos}(X,\ell_{\mathrm{past}}))\in\mathbb{R}^{B\times T\times d}.

wte has shape [n_vocab, n_embd] and is initialized with std (0.02). wpe has shape [n_ctx, n_embd] and std (0.01). positions_for emits (0,\ldots,T-1) on a cold start, or (\ell_{\mathrm{past}}+0,\ldots,\ell_{\mathrm{past}}+T-1) when a key/value cache is reused. Positions must stay strictly below n_ctx.

Each of the (L) blocks (block, scopes h0 … h{L-1}) is pre-norm residual attention followed by a pre-norm feed-forward:

A^{(\ell)},\;P^{(\ell)}=\mathrm{Attn}\!\left(\mathrm{LN}_1(H^{(\ell-1)})\right),\qquad H' = H^{(\ell-1)}+A^{(\ell)},

M^{(\ell)}=\mathrm{MLP}\!\left(\mathrm{LN}_2(H')\right),\qquad H^{(\ell)}=H'+M^{(\ell)}.

Attention uses a fused (QKV) projection (c_attn) to (3d), splits into (h) heads of width (d_h=d/h), applies a causal mask, and projects back with c_proj. The MLP expands to (4d) with GELU, then projects back to (d). Both sublayers return tensors of shape [B, T, d]. If past is supplied, it is unstacked per layer as [B, 2, h, T_past, d_h] and concatenated onto the new keys and values on the sequence axis; each layer’s new present is stacked into results['present'] with shape [B, L, 2, h, T, d_h].

After the last block, a final norm (ln_f) is applied, then the language-model head ties the token embedding:

\mathrm{logits}=\mathrm{LN}_f(H^{(L)})\,W_e^{\top}\in\mathbb{R}^{B\times T\times V}.

The official implementation flattens to [B·T, d], multiplies by wte transposed, and reshapes. sample.py then slices logits[:, :, :n_vocab] before sampling, which is a defensive clip if the embedding table was padded. Training loss is not computed inside model(); callers take logits and a next-token cross-entropy themselves. The paper’s modified residual initialization ((1/\sqrt{N})) is a training-time weight scale and does not appear as a runtime multiply in this graph.

End-to-end, one uncached pass is therefore: embed → (L) pre-norm decoder blocks → final LN → tied projection. Cached generation feeds only the newest token and the stacked present from the previous step, which is how sample_sequence avoids re-encoding the prefix.

Official code

All of this lives in src/model.py of openai/gpt-2. Read default_hparams, model, block, attn, mlp, norm, gelu, positions_for, and past_shape. Autoregressive use is in src/sample.py (sample_sequence), called from src/generate_unconditional_samples.py and src/interactive_conditional_samples.py. The repo is an inference release; it loads published checkpoints rather than training from scratch.

Watch-outs

Sources