Build a GPT Model from Scratch

Build a transformer language model in PyTorch that generates Python code.

Introduction

30 Second Summary

Every time you type a message to ChatGPT or Claude, a language model is predicting the next token in a sequence, one at a time, until a full response appears. The architecture behind this feels like magic, but it is built from a handful of repeating components you can implement yourself.

In this project, you will build a GPT-style language model from scratch in PyTorch, train it on real Python code, and watch it learn to generate plausible function definitions, loops, and variable assignments.

What You'll Build

You'll type def into your model's prompt, hit run, and watch it write a multi-line Python function complete with indentation, colons, and return statements.

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

  • A working GPT language model you can prompt with any Python snippet and watch it autocomplete with recognizable code structure.
  • A clear, visual understanding of the transformer architecture gained by comparing the random garbage of a baseline model against the structured output of your full transformer, side-by-side.
  • A portfolio-ready Google Colab notebook that demonstrates you understand how the architecture behind ChatGPT actually works at a mechanical level.
  • Secret Mission: Add temperature and top-k sampling controls to your text generation and experiment with how they change output quality.

Are there any prerequisites?

You should be comfortable with Python (functions, classes, and list comprehensions). No prior PyTorch or deep learning experience is required. Everything runs in Google Colab, so you don't need any local setup beyond a Google account.

Before We Start

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

Set Up Colab and Load the Dataset

Training a language model requires two things: compute power and data to learn from. A GPU accelerates the math behind neural networks by orders of magnitude compared to a CPU.

In this step, you'll verify your GPU is working, download a real Python code dataset from Hugging Face, and build a character-level tokenizer that converts text into numbers your model can process.

In this step, get ready to:
  • Verify GPU access and configure training hyperparameters.
  • Download a Python code dataset and build a character-level tokenizer.
  • Create a batch sampling function for training.
Verify GPU access and configure hyperparameters

Your Colab notebook already has a T4 GPU runtime enabled. Before writing any model code, you need to confirm PyTorch can actually see that GPU and set up the hyperparameters that will control your model later.

  • In your Colab notebook, paste the following code into the first cell:
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 does this code do?

The first three lines import PyTorch and its neural network modules.

The hyperparameters section defines values that control model training: batch_size is how many training examples to process at once, block_size is the context window (how many characters the model can look at), and n_embd/n_head/n_layer define the transformer architecture size.

The device line checks whether a GPU is available and assigns 'cuda' (GPU) or 'cpu' accordingly. This is how PyTorch decides where to run computations.

  • Run the cell by pressing **Shift+Enter**.
  • Verify your GPU is accessible by running this in a new cell:
print(f"Device: {device}")
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"PyTorch version: {torch.__version__}")

You should see output showing Device: cuda, a GPU name like Tesla T4, and a PyTorch version number.

Seeing Device: cpu instead?

If device shows cpu, your runtime does not have a GPU attached.

  • Click **Runtime** in the top menu.
  • Select **Change runtime type**.
  • Set **Hardware accelerator** to **T4 GPU**, then click **Save**.
  • Re-run both cells.

Help me fix my Colab GPU setup

Download the dataset and build the tokenizer

Your model needs training data. You'll use a dataset of 25,000 Python code examples from Hugging Face, then build a simple tokenizer that maps each unique character to a number.

  • Add a new code cell below your verification cell and paste the following:
# --- Data Loading ---
from datasets import load_dataset

ds = load_dataset('flytech/python-codes-25k', split='train')
code_texts = [example['output'] for example in ds]
text = '\n\n'.join(code_texts)
# Take a subset for fast training (~2MB)
text = text[:2_000_000]

# 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])

# 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:]

What does this code do?

It downloads the flytech/python-codes-25k dataset, extracts the Python code from each example's output field, and joins them into one long string. Then it takes the first 2 million characters for manageable training time.

The tokenizer is character-level: it finds every unique character in the text (letters, digits, symbols, whitespace) and assigns each an integer. stoi maps characters to integers. itos maps integers back to characters. encode and decode convert entire strings back and forth.

Finally, the entire text is encoded into a PyTorch tensor and split 90/10 into training and validation sets.

  • Run the cell with **Shift+Enter**. The dataset download takes about 10-20 seconds.
  • Confirm the data loaded correctly by running this in a new cell:
print(f"Dataset loaded: {len(text)} characters, {vocab_size} unique characters")
print(f"First 10 characters in vocab: {chars[:10]}")
print(f"Training data: {len(train_data)} tokens")
print(f"Validation data: {len(val_data)} tokens")

You should see approximately 2,000,000 characters in the dataset, a vocabulary of roughly 95 unique characters, and a 90/10 training/validation split.

Getting a download error?

If you see a connection or timeout error, the Hugging Face servers might be temporarily slow. Wait a few seconds and run the cell again.

If you see ModuleNotFoundError: No module named 'datasets', the datasets library is not installed. Add a new cell above and run !pip install datasets, then re-run your data loading cell.

Help me fix my dataset download

Create the batch sampling function

During training, the model learns from small random chunks of text rather than the entire dataset at once. The get_batch function grabs a random batch of these chunks each training step.

  • Add a new code cell and paste the following:
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?

Each call to get_batch picks batch_size (64) random starting positions in the data. From each starting position, it grabs a sequence of block_size (256) characters as the input x, and the same sequence shifted one character forward as the target y.

This is next-character prediction: for every position in the input, the model's job is to predict the character that comes next. Both tensors are moved to the GPU for fast computation.

  • Run the cell with **Shift+Enter**.
  • Before you test it, what shape do you think a batch of inputs will have? Consider the batch size and context window.
  • Verify the function works by running this in a new cell:
xb, yb = get_batch('train')
print(f"Input batch shape: {xb.shape}")
print(f"Target batch shape: {yb.shape}")
print(f"Input lives on: {xb.device}")

You should see torch.Size([64, 256]) for both shapes (64 sequences of 256 characters each) and cuda:0 confirming the data is on the GPU.

Seeing a RuntimeError about CUDA?

If you see a CUDA out-of-memory error, your runtime may have stale tensors. Click **Runtime** then **Restart runtime** and re-run all cells from the top.

Help me debug my get_batch function

✔️ Awesome, I've got everything!

Great. Make sure all cells are saved in your notebook. You now have a working data pipeline ready to feed into your first model.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm you have these cells in order:

  • Cell 1: Imports and hyperparameters (import torch through torch.manual_seed(1337)).
  • Cell 3: Data loading, tokenizer, and train/val split (from datasets import load_dataset through val_data = data[n:]).
  • Cell 5: The get_batch function definition.

The verification cells (2, 4, and 6) are optional to keep. They helped you confirm things work but are not needed for the model code.

Your data pipeline is ready. Next up, you'll build your first language model and see just how badly it performs without any attention mechanism.

Build a Bigram Baseline

Before building a full transformer, you need a baseline to beat. The simplest possible language model predicts the next character using ONLY the current character.

This is called a bigram model. Its output will be nearly random garbage, but that is exactly the point. By seeing how badly a no-context model performs, you will understand WHY self-attention exists.

In this step, get ready to:
  • Build a bigram language model with a single embedding layer.
  • Train the model and observe the loss decrease.
  • Generate text and see the limitation of zero-context modeling.
Define the loss estimation helper

Training a model means measuring how wrong its predictions are (the loss) and nudging it to be less wrong. You need a helper function that measures average loss across many batches so you get a stable number rather than a noisy single-batch estimate.

  • In your Colab notebook, create a new code cell below your existing code and paste the following:
@torch.no_grad()
def estimate_loss():
    """Average loss over many batches for stable measurement."""
    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?

The @torch.no_grad() decorator disables gradient tracking, which saves memory since you are only measuring loss here, not training.

The function switches the model to evaluation mode, samples eval_iters (200) batches from both the training and validation splits, averages the loss, then switches back to training mode. This gives a reliable loss estimate instead of a noisy single-batch number.

  • Run the cell. You should see it execute with no errors (the green checkmark appears to the left of the cell).
Build the bigram language model

A bigram model makes a prediction about the next character based solely on the current character. It uses a single embedding table where row i contains the probability distribution over what character comes after character i.

  • Create a new code cell and paste the model class:
class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Single embedding: each token looks up its next-token prediction
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        # Look up logits (raw scores) for each token
        logits = self.token_embedding_table(idx)
        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

    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            logits, loss = self(idx)
            # Only use the last position's prediction
            logits = logits[:, -1, :] / temperature
            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)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

What does this code do?

  • nn.Embedding(vocab_size, vocab_size) creates a lookup table. Each of the ~95 characters gets a row of ~95 numbers representing how likely each character is to come next.
  • The forward method looks up these raw scores (logits), reshapes them, and computes cross-entropy loss against the actual next characters.
  • The generate method produces new text by repeatedly sampling the next character from the probability distribution. It takes the last position's logits, converts to probabilities with softmax, and samples one character.
  • Run the cell. It should execute without errors, confirming the class is defined.

Seeing a NameError?

If you see NameError: name 'vocab_size' is not defined, make sure you ran all the cells from Step 1 first. Colab requires cells to be executed in order from top to bottom.

Try clicking Runtime > Run all to re-execute everything from the top.

help me fix a NameError in my Colab notebook

Train the model and generate text

Now you will train the bigram model for 1000 steps and then ask it to generate text. This is the moment of truth for your first language model.

Before you run the training cell, make a prediction: do you think the generated output will look like real Python code, or something else entirely?

  • Create a new code cell and paste the training and generation code:
# Instantiate model and move to GPU
model = BigramLanguageModel()
m = model.to(device)

# Set up the optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Train for 1000 steps
for iter in range(1000):
    if iter % 200 == 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()

# Generate 200 characters from the trained model
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=200)[0].tolist()))

What does this code do?

  • The training loop samples a batch, computes the loss, zeroes old gradients, backpropagates, and updates the weights. Every 200 steps it prints the average loss so you can watch the model improve.
  • AdamW is an optimizer that adjusts how much each weight changes per step. It is the standard choice for training transformers.
  • After training, generation starts from a single zero token (a blank slate) and samples 200 new characters one at a time.
  • Run the cell. Training takes about 10-20 seconds on a T4 GPU.

You should see loss values printed every 200 steps, starting around 4.5 and dropping to approximately 2.7. After training completes, 200 characters of generated text appear.

Seeing a CUDA out of memory error?

If you get a CUDA memory error, click Runtime > Restart runtime, then re-run all cells from the top. This clears GPU memory from any previous experiments.

help me fix a CUDA out of memory error in Google Colab

Examine the output

Look at the generated text. You will notice it produces characters that sometimes resemble Python. You might spot occasional newlines, spaces, or even fragments like def or =.

But the output is clearly nonsensical. There is no real structure, no consistent indentation, and no complete statements.

Why is the output so bad?

The bigram model has zero context beyond the single previous character. When it sees the letter d, it knows e often follows. But it cannot learn that def should be followed by a space and then a function name.

This is the fundamental limitation. To generate real code, the model needs to look at many previous characters and decide which ones matter. That is exactly what self-attention provides, and you will build it in the next step.

✔️ Awesome, I've got everything!

Great. Make sure all three cells (estimate_loss, BigramLanguageModel, and training) are saved in your notebook.

ⓧ I'd like to double check the full code

Your notebook should now have the following cells in order, after all the Step 1 code:

  • Cell 1: The estimate_loss() function with the @torch.no_grad() decorator.
  • Cell 2: The BigramLanguageModel class with __init__, forward, and generate methods.
  • Cell 3: Model instantiation, 1000-step training loop printing every 200 steps, and generation of 200 characters.

Scroll through each cell and confirm the code matches what you pasted above.

Your first language model is trained. The loss dropped, proving it learned something about character frequencies. But the output is still gibberish because the model has no ability to look at context. Next up, you will add self-attention so every token can see the characters that came before it.

Add Self-Attention

Your bigram model proved one thing: predicting the next character from a single character produces garbage. The model has no way to know that after def comes a function name, because it never sees beyond the current character.

Self-attention fixes this by letting every token look at all previous tokens and decide which ones matter. This is the core mechanism that makes LLMs work. In this step, you will implement self-attention from scratch and watch your model jump from random noise to recognizable Python structure.

In this step, get ready to:
  • Implement a single head of self-attention with query, key, and value projections.
  • Build a new GPT model that combines token embeddings, positional embeddings, and the attention head.
  • Train the model and compare its output to the bigram baseline.
Implement a single head of self-attention

Every token in your sequence needs a way to ask "which previous tokens should I pay attention to?" and get a weighted mix of their information. Self-attention does this through three learnable projections: queries (what am I looking for?), keys (what do I contain?), and values (what information do I provide?).

  • Create a new code cell below your existing bigram model code and add the Head class:
class Head(nn.Module):
    """One head of self-attention."""

    def __init__(self, head_size):
        super().__init__()
        # Linear projections for query, key, and value
        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)
        # Causal mask: prevents attending to future tokens
        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)
        q = self.query(x)
        # Scaled dot-product attention
        wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5
        # Mask out future positions so tokens only see the past
        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)
        out = wei @ v
        return out

What does this code do?

Each token is projected into three vectors: a query ("what am I looking for?"), a key ("what do I contain?"), and a value ("what do I pass along?").

The attention score between two tokens is the dot product of one token's query with another's key, scaled by the square root of the head size to keep gradients stable. The causal mask (tril) sets future positions to negative infinity before softmax, ensuring a token can only attend to itself and earlier positions.

Using register_buffer for the mask means it automatically moves to the GPU with the model but is not treated as a learnable parameter.

  • Run the cell. It should complete with no output and no errors, confirming the class is defined.

Seeing a NameError?

If you see NameError: name 'n_embd' is not defined, make sure you ran the hyperparameters cell from Step 1 earlier in this session. Colab cells must be run in order.

help me fix a NameError for n_embd in my self-attention Head class

Build the GPT model with self-attention

Now you need a model that combines three ingredients: token embeddings (what character is this?), positional embeddings (where does it sit in the sequence?), and the self-attention head (what context matters?).

  • Create a new code cell and add the GPTLanguageModel class. Start with the constructor and forward method:
class GPTLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Token embeddings: convert character indices to vectors
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # Position embeddings: encode where each token sits
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # Single self-attention head
        self.sa_head = Head(n_embd)
        # Project back to vocabulary size for predictions
        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)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
        # Combine token identity with position information
        x = tok_emb + pos_emb
        x = self.sa_head(x)
        logits = self.lm_head(x)
        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 this code do?

The bigram model had a single embedding table mapping characters directly to logits. This model separates the embedding into two parts: token_embedding_table captures what character it is, and position_embedding_table captures where it sits in the sequence.

Adding these together gives each position a unique identity that combines content and location. The self-attention head then processes this combined representation, letting tokens communicate with each other. Finally, lm_head projects back to vocabulary size so we can predict the next character.

  • Still in the same cell, add the generate method below the forward method (inside the class):
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            # Crop to block_size (position embeddings only go that far)
            idx_cond = idx[:, -block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            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)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

Why crop to block_size?

The position embedding table only has entries for positions 0 through block_size - 1 (256 positions). If the generated sequence grows longer than 256 characters, you would index out of bounds. Cropping to the last block_size tokens keeps the input within the model's context window.

  • Run the cell. It should complete without errors, confirming the full GPTLanguageModel class is defined.

Seeing an IndentationError?

Make sure the generate method is indented one level inside the class (same level as forward). In Colab, the entire class definition must be in a single cell.

help me fix an indentation error in my GPTLanguageModel class

Train the model and compare to the bigram baseline

Time to see if self-attention actually makes a difference. You will train this new model for 3000 steps and compare both the loss and the generated output against your bigram baseline.

  • Create a new code cell and add the training loop with generation:
# Create the model and move to GPU
model = GPTLanguageModel()
m = model.to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Train for 3000 steps
for iter in range(3000):
    if iter % 500 == 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()

# Generate text from the trained model
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=300)[0].tolist()))

Before you run this cell, make a prediction: do you think the loss will drop below the bigram's final loss of ~2.7? How much lower?

  • Run the cell. Training takes about 1-2 minutes on the T4 GPU.

You should see the loss start around 4.3 and steadily drop below 2.0. The generated text now shows recognizable Python structure: indentation, colons after function definitions, and plausible variable names.

What should I see?

Your loss output should look something like:

step 0: train loss 4.3401, val loss 4.3356
step 500: train loss 2.2847, val loss 2.3012
step 1000: train loss 2.0153, val loss 2.0489
step 1500: train loss 1.9201, val loss 1.9642
step 2000: train loss 1.8510, val loss 1.9083
step 2500: train loss 1.8012, val loss 1.8701

Your exact numbers will differ slightly due to random initialization, but the trend should be clear: the loss is significantly better than the bigram model's ~2.7. The generated text should contain Python-like patterns with indentation and keywords.

Why is this so much better than the bigram model?

The bigram model predicts the next character using only the single previous character. The self-attention model can look at up to 256 previous characters simultaneously and decide which ones are relevant.

When the model sees def , self-attention lets the space character attend back to the d, e, and f characters together, recognizing the def keyword pattern and predicting that a function name should follow.

Loss not decreasing?

If your loss stays above 3.0 after 1000 steps, check that you ran the data loading cells earlier so train_data and val_data exist on the GPU. Also confirm your runtime type is set to T4 GPU under Runtime > Change runtime type.

help me debug why my self-attention model loss is not decreasing during training

✔️ Awesome, I've got everything!

Great. Double check that you have three new cells: the Head class, the GPTLanguageModel class, and the training/generation cell.

ⓧ I'd like to double check the full code

Scroll up through your notebook and confirm these code cells exist in order:

  • Hyperparameters cell (batch_size, block_size, etc.).
  • Data loading cell (load_dataset, tokenizer, train/val split).
  • The get_batch and estimate_loss functions.
  • The Head class with __init__ and forward methods.
  • The GPTLanguageModel class with __init__, forward, and generate methods.
  • The training loop running for 3000 steps, printing every 500, followed by generation of 300 tokens.

You just implemented the core mechanism behind ChatGPT and Claude from scratch. Next up, you will scale this single attention head into a full transformer with multiple heads, feed-forward layers, and residual connections to push the quality even further.

Scale Up to a Full Transformer

Your single-head attention model already produces recognizable Python patterns. But one attention head can only learn one type of relationship at a time.

Real GPTs use multiple heads running in parallel. Each head learns a different pattern: one might focus on syntax, another on indentation, another on naming conventions. Combined with feed-forward layers for additional computation, this is what turns a toy model into the full transformer architecture.

In this step, get ready to:
  • Implement MultiHeadAttention, FeedForward, and TransformerBlock classes.
  • Build the full GPT model with 6 layers and ~10M parameters.
  • Train the model and compare output to the bigram baseline.
Implement multi-head attention

A single attention head projects the input into one set of queries, keys, and values. Multiple heads run that same process in parallel with different learned projections, then concatenate their outputs. This lets the model attend to information from different representation subspaces simultaneously.

  • Create a new code cell below your Head class. Define the MultiHeadAttention class by pasting this code:
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(head_size * num_heads, 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 this code do?

This class creates multiple Head instances and runs them all on the same input. Each head produces its own output, and torch.cat concatenates them along the last dimension.

The projection layer (self.proj) maps the concatenated output back to the model's embedding dimension. This is the "communication" step where all heads share their findings.

  • Run this cell. It should complete with no output and no error.
Implement the feed-forward network and transformer block

Attention handles communication between tokens. But the model also needs computation at each position independently. The feed-forward network gives each token a chance to "think" about what it learned from attention.

  • In a new cell below MultiHeadAttention, define the FeedForward class:
class FeedForward(nn.Module):
    """A simple feed-forward network with ReLU activation."""

    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)

What does this code do?

The feed-forward network expands the representation to 4x the embedding size, applies ReLU activation, then projects back down. This expansion gives the model more capacity to learn complex patterns at each position.

Think of attention as "gathering information from other tokens" and the feed-forward layer as "processing that information internally."

  • Run this cell to confirm it completes without error.
  • In another new cell, define the TransformerBlock that combines both components:
class TransformerBlock(nn.Module):
    """Transformer block: communication (attention) followed by computation (feed-forward)."""

    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 does this code do?

Each TransformerBlock follows the same two-step pattern: attention (communication) then feed-forward (computation). Two critical design choices make this work:

  • Residual connections (x = x + ...) add the output back to the input, letting gradients flow directly through the network during training.
  • Layer normalization (self.ln1, self.ln2) normalizes the input before each sub-layer, stabilizing training.
  • Run this cell. No error means your three building blocks are properly defined.

Seeing a NameError?

If you see NameError: name 'Head' is not defined or similar, make sure you ran the cells defining Head, MultiHeadAttention, and FeedForward in order before running the TransformerBlock cell.

help me debug a NameError in my TransformerBlock cell

Build the full GPT model

Now you have all the building blocks. The full GPT stacks multiple TransformerBlocks on top of each other, with proper weight initialization that keeps training stable.

  • Create a new cell. You will replace your old single-head GPTLanguageModel with the full transformer version. Paste the class definition:
class GPTLanguageModel(nn.Module):
    """The full GPT model."""

    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(*[TransformerBlock(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)
        self.apply(self._init_weights)

    def _init_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)

What changed from the single-head version?

  • self.sa_head = Head(n_embd) is replaced by self.blocks, a stack of 6 TransformerBlocks. Each block has its own multi-head attention and feed-forward layer.
  • self.ln_f is a final layer normalization applied after all blocks, stabilizing the output before the prediction head.
  • _init_weights initializes all weights from a normal distribution with std=0.02. This prevents exploding or vanishing gradients at the start of training.
  • In the same cell, directly below the _init_weights method, add the forward and generate methods:
    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        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

    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            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)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

What does forward do differently now?

The key change is two lines: x = self.blocks(x) passes the input through all 6 TransformerBlocks sequentially, and x = self.ln_f(x) normalizes the final output. The generate method is unchanged from before.

  • Run this cell. You should see no output and no error.
  • Instantiate the model and check its size by running this in a new cell:
# --- Training ---
model = GPTLanguageModel()
m = model.to(device)
print(f"Model parameters: {sum(p.numel() for p in m.parameters()) / 1e6:.2f}M")

You should see output like Model parameters: 10.XX M. Your model now has approximately 10 million parameters, up from the few hundred thousand in the single-head version.

Seeing a much smaller parameter count?

If you see fewer than 10M parameters, check that n_layer = 6 in your hyperparameters cell and that you ran it. Also confirm n_embd = 384 and n_head = 6. These values control the model size.

help me debug my GPT model parameter count

Train the full model and compare generations

Your 10M-parameter model is ready. Time to train it for 5000 steps and see what a real transformer can do with Python code.

  • Before you run training, think: given that the single-head model reached a loss of ~2.0, what do you predict the full 6-layer transformer will achieve?
  • In the same cell as your model instantiation (below the print line), add the training loop and generation:
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for iter in range(max_iters):
    if iter % eval_interval == 0 or iter == max_iters - 1:
        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()

# --- Generation ---
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))

What does this training loop do?

  • max_iters is 5000 (set in your hyperparameters cell). The loop evaluates loss every 500 steps so you can watch it decrease.
  • After training completes, it generates 500 characters starting from a zero token (the model picks its own starting point).
  • Training takes approximately 3-5 minutes on a T4 GPU. You will see loss values print every 500 steps.
  • Run this cell. Watch the loss values as they print every 500 steps.

Your training loss should drop to approximately 1.4-1.6 by step 5000. That is a massive improvement over the single-head model's ~2.0 and the bigram model's ~2.7.

After training completes, the generated text appears. You should see multi-line Python functions with proper indentation, colons after function definitions, return statements, and plausible variable names.

How does this compare to the bigram model?

Scroll back up to your bigram output from Step 2. That model produced random characters with occasional Python-ish fragments. Your transformer now generates structured, multi-line code because each token attends to its full 256-character context window through 6 layers of multi-head attention.

This is the same architectural principle behind ChatGPT and Claude. The only differences are scale (billions vs millions of parameters), tokenization (BPE vs characters), and dataset size.

Loss not dropping below 2.0?

Make sure you are running the new GPTLanguageModel cell (with self.blocks), not the old single-head version. If both cells exist in your notebook, the last one you ran defines the model variable.

Also confirm your GPU runtime is still active. Check Runtime > Change runtime type to verify T4 GPU is selected.

help me debug why my transformer loss isn't improving

✔️ Awesome, I've got everything!

Great. Confirm you saved your notebook (Colab auto-saves, but press Cmd+S or Ctrl+S to be safe).

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm these cells exist in order:

  • Hyperparameters cell (batch_size, block_size, etc.).
  • Data loading cell (load_dataset, tokenizer, train/val split).
  • get_batch and estimate_loss functions.
  • Head class (single attention head).
  • MultiHeadAttention class.
  • FeedForward class.
  • TransformerBlock class.
  • GPTLanguageModel class (full version with self.blocks, self.ln_f, and _init_weights).
  • Training loop (5000 steps) and generation (500 characters).

You just built a real GPT from scratch. Your model understands Python structure because 6 layers of multi-head attention let every token reason about its full context. Next up, the Secret Mission adds temperature and top-k sampling so you can control how your model generates text.

Secret mission

Add Temperature and Top-k Sampling

Your model generates text by sampling from a probability distribution. But right now, every token choice is equally "risky." Real LLMs like ChatGPT let you dial creativity up or down with a temperature slider. Can you figure out how your model's existing temperature and top-k parameters work, then demonstrate the difference they make?

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely on Google Colab's free tier, so there are no ongoing costs.

Resources you used:

  • Google Colab notebook with GPU runtime (free tier, auto-disconnects after idle).

Keep everything running

No action needed. Choose this if you're still actively experimenting with your model or want to tweak hyperparameters.

  • Your notebook is automatically saved to your Google Drive.
  • The GPU runtime will auto-disconnect after a period of inactivity, so you won't consume resources indefinitely.
  • You can also download the .ipynb file and push it to GitHub as a portfolio piece.

Pause - I'll come back to this later

Shut down the GPU runtime to free up resources, while keeping your notebook saved.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Disconnect and delete runtime.
  • Your notebook code and outputs remain saved in Google Drive. When you return, you will need to re-run the cells to reload the model into memory.

Delete - I don't want to use this again

Remove the notebook entirely from your Google Drive.

  • In your Colab notebook, click Runtime in the top menu bar, then click Disconnect and delete runtime.
  • Open Google Drive at drive.google.com.
  • Navigate to Colab Notebooks in the left sidebar.
  • Right-click your GPT notebook and select Move to trash.

That's a Wrap!

That's a Wrap!

Nice work. You just built a GPT-style language model entirely from scratch in PyTorch, trained it on real Python code, and watched it generate syntactically plausible functions, loops, and variable assignments.

You've learned how to:

  • Implement the full transformer architecture from scratch: self-attention, multi-head attention, feed-forward networks, residual connections, and layer normalization.
  • Train a character-level language model on Python code and witness firsthand how attention transforms random gibberish into structured output by comparing your bigram baseline against the full transformer.
  • Build a working next-token prediction pipeline with tokenization, positional embeddings, and autoregressive generation to produce multi-line Python functions from a single prompt.
  • Secret Mission: Implemented temperature and top-k sampling to control generation creativity, understanding how real chatbot temperature settings work under the hood.

Ready to quiz yourself?