nanoGPT — Build GPT From Scratch

import urllib.request
import torch
import torch.nn as nn
from torch.nn import functional as F

torch.manual_seed(1337)
device = "cuda" if torch.cuda.is_available() else "cpu"

# 1) Load raw text. This is the dataset the tiny GPT will imitate.
url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
text = urllib.request.urlopen(url).read().decode("utf-8")

# 2) Tokenize characters: each unique character gets an integer ID.
chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[ch] for ch in s]
decode = lambda ids: "".join(itos[i] for i in ids)

# 3) Split into train and validation text.
data = torch.tensor(encode(text), dtype=torch.long)
split = int(0.9 * len(data))
train_data, val_data = data[:split], data[split:]

# 4) Hyperparameters. n_embd=128 means each token becomes a 128-number vector.
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
dropout = 0.2

def get_batch(split_name):
    source = train_data if split_name == "train" else val_data
    # Pick 32 random starting points. Each example is 64 characters long.
    starts = torch.randint(len(source) - block_size, (batch_size,))
    x = torch.stack([source[i:i + block_size] for i in starts])
    y = torch.stack([source[i + 1:i + block_size + 1] for i in starts])
    return x.to(device), y.to(device)

# 5) Embeddings: token meaning-vector + position/order-vector.
token_embedding = nn.Embedding(vocab_size, n_embd).to(device)
position_embedding = nn.Embedding(block_size, n_embd).to(device)
xb, yb = get_batch("train")
positions = torch.arange(block_size, device=device)
x = token_embedding(xb) + position_embedding(positions)

print("device:", device)
print("characters:", len(text), "| vocabulary:", vocab_size)
print("input batch:", xb.shape)
print("combined embeddings:", x.shape)

class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer(
            "causal_mask",
            torch.tril(torch.ones(block_size, block_size)),
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        # key = what each token contains; query = what each token is looking for.
        k = self.key(x)
        q = self.query(x)

        # Query @ Key scores decide which earlier tokens matter.
        weights = q @ k.transpose(-2, -1)
        weights = weights * (k.shape[-1] ** -0.5)
        # Causal mask hides future tokens, so GPT cannot peek ahead.
        weights = weights.masked_fill(
            self.causal_mask[:T, :T] == 0,
            float("-inf"),
        )
        weights = F.softmax(weights, dim=-1)
        weights = self.dropout(weights)

        # value = information to mix together using the attention weights.
        v = self.value(x)
        return weights @ v


class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList(
            [Head(head_size) for _ in range(num_heads)]
        )
        self.projection = nn.Linear(num_heads * head_size, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        joined = torch.cat([head(x) for head in self.heads], dim=-1)
        return self.dropout(self.projection(joined))


# Shape check: attention changes the content, but keeps the same shape.
attention = MultiHeadAttention(n_head, n_embd // n_head).to(device)
print("attention output:", attention(x).shape)

class FeedForward(nn.Module):
    def __init__(self, embedding_size):
        super().__init__()
        # Attention communicates BETWEEN tokens. The FFN then processes each
        # token independently; it never reads a neighboring token directly.
        # 128 → 512 gives it room to learn richer features, then 512 → 128
        # returns to the original size so it can be added back to x.
        self.network = nn.Sequential(
            nn.Linear(embedding_size, 4 * embedding_size),
            nn.ReLU(),
            nn.Linear(4 * embedding_size, embedding_size),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.network(x)


class Block(nn.Module):
    def __init__(self, embedding_size, num_heads):
        super().__init__()
        head_size = embedding_size // num_heads
        self.attention = MultiHeadAttention(num_heads, head_size)
        self.feed_forward = FeedForward(embedding_size)
        # LayerNorm keeps each token vector at a stable numerical scale.
        # It changes values, not shape: [B, T, 128] stays [B, T, 128].
        self.layer_norm_1 = nn.LayerNorm(embedding_size)
        self.layer_norm_2 = nn.LayerNorm(embedding_size)

    def forward(self, x):
        # Residual 1: preserve the old token representation and add context.
        x = x + self.attention(self.layer_norm_1(x))
        # Residual 2: preserve that result and add the FFN's computation.
        x = x + self.feed_forward(self.layer_norm_2(x))
        return x


block = Block(n_embd, n_head).to(device)
print("transformer block output:", block(x).shape)

class GPTLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Convert token IDs and positions into vectors.
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        # Stack 4 Transformer blocks: attention + feed-forward, repeated.
        self.blocks = nn.Sequential(
            *[Block(n_embd, n_head) for _ in range(n_layer)]
        )
        self.final_layer_norm = nn.LayerNorm(n_embd)
        self.language_model_head = nn.Linear(n_embd, vocab_size)
        self.apply(self._initialize_weights)

    def _initialize_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        if T > block_size:
            raise ValueError(f"Sequence length {T} exceeds block_size {block_size}")

        # Add meaning (token) vectors and order (position) vectors.
        token_vectors = self.token_embedding(idx)
        positions = torch.arange(T, device=idx.device)
        position_vectors = self.position_embedding(positions)
        x = token_vectors + position_vectors

        # Transformer stack reads the context and produces one vector per token.
        x = self.blocks(x)
        x = self.final_layer_norm(x)
        # Logits are raw scores for the next character in the vocabulary.
        logits = self.language_model_head(x)

        loss = None
        if targets is not None:
            B, T, C = logits.shape
            # Cross-entropy asks: did the model score the true next token highly?
            loss = F.cross_entropy(logits.reshape(B * T, C), targets.reshape(B * T))

        return logits, loss

    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=40):
        self.eval()
        for _ in range(max_new_tokens):
            # Only the latest block_size tokens fit into the model's context window.
            context = idx[:, -block_size:]
            logits, _ = self(context)
            # Use only the final time step to choose the next character.
            logits = logits[:, -1, :] / temperature

            if top_k is not None:
                # Keep the top choices so sampling is creative but not totally wild.
                k = min(top_k, logits.shape[-1])
                threshold = torch.topk(logits, k).values[:, [-1]]
                logits = logits.masked_fill(logits < threshold, float("-inf"))

            probabilities = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probabilities, num_samples=1)
            idx = torch.cat((idx, next_id), dim=1)

        self.train()
        return idx


model = GPTLanguageModel().to(device)
parameter_count = sum(p.numel() for p in model.parameters())
print(f"parameters: {parameter_count / 1e6:.2f} million")

learning_rate = 3e-4
max_steps = 2000
eval_interval = 200
eval_batches = 50
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

@torch.no_grad()
def estimate_loss():
    # Average many random batches so the printed loss is not too noisy.
    model.eval()
    result = {}
    for split_name in ("train", "val"):
        losses = torch.zeros(eval_batches)
        for i in range(eval_batches):
            inputs, targets = get_batch(split_name)
            _, batch_loss = model(inputs, targets)
            losses[i] = batch_loss.item()
        result[split_name] = losses.mean().item()
    model.train()
    return result


for step in range(max_steps + 1):
    if step % eval_interval == 0:
        losses = estimate_loss()
        print(
            f"step {step:4d} | "
            f"train {losses['train']:.4f} | val {losses['val']:.4f}"
        )
    if step == max_steps:
        break

    # One real training step: forward, loss, backward, optimizer step.
    inputs, targets = get_batch("train")
    _, loss = model(inputs, targets)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()