Building a GPT

A language model from scratch.

⏱ 2 min read▶ Video at 1:09:10🎓 Free

Watch this lesson in the Deep Learning courseWatch on YouTube ↗Full Deep Learning course

This is it. G P T stands for generative pre-trained transformer. In the next few minutes, you'll write one yourself. It'll be small, but it's the real architecture. Scale the numbers up, and add a huge amount of data, and this is the same design behind the A I chatbots that changed the world.

Generative pre-trained transformer
Small, but the real design
Scale it up and it's the real thing

The whole model

Here's the full path. Text becomes token numbers. Each token gets a word vector and a position vector. These flow through a stack of transformer blocks, with a mask, so that no word can see the future. A final linear layer gives a score for every possible next token. That's all a GPT is. A machine that predicts the next token.

Tokenstext to numbers
Embeddingsword + position
Blocksmasked attention, stacked
Headscores for next token

TinyGPT in forty lines

Here it is, all in one place. The block that we built. Then the model. Two embedding tables, a stack of blocks, a final norm, and a linear head. In forward, we add the embeddings, build the mask, run the blocks, and return scores. Run it. Our Tiny GPT has about three hundred thousand weights. GPT four style models have hundreds of billions or more, which is around a million times bigger. But the shape of the code is very much the same. Read those forty lines again, slowly. You've just written the architecture that changed the world.

tiny_gpt.pypython

import torch
from torch import nn

class Block(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(),
                                 nn.Linear(4 * dim, dim))
        self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)

    def forward(self, x, mask):
        y = self.norm1(x)
        x = x + self.attn(y, y, y, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.norm2(x))

class TinyGPT(nn.Module):
    def __init__(self, vocab, block_size=64, dim=96, heads=4, layers=4):
        super().__init__()
        self.tok = nn.Embedding(vocab, dim)
        self.pos = nn.Embedding(block_size, dim)
        self.blocks = nn.ModuleList(Block(dim, heads) for _ in range(layers))
        self.norm = nn.LayerNorm(dim)
        self.head = nn.Linear(dim, vocab)

    def forward(self, ids):
        T = ids.size(1)
        x = self.tok(ids) + self.pos(torch.arange(T))
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool), 1)
        for block in self.blocks:
            x = block(x, mask)
        return self.head(self.norm(x))

model = TinyGPT(vocab=65)
print("Weights:", sum(p.numel() for p in model.parameters()))
print("Output:", list(model(torch.zeros(2, 32, dtype=torch.long)).shape))

Output

Weights: 466241
Output: [2, 32, 65]

Summary

In forty lines, you've built a GPT. It only has random weights so far, so it can only produce noise. Next, we feed it real text, train it, and listen to it speak for the first time.

Key takeaways

GPT: predict the next token
Forty lines, the real architecture
Scale changes, the design doesn't

Next up: Training Your GPT — teaching it to write

❮ The Transformer Block Training Your GPT ❯