Overview

Following Andrej Karpathy’s nanoGPT, we build a small GPT that can generate text.

Model Architecture

class GPT(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layer):
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.blocks = nn.Sequential(*[TransformerBlock(n_embd, n_head) for _ in range(n_layer)])
        self.lm_head = nn.Linear(n_embd, vocab_size)

Complete Implementing Transformer in PyTorch first — it covers the building blocks you’ll need. The theory foundations are in Transformer Architecture Explained.