How to build GPT from scratch, step by step in PyTorch

In this article

1. Before you start
2. Why interviewers ask this
3. Building GPT from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

Why interviewers ask this

A candidate who can write a small GPT shows that attention, masking, normalisation and the training loop are not black boxes, and that the candidate can reason about tensor shapes, memory and cost. In our question bank, 68 of 249 prompts in the machine learning coding cluster mention attention, masking, a feed-forward block or a transformer (counted 2026-10-24 by assets/data/bank-count.py, output in assets/data/bank-count.json). A weak answer pastes a library call. A strong answer names the shape at every step, counts the parameters and explains why the loss at step zero should equal a known number. The wider group of questions lives in the machine learning coding round hub.

Building GPT from first principles

Decoder block

A decoder block is a repeated unit of a GPT that applies masked self-attention and then a small feed-forward network, each wrapped in layer normalisation and a residual connection, and returns vectors of the same shape it received.

Weight tying

Weight tying is the practice of using one matrix both as the token-embedding table and as the output projection, so the model stores one copy of a vocabulary-sized matrix instead of two.

In plain words

Think of the autocomplete on a phone keyboard. It looks at the words typed so far and offers the next one. A GPT is a very large version of that idea. It reads a stretch of text and gives a probability to each possible next token. The analogy breaks in one place. A keyboard looks back a word or two. A GPT lets each token look at all earlier tokens in its window.

In one precise sentence, a GPT is a decoder-only transformer that maps token ids to a score for every vocabulary entry at every position, trained so the score at position t favours the token at position t plus one. The name stands for generative pre-trained transformer, as explained in what GPT stands for.

Step 1: Turn text into token ids

A model cannot read letters, so the first part maps text to integers. Our model uses characters, numbering each distinct one. The training text is the documentation strings of twenty Python standard-library modules, restricted to plain ASCII, 180,449 characters in all, with 94 distinct characters.¹ The vocabulary size V is therefore 94.

The input is a string and the output is a list of integers, one per character. Real GPTs use a larger vocabulary: GPT-2 uses byte-level byte pair encoding with 50,257 tokens (Radford et al., 2019). The rest of the build does not change, because everything after this step sees only integers, so swapping the tokenizer later is cheap.

Step 2: Embed tokens and positions

Each integer becomes a vector by looking up a row in an embedding table, a grid with V rows and d columns, where d is the model width. We use d = 96. Attention alone cannot tell the order of tokens, so a second table with one row per position, up to the context length T, adds position information. We use T = 64.

The vector for a token is the sum of its token row and its position row. The output is a tensor of shape batch by T by d. GPT-2 learns its position table in the same way, with a context of 1,024 tokens (Radford et al., 2019).

Step 3: Mix each token with the ones before it

Attention lets each position build its vector from a weighted mix of itself and earlier positions. A linear layer produces a query, a key and a value from each vector. A score is the dot product of one position's query with another's key, divided by the square root of the head size so that the scores keep a stable spread, as set out in Vaswani et al., 2017. The causal mask then sets scores for later positions to minus infinity, so softmax gives them zero weight. Each token sees itself and everything before it, and never the answer it is trying to produce. The first position can only see itself.

Our model splits d = 96 into 4 heads of 24, so each head works in a smaller space, then joins the results and applies one more linear layer. See self-attention and multi-head attention.

Step 4: Finish the block, stack it and add the head

After attention, each block applies a feed-forward network to every position separately: a linear layer that widens the vector to 4d, a nonlinearity called the Gaussian error linear unit (GELU), and a linear layer back to d. Each of the two parts is preceded by layer normalisation and followed by a residual connection that adds the part's input to its output. We stack four blocks. A final layer normalisation and the language-model head come next. The language-model head is the final linear layer that produces one score per vocabulary entry at each position, 94 in our case. It reuses the token table, which is weight tying.

The output is a tensor of shape batch by T by V. Here is the full block, constructor included:

class Block(nn.Module):
    def __init__(s, d, h, T):                      # d = width, h = heads, T = context
        super().__init__(); s.h = h
        s.ln1 = nn.LayerNorm(d); s.qkv = nn.Linear(d, 3 * d); s.proj = nn.Linear(d, d)
        s.ln2 = nn.LayerNorm(d); s.fc = nn.Linear(d, 4 * d); s.out = nn.Linear(4 * d, d)
        s.register_buffer("mask", torch.tril(torch.ones(T, T)).view(1, 1, T, T))  # 1 = may look
    def forward(s, x):
        B, T, C = x.shape
        q, k, v = s.qkv(s.ln1(x)).split(C, dim=2)           # pre-norm, then one projection
        q, k, v = [t.view(B, T, s.h, C // s.h).transpose(1, 2) for t in (q, k, v)]
        att = (q @ k.transpose(-2, -1)) / math.sqrt(C // s.h)
        att = att.masked_fill(s.mask[:, :, :T, :T] == 0, float("-inf")).softmax(-1)
        y = (att @ v).transpose(1, 2).reshape(B, T, C)
        x = x + s.proj(y)                                    # residual 1
        return x + s.out(F.gelu(s.fc(s.ln2(x))))             # residual 2

The second line of forward normalises before attention, the line with minus infinity is the mask (the lower-triangle matrix has ones on its diagonal, so a token may look at itself), and the last two lines are the two residual connections.

Step 5: Train on next-token loss and check the first number

Training shows the model 64-character windows and asks for the next character at every position. The loss is cross-entropy, the average of minus the log of the probability the model gave to the right token (cross-entropy loss). We use the AdamW optimiser, batches of 32 windows, a learning rate of 0.003 and 600 steps. The key check comes before training: a model that knows nothing gives every token the same probability, so its loss should be the natural log of V, which is 4.543 for V = 94.¹

Ours did not. With PyTorch's default initialisation the tied embedding table starts with unit-size values, the scores are huge, and the first loss was 59.588. Setting every weight to small random values (standard deviation 0.02) gave 4.511 on training text and 4.529 on held-out text. A wrong first number means a bug or a bad initialisation. The loop is plain gradient descent.

Here are the model around the blocks, the initialisation and the training step:

class GPT(nn.Module):
    def __init__(s, V, T, d, h, L):
        super().__init__()
        s.tok = nn.Embedding(V, d); s.pos = nn.Embedding(T, d)
        s.blocks = nn.ModuleList([Block(d, h, T) for _ in range(L)])
        s.lnf = nn.LayerNorm(d); s.head = nn.Linear(d, V, bias=False)
        s.head.weight = s.tok.weight                         # weight tying
    def forward(s, idx, targets=None):
        x = s.tok(idx) + s.pos(torch.arange(idx.shape[1]))   # B x T x d
        for b in s.blocks: x = b(x)
        logits = s.head(s.lnf(x))                            # B x T x V
        loss = None if targets is None else F.cross_entropy(
            logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

def init(m):                                                 # the "one initialisation line" fix
    if isinstance(m, (nn.Linear, nn.Embedding)):
        nn.init.normal_(m.weight, mean=0.0, std=0.02)
        if getattr(m, "bias", None) is not None: nn.init.zeros_(m.bias)

model = GPT(V=94, T=64, d=96, h=4, L=4); model.apply(init)
opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
for step in range(600):                                      # x, y: batch of 32 windows and their next characters
    x, y = get_batch(train)                                  # y is x shifted left by one
    _, loss = model(x, y)
    opt.zero_grad(); loss.backward(); opt.step()

get_batch picks 32 random windows of 64 characters and their one-step-shifted targets; the full script is assets/data/mini-gpt.py.

1 · Tokens2 · Embeddings3 · Attention4 · Blocks and head5 · LossReplay

Step-by-step description

1. Text becomes integer token ids.
2. Each id plus its position becomes a vector.
3. A causal mask hides later tokens; each token sees itself and earlier ones.
4. Blocks stack, and a head scores every token.
5. Next-token loss trains the weights.

• Text becomes token ids.
• Each token id becomes a vector, with position added.
• Causal attention mixes each token with itself and earlier ones.
• The language-model head produces one score per vocabulary entry.

Worked example

The script trained the model for 600 steps on one central processing unit (CPU) thread in about four minutes, with a fixed seed. Training loss fell from 4.511 to 1.862, and held-out loss on 10% of the text fell from 4.529 to 2.255, so the model learned the spelling patterns of the text but is far from fluent.¹ A sample from it begins with fragments like "intracDiong of lilosce be arguined Argument", which resembles the text's style without being words. The gap between 1.862 and 2.255 shows the model fits training windows better than unseen ones.

The GPT-2 figures for layers, width, context and vocabulary come from the GPT-2 paper, and the model card states 124M parameters for the smallest version, which matches the formula below.²

The parameter count has a formula. Embeddings contribute V times d plus T times d. Each block contributes 12 times d squared plus 13 times d, made of the attention projections, the feed-forward layers and two layer normalisations. The final layer normalisation adds 2 times d. For our model that is 462,720, and the PyTorch count matched it exactly.²¹ For GPT-2 small it gives 124,439,808, about 85 million in the blocks and 38.6 million in the token table.

Text description of the diagram

• Loss at step 0 is 4.511 on training text and 4.529 on held-out text.
• Loss at step 600 is 1.862 on training text and 2.255 on held-out text.

Trade-offs and edge cases

Weight tying saves a large share of parameters at a large vocabulary. In GPT-2 small the token table holds 38.6 million parameters, 31% of the model, and a separate output matrix would add the same number again, giving 163,037,184 in total.² If our 94-character model switched to GPT-2's 50,257-token vocabulary, the table alone would hold 4,824,672 of 5,278,368 parameters, 91.4%.²

Text description of the diagram

• The token embedding holds 38.6 million parameters.
• The twelve blocks hold 85.05 million parameters.

Pre-norm and post-norm train differently. The original transformer put layer normalisation after each residual addition. GPT-2 moved it to the input of each sub-block and added one after the last block (Radford et al., 2019). Xiong et al. (2020) argue that with the original placement the gradients near the output are large at initialisation, which makes a warm-up stage necessary, while the pre-norm placement has well-behaved gradients. Our block is pre-norm.

Attention cost grows with the square of the context. Each head in each layer computes T by T scores: 4,096 at T = 64 and 1,048,576 at T = 1,024.² That is why long contexts are costly and why generation caches earlier keys and values, as described in the key-value (KV) cache.

Generation is a loop around the same model. To sample text, feed in at most the last 64 characters, take the scores at the last position, turn them into probabilities with softmax, draw one character and append it. Our script repeats this 120 times.¹ Because the position table has only 64 rows, the model can never look back further than its context length, however long the output grows.

Training windows and held-out text must not overlap. We hold out the last 10% of the text, 18,045 characters, and never train on it, so the held-out loss measures generalisation instead of memory.¹

A toy is not a recipe for scale. Our run uses 180,449 characters and 600 steps. Training real models needs far more data, mixed precision and sharded optimiser state, topics covered in optimizer states. The author of the popular nanoGPT repository noted in November 2025 that it is now old and deprecated, left up for reference, so treat any single codebase as a snapshot.

How to answer it in an interview

The spoken structure is: the five parts in order, the shape that flows through them, then the check you run before trusting it.

Common mistakes

"The mask is applied after the softmax." (wrong) Zeroing weights afterwards leaves rows that do not add to one. (right) "I set future scores to minus infinity before the softmax, so the weights of allowed positions still sum to one."

"The first loss can be anything." (wrong) A random model gives every token about the same probability, so its loss should be near the log of the vocabulary size. (right) "At step zero I expect about 4.5 for a 94-token vocabulary, and a very different number means a bug."

"I forgot the square root in the attention scores." (wrong) Without scaling, scores grow with the head size and softmax saturates. (right) "I divide by the square root of the head size so the spread stays stable."

"Weight tying adds parameters." (wrong) It removes a vocabulary-sized matrix. (right) "Tying shares the token table with the output layer, which saves 38.6 million parameters in GPT-2 small."

Follow-up questions

How many parameters does a GPT model have?

Add V times d for tokens, T times d for positions, 12 times d squared plus 13 times d for each block, and 2 times d for the final norm. With 12 layers, width 768, 50,257 tokens and context 1,024 this gives 124,439,808, the 124M of the GPT-2 small model card.²

Should layer normalisation go before or after attention?

Modern GPTs put it before, inside the residual branch, which is the pre-norm design GPT-2 used. Xiong et al. report that this makes gradients well-behaved at initialisation and can remove the warm-up stage.

How do you test that your GPT implementation is correct?

Check the first loss against the log of the vocabulary size, confirm the mask by changing a late token and verifying that earlier outputs do not change, and overfit one tiny batch until the loss is near zero. Then compare your parameter count with the formula.

What changes if the vocabulary is 50,257 tokens instead of 94 characters?

Only the table sizes change. In our model the parameter count rises from 462,720 to 5,278,368, and the token table becomes 91.4% of it, so tying the head to the table matters much more.²

Sources

1. Attention Is All You Need (arXiv 1706.03762 v7) — arXiv (Vaswani et al.), 12 Jun 2017
2. Language Models are Unsupervised Multitask Learners (GPT-2 paper) — OpenAI (Radford et al.), 14 Feb 2019
3. openai-community/gpt2 model card — Hugging Face (date = read date), 24 Oct 2026
4. On Layer Normalization in the Transformer Architecture (arXiv 2002.04745 v2) — arXiv (Xiong et al.), 12 Feb 2020
5. nanoGPT README — Andrej Karpathy (GitHub, README update of November 2025), 1 Nov 2025