Build a GPT From Scratch With PyTorch

Build and train a GPT-style transformer from scratch using PyTorch.

Introduction

30 Second Summary

Every time you type a question into ChatGPT or Claude, a model is predicting the next word in the sentence, one word at a time. That prediction loop is the core of every large language model, and it runs on a surprisingly small set of building blocks.

In this project, you will build a miniature GPT from the ground up using PyTorch. You will code every component of the transformer architecture in a single Python file and train it to generate text right on your computer.

What You'll Build

You type a starting word into your terminal and your model completes it with a sequence of words that mirrors the patterns it learned during training.

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

  • A text-generating model that completes any starting word with text capturing patterns from its training data.
  • A real-time view of learning as the training loss drops from random guessing toward near-zero, printed live in your terminal.
  • Proof of why self-attention is the key innovation behind every LLM. Your model goes from generating word salad without it to producing coherent patterns with it.
  • Secret Mission: Add temperature and top-k sampling controls to tune how creative or conservative your model's output is.

Are there any prerequisites?

Basic Python programming and some familiarity with array operations (like numpy) is all you need. No prior PyTorch or deep learning experience is required.

Before We Start

Before you write a single line of code, let's lock in what you're about to build and why it matters.

Set Up PyTorch

Every component of your GPT model depends on PyTorch, a deep learning framework that handles tensor operations, automatic gradient computation, and neural network building blocks.

In this step, you'll install PyTorch, create your project file, and verify everything is working before you start building.

In this step, get ready to:
  • Install PyTorch using pip.
  • Create gpt.py with the core PyTorch imports.
  • Verify PyTorch works by printing a version number and a random tensor.
Install PyTorch

PyTorch deprecated its official Anaconda channel starting with version 2.6. pip is now the recommended installation method.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Terminal and press Enter to open it.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir gpt-from-scratch
cd gpt-from-scratch
  • Install PyTorch by running:
pip install torch

PyTorch is a large download, so this may take a minute or two. You should see download progress bars followed by a Successfully installed message when it finishes.

Install failing or hanging?

Make sure you have an active internet connection. pip downloads PyTorch from the Python Package Index, so it needs network access.

If you see a permissions error, try running pip install --user torch instead.

help me troubleshoot my PyTorch installation error

Create your project file and verify the setup

Your entire GPT model will live in a single Python file called gpt.py. Let's create it and add the three imports that every piece of your model will use.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type VS Code and press Enter to open it.
  • Click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop and select the gpt-from-scratch folder you just created, then click Open.
  • Create a new file by clicking the new file icon in the Explorer sidebar, or by pressing Cmd+N (macOS) or Ctrl+N (Windows).
  • Save it immediately as gpt.py by pressing Cmd+S (macOS) or Ctrl+S (Windows) and typing gpt.py as the filename.
  • Add the core PyTorch imports to gpt.py by copying the following code:
import torch
import torch.nn as nn
import torch.nn.functional as F

What do these imports do?

  • torch is the core library. It provides tensors (multi-dimensional arrays, like numpy arrays but with GPU support and automatic gradient tracking).
  • torch.nn (aliased as nn) contains the building blocks for neural networks: layers, embeddings, and loss functions.
  • torch.nn.functional (aliased as F) provides stateless functions like softmax and cross_entropy that you'll use during training and generation.
  • Save gpt.py.

Before you run the verification, what do you think the output will show if PyTorch installed correctly?

  • Switch back to your terminal and verify that PyTorch is working by running:
python3 -c "import torch; print(torch.__version__); print(torch.rand(3, 3))"

You should see a version number (like 2.13.0) followed by a 3x3 grid of random decimal numbers. That grid is a tensor, the fundamental data structure your model will use for everything.

Seeing ModuleNotFoundError: No module named 'torch'?

This means the python3 command is pointing to a different Python installation than the one where pip installed PyTorch.

  • Check which Python pip is linked to by running pip --version. It should show your Anaconda Python 3.13 path.
  • If the paths don't match, try running python3 -m pip install torch to ensure the install targets the correct Python.

help me fix the ModuleNotFoundError for torch in my Anaconda Python 3.13 setup

✔️ Awesome, I've got everything!

Great. Double check you've saved gpt.py with all three imports.

ⓧ I'd like to double check the full code

Here is the complete gpt.py file so far. Make sure yours matches exactly.

import torch
import torch.nn as nn
import torch.nn.functional as F

PyTorch is ready to go. Next up, you'll build your first language model and watch it try to generate text.

Build a Simple Language Model

With PyTorch installed, you have everything you need to build a language model. Before stacking on complexity, you are going to build the simplest possible next-word predictor so you can see the entire pipeline working end-to-end.

This first model will be intentionally limited. Each word position is processed in complete isolation, with no way to look at surrounding context. It will run, it will train, and it will generate text. The question is: will that text make any sense?

In this step, get ready to:
  • Build a data pipeline that turns raw text into numbers a model can learn from.
  • Create a SimpleGPT model with embeddings and a linear projection.
  • Train the model and generate text to see what it produces.
Define your training corpus and hyperparameters

Every language model starts with data. You need a text corpus for the model to learn from, and a set of hyperparameters that control how the model is built and trained.

  • In gpt.py, add this code below your three import lines to define the hyperparameters and a small training corpus:
# ---- Hyperparameters ----
block_size = 32
embedding_dim = 64
n_heads = 4
n_layers = 3
learning_rate = 1e-3
epochs = 500
batch_size = 32

# ---- Data Preparation ----
corpus = """the cat sat on the mat and looked out the window
the dog ran through the park and chased the ball
the bird flew over the tree and sang a beautiful song
the cat chased the mouse around the old house
the dog and the cat sat together by the warm fire
the bird built a nest in the tall tree near the river
the mouse ran under the mat to hide from the cat
the cat climbed up the tree to watch the bird fly
the dog played in the park with the other dogs
the bird sang a song every morning at dawn
the cat slept on the warm mat by the fire all day
the dog fetched the ball from the river and ran back
the mouse found cheese under the old table in the house
the cat and the dog played together in the yard
the bird flew from the tree to the river for water"""

What does this code do?

The hyperparameters control the shape and training of your model. block_size sets how many words of context the model can see at once (32 words). embedding_dim sets the size of each word's numerical representation (64 dimensions).

The corpus is a small collection of simple sentences about cats, dogs, birds, and mice. It is tiny on purpose. A ~150 word corpus trains in seconds on your CPU, and the model will memorize its patterns quickly so you can clearly see whether the architecture is working.

Build the tokenizer and batch function

Neural networks cannot read words directly. Tokenization converts each word into a unique number so the model can process it mathematically.

  • Add this code below the corpus to build the tokenizer, convert the text to numbers, and create a function that produces training batches:
words = corpus.lower().split()
vocab = sorted(set(words))
vocab_size = len(vocab)
word_to_idx = {word: idx for idx, word in enumerate(vocab)}
idx_to_word = {idx: word for word, idx in word_to_idx.items()}

data = torch.tensor([word_to_idx[w] for w in words], dtype=torch.long)

print(f"Corpus: {len(words)} words, {vocab_size} unique")


def get_batch():
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    return x, y

What does this code do?

The first section splits the corpus into individual words and builds a vocabulary. word_to_idx maps each word to a unique integer (like "cat" to 3). idx_to_word does the reverse so you can convert the model's number outputs back to readable words.

The get_batch function creates random training pairs. For each pair, x is a sequence of 32 word indices, and y is the same sequence shifted one position forward. The model's job is to predict each next word from the words before it.

  • Save gpt.py.
  • Run your script to confirm the data pipeline works by running this command in your terminal:
python3 gpt.py

You should see a line like Corpus: 148 words, 42 unique printed. The exact numbers depend on the corpus, but you should see both a total word count and a unique vocabulary count.

Seeing a SyntaxError?

Check that the triple-quoted string for corpus starts and ends with three double-quotes ("""). A missing closing set of quotes is the most common cause.

help me debug a SyntaxError in my gpt.py corpus definition

Build the SimpleGPT model class

Now you need a model that turns those word numbers into predictions. This first version uses word embeddings to convert each word index into a dense vector of numbers, adds positional embeddings so the model knows where each word sits in the sequence, then projects straight to output predictions.

  • Add this class below the get_batch function in gpt.py:
# ---- Simple GPT Model (no attention) ----
class SimpleGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, embedding_dim)
        self.position_embedding = nn.Embedding(block_size, embedding_dim)
        self.head = nn.Linear(embedding_dim, vocab_size)

    def forward(self, idx):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T))
        x = tok_emb + pos_emb
        logits = self.head(x)
        return logits

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_crop = idx[:, -block_size:]
            logits = self(idx_crop)
            logits = logits[:, -1, :]
            next_idx = torch.argmax(logits, dim=-1, keepdim=True)
            idx = torch.cat([idx, next_idx], dim=1)
        return idx

What does this model do?

The __init__ method creates three layers. token_embedding converts each word index into a 64-dimensional vector. position_embedding adds position information so the model knows word order. head projects from embedding space directly to vocabulary scores (logits).

The forward method runs the model: embeddings in, logits out. The generate method produces text one word at a time using greedy decoding. torch.argmax picks whichever word has the highest score at each step.

What is missing from this model?

Notice that forward goes straight from embeddings to the output projection. There is nothing in between that lets one word look at the other words around it. Each word position is processed completely independently. The model knows what word is at position 5, but it has no idea what word is at position 4.

Train the model and generate text

The final piece is a training loop that teaches the model to predict next words, followed by a generation section that puts it to the test.

  • Add this training and generation code at the bottom of gpt.py, below the SimpleGPT class:
# ---- Training ----
model = SimpleGPT()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

print("Training...")
for epoch in range(epochs):
    x, y = get_batch()
    logits = model(x)
    B, T, C = logits.shape
    loss = F.cross_entropy(logits.view(B * T, C), y.view(B * T))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print(f"  Epoch {epoch:>4d} | Loss: {loss.item():.4f}")

print("Training complete!")

# ---- Generation ----
print("\nGenerating text...")
start_word = "the"
start_idx = word_to_idx[start_word]
prompt = torch.tensor([[start_idx]], dtype=torch.long)
generated = model.generate(prompt, max_new_tokens=20)
result = ' '.join([idx_to_word[i.item()] for i in generated[0]])
print(f"\n{result}")

What does this code do?

The training loop runs for 500 epochs. Each epoch grabs a random batch, runs it through the model, calculates cross-entropy loss (how wrong the predictions are), then updates the model's weights using the Adam optimizer. Loss is printed every 100 epochs so you can watch the model learn.

The generation section feeds the word "the" as a starting prompt and asks the model to predict 20 words, one at a time. Each predicted word gets appended to the sequence, and the model uses the growing sequence to pick the next word.

  • Save gpt.py.

Before you run this, think about what you expect to see. The model has learned which words tend to follow other words, but each position is processed completely independently. Do you think the generated text will form coherent sentences?

  • Run your script by running this command:
python3 gpt.py

You should see the training loss printed every 100 epochs, dropping from around 4.0 down toward 2.0 or lower. Then the generated text appears at the bottom.

Look at the generated text. The words are all real words from your corpus, but the output is repetitive nonsense. You might see the same word repeated over and over, or word combinations that never appeared in the training data. That is exactly the expected result.

Why is the output nonsensical?

This model processes each word position in complete isolation. When it predicts the word at position 5, it has no way to check what word is at position 4. It knows that "the" often appears in the corpus, but it cannot learn that "the cat" should be followed by "sat" because it never sees those words together.

This is the fundamental limitation that self-attention was invented to solve. In the next step, you will add attention so each word can look at all the words before it and decide which ones matter most.

Seeing a NameError or RuntimeError?

If you see NameError: name 'SimpleGPT' is not defined, make sure the SimpleGPT class is defined above the training section, not below it. Python reads files top to bottom.

If you see a size mismatch error, double check that embedding_dim is set to 64 in your hyperparameters and that all three layers in __init__ use embedding_dim.

help me debug an error when running my SimpleGPT training script

✔️ Awesome, I've got everything!

Great. Double check you have saved gpt.py and that the generated text appeared at the bottom of your terminal output.

ⓧ I'd like to double check the full code

Scroll through your gpt.py file and confirm the following sections appear in this order, top to bottom:

  • Three imports (torch, torch.nn, torch.nn.functional).
  • Hyperparameters section (block_size through batch_size).
  • Data preparation section (corpus, tokenization, get_batch function).
  • SimpleGPT class with __init__, forward, and generate methods.
  • Training section (model creation, optimizer, 500-epoch loop).
  • Generation section (starting from "the", generating 20 words).

Your model trains and generates text, but the output is a mess. Every word is processed in isolation, so the model has no ability to look at context. In the next step, you will add self-attention, the mechanism that lets each word look at the words before it and decide which ones matter most for predicting what comes next.

Add Self-Attention

In the last step, your model ran and generated text. But the output was repetitive nonsense. The problem is that each word position is processed in complete isolation. Your model has no way to look at surrounding words for context.

Self-attention solves this. It lets each word look at every word that came before it and decide which ones matter most for predicting the next word. This is the core mechanism that makes transformers work.

In this step, get ready to:
  • Build an AttentionHead class that implements scaled dot-product attention with Query, Key, and Value projections.
  • Build a MultiHeadAttention class that runs 4 attention heads in parallel and combines their outputs.
  • Wire attention into your SimpleGPT model and see the improvement in generated text.
Build the AttentionHead class

Every word in a sentence matters differently depending on what you are trying to predict. When your model sees "the cat sat on the", the word "cat" is far more useful for predicting the next word than "the" or "on".

An attention head learns to figure out which words matter most. It does this by computing three things for each word: a Query ("what am I looking for?"), a Key ("what do I contain?"), and a Value ("what information do I carry?").

  • In gpt.py, find the get_batch function.
  • Add the following AttentionHead class below get_batch, and above the SimpleGPT class:
# ---- Attention Head ----
class AttentionHead(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.query = nn.Linear(embedding_dim, head_size, bias=False)
        self.key = nn.Linear(embedding_dim, head_size, bias=False)
        self.value = nn.Linear(embedding_dim, head_size, bias=False)
        self.register_buffer(
            'mask', torch.tril(torch.ones(block_size, block_size))
        )

    def forward(self, x):
        B, T, C = x.shape
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        weights = q @ k.transpose(-2, -1) * (k.shape[-1] ** -0.5)
        weights = weights.masked_fill(self.mask[:T, :T] == 0, float('-inf'))
        weights = F.softmax(weights, dim=-1)
        return weights @ v

What does this code do?

  • self.query, self.key, and self.value are three separate linear projections. Each word's embedding gets transformed into a Query, Key, and Value vector.
  • register_buffer('mask', torch.tril(...)) creates a causal mask. This is a lower-triangular matrix of ones that prevents each word from "cheating" by looking at future words.
  • The forward method computes the attention formula: Q times K-transpose (how much does each word care about every other word?), scaled by the square root of the head size to keep values stable, masked so future words are invisible, softmaxed into probabilities, then multiplied by V to get the final output.

Why scale by the square root?

When Q and K are multiplied together, the resulting values can get very large. Large values push softmax into regions where it produces near-zero gradients, which stalls learning. Dividing by the square root of the head size keeps the values in a range where softmax behaves well.

  • Save gpt.py.
Build the MultiHeadAttention class

A single attention head can only learn one kind of relationship between words. Multi-head attention runs several heads in parallel, each learning a different pattern. One head might learn that adjectives relate to their nouns. Another might learn that verbs relate to their subjects.

  • Add the MultiHeadAttention class directly below the AttentionHead class you just wrote:
# ---- Multi-Head Attention ----
class MultiHeadAttention(nn.Module):
    def __init__(self):
        super().__init__()
        head_size = embedding_dim // n_heads
        self.heads = nn.ModuleList(
            [AttentionHead(head_size) for _ in range(n_heads)]
        )
        self.projection = nn.Linear(embedding_dim, embedding_dim)

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

What does this code do?

  • head_size = embedding_dim // n_heads splits the 64-dimensional embedding into 4 equal slices of 16 dimensions each. Each head works on its own 16-dimensional slice.
  • nn.ModuleList creates 4 independent AttentionHead instances that all run in parallel on the same input.
  • torch.cat glues the 4 outputs (each 16-dimensional) back into a single 64-dimensional vector. The final self.projection linear layer mixes the combined output so the heads can learn from each other.
  • Save gpt.py.
  • Confirm your new classes have no syntax errors by running your script:
python3 gpt.py

The script trains and generates text exactly as before. Your new classes are defined but not wired into the model yet. The important thing is that the script runs without any errors.

Seeing a syntax error?

  • Check that the AttentionHead class is placed between get_batch and SimpleGPT, not inside either of them.
  • Make sure MultiHeadAttention is placed directly after AttentionHead. It depends on AttentionHead, so it must come after it in the file.
  • Double check that register_buffer( is followed by a closing ) on the next line. The parentheses span two lines.

Help me debug a syntax error in my AttentionHead or MultiHeadAttention class.

Wire attention into your SimpleGPT model

Your attention classes are defined but sitting unused. Now you will plug MultiHeadAttention into your SimpleGPT model so every word can look at the words before it.

  • In gpt.py, scroll down to the SimpleGPT class. Update the comment above it from (no attention) to (with attention).
  • Inside __init__, find these two lines:
self.position_embedding = nn.Embedding(block_size, embedding_dim)
self.head = nn.Linear(embedding_dim, vocab_size)
  • Add a new line between them so the three lines look like this:
self.position_embedding = nn.Embedding(block_size, embedding_dim)
self.attention = MultiHeadAttention()
self.head = nn.Linear(embedding_dim, vocab_size)
  • Now inside the forward method, find these two lines:
x = tok_emb + pos_emb
logits = self.head(x)
  • Add a new line between them so the three lines look like this:
x = tok_emb + pos_emb
x = self.attention(x)
logits = self.head(x)

What changed?

Before, the embeddings went straight to the output projection. Each word position was processed in complete isolation. Now the embeddings pass through MultiHeadAttention first, which lets every word attend to all the words before it. The model can now learn that "cat sat" goes together and that "on the" often follows.

  • Save gpt.py.

Before you run the script, think about what you expect to be different this time. In Step 2, the loss hovered around a certain value and the generated text was repetitive. What do you think attention will change?

  • Run the script to train your attention-powered model:
python3 gpt.py

Watch the training loss. It drops faster and reaches lower values than before. The generated text at the end now captures multi-word patterns from the training data instead of repeating the same words over and over.

Why is this such a big difference?

Without attention, the model saw each position as an island. It could learn that "the" is a common word, but it could never learn that "the cat" tends to be followed by "sat" or "chased". Attention gives the model the ability to build relationships between words, which is exactly what language requires.

This is the breakthrough that makes transformers work. Every modern LLM, from GPT to Claude, is built on this same self-attention mechanism.

Not seeing improvement in the output?

  • Make sure you added self.attention = MultiHeadAttention() inside __init__. If you only added it in forward, the model creates a new random attention layer every forward pass.
  • Confirm that x = self.attention(x) appears between x = tok_emb + pos_emb and logits = self.head(x). The order matters.
  • If the loss is not dropping at all, check that your AttentionHead and MultiHeadAttention classes are defined above SimpleGPT in the file. Python reads top to bottom, so the classes must exist before they are used.

Help me debug why my SimpleGPT model with attention is not producing better output than without attention.

Here is what your complete gpt.py should look like at this point.

✔️ Awesome, I've got everything!

Great. Double check you have saved gpt.py.

ⓧ I'd like to double check the full code

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

  • Imports (torch, torch.nn, torch.nn.functional).
  • Hyperparameters (block_size, embedding_dim, n_heads, etc.).
  • Data Preparation (corpus, tokenization, get_batch).
  • AttentionHead class (the one you added in this step).
  • MultiHeadAttention class (the one you added in this step).
  • SimpleGPT class with self.attention = MultiHeadAttention() in __init__ and x = self.attention(x) in forward.
  • Training section (creates model, optimizer, runs 500 epochs).
  • Generation section (generates 20 words from "the").

Your model can now look at context. The generated text is noticeably better, but there is still room to grow. Next up, you will stack attention with feed-forward networks, layer normalization, and residual connections to build a complete GPT.

Assemble the Full GPT

Your self-attention mechanism is working. The model now looks at surrounding context when predicting the next word, and the generated text captures multi-word patterns from the corpus.

Attention alone is a huge improvement, but production transformers like GPT and Claude go further. They stack attention with feed-forward networks, layer normalization, and residual connections into repeated blocks. These components let the model learn deeper, more abstract patterns and train stably even as the network grows.

In this step, get ready to:
  • Build the feed-forward network and transformer block components.
  • Assemble the full GPT model with stacked transformer blocks and probabilistic generation.
  • Train the complete model and verify it generates text that mirrors the training corpus.
Build the feed-forward network and transformer block

Each transformer block has two sub-layers: the attention layer you already built and a feed-forward network. The feed-forward network gives each token a chance to process its attended information through a non-linear transformation.

  • In gpt.py, find the MultiHeadAttention class. Add the FeedForward class below it by copying this code:
# ---- Feed-Forward Network ----
class FeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embedding_dim, 4 * embedding_dim),
            nn.ReLU(),
            nn.Linear(4 * embedding_dim, embedding_dim),
        )

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

What does this code do?

  • The first nn.Linear expands each token's 64-dimensional embedding to 256 dimensions (4x). This gives the network more room to find patterns.
  • nn.ReLU() is the non-linear activation. Without it, stacking linear layers would collapse into a single linear layer.
  • The second nn.Linear compresses back to 64 dimensions so the output fits the rest of the model.

Now you have the two building blocks: attention and feed-forward. The transformer block wraps them together with normalization and residual connections for stable training.

  • Add the TransformerBlock class directly below the FeedForward class you just added:
# ---- Transformer Block ----
class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = MultiHeadAttention()
        self.feed_forward = FeedForward()
        self.norm1 = nn.LayerNorm(embedding_dim)
        self.norm2 = nn.LayerNorm(embedding_dim)

    def forward(self, x):
        x = x + self.attention(self.norm1(x))
        x = x + self.feed_forward(self.norm2(x))
        return x

What does this code do?

  • nn.LayerNorm normalizes the input before each sub-layer. This is called pre-norm and is the pattern used in modern transformers like GPT-2 and LLaMA. It keeps training stable as the network gets deeper.
  • The x = x + ... pattern is a residual connection. By adding the original input back to the sub-layer's output, gradients can flow directly through the addition during backpropagation. This prevents the "vanishing gradient" problem.
  • The forward pass runs: normalize, attend, add residual, then normalize, feed-forward, add residual. Two sub-layers, each with its own norm and skip connection.
  • Save gpt.py.
  • Confirm the new classes have valid syntax by running the script:
python3 gpt.py

The script still uses your SimpleGPT model, so the output looks the same as Step 3. The important thing is that it runs without errors, which means your new classes are syntactically correct.

Seeing a syntax error?

  • Check that FeedForward is defined after MultiHeadAttention and before TransformerBlock. The TransformerBlock class references FeedForward, so it must come after it in the file.
  • Make sure each class has correct indentation. All method definitions should be indented one level inside the class.
  • help me debug a syntax error in my FeedForward or TransformerBlock class in PyTorch
Replace SimpleGPT with the full GPT model

The SimpleGPT model passed embeddings straight through attention to the output projection. The full GPT model stacks multiple transformer blocks in sequence, giving the network depth to learn increasingly abstract patterns.

  • In gpt.py, find the entire SimpleGPT class (starting with # ---- Simple GPT Model (with attention) ----).
  • Delete the entire SimpleGPT class, including the comment above it.
  • In its place, add the new GPT class:
# ---- GPT Model ----
class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, embedding_dim)
        self.position_embedding = nn.Embedding(block_size, embedding_dim)
        self.blocks = nn.Sequential(
            *[TransformerBlock() for _ in range(n_layers)]
        )
        self.norm = nn.LayerNorm(embedding_dim)
        self.head = nn.Linear(embedding_dim, vocab_size)

    def forward(self, idx):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.norm(x)
        logits = self.head(x)
        return logits

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_crop = idx[:, -block_size:]
            logits = self(idx_crop)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            next_idx = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_idx], dim=1)
        return idx

What changed from SimpleGPT?

  • self.blocks = nn.Sequential(*[TransformerBlock() for _ in range(n_layers)]) stacks 3 transformer blocks in sequence. Data flows through all 3, so the model processes information at three levels of depth.
  • A final nn.LayerNorm stabilizes the output after all blocks, right before the vocabulary projection.
  • The generate method now uses torch.multinomial for probabilistic sampling instead of torch.argmax. Instead of always picking the single most likely word, it samples from the probability distribution. This adds natural variation to the output.

The training section still references SimpleGPT, so you need to update it before you can run the script.

  • In the # ---- Training ---- section, find this line:
model = SimpleGPT()
  • Replace SimpleGPT() with GPT() so the line looks like this:
model = GPT()
  • Save gpt.py and run the script to see your full GPT model in action:
python3 gpt.py

You should see training loss drop much faster than before. The generated text now captures longer patterns from the corpus.

Seeing NameError: name 'SimpleGPT' is not defined?

  • Make sure you changed model = SimpleGPT() to model = GPT() in the training section at the bottom of the file.
  • help me fix a NameError in my GPT training script
Add parameter counting and finalize generation

Production models always report their parameter count. This number tells you how many learnable values the model has. Your miniature GPT will have roughly 80,000 to 100,000 parameters. GPT-2 has 124 million.

  • In the training section, find these lines:
model = GPT()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

print("Training...")
  • Add the parameter count lines between the optimizer and the print("Training...") line, so the section now looks like this:
model = GPT()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

param_count = sum(p.numel() for p in model.parameters())
print(f"Model parameters: {param_count:,}")
print("Training...")

What does this code do?

  • p.numel() counts the number of individual values in each parameter tensor. Summing across all parameters gives the total learnable values in the model.
  • The :, format adds commas for readability (e.g. 89,475 instead of 89475).
  • Scroll down to the generation section at the bottom. Find this line:
generated = model.generate(prompt, max_new_tokens=20)
  • Change max_new_tokens=20 to max_new_tokens=30 so the model generates a longer passage:
generated = model.generate(prompt, max_new_tokens=30)
  • Save gpt.py.

Before you run the script, think about what you expect to see. The model now has 3 stacked transformer blocks, layer normalization, and residual connections. How do you think the training loss and generated text will compare to the simpler model?

  • Run the final version of your GPT:
python3 gpt.py

You should see the total parameter count printed first, followed by training loss that drops to near zero. The generated text closely mirrors the style and structure of the training corpus.

What should I see?

The model prints its parameter count (roughly 80,000 to 100,000). Training loss starts around 4.0 and drops to near 0.0 by epoch 400. The generated 30-word passage reads like a remixed version of the training sentences, with recognizable phrases like "the cat sat on the mat" or "the bird flew over the tree."

Because the model now uses probabilistic sampling with torch.multinomial, your output will be slightly different each time you run it. That is expected and intentional.

Loss not dropping to near zero?

  • Make sure you have all 3 transformer blocks. Check that n_layers = 3 is set in your hyperparameters at the top of the file.
  • Verify your TransformerBlock has residual connections. The forward method should use x = x + ... (not just x = ...). Without residuals, gradients vanish and the model cannot learn.
  • help me debug why my GPT model's loss isn't dropping during training

✔️ Awesome, I've got everything!

Great. Double check you have saved gpt.py.

ⓧ I'd like to double check the full code

The full file is too long to display in one block. Scroll up through your gpt.py and confirm each of these sections exists, in this order:

  • Imports: torch, torch.nn, torch.nn.functional.
  • Hyperparameters: block_size, embedding_dim, n_heads, n_layers, learning_rate, epochs, batch_size.
  • Data Preparation: corpus, vocabulary, word_to_idx, idx_to_word, get_batch().
  • AttentionHead class with query, key, value projections and causal mask.
  • MultiHeadAttention class with 4 heads and a projection layer.
  • FeedForward class with 4x expansion MLP.
  • TransformerBlock class with pre-norm LayerNorm and residual connections.
  • GPT class with token/position embeddings, 3 stacked blocks, final norm, linear head, and generate using torch.multinomial.
  • Training section: model = GPT(), parameter count print, 500-epoch loop.
  • Generation section: prompt with "the", max_new_tokens=30.

You just built a complete GPT-style transformer from scratch. Every component inside production LLMs from GPT to Claude is now something you have coded yourself. Next up, the Secret Mission challenges you to add the sampling controls that production LLMs use to tune creativity.

Secret mission

Add Temperature and Top-k Sampling

Your GPT generates text by sampling from the full probability distribution. Production LLMs like Claude and GPT-4 use two powerful controls to tune their output: temperature (how creative vs. conservative) and top-k sampling (how many word candidates to consider). Can you add both to your model and see the difference?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • gpt.py file (your GPT model script).
  • PyTorch package (installed via pip).

Keep everything running

No action needed. Choose this if you want to keep experimenting with your GPT model or extend it further.

  • Your gpt.py file stays wherever you saved it.
  • PyTorch remains installed in your Python environment and is available for future projects.

Pause - I'll come back to this later

There are no running processes to stop. Your script runs once and exits, so there is nothing to pause.

  • Your gpt.py file and PyTorch installation will be ready whenever you come back.
  • No memory or CPU is used while the script is not running.

Delete - I don't want to use this again

Remove all project resources and start fresh.

  • Delete your gpt.py file. In your terminal, navigate to the folder where you saved it and run:
rm gpt.py

On Windows, run this instead:

del gpt.py
  • Uninstall PyTorch by running:
pip uninstall torch -y

Nice Work!

Nice Work!

You just built a working GPT-style language model from the ground up in a single Python file. You coded every component that powers models like GPT and Claude, trained it on your own machine, and watched it learn to generate text.

You've learned how to:

  • Build every core component of the transformer architecture from scratch: tokenization, token and positional embeddings, scaled dot-product attention with causal masking, multi-head attention, feed-forward networks, layer normalization, and residual connections.
  • See firsthand why self-attention is the breakthrough that makes LLMs work, by comparing a model without it (repetitive nonsense) to one with it (coherent, pattern-aware text).
  • Train a model using cross-entropy loss and the Adam optimizer, then generate text autoregressively, one token at a time, the same way every modern LLM produces language.
  • Secret Mission: Implement temperature and top-k sampling, the exact controls used in production LLMs to tune creativity versus coherence.

Ready to quiz yourself?