LM4-T5 — Output Logits, Parameter Counts and Shapes

Learning objectives

After this lesson, you should be able to:

• Explain how the output head maps the final hidden states to logits and what weight tying changes.
• Derive the parameter count of a GPT from its configuration and reproduce 163,009,536 and 124,412,160.
• Convert parameters to memory (weights, gradients, optimizer state) and to FLOPs per token.
• State the shape of every tensor from input to logits, and the expected initial loss.

The idea: the head turns vectors into votes for every token

After the last block and the final layer norm, each position has a vector of width d. The output head is a linear layer without bias that multiplies this vector by a matrix of shape (d, V). The result is a vector of V numbers, one per vocabulary entry: the logits. A larger logit means the model considers that token more likely to come next. Softmax turns the logits into probabilities (Lesson LM1-T2); training compares them with the actual next token (Module 5).

Weight tying

The head matrix has the same shape as the transpose of the token embedding table. Weight tying uses the same matrix for both: the embedding row of a token is also the vector that is compared with the final hidden state to score that token. It saves V · d parameters, and the idea is natural: if a hidden state is close to the embedding of “box”, then “box” is a likely next token. The original GPT-2 uses tying, so its 124-million-parameter count includes the table only once. The book’s implementation keeps the two separate for clarity, which gives 163 million parameters and is why you will see both numbers.

Logits are big

The logits tensor has shape (B, n, V) and, for a real vocabulary, is larger than any other activation: with B = 8, n = 1,024 and V = 50,257 it has 411,705,344 entries, 1.6 GB in float32. Training code therefore often computes the loss in chunks or in lower precision.

The starting loss

A freshly initialised model has no knowledge, so its predicted distribution is roughly uniform: each token gets probability 1/V. The cross-entropy of a uniform guess is ln V, which is 10.82 for V = 50,257. If the first training loss is far from ln V (much larger, or already small), something is wrong in the initialisation, the data or the loss. This is one of the cheapest and best sanity checks in the field.

The mathematics: counting parameters, memory and compute

Let V be the vocabulary size, L the context length, d the width, N the number of blocks, with feed-forward expansion 4 and no biases on Q, K and V. The parameters are:

V, L, d, N

vocabulary size, context length, width, number of blocks.

12d² + 10d

parameters of one block: attention (4d² + d), feed-forward (8d² + 5d) and two layer norms (4d).

P

the total. Q, K and V biases would add 3d per block.

For GPT-2 small: 38,597,376 + 786,432 + 85,026,816 + 1,536 = 124,412,160 with a tied head, and 124,412,160 + 38,597,376 = 163,009,536 with a separate head. (The released GPT-2 weights have biases on Q, K and V, so their file contains 124,439,808 parameters: 27,648 more, which is 12 blocks × 3 × 768.)

For large d and N, P ≈ 12 · N · d² plus embeddings. Check with GPT-2 XL: 12 · 48 · 1,600² = 1.475 billion, plus about 80 million for embeddings, close to its 1.558 billion.

Memory. The weights take P · bytes: 124.4 million parameters need 498 MB in float32 (4 bytes) and 249 MB in 16-bit formats. Training with the Adam optimizer keeps, per parameter, the weight (4 bytes), its gradient (4) and two optimizer moments (8): about 16 bytes, so 124.4 million parameters need about 2.0 GB before counting activations. For a 7-billion-parameter model that is 112 GB, which is why large models are trained across many accelerators.

Compute. A forward pass uses about 2 FLOPs per parameter per token (one multiply and one add per weight), ignoring the attention products that depend on n. For 124.4 million parameters, about 249 million FLOPs per token. Training costs about three times a forward pass (forward, backward with respect to activations, backward with respect to weights), so 6 FLOPs per parameter per token, the rule used in Module 5 to estimate training cost.

A worked example: from configuration to bytes

Configuration: V = 50,257, L = 1,024, d = 768, N = 12.

1. Per block: 12 · 768² + 10 · 768 = 7,077,888 + 7,680 = 7,085,568.
2. All blocks: 12 · 7,085,568 = 85,026,816.
3. Embeddings: 50,257 · 768 = 38,597,376 and 1,024 · 768 = 786,432, total 39,383,808.
4. Final norm: 1,536. Sum with a tied head: 85,026,816 + 39,383,808 + 1,536 = 124,412,160.
5. Untied head adds another 38,597,376: 163,009,536.
6. Memory in float32: 163,009,536 · 4 = 652 MB for weights; training with Adam roughly 16 · 163,009,536 ≈ 2.6 GB.

Forward FLOPs per token for the tied model: 2 · 124,412,160 ≈ 249 million, plus the attention products 4 · n · d per layer per token (Lesson LM3-T6), which at n = 1,024 adds 4 · 1,024 · 768 · 12 ≈ 37.7 million, about 15% more.

import math

def gpt_params(V, L, d, N, tied=False, qkv_bias=False):
    block = 4 * d * d + d + (3 * d if qkv_bias else 0)      # attention: Q, K, V, output projection
    block += 8 * d * d + 5 * d                               # feed-forward (hidden width 4d, with biases)
    block += 4 * d                                           # two layer norms (scale and shift each)
    total = V * d + L * d + N * block + 2 * d                # embeddings, blocks, final layer norm
    return total if tied else total + V * d                  # output head

print(gpt_params(50257, 1024, 768, 12))                      # 163009536
print(gpt_params(50257, 1024, 768, 12, tied=True))           # 124412160
print(gpt_params(50257, 1024, 768, 12, tied=True, qkv_bias=True))   # 124439808
print(round(math.log(50257), 3))                             # 10.825  (loss of a uniform guess)

p = gpt_params(50257, 1024, 768, 12, tied=True)
print(round(p * 4 / 1e6), "MB of float32 weights;", round(p * 16 / 1e9, 2), "GB for Adam training state")   # 498 MB; 1.99 GB
print(2 * p, "forward FLOPs per token (without the attention products)")   # 248824320

The three counts agree with the published figures, which is the check to run on any reimplementation.

Shapes, once more

Engineering implications and limitations

Always print the parameter count. Compare it with the formula and, for known models, with the published value. It is the cheapest test that catches a missing norm, a wrong expansion or an extra bias.

Estimate memory before running. Weights, gradients and optimizer state scale with P; activations scale with B · n. If a run runs out of memory, reduce the batch size or sequence length first; the parameter memory is fixed by the model.

Check the initial loss. About ln V. Much higher suggests too large an initialisation; the book’s configuration gives a number close to ln V within a small margin.

Tie weights for small models. If the budget is a fixed number of parameters, tying lets you spend the saved V · d elsewhere. When loading pretrained weights, respect the model’s choice, or logits will not match.

Vocabulary size affects efficiency. Rounding V up to a multiple of 64 can speed up matrix products on accelerators. The extra rows are never used.

Mixed precision. Training in 16-bit formats halves activation memory and speeds up matrix products, with a float32 copy of the weights kept by the optimizer. Loss and softmax are computed in higher precision.

Meta device for planning. PyTorch’s meta device creates a model without allocating weights, so you can count parameters of a configuration of any size without having the memory (the notebook uses it for GPT-2 small).

Common misconceptions

Engineering takeaways

1. The output head is a (d, V) matrix without bias; logits have shape (B, n, V); tying it to the token embedding saves V·d parameters.
2. Total parameters: V·d + L·d + N·(12d² + 10d) + 2d, plus V·d for an untied head.
3. GPT-2 small: 163,009,536 untied, 124,412,160 tied, 124,439,808 with Q, K, V biases.
4. Memory: 4 bytes per weight in float32, about 16 bytes per parameter for Adam training state; compute: about 2 FLOPs per parameter per token forward, 6 for training.
5. An untrained model’s loss should be near ln V (10.82 for V = 50,257).

References and further reading

See Chapter 4 of the book (computing the number of parameters and the memory requirements of the model) for the same counts and for the discussion of weight tying.

• Radford, A. et al. (2019). Language Models are Unsupervised Multitask Learners (GPT-2). OpenAI technical report.
• Build a Large Language Model (From Scratch), Chapter 4: Implementing a GPT Model from Scratch (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

That completes Module 4. Next, Module 5 trains this model: the loss, the training loop and text generation.

Related content

Nothing below is required to finish this lesson.

• Before
  LM4-T4: The Transformer Block and Assembling GPT

  Combine attention, feed-forward, normalization and residuals into a block and stack blocks into a model.

• Lab
  LM4-L1: Count GPT Parameters and Tensor Shapes

  Write a parameter counter and shape tracer for a GPT configuration and check it against known sizes.

• Lab
  LM4-L3: Assemble a GPT Model in PyTorch

  Implement the GPT components, run a forward pass, count parameters and generate text from random weights.