Build Mini GPT-2 in Colab

Introduction

30 Second Summary

Text suggestions appear while you type, often before you have decided how to finish the sentence. The steps behind each suggestion are usually hidden.

In this project, you will build a compact character-level GPT-2-style language model in Google Colab using PyTorch. You will train it from random weights after proving that a causal mask keeps each prediction from reading future characters.

What You'll Build

Your finished notebook lowers its validation loss during training before turning ROMEO: into a Shakespeare-style continuation from the model you assembled yourself.

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

• A character tokenizer that round-trips ROMEO: through character IDs. Its paired bigram baseline produces an immediate continuation that exposes the limits of one-character context.
• A causality test that catches future-token leakage numerically. Earlier outputs stay fixed once causal attention is working.
• A three-layer mini GPT-2 trained from random weights. It reports its initial and final validation loss before producing a prompt-conditioned sample.
• Secret Mission: Temporarily break the causal mask to trigger future-token leakage. Use the numerical test to diagnose the bug before restoring the mask.

Am I ready for this project?

You'll need a Google account with access to Google Colab plus confidence with Python classes, tensors, and simple neural networks.

All computation runs in the hosted runtime, so your local machine needs no installation.

Before We Start

This first checkpoint commits you to building a compact character-level decoder-only GPT-2-style Transformer from PyTorch primitives. Implementing causal attention yourself reveals how the model prevents future-token leakage.

Pin and Verify the Colab Runtime

A managed Google Colab runtime can change its preinstalled libraries between sessions. Pinning the runtime protects this project from version drift before you write version-sensitive PyTorch code.

In this step, you'll create the notebook that holds your mini GPT-2 project. You'll pin its software versions before checking which compute device is available.

In this step, get ready to:

• Create the named Google Colab notebook.
• Pin the Colab runtime version with a GPU when available.
• Run the first Python cell to verify PyTorch plus the compute device.

Create and name the notebook

Your notebook keeps every code cell plus its visible output in one shareable file. Colab stores the notebook in Google Drive so the file persists after the hosted runtime disconnects.

• Open Google Colab in your browser.
• Create a new notebook from Colab's opening page.
• Click the current notebook name in the top left corner.
• Replace the current name with mini_gpt2_from_scratch.ipynb.
• Press Enter to save the new name.

You'll see mini_gpt2_from_scratch.ipynb in the top left corner. That name confirms you are working in the project notebook stored in Google Drive.

Pin the runtime version

A Colab runtime bundles Python plus its installed libraries. Runtime version 2026.07 provides Python 3.12.13 with PyTorch 2.11.0.

A GPU can speed up the later training step. Colab may provide a CPU because free accelerator availability varies.

• Click Runtime in the top menu.
• Select Change runtime type.
• Select 2026.07 from the Runtime Version dropdown.
• Choose a GPU hardware accelerator if one is offered.
• Keep the standard accelerator when a GPU is unavailable.
• Click Save to apply the runtime settings.
• Click Connect in the top right corner if the notebook is not already connected.

The notebook reconnects with the pinned software image. The CPU fallback keeps every later cell runnable when Colab does not offer a GPU.

Why Pin This Runtime?

Colab updates its default runtime over time. Pinning version 2026.07 gives this notebook the Python and PyTorch versions used to test the project.

That stable environment makes later tensor operations reproducible. It also keeps the attention code aligned with the APIs used throughout the notebook.

Add and run the runtime check

The first cell establishes the constants used by every later model component. It also selects cuda when PyTorch can access it.

• Click the first code cell in mini_gpt2_from_scratch.ipynb.
• Replace the cell contents with this runtime configuration code:

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

SEED = 1337
BATCH_SIZE = 32
BLOCK_SIZE = 64
N_EMBD = 96
N_HEAD = 4
N_LAYER = 3
EVAL_ITERS = 10
BIGRAM_STEPS = 100
GPT_STEPS = 400
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_URL = "https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"

assert N_EMBD % N_HEAD == 0
torch.manual_seed(SEED)
print(f"PyTorch: {torch.__version__}")
print(f"Device: {DEVICE}")

What Does This Code Do?

• The imports provide tensor operations plus the neural-network building blocks used throughout the notebook.
• The urlopen import prepares the notebook to download the Shakespeare corpus in the next step.
• The model constants define the batch size plus the maximum context length. They also set the embedding size plus the number of attention heads.
• The training constants set the evaluation count plus the training duration for both models.
• DEVICE selects cuda when it is available. Its fallback value is cpu.
• The assertion confirms that the embedding dimension divides evenly across the four attention heads.
• torch.manual_seed(SEED) establishes the random seed used by later model initialization plus sampling.
• The final two lines print the installed PyTorch version plus the selected device.

Before you run the cell, which device do you think Colab will report?

• Run the cell by clicking its play button on the left.

✔️ I see PyTorch 2.11.0

The first line reports PyTorch 2.11.0. The second line reports either cuda or cpu as the device.

Your runtime is ready. The notebook now has a verified software environment plus a working compute target.

ⓧ I see another PyTorch version

Another version means the notebook connected to a different Colab runtime image.

• Click Runtime in the top menu.
• Select Change runtime type.
• Choose 2026.07 from the Runtime Version dropdown.
• Click Save to apply the corrected version.
• Run the first cell again after Colab reconnects.

ⓧ The cell does not run

The cell needs an active hosted runtime before Python can execute.

• Click Connect in the top right corner.
• Wait for Colab to finish connecting the runtime.
• Run the first cell again.

Cell Stops Before Printing?

• Check that N_EMBD = 96 appears exactly as shown if the assertion stops the cell.
• Check that N_HEAD = 4 appears exactly as shown if the assertion stops the cell.
• Reconnect the runtime if the cell remains queued without producing output.

Still stuck? Help me diagnose why my Colab runtime verification cell does not print the expected PyTorch version and device.

✔️ Awesome, I've got everything!

Your first cell now contains the complete runtime configuration. Its output confirms the pinned PyTorch version plus the selected device.

ⓧ I'd like to double check the full code

• Compare your first cell with this complete reference:

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

SEED = 1337
BATCH_SIZE = 32
BLOCK_SIZE = 64
N_EMBD = 96
N_HEAD = 4
N_LAYER = 3
EVAL_ITERS = 10
BIGRAM_STEPS = 100
GPT_STEPS = 400
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_URL = "https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"

assert N_EMBD % N_HEAD == 0
torch.manual_seed(SEED)
print(f"PyTorch: {torch.__version__}")
print(f"Device: {DEVICE}")

Your Colab runtime is pinned with a verified PyTorch installation. Next, you'll train a context-blind generator on character-encoded Shakespeare text.

Train a Context-Blind Text Generator

Your Google Colab runtime is pinned to the expected environment. You can now turn the verified notebook into a working text generator.

A bigram model predicts the next character using only the latest character. Its limited continuation gives you an immediate result while exposing why language models need longer context.

In this step, get ready to:

• Download the Shakespeare corpus and convert its characters into token IDs.
• Build training batches and a context-blind bigram model.
• Train the baseline and generate a continuation beginning with ROMEO:.

Prepare the Shakespeare text

Character tokenization assigns one integer to every distinct character in the corpus. The model trains on those integers before decoding its predictions back into readable text.

• Scroll to the bottom of the existing Python cell.
• Append the corpus and tokenizer pipeline by copying this code beneath the device verification prints:

with urlopen(DATA_URL, timeout=30) as response:
    text = response.read().decode("utf-8")

chars = sorted(set(text))
vocab_size = len(chars)
stoi = {char: index for index, char in enumerate(chars)}
itos = {index: char for index, char in enumerate(chars)}


def encode(value):
    return [stoi[char] for char in value]


def decode(token_ids):
    return "".join(itos[token_id] for token_id in token_ids)


data = torch.tensor(encode(text), dtype=torch.long)
split_index = int(0.9 * len(data))
train_data = data[:split_index]
val_data = data[split_index:]

print(f"Corpus characters: {len(text):,}")
print(f"Vocabulary size: {vocab_size}")
print(f"Round trip: {decode(encode('ROMEO:'))}")

What Does This Code Do?

• The download reads the verified Shakespeare text into text.
• The stoi mapping converts characters into integer IDs. The itos mapping reverses that conversion.
• The encode() function prepares text for the model. The decode() function turns predicted IDs back into characters.
• The split reserves the final ten percent of the token IDs as validation data. This gives training a separate sample for measuring progress.

• Run the complete Python cell using the play button beside it.

You should see the corpus character count followed by the vocabulary size. The final line should show Round trip: ROMEO:.

Corpus Cell Not Finishing?

• Confirm that the Colab runtime still shows as connected.
• Check that DATA_URL still matches the value from the runtime setup cell.
• Check that the new code sits below the imports and constants in the same Python cell.

Still stuck? Help me troubleshoot the Shakespeare download or ROMEO: tokenization round trip.

Build batches and the bigram model

Next-character training needs input sequences paired with the character that follows each position. Random batches expose the model to different sections of the corpus during every update.

• Append the batch sampler beneath the tokenizer output prints by copying this code:

def get_batch(split_name):
    source = train_data if split_name == "train" else val_data
    starts = torch.randint(0, len(source) - BLOCK_SIZE - 1, (BATCH_SIZE,))
    inputs = torch.stack([source[i : i + BLOCK_SIZE] for i in starts])
    targets = torch.stack([source[i + 1 : i + BLOCK_SIZE + 1] for i in starts])
    return inputs.to(DEVICE), targets.to(DEVICE)

How Do the Batches Work?

• The split_name argument selects the training or validation token tensor.
• Each random starting position produces one sequence containing BLOCK_SIZE characters.
• Each target sequence is shifted forward by one character. Every input position therefore has a next-character answer.
• The final line moves both tensors onto the selected DEVICE.

• Run the complete Python cell again.

The cell should complete without a tensor shape error. You should still see the corpus statistics followed by the ROMEO: round trip.

Seeing a Batch Error?

• Check that get_batch() appears after train_data and val_data are defined.
• Confirm that every line inside get_batch() uses four spaces of indentation.

Need another pair of eyes? Help me debug the get_batch() function in my Colab notebook.

The baseline uses an embedding table as a next-character lookup. Each current character selects one row containing scores for every possible next character.

• Append the bigram model beneath get_batch() by copying this code:

class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.table = nn.Embedding(vocab_size, vocab_size)
        nn.init.zeros_(self.table.weight)

    def forward(self, idx, targets=None):
        logits = self.table(idx)
        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss

What Does the Bigram Model Learn?

• The embedding table has one row for each possible current character. Every row contains one score for each possible next character.
• Zero initialization gives every next character an equal starting score. Training reshapes those scores around patterns in the corpus.
• Cross-entropy loss compares the predicted scores with the real next-character IDs.
• The returned loss remains None during generation because sampled text has no target answers.

• Run the complete Python cell again.

The cell should finish with the same three corpus checks. This confirms that Colab accepted the model class without a syntax or indentation error.

Bigram Class Not Loading?

• Confirm that BigramLanguageModel inherits from nn.Module.
• Check that return logits, loss lines up with the first statement inside forward().
• Compare the two reshape() calls with the dimensions unpacked from logits.shape.

Still seeing a class error? Help me debug my BigramLanguageModel definition.

Train and generate the baseline

Training repeatedly compares predictions with real next characters. The optimizer uses those losses to update the embedding table.

• Append the evaluation and training helpers beneath BigramLanguageModel by copying this code:

@torch.no_grad()
def evaluate(model):
    model.eval()
    losses = []
    for _ in range(EVAL_ITERS):
        inputs, targets = get_batch("val")
        _, loss = model(inputs, targets)
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)


def train_model(model, steps, learning_rate, label):
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    initial_loss = evaluate(model)

    for step in range(steps):
        inputs, targets = get_batch("train")
        _, loss = model(inputs, targets)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    final_loss = evaluate(model)
    print(f"{label} validation loss: {initial_loss:.3f} -> {final_loss:.3f}")

How Does Training Measure Progress?

• The evaluate() helper averages validation loss across several random batches.
• The AdamW optimizer updates the embedding table after backpropagation calculates its gradients.
• The initial and final validation losses measure progress on held-out text. A lower final value shows that the table learned next-character patterns.

• Run the complete Python cell again.

The cell should complete with the corpus checks. The two training helpers are now ready to evaluate any compatible model.

Training Helpers Not Defined?

• Check that evaluate() appears before train_model().
• Confirm that the decorator sits directly above evaluate().
• Check that optimizer.step() remains inside the training loop.

Need help with the loop? Help me debug evaluate() and train_model().

Autoregressive sampling feeds each sampled character back into the model. Repeating that cycle grows one prompt into a longer continuation.

• Append the generation helper beneath train_model() by copying this code:

@torch.no_grad()
def generate(model, prompt, max_new_tokens, temperature=1.0):
    model.eval()
    token_ids = torch.tensor([encode(prompt)], dtype=torch.long, device=DEVICE)

    for _ in range(max_new_tokens):
        context = token_ids[:, -BLOCK_SIZE:]
        logits, _ = model(context)
        next_logits = logits[:, -1, :] / temperature
        probabilities = F.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probabilities, num_samples=1)
        token_ids = torch.cat((token_ids, next_token), dim=1)

    return decode(token_ids[0].tolist())

How Does Generation Work?

• The prompt becomes a tensor of character IDs on the selected device.
• Each loop keeps at most BLOCK_SIZE tokens as context.
• Softmax converts the latest logits into sampling probabilities.
• The sampled token joins the sequence before the next loop begins.

• Run the complete Python cell again.

The cell should complete with the corpus checks. The notebook can now train a model and sample characters from its predictions.

Generation Helper Failing?

• Confirm that generate() appears after encode() and decode().
• Check that next_logits selects the final time position.
• Confirm that torch.cat() uses dimension 1.

Still blocked? Help me debug autoregressive generation in my bigram notebook.

The final lines create a freshly seeded model before running 100 optimizer updates. They then request 200 new characters from the prompt ROMEO:.

• Append the baseline training and generation calls beneath generate() by copying this code:

torch.manual_seed(SEED)
bigram = BigramLanguageModel().to(DEVICE)
train_model(bigram, BIGRAM_STEPS, learning_rate=0.02, label="Bigram")
print("\nBigram sample:\n")
print(generate(bigram, "ROMEO:", max_new_tokens=200))

What Runs in This Final Chunk?

• Resetting the seed makes the baseline run reproducible within the pinned environment.
• The model moves to the selected compute device before training begins.
• The training call uses the configured BIGRAM_STEPS value and a learning rate of 0.02.
• The final call preserves the prompt before adding sampled characters.

Before you run the cell, consider whether a model that sees only the latest character can produce a coherent Shakespeare continuation.

• Run the complete Python cell to train the baseline and reveal its continuation.

You should see a Bigram validation loss line containing initial and final values. The final value should be lower.

Below Bigram sample: you should see a continuation that begins with ROMEO:. Its character patterns may resemble the corpus while its words remain mostly incoherent.

That is your first visible language-model result. The weak continuation gives you a concrete reason to add longer context.

No Bigram Sample?

• Confirm that the runtime remains connected while the optimizer updates the model.
• Check that the training calls appear after every helper and class definition.
• Confirm that bigram moves to DEVICE before training.

Need help finding the failing line? Help me troubleshoot why the bigram model does not print its validation loss or ROMEO: sample.

✔️ Awesome, I've got everything!

Great work. Your notebook now trains the context-blind baseline and prints a generated continuation beginning with ROMEO:.

ⓧ I'd like to double check the full code

Compare your single Python cell with this cumulative version. Keep every section in the same order.

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

SEED = 1337
BATCH_SIZE = 32
BLOCK_SIZE = 64
N_EMBD = 96
N_HEAD = 4
N_LAYER = 3
EVAL_ITERS = 10
BIGRAM_STEPS = 100
GPT_STEPS = 400
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_URL = "https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"

assert N_EMBD % N_HEAD == 0
torch.manual_seed(SEED)
print(f"PyTorch: {torch.__version__}")
print(f"Device: {DEVICE}")

with urlopen(DATA_URL, timeout=30) as response:
    text = response.read().decode("utf-8")

chars = sorted(set(text))
vocab_size = len(chars)
stoi = {char: index for index, char in enumerate(chars)}
itos = {index: char for index, char in enumerate(chars)}


def encode(value):
    return [stoi[char] for char in value]


def decode(token_ids):
    return "".join(itos[token_id] for token_id in token_ids)


data = torch.tensor(encode(text), dtype=torch.long)
split_index = int(0.9 * len(data))
train_data = data[:split_index]
val_data = data[split_index:]

print(f"Corpus characters: {len(text):,}")
print(f"Vocabulary size: {vocab_size}")
print(f"Round trip: {decode(encode('ROMEO:'))}")


def get_batch(split_name):
    source = train_data if split_name == "train" else val_data
    starts = torch.randint(0, len(source) - BLOCK_SIZE - 1, (BATCH_SIZE,))
    inputs = torch.stack([source[i : i + BLOCK_SIZE] for i in starts])
    targets = torch.stack([source[i + 1 : i + BLOCK_SIZE + 1] for i in starts])
    return inputs.to(DEVICE), targets.to(DEVICE)


class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.table = nn.Embedding(vocab_size, vocab_size)
        nn.init.zeros_(self.table.weight)

    def forward(self, idx, targets=None):
        logits = self.table(idx)
        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss


@torch.no_grad()
def evaluate(model):
    model.eval()
    losses = []
    for _ in range(EVAL_ITERS):
        inputs, targets = get_batch("val")
        _, loss = model(inputs, targets)
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)


def train_model(model, steps, learning_rate, label):
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    initial_loss = evaluate(model)

    for step in range(steps):
        inputs, targets = get_batch("train")
        _, loss = model(inputs, targets)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    final_loss = evaluate(model)
    print(f"{label} validation loss: {initial_loss:.3f} -> {final_loss:.3f}")


@torch.no_grad()
def generate(model, prompt, max_new_tokens, temperature=1.0):
    model.eval()
    token_ids = torch.tensor([encode(prompt)], dtype=torch.long, device=DEVICE)

    for _ in range(max_new_tokens):
        context = token_ids[:, -BLOCK_SIZE:]
        logits, _ = model(context)
        next_logits = logits[:, -1, :] / temperature
        probabilities = F.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probabilities, num_samples=1)
        token_ids = torch.cat((token_ids, next_token), dim=1)

    return decode(token_ids[0].tolist())


torch.manual_seed(SEED)
bigram = BigramLanguageModel().to(DEVICE)
train_model(bigram, BIGRAM_STEPS, learning_rate=0.02, label="Bigram")
print("\nBigram sample:\n")
print(generate(bigram, "ROMEO:", max_new_tokens=200))

How to Use This Reference

This cell combines the verified runtime setup with the complete bigram pipeline. Compare definitions and indentation before rerunning your own cell.

Your baseline can now learn next-character frequencies and generate text. Next, you will test whether attention can read characters that should still be hidden.

Catch Attention Looking Into the Future

The bigram baseline in your Google Colab runtime now produces a Shakespeare continuation from ROMEO:. Its context still stops at the latest character.

PyTorch multi-head self-attention can connect each position to every other position. You will build that unrestricted path before using a deterministic test to expose future-token leakage.

In this step, get ready to:

• Build four-head self-attention from query-key-value projections.
• Create a causality test that changes only future embeddings.
• Run an unmasked probe that measures changes in an earlier prefix.

Build unmasked multi-head attention

Each attention head compares positions through scaled query-key scores. The four attended views return to one embedding for each sequence position.

• In the existing Python cell, find the @torch.no_grad() decorator directly above evaluate(model).
• Add the attention class above that decorator by copying the code below:

class CausalSelfAttention(nn.Module):
    def __init__(self, causal=True):
        super().__init__()
        self.causal = causal
        self.head_size = N_EMBD // N_HEAD
        self.qkv = nn.Linear(N_EMBD, 3 * N_EMBD, bias=False)
        self.projection = nn.Linear(N_EMBD, N_EMBD, bias=False)

    def forward(self, x):
        batch, time, channels = x.shape
        queries, keys, values = self.qkv(x).chunk(3, dim=-1)

        queries = queries.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        keys = keys.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        values = values.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)

        scores = (queries @ keys.transpose(-2, -1)) * (self.head_size ** -0.5)
        weights = F.softmax(scores, dim=-1)
        attended = weights @ values
        attended = attended.transpose(1, 2).contiguous().view(batch, time, channels)
        return self.projection(attended)

What does this code do?

• The qkv projection creates three views of every input embedding.
• The view() operations divide each view across N_HEAD attention heads. Each head receives head_size channels.
• The scaled query-key product measures how strongly each position attends to every sequence position.
• The projection layer combines the attended heads into the original embedding width.

• Run the updated notebook cell to load CausalSelfAttention into the connected runtime.

You should see the existing corpus details plus the bigram validation output complete without an exception. That is the attention plumbing in place: your runtime can calculate four heads while preserving the input shape.

Seeing a tensor shape error?

• Check that N_EMBD divides evenly across N_HEAD.
• Confirm every head reshape uses self.head_size.
• Compare the final view(batch, time, channels) with the shape unpacked at the start of forward().

Still stuck? Help me trace the tensor shapes in my attention class.

Create the causality test

A causality test treats prefix stability as an invariant. It preserves the first four embeddings before changing every later embedding.

• Find the final return decode(token_ids[0].tolist()) line inside generate().
• Add the test function below generate() by copying this code:

@torch.no_grad()
def causality_gap(attention):
    attention.eval()
    torch.manual_seed(SEED + 1)
    embeddings = torch.randn(1, 8, N_EMBD, device=DEVICE)
    changed = embeddings.clone()
    changed[:, 4:, :] = torch.randn_like(changed[:, 4:, :])
    original_prefix = attention(embeddings)[:, :4, :]
    changed_prefix = attention(changed)[:, :4, :]
    return (original_prefix - changed_prefix).abs().max().item()

How does the test work?

• The fixed random seed makes repeated runs use the same initial embeddings.
• The changed tensor preserves positions zero through three.
• The assignment to changed[:, 4:, :] replaces only the future positions.
• The return value captures the largest absolute change across the preserved prefix outputs.

• Run the updated notebook cell to load causality_gap() into the connected runtime.

You should see the existing notebook output complete again without an exception. Your test harness is ready to detect whether future positions influence an earlier prefix.

Test function not loading?

• Confirm the @torch.no_grad() decorator sits immediately above causality_gap(attention).
• Check that both attention results use the prefix slice [:, :4, :].
• Rerun the runtime setup if SEED or N_EMBD is missing from memory.

Need help checking the experiment? Help me debug my causality test.

Expose future-token leakage

The probe selects the unrestricted attention path. Its output gives you one number for the effect of changing future embeddings.

• Add the unmasked probe directly below causality_gap() by copying this code:

probe = CausalSelfAttention(causal=False).to(DEVICE)
unmasked_gap = causality_gap(probe)
print(f"Unmasked future-token gap: {unmasked_gap:.8f}")

What does this probe measure?

• The causal=False argument creates the unmasked experiment.
• The attention softmax receives every scaled score from the sequence.
• The unmasked_gap value records the largest change found in the preserved prefix.

Before you run the cell, do you think replacing only positions four onward will change the outputs for positions zero through three?

• Run the updated notebook cell to test the unmasked attention path.

You should see Unmasked future-token gap followed by a value greater than 0. The result appears before the bigram validation output.

You caught the leak. The earlier outputs changed even though the experiment replaced only future embeddings.

Gap prints zero?

• Verify that the probe uses CausalSelfAttention(causal=False).
• Confirm changed[:, 4:, :] receives new random values.
• Rerun the full notebook cell so the runtime uses the latest class definition.

If the gap remains zero, help me diagnose the unmasked causality probe.

✔️ Awesome, I've got everything!

Your notebook now contains four-head unmasked attention plus a deterministic leakage test. Keep the non-zero gap visible for the checkpoint.

ⓧ I'd like to double check the full code

• Compare the cumulative Python cell in mini_gpt2_from_scratch.ipynb with this reference:

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

SEED = 1337
BATCH_SIZE = 32
BLOCK_SIZE = 64
N_EMBD = 96
N_HEAD = 4
N_LAYER = 3
EVAL_ITERS = 10
BIGRAM_STEPS = 100
GPT_STEPS = 400
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_URL = "https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"

assert N_EMBD % N_HEAD == 0
torch.manual_seed(SEED)
print(f"PyTorch: {torch.__version__}")
print(f"Device: {DEVICE}")

with urlopen(DATA_URL, timeout=30) as response:
    text = response.read().decode("utf-8")

chars = sorted(set(text))
vocab_size = len(chars)
stoi = {char: index for index, char in enumerate(chars)}
itos = {index: char for index, char in enumerate(chars)}


def encode(value):
    return [stoi[char] for char in value]


def decode(token_ids):
    return "".join(itos[token_id] for token_id in token_ids)


data = torch.tensor(encode(text), dtype=torch.long)
split_index = int(0.9 * len(data))
train_data = data[:split_index]
val_data = data[split_index:]

print(f"Corpus characters: {len(text):,}")
print(f"Vocabulary size: {vocab_size}")
print(f"Round trip: {decode(encode('ROMEO:'))}")


def get_batch(split_name):
    source = train_data if split_name == "train" else val_data
    starts = torch.randint(0, len(source) - BLOCK_SIZE - 1, (BATCH_SIZE,))
    inputs = torch.stack([source[i : i + BLOCK_SIZE] for i in starts])
    targets = torch.stack([source[i + 1 : i + BLOCK_SIZE + 1] for i in starts])
    return inputs.to(DEVICE), targets.to(DEVICE)


class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.table = nn.Embedding(vocab_size, vocab_size)
        nn.init.zeros_(self.table.weight)

    def forward(self, idx, targets=None):
        logits = self.table(idx)
        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss


class CausalSelfAttention(nn.Module):
    def __init__(self, causal=True):
        super().__init__()
        self.causal = causal
        self.head_size = N_EMBD // N_HEAD
        self.qkv = nn.Linear(N_EMBD, 3 * N_EMBD, bias=False)
        self.projection = nn.Linear(N_EMBD, N_EMBD, bias=False)

    def forward(self, x):
        batch, time, channels = x.shape
        queries, keys, values = self.qkv(x).chunk(3, dim=-1)

        queries = queries.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        keys = keys.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        values = values.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)

        scores = (queries @ keys.transpose(-2, -1)) * (self.head_size ** -0.5)
        weights = F.softmax(scores, dim=-1)
        attended = weights @ values
        attended = attended.transpose(1, 2).contiguous().view(batch, time, channels)
        return self.projection(attended)


@torch.no_grad()
def evaluate(model):
    model.eval()
    losses = []
    for _ in range(EVAL_ITERS):
        inputs, targets = get_batch("val")
        _, loss = model(inputs, targets)
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)


def train_model(model, steps, learning_rate, label):
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    initial_loss = evaluate(model)

    for step in range(steps):
        inputs, targets = get_batch("train")
        _, loss = model(inputs, targets)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    final_loss = evaluate(model)
    print(f"{label} validation loss: {initial_loss:.3f} -> {final_loss:.3f}")


@torch.no_grad()
def generate(model, prompt, max_new_tokens, temperature=1.0):
    model.eval()
    token_ids = torch.tensor([encode(prompt)], dtype=torch.long, device=DEVICE)

    for _ in range(max_new_tokens):
        context = token_ids[:, -BLOCK_SIZE:]
        logits, _ = model(context)
        next_logits = logits[:, -1, :] / temperature
        probabilities = F.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probabilities, num_samples=1)
        token_ids = torch.cat((token_ids, next_token), dim=1)

    return decode(token_ids[0].tolist())


@torch.no_grad()
def causality_gap(attention):
    attention.eval()
    torch.manual_seed(SEED + 1)
    embeddings = torch.randn(1, 8, N_EMBD, device=DEVICE)
    changed = embeddings.clone()
    changed[:, 4:, :] = torch.randn_like(changed[:, 4:, :])
    original_prefix = attention(embeddings)[:, :4, :]
    changed_prefix = attention(changed)[:, :4, :]
    return (original_prefix - changed_prefix).abs().max().item()


probe = CausalSelfAttention(causal=False).to(DEVICE)
unmasked_gap = causality_gap(probe)
print(f"Unmasked future-token gap: {unmasked_gap:.8f}")

torch.manual_seed(SEED)
bigram = BigramLanguageModel().to(DEVICE)
train_model(bigram, BIGRAM_STEPS, learning_rate=0.02, label="Bigram")
print("\nBigram sample:\n")
print(generate(bigram, "ROMEO:", max_new_tokens=200))

What should this reference prove?

The cumulative cell preserves the runtime setup plus the trained bigram baseline. It adds the unmasked attention class before the existing helpers.

The causality test appears below generate(). Its probe prints the non-zero gap before the bigram training output.

You have turned a silent data leak into a measured failure. Next, a lower-triangular causal mask will restrict every position to available context.

Add Causal Mask and Stack GPT-2 Blocks

Your unmasked attention test exposed a future-token leak. Earlier positions changed when you replaced later embeddings, which means training could use answers hidden in the target sequence.

A causal mask closes that leak by limiting each position to its current token plus earlier tokens. You will use the repaired attention inside a three-layer Transformer that follows the core GPT-2 block pattern.

In this step, get ready to:

• Block future-token leakage with a lower-triangular attention mask.
• Build a GELU feed-forward path inside a pre-normalized residual block.
• Assemble MiniGPT2 with tied output weights and run an untrained generation test.

Repair attention with a causal mask

The attention scores form a square matrix of source positions and destination positions. A lower-triangular mask keeps the current position plus its history visible while hiding every future position.

• In the main code cell of mini_gpt2_from_scratch.ipynb, find the line that calculates scores inside CausalSelfAttention.forward().
• Insert this mask logic immediately below that line:

        if self.causal:
            mask = torch.tril(
                torch.ones(time, time, device=x.device, dtype=torch.bool)
            )
            scores = scores.masked_fill(~mask, float("-inf"))

How does the mask stop leakage?

• The Boolean mask contains a lower triangle of allowed attention links.
• The inverted mask selects links that point into the future.
• Filling those scores with negative infinity gives them zero probability after softmax.
• The self.causal condition keeps the earlier unmasked probe available for comparison.

• Save mini_gpt2_from_scratch.ipynb.

The existing probe still uses causal=False, so its non-zero gap remains your control result. Changing that same probe to causal mode lets you test the repair without introducing another random attention module.

Before you rerun the cell, what do you expect the prefix gap to become when future positions are completely hidden?

• Find the line that prints Unmasked future-token gap.
• Append this second probe check directly below it:

probe.causal = True
masked_gap = causality_gap(probe)
print(f"Masked future-token gap: {masked_gap:.8f}")

What does this test prove?

• Setting probe.causal to True activates the mask inside the same attention module.
• The repeated causality_gap() call changes future embeddings while preserving the prefix.
• A zero maximum difference proves that later positions cannot influence earlier outputs.

• Save mini_gpt2_from_scratch.ipynb.
• Run the notebook's main code cell.

You should see Masked future-token gap: 0.00000000. That zero confirms the prefix is invariant to every change made after position 3.

Still seeing a non-zero masked gap?

• Check that scores.masked_fill() uses ~mask so it targets the upper triangle.
• Confirm the masked scores replace the original scores before F.softmax() runs.
• Confirm the test sets probe.causal = True before calculating masked_gap.

Still stuck? Help me diagnose why my causal attention gap is not zero.

Build the feed-forward and residual block

Attention mixes information across allowed positions. Each decoder block also needs a per-position feed-forward network plus residual connections that preserve a direct path through both transformations.

This block uses pre-layer normalization. Each sublayer receives normalized input before its result is added back to the residual stream.

• Find return self.projection(attended) at the end of CausalSelfAttention.
• Add the MLP class directly below the attention class:

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.expand = nn.Linear(N_EMBD, 4 * N_EMBD)
        self.activation = nn.GELU(approximate="tanh")
        self.project = nn.Linear(4 * N_EMBD, N_EMBD)

    def forward(self, x):
        return self.project(self.activation(self.expand(x)))

What does the feed-forward path do?

• The first linear layer expands each position from N_EMBD features to four times that width.
• The GELU activation introduces a nonlinear transformation using its tanh approximation.
• The final linear layer projects each position back to N_EMBD features so it can rejoin the residual stream.

• Save mini_gpt2_from_scratch.ipynb.
• Run the notebook's main code cell.

You should see the existing leakage checks and Bigram sample complete again without an MLP class-definition error.

Seeing an MLP definition error?

• Check that forward() is indented inside MLP.
• Confirm both linear layers use N_EMBD as the residual-stream width.
• Check that the GELU approximation value is tanh.

Need another pair of eyes? Help me fix my MLP class definition.

• Add the Block class directly below MLP.

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm_1 = nn.LayerNorm(N_EMBD)
        self.attention = CausalSelfAttention(causal=True)
        self.norm_2 = nn.LayerNorm(N_EMBD)
        self.mlp = MLP()

    def forward(self, x):
        x = x + self.attention(self.norm_1(x))
        x = x + self.mlp(self.norm_2(x))
        return x

How does the decoder block work?

• The first layer normalization prepares the residual stream for causal attention.
• The attention result is added to the original input through the first residual path.
• The second layer normalization prepares that updated stream for the feed-forward network.
• The MLP result returns through the second residual path without changing the tensor width.

• Save mini_gpt2_from_scratch.ipynb.
• Run the notebook's main code cell.

The cell should complete through both causality checks and the existing Bigram output. This confirms the pre-normalized residual block can be defined with the repaired attention module.

Block definition not running?

• Confirm both normalization layers receive N_EMBD.
• Check that the block creates CausalSelfAttention(causal=True).
• Make sure each residual expression assigns its result back to x.

If the block still fails, help me trace the tensor flow through my Block class.

Assemble and smoke-test MiniGPT2

The full model converts character IDs into token embeddings. Learned position embeddings add sequence order before three decoder blocks process the combined representation.

The output calculation reuses the token-embedding weights as its vocabulary projection. The untrained generation test checks that these components agree on every tensor shape before training begins.

• Add the first part of MiniGPT2 directly below Block.

class MiniGPT2(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, N_EMBD)
        self.position_embedding = nn.Embedding(BLOCK_SIZE, N_EMBD)
        self.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)])
        self.final_norm = nn.LayerNorm(N_EMBD)
        self.apply(self._init_weights)

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

What does the model setup create?

• The token embedding maps every character ID to an N_EMBD-wide vector.
• The position embedding provides a learned vector for every location up to BLOCK_SIZE.
• The module list registers the three decoder blocks defined by N_LAYER.
• The initialization helper fills linear and embedding weights with small random values while clearing available linear biases.

• Save mini_gpt2_from_scratch.ipynb.
• Run the notebook's main code cell.

You should see the existing outputs complete without a class-definition error. The model container is now ready for its forward pass.

Model setup failing?

• Check that self.blocks uses nn.ModuleList so every decoder block is registered.
• Confirm the position embedding uses BLOCK_SIZE as its number of entries.
• Check that self.apply(self._init_weights) appears after every model component is created.

Need help locating the mismatch? Help me debug the MiniGPT2 model setup.

• Add the forward pass directly below _init_weights() inside MiniGPT2.

    def forward(self, idx, targets=None):
        batch, time = idx.shape
        if time > BLOCK_SIZE:
            raise ValueError("Sequence is longer than BLOCK_SIZE.")

        positions = torch.arange(time, device=idx.device)
        x = self.token_embedding(idx) + self.position_embedding(positions)
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        logits = x @ self.token_embedding.weight.transpose(0, 1)

        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss

How does the forward pass connect the model?

• The sequence-length check protects the learned position-embedding range.
• Token vectors and position vectors combine before passing through all three blocks.
• The final normalization prepares the residual stream for vocabulary prediction.
• Multiplying by the transposed token-embedding weights creates tied output logits for every character.
• The optional target path calculates the same next-token cross-entropy loss used by the baseline.

• Save mini_gpt2_from_scratch.ipynb.
• Run the notebook's main code cell.

You should see the existing outputs complete without an indentation or class-definition error. The next run exercises the full model with actual character IDs.

Forward pass not defined correctly?

• Confirm forward() has the same indentation as _init_weights().
• Check that positions uses the same device as idx.
• Confirm the output projection uses self.token_embedding.weight.transpose(0, 1).

If the method is still failing, help me inspect my MiniGPT2 forward pass.

Before you run the architecture smoke test, do you expect random weights to generate meaningful dialogue or simply prove that generation works?

• Append the untrained model smoke test after the existing Bigram sample code:

torch.manual_seed(SEED)
mini_gpt2 = MiniGPT2().to(DEVICE)
print("\nUntrained Mini GPT-2 smoke test:\n")
print(generate(mini_gpt2, "ROMEO:", max_new_tokens=40))

What does the smoke test check?

• Resetting the seed makes the random initialization reproducible.
• Moving mini_gpt2 to DEVICE keeps its parameters beside the generated token tensor.
• The shared generate() helper repeatedly runs the model and samples the next character.
• Completing 40 sampling iterations proves that embeddings, attention heads, residual paths, and tied logits have compatible shapes.

• Save mini_gpt2_from_scratch.ipynb.
• Run the notebook's main code cell.

You should see Untrained Mini GPT-2 smoke test followed by a sample that begins with ROMEO: and contains 40 newly sampled characters. The text is random at this stage, while the completed generation proves that the architecture runs without a shape error.

Seeing a shape error?

• Check that each attention head returns to N_EMBD channels before the projection.
• Confirm each block preserves the input shape across both residual additions.
• Check that the output projection transposes the token-embedding weight before matrix multiplication.

Still seeing incompatible tensors? Help me trace the MiniGPT2 smoke-test shapes.

✔️ Awesome, I've got everything!

Great work. Your causal mask passes its invariant test, while the untrained model can generate a complete sample from ROMEO:.

ⓧ I'd like to double check the full code

Compare the cumulative main code cell in mini_gpt2_from_scratch.ipynb with this reference:

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

SEED = 1337
BATCH_SIZE = 32
BLOCK_SIZE = 64
N_EMBD = 96
N_HEAD = 4
N_LAYER = 3
EVAL_ITERS = 10
BIGRAM_STEPS = 100
GPT_STEPS = 400
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_URL = "https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"

assert N_EMBD % N_HEAD == 0
torch.manual_seed(SEED)
print(f"PyTorch: {torch.__version__}")
print(f"Device: {DEVICE}")

with urlopen(DATA_URL, timeout=30) as response:
    text = response.read().decode("utf-8")

chars = sorted(set(text))
vocab_size = len(chars)
stoi = {char: index for index, char in enumerate(chars)}
itos = {index: char for index, char in enumerate(chars)}


def encode(value):
    return [stoi[char] for char in value]


def decode(token_ids):
    return "".join(itos[token_id] for token_id in token_ids)


data = torch.tensor(encode(text), dtype=torch.long)
split_index = int(0.9 * len(data))
train_data = data[:split_index]
val_data = data[split_index:]

print(f"Corpus characters: {len(text):,}")
print(f"Vocabulary size: {vocab_size}")
print(f"Round trip: {decode(encode('ROMEO:'))}")


def get_batch(split_name):
    source = train_data if split_name == "train" else val_data
    starts = torch.randint(0, len(source) - BLOCK_SIZE - 1, (BATCH_SIZE,))
    inputs = torch.stack([source[i : i + BLOCK_SIZE] for i in starts])
    targets = torch.stack([source[i + 1 : i + BLOCK_SIZE + 1] for i in starts])
    return inputs.to(DEVICE), targets.to(DEVICE)


class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.table = nn.Embedding(vocab_size, vocab_size)
        nn.init.zeros_(self.table.weight)

    def forward(self, idx, targets=None):
        logits = self.table(idx)
        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss


class CausalSelfAttention(nn.Module):
    def __init__(self, causal=True):
        super().__init__()
        self.causal = causal
        self.head_size = N_EMBD // N_HEAD
        self.qkv = nn.Linear(N_EMBD, 3 * N_EMBD, bias=False)
        self.projection = nn.Linear(N_EMBD, N_EMBD, bias=False)

    def forward(self, x):
        batch, time, channels = x.shape
        queries, keys, values = self.qkv(x).chunk(3, dim=-1)

        queries = queries.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        keys = keys.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        values = values.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)

        scores = (queries @ keys.transpose(-2, -1)) * (self.head_size ** -0.5)
        if self.causal:
            mask = torch.tril(
                torch.ones(time, time, device=x.device, dtype=torch.bool)
            )
            scores = scores.masked_fill(~mask, float("-inf"))

        weights = F.softmax(scores, dim=-1)
        attended = weights @ values
        attended = attended.transpose(1, 2).contiguous().view(batch, time, channels)
        return self.projection(attended)


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.expand = nn.Linear(N_EMBD, 4 * N_EMBD)
        self.activation = nn.GELU(approximate="tanh")
        self.project = nn.Linear(4 * N_EMBD, N_EMBD)

    def forward(self, x):
        return self.project(self.activation(self.expand(x)))


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm_1 = nn.LayerNorm(N_EMBD)
        self.attention = CausalSelfAttention(causal=True)
        self.norm_2 = nn.LayerNorm(N_EMBD)
        self.mlp = MLP()

    def forward(self, x):
        x = x + self.attention(self.norm_1(x))
        x = x + self.mlp(self.norm_2(x))
        return x


class MiniGPT2(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, N_EMBD)
        self.position_embedding = nn.Embedding(BLOCK_SIZE, N_EMBD)
        self.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)])
        self.final_norm = nn.LayerNorm(N_EMBD)
        self.apply(self._init_weights)

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

    def forward(self, idx, targets=None):
        batch, time = idx.shape
        if time > BLOCK_SIZE:
            raise ValueError("Sequence is longer than BLOCK_SIZE.")

        positions = torch.arange(time, device=idx.device)
        x = self.token_embedding(idx) + self.position_embedding(positions)
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        logits = x @ self.token_embedding.weight.transpose(0, 1)

        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss


@torch.no_grad()
def evaluate(model):
    model.eval()
    losses = []
    for _ in range(EVAL_ITERS):
        inputs, targets = get_batch("val")
        _, loss = model(inputs, targets)
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)


def train_model(model, steps, learning_rate, label):
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    initial_loss = evaluate(model)

    for step in range(steps):
        inputs, targets = get_batch("train")
        _, loss = model(inputs, targets)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    final_loss = evaluate(model)
    print(f"{label} validation loss: {initial_loss:.3f} -> {final_loss:.3f}")


@torch.no_grad()
def generate(model, prompt, max_new_tokens, temperature=1.0):
    model.eval()
    token_ids = torch.tensor([encode(prompt)], dtype=torch.long, device=DEVICE)

    for _ in range(max_new_tokens):
        context = token_ids[:, -BLOCK_SIZE:]
        logits, _ = model(context)
        next_logits = logits[:, -1, :] / temperature
        probabilities = F.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probabilities, num_samples=1)
        token_ids = torch.cat((token_ids, next_token), dim=1)

    return decode(token_ids[0].tolist())


@torch.no_grad()
def causality_gap(attention):
    attention.eval()
    torch.manual_seed(SEED + 1)
    embeddings = torch.randn(1, 8, N_EMBD, device=DEVICE)
    changed = embeddings.clone()
    changed[:, 4:, :] = torch.randn_like(changed[:, 4:, :])
    original_prefix = attention(embeddings)[:, :4, :]
    changed_prefix = attention(changed)[:, :4, :]
    return (original_prefix - changed_prefix).abs().max().item()


probe = CausalSelfAttention(causal=False).to(DEVICE)
unmasked_gap = causality_gap(probe)
print(f"Unmasked future-token gap: {unmasked_gap:.8f}")

probe.causal = True
masked_gap = causality_gap(probe)
print(f"Masked future-token gap: {masked_gap:.8f}")

torch.manual_seed(SEED)
bigram = BigramLanguageModel().to(DEVICE)
train_model(bigram, BIGRAM_STEPS, learning_rate=0.02, label="Bigram")
print("\nBigram sample:\n")
print(generate(bigram, "ROMEO:", max_new_tokens=200))

torch.manual_seed(SEED)
mini_gpt2 = MiniGPT2().to(DEVICE)
print("\nUntrained Mini GPT-2 smoke test:\n")
print(generate(mini_gpt2, "ROMEO:", max_new_tokens=40))

How to use this reference

Check the class order, method indentation, mask placement, and smoke-test calls against your main code cell. The reference represents the complete notebook state for this step.

Your model now blocks future-token leakage and carries character representations through three GPT-2-style decoder blocks. Next, you will train those random weights and compare the model's initial loss with its final loss.

Train Mini GPT-2 and Generate a Continuation

Your causal decoder now blocks future information. Its untrained continuation proves that every tensor shape fits.

Random weights only prove that the architecture runs. Next-token prediction gives the model a learning signal from the Shakespeare corpus.

You will track cross-entropy loss with PyTorch. The final sample shows whether 400 updates produced visible structure.

In this step, get ready to:

• Add validation-loss evaluation plus an AdamW training loop.
• Train Mini GPT-2 while comparing its initial loss with its final loss.
• Generate a 300-character continuation from ROMEO:.

Add evaluation and training loops

Validation loss measures the model on text excluded from its parameter updates. Comparing the loss before training with the loss after training gives you direct evidence that the model learned.

• In the existing notebook cell, find the @torch.no_grad() decorator immediately above generate().
• Insert the following training utilities above that decorator:

@torch.no_grad()
def evaluate(model):
    model.eval()
    losses = []
    for _ in range(EVAL_ITERS):
        inputs, targets = get_batch("val")
        _, loss = model(inputs, targets)
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)


def train_model(model, steps, learning_rate, label):
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    initial_loss = evaluate(model)

    for step in range(steps):
        inputs, targets = get_batch("train")
        _, loss = model(inputs, targets)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    final_loss = evaluate(model)
    print(f"{label} validation loss: {initial_loss:.3f} -> {final_loss:.3f}")

How does training work?

• The @torch.no_grad() decorator keeps evaluation from building gradient data that the model does not need.
• The evaluate() function averages loss across EVAL_ITERS validation batches.
• The train_model() function uses AdamW to update every trainable parameter.
• The loss.backward() call calculates gradients from the current batch.
• The final print compares the two validation measurements with one concise result.

Run the complete training pass

The final execution sequence resets the random seed before creating the training result. This keeps the run reproducible within the pinned runtime.

• Move to the bottom of the existing notebook cell.
• Append the final training sequence beneath the untrained Mini GPT-2 smoke test:

torch.manual_seed(SEED)
parameter_count = sum(parameter.numel() for parameter in mini_gpt2.parameters())
print(f"\nMini GPT-2 parameters: {parameter_count:,}")
train_model(mini_gpt2, GPT_STEPS, learning_rate=0.002, label="Mini GPT-2")
print("\nMini GPT-2 sample:\n")
print(generate(mini_gpt2, "ROMEO:", max_new_tokens=300, temperature=0.9))

What does this sequence do?

• The parameter count shows the total number of values that training can adjust.
• The train_model() call performs GPT_STEPS updates with a learning rate of 0.002.
• The final generate() call performs autoregressive sampling for 300 new characters.
• The temperature of 0.9 shapes the probability distribution used to sample each character.

The training loop reports its result after all 400 updates. A busy cell during that pass is expected on the cpu path.

Before you run the cell, do you expect the final validation loss to finish above or below its initial value?

• Run the complete notebook cell with its run control.
• Keep the hosted runtime connected until the cell finishes.

You should see a line beginning with Mini GPT-2 parameters:. You should then see Mini GPT-2 validation loss: followed by an initial value plus a lower final value.

The output ends with Mini GPT-2 sample: plus a continuation beginning with ROMEO:.

Training did not finish?

• If the cell reports an undefined name, confirm that evaluate() plus train_model() appear above generate().
• If the runtime runs out of accelerator memory, reconnect with the standard accelerator so DEVICE selects cpu.
• If the final loss does not decrease, compare SEED plus GPT_STEPS with the completed code below.

Still stuck? Help me diagnose my Mini GPT-2 training run.

Inspect the learned continuation

A short educational run teaches local character patterns before it learns strong semantic coherence. The sample should be judged by visible structure from the corpus.

• Find the generated text beneath Mini GPT-2 sample:.
• Confirm that the continuation begins with ROMEO:.
• Look for recognizable character sequences in the continuation.
• Look for word-like spacing in the continuation.
• Look for line breaks in the continuation.
• Look for speaker-name patterns in the continuation.
• Treat the generated passage as an educational sample instead of factual text.

You made the full learning loop work. Random parameters now carry patterns learned from the Shakespeare corpus.

✔️ Awesome, I've got everything!

Great. Your notebook now measures validation loss before training plus after training. It also generates a prompt-conditioned continuation from the trained model.

ⓧ I'd like to double check the full code

Compare your complete notebook cell with this final version:

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

SEED = 1337
BATCH_SIZE = 32
BLOCK_SIZE = 64
N_EMBD = 96
N_HEAD = 4
N_LAYER = 3
EVAL_ITERS = 10
BIGRAM_STEPS = 100
GPT_STEPS = 400
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_URL = "https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"

assert N_EMBD % N_HEAD == 0
torch.manual_seed(SEED)
print(f"PyTorch: {torch.__version__}")
print(f"Device: {DEVICE}")

with urlopen(DATA_URL, timeout=30) as response:
    text = response.read().decode("utf-8")

chars = sorted(set(text))
vocab_size = len(chars)
stoi = {char: index for index, char in enumerate(chars)}
itos = {index: char for index, char in enumerate(chars)}


def encode(value):
    return [stoi[char] for char in value]


def decode(token_ids):
    return "".join(itos[token_id] for token_id in token_ids)


data = torch.tensor(encode(text), dtype=torch.long)
split_index = int(0.9 * len(data))
train_data = data[:split_index]
val_data = data[split_index:]

print(f"Corpus characters: {len(text):,}")
print(f"Vocabulary size: {vocab_size}")
print(f"Round trip: {decode(encode('ROMEO:'))}")


def get_batch(split_name):
    source = train_data if split_name == "train" else val_data
    starts = torch.randint(0, len(source) - BLOCK_SIZE - 1, (BATCH_SIZE,))
    inputs = torch.stack([source[i : i + BLOCK_SIZE] for i in starts])
    targets = torch.stack([source[i + 1 : i + BLOCK_SIZE + 1] for i in starts])
    return inputs.to(DEVICE), targets.to(DEVICE)


class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.table = nn.Embedding(vocab_size, vocab_size)
        nn.init.zeros_(self.table.weight)

    def forward(self, idx, targets=None):
        logits = self.table(idx)
        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss


class CausalSelfAttention(nn.Module):
    def __init__(self, causal=True):
        super().__init__()
        self.causal = causal
        self.head_size = N_EMBD // N_HEAD
        self.qkv = nn.Linear(N_EMBD, 3 * N_EMBD, bias=False)
        self.projection = nn.Linear(N_EMBD, N_EMBD, bias=False)

    def forward(self, x):
        batch, time, channels = x.shape
        queries, keys, values = self.qkv(x).chunk(3, dim=-1)

        queries = queries.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        keys = keys.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)
        values = values.view(batch, time, N_HEAD, self.head_size).transpose(1, 2)

        scores = (queries @ keys.transpose(-2, -1)) * (self.head_size ** -0.5)
        if self.causal:
            mask = torch.tril(
                torch.ones(time, time, device=x.device, dtype=torch.bool)
            )
            scores = scores.masked_fill(~mask, float("-inf"))

        weights = F.softmax(scores, dim=-1)
        attended = weights @ values
        attended = attended.transpose(1, 2).contiguous().view(batch, time, channels)
        return self.projection(attended)


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.expand = nn.Linear(N_EMBD, 4 * N_EMBD)
        self.activation = nn.GELU(approximate="tanh")
        self.project = nn.Linear(4 * N_EMBD, N_EMBD)

    def forward(self, x):
        return self.project(self.activation(self.expand(x)))


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm_1 = nn.LayerNorm(N_EMBD)
        self.attention = CausalSelfAttention(causal=True)
        self.norm_2 = nn.LayerNorm(N_EMBD)
        self.mlp = MLP()

    def forward(self, x):
        x = x + self.attention(self.norm_1(x))
        x = x + self.mlp(self.norm_2(x))
        return x


class MiniGPT2(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, N_EMBD)
        self.position_embedding = nn.Embedding(BLOCK_SIZE, N_EMBD)
        self.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)])
        self.final_norm = nn.LayerNorm(N_EMBD)
        self.apply(self._init_weights)

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

    def forward(self, idx, targets=None):
        batch, time = idx.shape
        if time > BLOCK_SIZE:
            raise ValueError("Sequence is longer than BLOCK_SIZE.")

        positions = torch.arange(time, device=idx.device)
        x = self.token_embedding(idx) + self.position_embedding(positions)
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        logits = x @ self.token_embedding.weight.transpose(0, 1)

        loss = None
        if targets is not None:
            batch, time, classes = logits.shape
            loss = F.cross_entropy(
                logits.reshape(batch * time, classes),
                targets.reshape(batch * time),
            )
        return logits, loss


@torch.no_grad()
def evaluate(model):
    model.eval()
    losses = []
    for _ in range(EVAL_ITERS):
        inputs, targets = get_batch("val")
        _, loss = model(inputs, targets)
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)


def train_model(model, steps, learning_rate, label):
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    initial_loss = evaluate(model)

    for step in range(steps):
        inputs, targets = get_batch("train")
        _, loss = model(inputs, targets)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    final_loss = evaluate(model)
    print(f"{label} validation loss: {initial_loss:.3f} -> {final_loss:.3f}")


@torch.no_grad()
def generate(model, prompt, max_new_tokens, temperature=1.0):
    model.eval()
    token_ids = torch.tensor([encode(prompt)], dtype=torch.long, device=DEVICE)

    for _ in range(max_new_tokens):
        context = token_ids[:, -BLOCK_SIZE:]
        logits, _ = model(context)
        next_logits = logits[:, -1, :] / temperature
        probabilities = F.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probabilities, num_samples=1)
        token_ids = torch.cat((token_ids, next_token), dim=1)

    return decode(token_ids[0].tolist())


@torch.no_grad()
def causality_gap(attention):
    attention.eval()
    torch.manual_seed(SEED + 1)
    embeddings = torch.randn(1, 8, N_EMBD, device=DEVICE)
    changed = embeddings.clone()
    changed[:, 4:, :] = torch.randn_like(changed[:, 4:, :])
    original_prefix = attention(embeddings)[:, :4, :]
    changed_prefix = attention(changed)[:, :4, :]
    return (original_prefix - changed_prefix).abs().max().item()


probe = CausalSelfAttention(causal=False).to(DEVICE)
unmasked_gap = causality_gap(probe)
print(f"Unmasked future-token gap: {unmasked_gap:.8f}")

probe.causal = True
masked_gap = causality_gap(probe)
print(f"Masked future-token gap: {masked_gap:.8f}")

torch.manual_seed(SEED)
bigram = BigramLanguageModel().to(DEVICE)
train_model(bigram, BIGRAM_STEPS, learning_rate=0.02, label="Bigram")
print("\nBigram sample:\n")
print(generate(bigram, "ROMEO:", max_new_tokens=200))

torch.manual_seed(SEED)
mini_gpt2 = MiniGPT2().to(DEVICE)
print("\nUntrained Mini GPT-2 smoke test:\n")
print(generate(mini_gpt2, "ROMEO:", max_new_tokens=40))

torch.manual_seed(SEED)
parameter_count = sum(parameter.numel() for parameter in mini_gpt2.parameters())
print(f"\nMini GPT-2 parameters: {parameter_count:,}")
train_model(mini_gpt2, GPT_STEPS, learning_rate=0.002, label="Mini GPT-2")
print("\nMini GPT-2 sample:\n")
print(generate(mini_gpt2, "ROMEO:", max_new_tokens=300, temperature=0.9))

Your notebook now trains a decoder-only Transformer from random weights. Its loss comparison plus prompt-conditioned sample show the complete path from architecture to learned behavior.

Secret mission

Break and Repair Causality

Temporarily disable the causal mask while keeping causal mode enabled. Use the leakage test to catch the bug. Then restore the mask and document why a low training loss cannot prove genuine next-token learning when future characters are visible.

Clean Up Your Resources

Clean Up Your Resources

Google Colab keeps this project's notebook in Google Drive at no ongoing cost. Decide whether to keep the notebook, disconnect its temporary runtime, or delete the notebook entirely.

Resources you used:

• The mini_gpt2_from_scratch.ipynb notebook stored in Google Drive contains the completed code, generated outputs, and causality debugging note.
• A temporary hosted runtime contains the downloaded corpus, baseline model, trained in-memory mini_gpt2, and repaired causal-mask test.

Keep everything running

No action is needed. Choose this if you want to keep the completed notebook ready for another walkthrough.

• Keep mini_gpt2_from_scratch.ipynb in Google Drive.
• Leave the generated text, validation losses, and causality results visible in the saved notebook.
• Rerun the notebook from the first cell whenever a fresh runtime needs to rebuild the models.

Pause - I'll come back to this later

Pausing releases the temporary runtime while preserving your notebook in Google Drive.

• Switch back to the connected mini_gpt2_from_scratch.ipynb notebook.
• Use the connection control in the top-right corner to disconnect the hosted runtime.
• Confirm that the notebook shows no connected runtime.

You are paused cleanly. The saved notebook and its visible output remain available in Google Drive.

A future runtime starts without the trained in-memory weights. Rerunning the notebook retrains both models.

Delete - I don't want to use this again

Deletion can feel drastic. This cleanup targets only mini_gpt2_from_scratch.ipynb plus its temporary runtime.

• Switch back to the connected mini_gpt2_from_scratch.ipynb notebook.
• Use the connection control in the top-right corner to disconnect the hosted runtime.
• Confirm that the notebook shows no connected runtime.
• Open Google Drive in a browser tab.
• Use the notebook's context menu to move mini_gpt2_from_scratch.ipynb to Google Drive's trash.
• Open the trash section in Google Drive.
• Permanently delete mini_gpt2_from_scratch.ipynb from the trash.
• Confirm that mini_gpt2_from_scratch.ipynb no longer appears in Google Drive's trash.

Cleanup complete. The notebook is gone from Google Drive.

No project runtime remains connected.

Nice Work!

Nice Work!

You did it! You built a character-level GPT-2-style language model from random weights in Google Colab. The completed notebook turns the prompt ROMEO: into a learned Shakespeare-style continuation.

You've learned how to:

• Turn Shakespeare text into character IDs with a character tokenizer. Train a context-blind bigram baseline that reveals the limits of single-character context.
• Expose future-token leakage with the deterministic causality_gap() test. Repair the leak with a causal mask that keeps earlier outputs invariant when later inputs change.
• Build a three-layer decoder-only Transformer from PyTorch primitives. Combine multi-head self-attention with learned position embeddings. Add pre-layer normalization around residual paths. Use GELU feed-forward layers inside each block. Tie the output projection to the token embeddings. Train the model from random weights. Compare its initial loss with its final validation loss. Generate a prompt-conditioned continuation from ROMEO:.
• Complete the optional Secret Mission by disabling the restored mask on purpose. Use the zero-gap invariant to diagnose the future-token leak. Repair the mask until the test returns 0.00000000. Document why low training loss becomes misleading when future targets are visible.

Ready to quiz yourself?