Build a GPT Language Model from Scratch

Build a ~10M parameter GPT in PyTorch and train it to write Shakespeare.

Introduction

30 Second Summary

Every time you ask ChatGPT a question or watch Claude write code, a transformer neural network is predicting the next word in the sequence.

In this project, you will build a GPT language model from scratch using PyTorch, train it on Shakespeare's complete works, and generate new text that reads like Shakespeare wrote it.

What You'll Build

Feed your trained model a blank prompt and watch it write grammatically correct Shakespearean dialogue, complete with character names, stage directions, and iambic-style phrasing.

By the end of this project, you'll have:

  • A working GPT-style language model (~10 million parameters) that you can feed any text prompt and watch it continue writing in Shakespearean style.
  • A side-by-side comparison proving why self-attention matters: your simple bigram model produces nonsense, while your transformer produces coherent English sentences.
  • First-hand experience building every component of the transformer architecture that powers GPT-4 and Claude, just scaled down to fit on your GPU.
  • Secret Mission: Add temperature and top-k sampling controls to shape your model's creativity.

Are there any prerequisites?

You should be comfortable writing Python (functions, classes, loops). Basic linear algebra helps too (matrices and dot products).

You also need PyTorch installed with CUDA support and an NVIDIA GPU. If you already have a working PyTorch + CUDA setup, you're good to go.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters to you.

Set Up Data and Tokenizer

Neural networks only understand numbers. Before you can train a language model, every character in your training text needs to be converted into an integer that the network can process.

In this step, you will download Shakespeare's complete works as training data, build a character-level tokenizer to convert text into numbers and back, and create a data loader that feeds batches of sequences to your model during training.

In this step, get ready to:
  • Download the Tiny Shakespeare dataset and load it into Python.
  • Build a character-level tokenizer with encode and decode functions.
  • Create a batch loader function that samples training sequences for the model.
Download the Tiny Shakespeare dataset

Your model needs a large body of text to learn from. The Tiny Shakespeare dataset contains 1.1 million characters of Shakespeare's plays, which is small enough to train quickly but large enough to learn real English patterns.

  • Open bigram.py in your code editor and add the following imports and hyperparameters at the top of the file:
import torch
import torch.nn as nn
from torch.nn import functional as F

# hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 1e-3
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200

torch.manual_seed(1337)

What does this code do?

These lines import PyTorch and set the training configuration values your model will use later. The device variable automatically picks your GPU if CUDA is available, falling back to CPU otherwise.

The block_size=256 means the model will look at up to 256 characters of context when predicting the next character. The manual seed ensures reproducible results.

  • Below the seed line, add the data download logic that fetches Shakespeare's text from the internet:
# download and read data
import os
import requests

data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
if not os.path.exists('input.txt'):
    with open('input.txt', 'w', encoding='utf-8') as f:
        f.write(requests.get(data_url).text)

with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

What does this code do?

The script checks if input.txt already exists before downloading. This prevents re-downloading every time you run the script. The requests.get() call fetches the dataset from GitHub, and .text extracts the response body as a string.

Install requests if needed

The requests library is not part of Python's standard library. If you get an import error when running the script, install it first with pip install requests.

  • Save bigram.py and run it to download the dataset:
python bigram.py
  • Confirm the dataset file was created by running:
ls input.txt

You should see input.txt listed in the output. The script completed without errors and downloaded 1.1 million characters of Shakespeare.

Seeing a connection error or ModuleNotFoundError?

  • If you see ModuleNotFoundError: No module named 'requests', run pip install requests in your terminal first.
  • If you see a connection error, check your internet connection. The script needs to reach GitHub to download the dataset.

help me fix a connection or import error when downloading the Tiny Shakespeare dataset

Build the character-level tokenizer

Language models work with numbers, not letters. A tokenizer converts each character to a unique integer and back again. Your tokenizer will map all 65 unique characters in Shakespeare (letters, punctuation, spaces) to integers 0 through 64.

  • Add the tokenizer code below the data-reading block in bigram.py:
# character-level tokenizer
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[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

What does this code do?

  • chars extracts every unique character from the text and sorts them alphabetically.
  • stoi (string-to-integer) maps each character to a number. itos (integer-to-string) maps back the other way.
  • encode() converts a string into a list of integers. decode() converts a list of integers back into readable text.
  • Now add the train/val split and verification prints below the tokenizer code:
# train/val split
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

print(f"vocab size: {vocab_size}")
print(f"train has {len(train_data):,} tokens")
print(f"val has {len(val_data):,} tokens")
print(f"CUDA available: {torch.cuda.is_available()}")

What does this code do?

The encode() function converts the entire Shakespeare text into a tensor of integers. The 90/10 split gives you 1,003,854 tokens for training and 111,540 for validation. Validation data is text the model never trains on, so you can check if it's actually learning general patterns rather than memorizing.

  • Save bigram.py and run it to verify your tokenizer and split are working:
python bigram.py

You should see four lines printed in your terminal:

vocab size: 65
train has 1,003,854 tokens
val has 111,540 tokens
CUDA available: True

What do these numbers mean?

65 unique characters is your model's entire vocabulary. That includes all uppercase and lowercase letters, punctuation, spaces, and newlines. The model will learn to predict which of these 65 characters comes next at every position.

CUDA available: True confirms PyTorch can see your NVIDIA GPU. Training will use the GPU for massive speed gains.

Seeing CUDA available: False?

  • Make sure you installed the CUDA-enabled version of PyTorch. Visit the PyTorch install page and select your CUDA version to get the correct install command.
  • If you have an AMD GPU or no dedicated GPU, the model will still train on CPU but much more slowly.

help me get PyTorch to recognize my NVIDIA GPU with CUDA

Build the data loader function

During training, your model needs random batches of text sequences to learn from. The get_batch() function grabs 64 random sequences of 256 characters each, along with their target outputs (the same sequences shifted by one character).

  • Add the get_batch() function at the bottom of bigram.py, below the print statements:
def get_batch(split):
    data_split = train_data if split == 'train' else val_data
    ix = torch.randint(len(data_split) - block_size, (batch_size,))
    x = torch.stack([data_split[i:i+block_size] for i in ix])
    y = torch.stack([data_split[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

What does this code do?

  • ix picks 64 random starting positions from the text.
  • x is the input: 64 sequences of 256 characters each (what the model sees).
  • y is the target: the same sequences shifted one position forward (what the model should predict).
  • .to(device) moves the tensors to your GPU for fast computation.
  • Save bigram.py.

Before you run the script again, what do you expect to see? Adding a function definition should not change the printed output.

  • Run the script one final time to confirm everything still works:
python bigram.py

You should see the same four lines as before. No errors means your get_batch() function is defined correctly and ready for the training loop.

Seeing a NameError or IndentationError?

  • Check that get_batch() is defined AFTER the train_data and val_data variables. Python reads top-to-bottom, so the function body must reference variables that already exist at module level.
  • Make sure all lines inside the function are indented with 4 spaces consistently.

help me fix an error in my get_batch function in bigram.py

✔️ Awesome, I've got everything!

Great. Double check you have saved bigram.py and that running it prints all four verification lines with no errors.

ⓧ I'd like to double check the full code

Scroll up and confirm your bigram.py contains all five code blocks in this order:

  • Imports and hyperparameters (starting with import torch, ending with torch.manual_seed(1337)).
  • Data download logic (starting with # download and read data).
  • Character-level tokenizer (starting with # character-level tokenizer).
  • Train/val split and print statements (starting with # train/val split).
  • The get_batch() function definition at the bottom.

Your data pipeline is ready. You have 1 million characters of Shakespeare converted into GPU-ready tensors, split into training and validation sets, with a batch loader that can feed random sequences to a model. Next up, you will build and train a bigram language model to see what happens when a model predicts the next character with zero context.

Build and Train the Bigram Model

Your data is tokenized and sitting on the GPU, ready to feed to a neural network. The next question is: what kind of model should predict the next character?

The simplest possible approach is a bigram model. It only looks at the single previous character to decide what comes next. You will build it, train it, and see exactly where it falls short.

In this step, get ready to:
  • Build the BigramLanguageModel class with an embedding layer and loss computation.
  • Write an autoregressive generate method that produces text one character at a time.
  • Train the model for 5000 steps and observe the generated output.
Build the BigramLanguageModel class

A bigram model predicts the next character based solely on the current character. It uses an embedding table where each row stores the probability distribution over what character is likely to follow.

What is a bigram model?

Imagine you only ever look at the letter you just typed to guess the next one. After the letter "q", you'd almost always guess "u". That single-character lookup is all a bigram model does.

The model stores these lookup probabilities in a 65x65 table (one row per character in the vocabulary, one column per possible next character). Training adjusts these probabilities to match real Shakespeare patterns.

  • First, add a loss estimation helper function below your get_batch() function in bigram.py. This averages loss over many batches to give you a stable measurement during training:
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

What does this code do?

  • @torch.no_grad() disables gradient tracking during evaluation, saving memory and speed.
  • model.eval() switches the model to evaluation mode (disables dropout and similar training-only behaviors).
  • The function averages loss over 200 random batches per split, giving a stable estimate instead of a single noisy measurement.
  • Now add the model class itself below estimate_loss(), followed by the lines that instantiate it and print the parameter count:
class BigramLanguageModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        logits = self.token_embedding_table(idx)  # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss


model = BigramLanguageModel()
model = model.to(device)

print(f"\nBigram model parameters: {sum(p.numel() for p in model.parameters()):,}")

What does this code do?

  • nn.Embedding(vocab_size, vocab_size) creates a 65x65 lookup table. Given character index i, the model looks up row i and uses those 65 numbers as predictions for the next character.
  • The forward() method looks up embeddings and reshapes the tensors so F.cross_entropy can compare predictions against actual next characters.
  • model.to(device) moves the model to your GPU for fast computation.
  • Save bigram.py and run it to confirm the model instantiates correctly:
python bigram.py

You should see the same vocab and token output as before, plus a new line showing the parameter count:

Bigram model parameters: 4,225 confirms the 65x65 embedding table loaded onto your GPU.

Seeing a NameError or SyntaxError?

  • Make sure estimate_loss() is placed after get_batch() and before the class definition. Order matters in a top-to-bottom script.
  • Check that the class definition starts at column 0 (no accidental indentation) and that the model = BigramLanguageModel() line is outside and below the class.
  • help me debug a NameError or SyntaxError in my bigram model code
Write the generate method

The model can compute loss, but it also needs to produce new text. The generation process is autoregressive: predict the next character, append it, then use the extended sequence to predict the character after that. Repeat.

  • In bigram.py, find the forward() method inside your BigramLanguageModel class.
  • Add this generate() method below forward(), still inside the class (indented one level):
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            logits, loss = self(idx)
            logits = logits[:, -1, :]  # (B, C)
            probs = F.softmax(logits, dim=-1)  # (B, C)
            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)
            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)
        return idx

What does this code do?

  • logits[:, -1, :] grabs only the prediction for the last character in the sequence (the one we need to extend from).
  • F.softmax converts raw scores into a probability distribution that sums to 1.
  • torch.multinomial randomly samples from that distribution. This means the model does not always pick the highest-probability character, which adds natural variation to the output.
  • The new character is appended to the sequence with torch.cat, and the loop continues.
  • Save bigram.py. The generate method is ready to use once training begins in the next substep.
Train the model and generate text

Time to train your bigram model and see what it produces. You will use the AdamW optimizer for 5000 steps, printing loss every 500 steps so you can watch the model learn.

What loss should you expect at the start?

Before training, the model guesses randomly across 65 characters. The expected loss of a uniform random guess is -ln(1/65) = ln(65), which is approximately 4.17. If you see a starting loss near 4.17, that confirms everything is wired up correctly.

  • Add this training and generation code at the very bottom of bigram.py, below the print(f"\nBigram model parameters:...") line:
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

print("\nTraining bigram model...")
for iter in range(max_iters):
    if iter % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

losses = estimate_loss()
print(f"step {max_iters}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

print("\n--- Bigram Model Output ---")
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))

What does this code do?

  • AdamW is an optimizer that adjusts the model's embedding table to reduce prediction error on each batch.
  • Every 500 steps, estimate_loss() evaluates both training and validation loss so you can watch improvement over time.
  • After training, the script generates 500 characters starting from a single zero token (a newline character in the vocabulary).
  • Save bigram.py.

Before you run this, think: do you expect the generated text to read like actual Shakespeare? Or something less coherent?

  • Run the script:
python bigram.py

You should see the loss drop from approximately 4.17 (random guessing) down to around 2.5 over 5000 steps. Then the model generates 500 characters.

Look at the generated text. You will see occasional recognizable English words, but mostly nonsensical fragments with no sentence structure. Character names appear broken, dialogue makes no sense, and there is no coherent flow from one word to the next.

Why is the output so bad?

The bigram model only sees ONE character of context. It knows that "q" is usually followed by "u", but it has no idea what word it is in the middle of, what sentence it is writing, or who is speaking. A loss of 2.5 is better than random (4.17), but still far from good.

This is the exact limitation you will solve in the next step. The self-attention mechanism lets every character look back at ALL previous characters to decide what comes next.

Training seems stuck or loss isn't dropping?

  • If the first loss is much higher or lower than 4.17, check that vocab_size equals 65 (printed earlier). A wrong value means the tokenizer or embedding size is misconfigured.
  • If you see a CUDA out-of-memory error, confirm that no other GPU processes are running. Close other notebooks or scripts that might be using your GPU.
  • If loss stays flat at 4.17, check that optimizer.zero_grad(set_to_none=True) and loss.backward() are both inside the training loop.
  • help me debug my bigram model training loop

✔️ Awesome, I've got everything!

Great. Double check you have saved your file, then move on to the next step.

ⓧ I'd like to double check the full code

Scroll up and confirm each of these sections is present in your bigram.py file, in this order:

  • Imports, hyperparameters, and torch.manual_seed(1337) at the top.
  • Data download, reading, and the character-level tokenizer (chars, stoi, itos, encode, decode).
  • Train/val split and the print statements for vocab size and token counts.
  • get_batch() function.
  • estimate_loss() function with the @torch.no_grad() decorator.
  • BigramLanguageModel class with __init__, forward, and generate methods.
  • Model instantiation (model = BigramLanguageModel()) and parameter count print.
  • Optimizer, training loop (5000 iterations), final loss print, and 500-character generation.

Your bigram model proves a single embedding table can learn basic character-pair statistics. But a loss of 2.5 and incoherent text is not good enough. In the next step, you will build the transformer architecture that lets the model look at hundreds of characters of context at once, and the difference in output quality will be dramatic.

Build the Transformer Architecture

Your bigram model can generate text, but the output is incoherent because it only looks at the previous single character to predict the next one. It has zero understanding of context.

The self-attention mechanism solves this by letting every character in a sequence look back at all previous characters and decide which ones are relevant for predicting what comes next. This is the core innovation behind GPT.

In this step, get ready to:
  • Build single-head and multi-head self-attention classes.
  • Build feedforward and transformer block classes with residual connections.
  • Assemble the full GPT model and verify it has approximately 10.7 million parameters.
Build single-head and multi-head self-attention

You need a new file for your transformer model. Much of the foundation (data loading, batching) is reused from bigram.py, but the architecture is entirely different.

  • Create a new file called gpt.py in your project folder.
  • Add the imports, hyperparameters, and random seed by pasting this code:
import torch
import torch.nn as nn
from torch.nn import functional as F

# hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2

torch.manual_seed(1337)

What's different from bigram.py?

The learning rate is 3e-4 instead of 1e-3 because transformers train better with a smaller learning rate.

The four new hyperparameters control the transformer's size: n_embd=384 is the embedding dimension, n_head=6 is the number of attention heads, n_layer=6 is how many transformer blocks to stack, and dropout=0.2 prevents overfitting.

  • Add the data loading and tokenizer below the seed. This is the same code from bigram.py (without the print statements since you already verified the data):
# load data
import os
import requests

data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
if not os.path.exists('input.txt'):
    with open('input.txt', 'w', encoding='utf-8') as f:
        f.write(requests.get(data_url).text)

with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

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[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
  • Save gpt.py and verify the foundation loads without errors by running:
python gpt.py

You should see no output and no errors. The script loaded the data silently.

Seeing a ModuleNotFoundError?

If you see an error about requests not being found, install it with pip install requests. Since you already ran bigram.py successfully, this should not happen.

help me debug a ModuleNotFoundError in my gpt.py file

  • Add the get_batch() and estimate_loss() helper functions below the data loading:
def get_batch(split):
    data_split = train_data if split == 'train' else val_data
    ix = torch.randint(len(data_split) - block_size, (batch_size,))
    x = torch.stack([data_split[i:i+block_size] for i in ix])
    y = torch.stack([data_split[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y


@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

What do these functions do?

These are identical to the helpers in bigram.py. get_batch() randomly samples training sequences. estimate_loss() averages loss over 200 batches for a stable measurement.

Now for the core innovation. The Head class implements single-head causal self-attention. This is the mechanism that lets each character "look back" at all previous characters and decide which ones matter for its prediction.

  • Add the Head class below estimate_loss():
class Head(nn.Module):
    """One head of self-attention."""

    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('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)    # (B, T, head_size)
        q = self.query(x)  # (B, T, head_size)
        # compute attention scores
        wei = q @ k.transpose(-2, -1) * C**-0.5  # (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)
        # weighted aggregation of values
        v = self.value(x)  # (B, T, head_size)
        out = wei @ v      # (B, T, head_size)
        return out

How does self-attention work?

Every token produces three vectors: a query ("what am I looking for?"), a key ("what do I contain?"), and a value ("what do I communicate?").

The attention score between two tokens is the dot product of the query and key, scaled by the square root of the head size to prevent gradients from exploding. The tril mask sets future positions to negative infinity so each token can only attend to tokens that came before it. This is what makes it "causal" (no peeking at the future).

  • Save gpt.py and run it to verify the Head class has no syntax errors:
python gpt.py

No output and no errors means your self-attention head is syntactically valid.

A single attention head can only focus on one type of relationship at a time. MultiHeadAttention runs multiple heads in parallel so the model can attend to different patterns simultaneously (one head might track grammar, another might track meaning).

  • Add the MultiHeadAttention class below Head:
class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""

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

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

What does the projection layer do?

Each head produces a partial view of size head_size (384 / 6 = 64). Concatenating all 6 heads gives a 384-dimensional vector. The projection layer (self.proj) lets the model mix information across heads before passing it forward.

Build the feedforward and transformer block classes

Self-attention handles communication between tokens. But the model also needs a place to "think" about what it learned. The feedforward network gives each token time to process the attention output independently.

  • Add the FeedForward class below MultiHeadAttention:
class FeedForward(nn.Module):
    """A simple linear layer followed by a non-linearity."""

    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

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

Why the 4x expansion?

The inner layer expands from 384 to 1536 dimensions (4x), then compresses back to 384. This wider inner layer gives the network more capacity to learn complex patterns before projecting back to the residual stream. This 4x ratio is standard across GPT-2, GPT-3, and GPT-4.

  • Save gpt.py and run it to check for errors:
python gpt.py

No errors means your feedforward layer is valid. Now you can combine attention and feedforward into a single transformer block.

The Block class is the fundamental repeating unit in every transformer. It runs self-attention (communication), then feedforward (computation), with residual connections and LayerNorm for training stability.

  • Add the Block class below FeedForward:
class Block(nn.Module):
    """Transformer block: communication followed by computation."""

    def __init__(self, n_embd, n_head):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

What are residual connections?

The x = x + self.sa(...) pattern is a residual connection. Instead of replacing x, the attention output is added to it. This creates a "highway" for gradients to flow directly through the network during training, which prevents the vanishing gradient problem in deep networks.

LayerNorm is applied before each sub-layer (pre-norm architecture). This is more stable for training than the original post-norm design and is what modern GPT implementations use.

Assemble the full GPT model and verify the parameter count

With all the building blocks ready, you can now assemble the full GPT model. It stacks 6 transformer blocks, adds token and position embeddings, and projects the output back to vocabulary-sized predictions.

  • Add the GPT class with its __init__ and forward methods below Block:
class GPT(nn.Module):

    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.token_embedding_table(idx)       # (B, T, n_embd)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T, n_embd)
        x = tok_emb + pos_emb                           # (B, T, n_embd)
        x = self.blocks(x)                              # (B, T, n_embd)
        x = self.ln_f(x)                                # (B, T, n_embd)
        logits = self.lm_head(x)                        # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

What does the forward pass do?

  • token_embedding_table converts each character ID into a 384-dimensional vector that represents its meaning.
  • position_embedding_table adds position information so the model knows where each token sits in the sequence.
  • The 6 transformer blocks (self.blocks) apply self-attention and feedforward processing repeatedly, building deeper understanding.
  • lm_head projects the 384-dimensional output back to 65 dimensions (one per character in the vocabulary) to produce predictions.
  • Now add the generate() method inside the GPT class (indented under it), followed by the model instantiation outside the class:
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            # crop context to block_size
            idx_cond = idx[:, -block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :]  # (B, C)
            # apply temperature
            logits = logits / temperature
            # apply top-k filtering
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            probs = F.softmax(logits, dim=-1)  # (B, C)
            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)
            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)
        return idx


model = GPT()
model = model.to(device)

print(f"GPT model parameters: {sum(p.numel() for p in model.parameters()):,}")

How is this generate() different from the bigram's?

The key difference is idx_cond = idx[:, -block_size:]. The bigram model could take unlimited context (it ignored everything except the last token anyway). The transformer processes all block_size=256 characters simultaneously through self-attention, so you must crop the context to fit.

The temperature and top_k parameters are included for the Secret Mission later. At their defaults (1.0 and None), they have no effect on generation.

  • Save gpt.py.

Before you run this, what do you think the parameter count will be? The bigram model had only 4,225 parameters (a 65x65 embedding table). This model has 384-dimensional embeddings, 6 attention heads, and 6 transformer blocks.

  • Run gpt.py to see your model's parameter count:
python gpt.py

You should see approximately 10.7 million parameters printed. That is over 2,500x more parameters than your bigram model, all dedicated to understanding context through self-attention.

Seeing a NameError or IndentationError?

Make sure the generate() method is indented inside the GPT class (one level of indentation, matching forward()). The model = GPT() line must be at the top level (no indentation), outside the class.

If you see a NameError about n_embd or n_head, check that your hyperparameters section at the top of the file includes all four new variables.

help me fix an error when running gpt.py

✔️ Awesome, I've got everything!

Great. Double check you have saved gpt.py and the parameter count printed successfully.

ⓧ I'd like to double check the full code

Scroll up through your gpt.py file and confirm each section is present in this order:

  • Imports (torch, nn, F) and hyperparameters (including n_embd, n_head, n_layer, dropout) with torch.manual_seed(1337).
  • Data loading, tokenizer, train/val split (no print statements).
  • get_batch() and estimate_loss() functions.
  • Head class with key, query, value projections, tril buffer, and scaled dot-product attention.
  • MultiHeadAttention class with nn.ModuleList of heads and a projection layer.
  • FeedForward class with 4x expansion and ReLU.
  • Block class combining attention + feedforward with residual connections and LayerNorm.
  • GPT class with __init__, forward, and generate methods.
  • Model instantiation (model = GPT()) and parameter count print.

You have built the exact same architecture used in GPT-2, GPT-3, and GPT-4, just scaled down to fit on a single GPU. Next up, you will train this model on Shakespeare and watch the loss drop far below what the bigram could achieve.

Train the GPT and Generate Text

Your transformer architecture is built and sitting on the GPU with ~10.7 million parameters ready to learn.

The bigram model from earlier only managed to push its loss down to ~2.5 and produced incoherent text because it sees just one character of context. Now you will train your GPT on the same Shakespeare data and find out whether self-attention, which lets every character look back at 256 characters of history, actually makes the output coherent.

In this step, get ready to:
  • Add the training loop with an AdamW optimizer and periodic loss evaluation.
  • Run training and watch the loss plummet far below the bigram baseline.
  • Generate 1000 characters and compare them side-by-side with the bigram output.
Add the training loop

The training loop feeds batches of Shakespeare through the model, measures how wrong its predictions are (the loss), and adjusts the weights to reduce that error over 5000 iterations.

  • In gpt.py, find the line that says print(f"GPT model parameters: {sum(p.numel() for p in model.parameters()):,}").
  • Add the following training loop code directly below that line:
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

print("\nTraining GPT model...")
for iter in range(max_iters):
    if iter % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

losses = estimate_loss()
print(f"step {max_iters}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

What does this code do?

  • AdamW is the optimizer that adjusts the model's 10.7 million weights each iteration, using a learning rate of 3e-4 (slower than the bigram's 1e-3 because deeper networks need gentler updates).
  • Every 500 steps, estimate_loss() averages loss over 200 batches from both train and val sets. This gives a stable reading instead of noisy single-batch numbers.
  • The core loop samples a batch, runs the forward pass, computes the loss, zeros old gradients, backpropagates, and updates weights.
Add text generation

After training completes, you want the model to prove it learned something by generating brand new text from scratch.

  • Still in gpt.py, add the following generation code at the very end of the file, below the training loop:
# generate text
print("\n--- GPT Model Output (temperature=1.0) ---")
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(model.generate(context, max_new_tokens=1000)[0].tolist()))

What does this code do?

  • context = torch.zeros((1, 1), ...) creates a starting token of index 0 (the newline character). The model generates from this blank slate.
  • The generate() method autoregressively produces 1000 new characters, each time feeding its full context (up to 256 characters) back through the transformer to predict the next one.
  • Save gpt.py.

✔️ Awesome, I've got everything!

Great. Make sure you have saved the file before running it.

ⓧ I'd like to double check the full code

Your gpt.py is too long to show in one block, so confirm these sections are present in order:

  • Hyperparameters at the top (batch_size through dropout).
  • Data loading, tokenizer, and train/val split.
  • get_batch() and estimate_loss() functions.
  • Head, MultiHeadAttention, FeedForward, and Block classes.
  • GPT class with forward() and generate() methods.
  • Model instantiation and parameter count print.
  • AdamW optimizer and the 5000-iteration training loop with loss printing.
  • Final loss evaluation and 1000-character text generation at the bottom.
Run training and compare results

Before you run this, think about what you expect. The bigram model reached a training loss of ~2.5. Do you think 10.7 million parameters with self-attention will do better or worse on the same data?

  • Run your GPT training script:
python gpt.py

How long does training take?

On an NVIDIA GPU, each iteration takes roughly 10-20ms. The full 5000 iterations should complete in about 1-2 minutes. You will see loss printed every 500 steps so you can watch it drop in real time.

Watch the loss output as it prints. You should see something like this:

GPT model parameters: 10,788,929

Training GPT model...
step 0: train loss 4.1743, val loss 4.1754
step 500: train loss 1.7896, val loss 1.8639
step 1000: train loss 1.5353, val loss 1.6285
step 1500: train loss 1.4302, val loss 1.5628
...

Why is this loss so much better?

The bigram model plateaued at ~2.5 because it could only look at one character. Your GPT starts at the same ~4.17 (random chance for 65 characters) but drops to ~1.4-1.5 on the training set. That gap from 2.5 to 1.5 is entirely due to self-attention giving the model 256 characters of context.

The validation loss should settle around 1.5-1.6. The small gap between train and val shows the model is generalizing, not just memorizing.

After training finishes, the script generates 1000 characters. Compare this output with what the bigram produced in Step 2.

The transformer output should contain proper character names followed by colons, dialogue that sounds like Shakespeare, stage directions like "Enter" and "Exit", and mostly grammatical English sentences.

What should I look for in the output?

  • Character names in ALL CAPS followed by colons (like KING HENRY: or ROMEO:).
  • Grammatical English sentences with subjects, verbs, and objects.
  • Stage directions in the style of Shakespearean plays.
  • Vocabulary and phrasing that sounds Elizabethan (thee, thou, hath, wherefore).

The text is not perfect Shakespeare. Some sentences trail off or mix up characters. But compared to the bigram's random character soup, the difference is dramatic.

Seeing an out-of-memory error?

  • If you see a CUDA out of memory error, your GPU may not have enough VRAM for these hyperparameters. Try reducing batch_size from 64 to 32, or reduce n_embd from 384 to 256 and n_layer from 6 to 4.
  • Make sure no other processes are using your GPU. Close other notebooks or training scripts that might be holding VRAM.

Help me debug my CUDA out of memory error when training my GPT model.

Your GPT is generating Shakespeare. You built the exact same architecture that powers GPT-2, GPT-3, GPT-4, and Claude, just at a smaller scale. In the Secret Mission, you will add temperature and top-k sampling to give yourself fine-grained control over how creative or focused the output is.

Secret mission

Temperature and Top-k Sampling

Your model generates text using raw probabilities, which produces decent output but gives you no control over creativity. Professional LLMs like GPT-4 and Claude let users dial creativity up or down. Can you add that same control to your model?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources to experiment further, or delete them entirely.

Resources you used:

  • bigram.py - your bigram language model script.
  • gpt.py - your full GPT transformer model with temperature and top-k sampling.
  • input.txt - the downloaded Tiny Shakespeare dataset (1.1 million characters).

Keep everything running

No action needed. Choose this if you want to keep experimenting with longer training runs, different hyperparameters, or new datasets.

  • Your project files are all local and take minimal disk space (a few megabytes total).
  • You can re-run gpt.py anytime to retrain or generate more text.
  • Try swapping input.txt with a different text file to train on new styles.

Pause - I'll come back to this later

Since all resources are local files with no running processes, there is nothing to pause. Your files stay on disk until you delete them.

  • No servers or background processes are running.
  • Your trained model state is embedded in the script (it retrains each time you run it). If you want to save a trained model for later, add torch.save(model.state_dict(), 'gpt_model.pt') to the end of gpt.py.

Delete - I don't want to use this again

Remove all project files from your machine.

  • Delete your project folder by running the appropriate command for your operating system.

On macOS or Linux, run:

cd ~/Desktop
rm -rf gpt-from-scratch

On Windows, run:

cd ~\Desktop
Remove-Item -Recurse -Force gpt-from-scratch
  • This removes bigram.py, gpt.py, input.txt, and any other files in the project folder.
  • Python and PyTorch remain installed on your system. Uninstall them separately through your package manager if you no longer need them.

That's a Wrap!

That's a Wrap!

Nice work! You just built a GPT-style language model entirely from scratch and trained it to generate Shakespeare on demand.

You've learned how to:

  • Build the complete transformer architecture from scratch: self-attention, multi-head attention, feedforward networks, residual connections, and layer normalization.
  • Train a ~10 million parameter GPT on Shakespeare and watch the loss drop from 4.17 (random chance) to ~1.5, generating coherent prose with character names, dialogue, and stage directions.
  • Prove why context matters by comparing a bigram model (nonsensical fragments) against your transformer (grammatical sentences), demonstrating exactly what self-attention buys you.
  • Secret Mission: Added temperature and top-k sampling controls to shape your model's creativity, from conservative and repetitive to wild and chaotic.

Ready to quiz yourself?