LearnAdvanced Training & AdaptationBuild GPT from Scratch Lab

⚡HardFine-Tuning & Training

Build GPT from Scratch Lab

Build and train a tiny GPT end to end on Shakespeare: tokenize with GPT-style subwords, remap active token IDs, run causal self-attention, track validation loss, save a checkpoint, and sample text.

15 min read

Learning path

Step 87 of 143 in the full curriculum

Pre-training Data at Scale Continued Pretraining for Domain Shift

Build GPT from Scratch Lab

The data-pipeline chapter ended with token shards ready for training. This lab shows what happens next: a tiny decoder-only GPT consumes token blocks, predicts next tokens, writes a checkpoint, and produces sample text you can judge.^[1]

Reference implementation inspiration matters here. Karpathy's nanoGPT is still one of the clearest end-to-end GPT training codebases, but its fastest Shakespeare on-ramp is character-level. For this lab we keep same spirit and tiny scale, while switching to GPT-2 BPE subwords because that is closer to how real GPT-family models tokenize text.^[2]^[3]

Corpus is bundled with lesson so readers can download same raw text and rerun experiment locally: Download shakespeare.txt. That file is Princeton's course mirror of Shakespeare text. We train on full bundled corpus, not single play excerpt. Run still stays CPU-friendly because model is tiny and update budget is tiny.^[4]

Why use Shakespeare in curriculum that normally stays anchored in orders, refunds, shipping, and customer-support flows? Because here learning target is pure next-token mechanics. Once you understand tiny GPT loop on public corpus, same pipeline transfers back to ecommerce support logs, delivery updates, and other domain text where you would usually fine-tune or continue pretrain.

Tiny GPT-from-scratch lab: full bundled Shakespeare is tokenized with GPT-2 BPE, active token ids are remapped into compact local vocabulary, packed windows feed causal attention, validation is tracked, and sampled text is checked with novel prompt. — Same control flow as bigger GPT runs, stripped down until you can inspect every step. Raw text becomes GPT-style token IDs, fixed windows feed causal self-attention, validation checks whether training is learning anything beyond memorization, and sampled text is fastest sanity check on checkpoint quality.

Why this lab feels different

Small GPT runs still have one non-negotiable contract:

tokenize raw text into IDs
let position t attend only to positions <= t
predict token t + 1
repeat across many packed windows
judge checkpoint with both validation loss and sampled text

If any piece is wrong, loop still runs and still prints numbers. That is why this lesson is valuable. It forces you to make tokenizer, batching, model, and generation agree with each other instead of studying them as isolated concepts.

Five moving parts

Piece	What it does	Failure mode
Corpus	supplies raw language distribution	using too much text for CPU lab or too little text for recognizable output
Tokenizer	turns text into GPT-style subword IDs	forgetting that token IDs depend on exact tokenizer
Active-vocab remap	shrinks sparse GPT-2 token IDs into small local range	keeping 50k-way output head in toy CPU lab
Block packer	slices token stream into fixed contexts and shifted labels	off-by-one windows or broken train/val split
Decoder loop	runs masked self-attention, loss, checkpoint, and generation	missing causal mask or trusting loss without sampling

That middle row matters. GPT-2 BPE can emit IDs anywhere in a 50k vocabulary. Our tiny lab still only activates subset of those IDs, so we remap active token IDs into contiguous local IDs like 0..17484. Same subword tokenization. Much smaller embedding table and output head.

Runnable lab

This lab is more useful if first result is still half-baked. So notebook flow is staged:

train short warmup run
sample rough checkpoint
keep training same model longer
sample again and compare

First cell builds whole pipeline in plain PyTorch, then only trains for 80 steps. That is enough to learn something, but not enough to look good.^[5]

build_gpt_from_scratch.py

from pathlib import Path
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import tiktoken

# Fix random seed so article output is reproducible.
torch.manual_seed(7)

# Load raw text exactly as learner will download it.
all_text = Path("assets/shakespeare.txt").read_text(encoding="utf-8")
corpus = all_text

# Use GPT-2 BPE so tokenization matches real GPT-style subword models.
encoding_name = "gpt2"
encoder = tiktoken.get_encoding(encoding_name)
corpus_token_ids = encoder.encode(corpus)

# GPT-2 token ids are sparse across a 50k vocabulary. This run only needs
# tokens that actually appear inside bundled Shakespeare corpus, so remap them to 0..N-1.
# That keeps embedding table and output head much smaller without changing
# which subword pieces the tokenizer produced.
active_token_ids = sorted(set(corpus_token_ids))
token_to_local = {token_id: idx for idx, token_id in enumerate(active_token_ids)}
local_to_token = {idx: token_id for token_id, idx in token_to_local.items()}
ids = [token_to_local[token_id] for token_id in corpus_token_ids]

# Each training example needs block_size input tokens plus 1 next-token label.
block_size = 48
chunk_size = block_size + 1
train_flat = []
val_flat = []

# Walk through token stream in fixed windows. Every 10th chunk goes to
# validation so train and val still come from same overall distribution.
for chunk_index, start in enumerate(range(0, len(ids) - chunk_size, chunk_size)):
    chunk = ids[start:start + chunk_size]
    target = val_flat if chunk_index % 10 == 0 else train_flat
    target.extend(chunk)

# Convert Python lists into tensors once so batch sampling is cheap.
train_ids = torch.tensor(train_flat, dtype=torch.long)
val_ids = torch.tensor(val_flat, dtype=torch.long)
batch_size = 12

print(
    f"tokens={len(ids)} active_vocab={len(active_token_ids)} "
    f"train={len(train_ids)} val={len(val_ids)}"
)

def sample_batch(source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    # Pick random starting positions from long token stream.
    starts = torch.randint(0, len(source) - block_size - 1, (batch_size,))

    # x is current context window.
    x = torch.stack([source[s:s + block_size] for s in starts])

    # y is same window shifted one token to left. This is whole learning target.
    y = torch.stack([source[s + 1:s + block_size + 1] for s in starts])
    return x, y

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int = 96, n_heads: int = 4):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        # One linear layer projects each position into query, key, and value vectors.
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)

        # Lower-triangular mask blocks attention to future positions.
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(block_size, block_size, dtype=torch.bool)),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seqlen, width = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)

        def split_heads(tensor: torch.Tensor) -> torch.Tensor:
            # Turn [batch, time, width] into [batch, heads, time, head_dim].
            return tensor.view(batch_size, seqlen, self.n_heads, self.head_dim).transpose(1, 2)

        q = split_heads(q)
        k = split_heads(k)
        v = split_heads(v)

        # Attention score = query-key similarity, scaled to keep softmax stable.
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # Any position above diagonal is future token, so hide it from model.
        attn = attn.masked_fill(~self.mask[:seqlen, :seqlen], float("-inf"))
        attn = attn.softmax(dim=-1)

        # Weighted sum of value vectors produces contextualized representation.
        out = attn @ v
        out = out.transpose(1, 2).contiguous().view(batch_size, seqlen, width)
        return self.proj(out)

class Block(nn.Module):
    def __init__(self, d_model: int = 96, n_heads: int = 4):
        super().__init__()

        # Pre-LN transformer block: normalize, attend, add residual, then MLP.
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

class TinyGPT(nn.Module):
    def __init__(self, vocab_size: int, d_model: int = 96, n_heads: int = 4, n_layers: int = 2):
        super().__init__()

        # Token embeddings say "which subword is this?".
        self.token_emb = nn.Embedding(vocab_size, d_model)

        # Position embeddings say "where is this token inside current window?".
        self.pos_emb = nn.Embedding(block_size, d_model)
        self.blocks = nn.ModuleList([Block(d_model, n_heads) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d_model)

        # Final linear layer turns hidden state back into next-token logits.
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, seqlen = x.shape
        positions = torch.arange(seqlen)

        # GPT adds token meaning and position meaning before any attention happens.
        h = self.token_emb(x) + self.pos_emb(positions)[None, :, :]
        for block in self.blocks:
            h = block(h)
        h = self.ln_f(h)
        return self.head(h)

# Build model and optimizer.
model = TinyGPT(vocab_size=len(active_token_ids))
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3)

def evaluate(source: torch.Tensor) -> tuple[float, float]:
    # Average across a few validation batches so accuracy is less noisy.
    model.eval()
    losses = []
    accuracies = []
    with torch.no_grad():
        for _ in range(8):
            x, y = sample_batch(source)
            logits = model(x)
            losses.append(F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item())
            accuracies.append((logits.argmax(dim=-1) == y).float().mean().item())
    model.train()
    return sum(losses) / len(losses), sum(accuracies) / len(accuracies)

def sample_completion(prompt: str, steps: int = 80) -> str:
    # Reset sampling seed so article output stays reproducible.
    torch.manual_seed(17)

    prompt_token_ids = encoder.encode(prompt)
    prompt_local_ids = [token_to_local[token_id] for token_id in prompt_token_ids]
    context = torch.tensor([prompt_local_ids], dtype=torch.long)

    model.eval()
    with torch.no_grad():
        for _ in range(steps):
            # If sample gets longer than block size, GPT only sees most recent window.
            x = context[:, -block_size:]
            logits = model(x)

            # Only final position matters for next-token sampling.
            next_logits = logits[:, -1, :]

            # Restrict to top candidates so toy model doesn't wander too wildly.
            top_values, top_indices = torch.topk(next_logits, k=8, dim=-1)
            probs = torch.softmax(top_values / 0.9, dim=-1)
            sampled_index = torch.multinomial(probs, num_samples=1)
            next_local_id = top_indices.gather(-1, sampled_index)

            # Append sampled token and continue autoregressive loop.
            context = torch.cat([context, next_local_id], dim=1)

    # Convert local ids back to original GPT-2 token ids, then decode to text.
    return encoder.decode([local_to_token[int(idx)] for idx in context[0]])

for step in range(81):
    # 1. Draw random training batch.
    x, y = sample_batch(train_ids)

    # 2. Predict next-token logits for every position.
    logits = model(x)

    # 3. Compare logits against shifted targets.
    loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))

    # 4. Backpropagate and update weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Print occasional train/val snapshots so learner can watch first useful learning happen.
    if step % 40 == 0:
        val_loss, val_acc = evaluate(val_ids)
        print(
            f"step={step:03d} train={loss.item():.3f} "
            f"val={val_loss:.3f} val_acc={val_acc:.3f}"
        )

# Save first-checkpoint metrics so later cell can compare improvement directly.
early_val_loss = val_loss
early_val_acc = val_acc

# Save enough state to keep sampling compatible with trained weights.
checkpoint = {
    "encoding_name": encoding_name,
    "active_token_ids": active_token_ids,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
}

Warmup training log

tokens=1255253 active_vocab=17485 train=1129695 val=125538
step=000 train=9.887 val=9.799 val_acc=0.000
step=040 train=6.522 val=6.415 val_acc=0.145
step=080 train=6.259 val=6.140 val_acc=0.154

Second cell samples that early checkpoint. Prompt is chosen to be not present verbatim in training corpus, so sample is less about copied heading and more about genuine recombination.

build_gpt_from_scratch_warmup_sample.py

# Prompt is intentionally not copied from training corpus verbatim.
prompt = "Good sir,\nSpeak plain.\n"
assert prompt not in corpus

# This first sample should still look rough and undertrained.
sample = sample_completion(prompt)
print(f"prompt_seen_verbatim={prompt in corpus}")
print("sample:")
print("\n".join(line.rstrip() for line in sample.splitlines()))

Warmup sample

prompt_seen_verbatim=False
sample:
Good sir,
Speak plain.

I , a to , and , a . and

And .

And of ,

And

That .
But .

That .
The the , and ,
And .

To ,
And I , and the the , and , and ,

And ,
O

That checkpoint is still half-baked. It knows Shakespeare-ish punctuation and function words, but it doesn't yet hold a stable thought.

Third cell keeps exact same model and optimizer state, trains longer, and prints whether held-out metrics improved.

build_gpt_from_scratch_continue_training.py

# Continue from exact same checkpoint instead of restarting from scratch.
for step in range(81, 321):
    x, y = sample_batch(train_ids)
    logits = model(x)
    loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 80 == 0:
        late_val_loss, late_val_acc = evaluate(val_ids)
        print(
            f"step={step:03d} train={loss.item():.3f} "
            f"val={late_val_loss:.3f} val_acc={late_val_acc:.3f}"
        )

print(f"val_loss_improved_by={early_val_loss - late_val_loss:.3f}")
print(f"val_acc_improved_by={late_val_acc - early_val_acc:.3f}")

checkpoint = {
    "encoding_name": encoding_name,
    "active_token_ids": active_token_ids,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
}

Longer training log

step=160 train=5.871 val=5.883 val_acc=0.166
step=240 train=6.057 val=5.718 val_acc=0.175
step=320 train=5.614 val=5.534 val_acc=0.179
val_loss_improved_by=0.606
val_acc_improved_by=0.025

Fourth cell samples again from same prompt so you can compare rough checkpoint against longer-trained checkpoint.

build_gpt_from_scratch_late_sample.py

# Same prompt, same sampling settings. Only model weights changed.
sample = sample_completion(prompt)
print(f"prompt_seen_verbatim={prompt in corpus}")
print("sample:")
print("\n".join(line.rstrip() for line in sample.splitlines()))

Longer-training sample

prompt_seen_verbatim=False
sample:
Good sir,
Speak plain.
What , he , my love , and the other man ?

He in my heart ,
And to her , and not in my life

And to you are his own time ,
That ,
What !

A father .

I am my lord ?

I have a good heart ;
I am the king .

He ?

That later checkpoint is still not good writing, but it is clearly more accurate than warmup checkpoint:

validation loss drops from 6.140 to 5.534
validation next-token accuracy rises from 0.154 to 0.179
sample shifts from punctuation soup toward dialogue-like clauses and role words

This is also where many readers misread toy-model output. You aren't asking, "is this fluent enough to publish?" You're asking, "did tokenization, masking, and training teach model enough local structure that output now looks Shakespeare-shaped instead of random?" Later run answers yes much more clearly:

Guide to reading sample GPT output: verify prompt is novel, look for speaker-like Shakespeare reply, notice court and family-role words, and expect broken syntax because the toy model is still small. — Read sample like debugging artifact, not polished writing. Novel prompt plus speaker-like reply, court words, and clause-shaped phrasing are positive signs. Repetition and weird syntax are still expected because model is tiny and undertrained.

Why subword IDs are closer to real GPT

Character tokens are useful for first contact because vocabulary is tiny and code is easy. But modern GPT-family models usually train on subword tokenizers, not raw characters.^[3]

That changes three practical things:

one token can span part of word, whole word, punctuation, or whitespace prefix
active vocabulary is much bigger than a character alphabet
generation quality becomes more word-shaped much earlier

This is why lab uses GPT-2 BPE plus local remap:

GPT-2 BPE keeps tokenizer behavior close to real GPTs
local remap keeps tiny model small enough for laptop CPU

If you skip remap and keep full 50k-way vocabulary head, toy model becomes needlessly heavy without teaching anything extra.

Why one-token shift still matters

Nothing about subwords changes core causal objective. Loss still compares logits at position t against ground-truth token at t + 1.

If you forget shift, model stops learning continuation and starts learning reconstruction. Loss may still fall, but checkpoint is optimizing wrong problem.

For packed token block:

text

x: [15496, 995, 11, 262]
y: [  995, 11, 262, 995]

Those are still "predict next token" pairs, even though each integer now stands for BPE token instead of character.

Packing and splitting data

Real GPT pretraining treats tokenized text like one long stream. Fixed windows are cut from stream, not from tidy sentence rows.

This lab keeps same idea:

encode text once
flatten token stream
slice fixed block_size windows
shift targets by one token

Train/val split is chunk-based instead of one big tail split. That keeps validation text closer in distribution to training text for tiny corpus, while still holding out separate windows.

Validation loss versus sample quality

Notice progression:

short warmup checkpoint is mostly punctuation soup
longer run lowers validation loss and raises validation accuracy
later sample has more dialogue-like clauses and role words

That is point of staged notebook. You can see model move from "barely shaped" to "somewhat useful" instead of pretending first checkpoint was already good.

In bigger runs you still don't choose between metrics and generations. You need both.

Common mistakes

Treating pieces as separate chores

Tokenizer, batch packer, causal mask, loss, checkpoint, and sampler have to agree on same ID space and same context contract. Debugging them separately can miss bugs that only appear when checkpoint is sampled.

Changing tokenizer after checkpoint exists

Tokenizer defines meaning of every token ID. If you change tokenizer, embedding rows and output logits stop matching intended text units.

Keeping sparse token IDs in tiny lab

GPT-2 BPE token 50256 is valid, but tiny lab doesn't need to allocate all rows between 0 and 50256 if corpus only touched 17,485 distinct tokens.

Trusting train loss alone

Train loss can keep improving while generations become repetitive or validation drifts up. That is overfitting, not success.

Choosing context window by convenience

block_size is a learning constraint, not only a speed knob. If window is too short for patterns you want model to learn, sample can stay locally plausible while losing longer dialogue structure.

What you should be able to defend

Skill	What strong answer includes
Coherent GPT loop	connects corpus, tokenizer, packed blocks, causal mask, shifted loss, checkpoint, and generation as one contract
Runnable small run	explains GPT-2 BPE tokenization, active-token remap, shifted targets, and held-out validation checks
Failure diagnosis	catches unshifted labels, train-only evaluation, tokenizer/checkpoint mismatch, tiny-context limits, and bad sample behavior

Practice answers

Prompt	Answer sketch
Why is this chapter different from earlier language-modeling and Transformer chapters?	Earlier chapters taught ingredients in isolation. This lab forces them to cooperate in a real GPT-style loop where tokenizer IDs, packed blocks, causal attention, validation, checkpointing, and generation all have to match.
What is most common silent bug in causal language-model training code?	Forgetting one-token shift between logits and labels. If logits at position `t` train against token `t` instead of token `t + 1`, loss can move while objective is wrong.
Why sample text after checkpoint instead of trusting loss alone?	Lower loss doesn't guarantee healthy generations. Sampling exposes repetition, collapse, context-length bugs, and tokenizer mistakes that scalar loss can hide.

What to remember

Real GPTs are token-ID machines, not string machines.
Subword BPE is closer to production GPT tokenization than character-level alphabets.
Tiny labs can still use real subword tokenization if they remap active tokens into compact local IDs.
GPT training loop is still corpus -> token IDs -> packed windows -> masked self-attention -> shifted loss -> checkpoint -> generation.
Validation loss and sampled text should both move in right direction before you trust checkpoint.

Next Step

Continue to Continued Pretraining and Domain Adaptation

This lab built a GPT-style base-model training loop from raw text to checkpoint and sample. Next chapter asks what changes when model already knows general language and you want to adapt it to narrower domain data instead of starting from zero.

PreviousPre-training Data at Scale

Share this article

X Facebook LinkedIn Bluesky Reddit Hacker News Email

References

CS336: Language Modeling from Scratch.

Stanford University · 2026

nanoGPT.

Karpathy, A. · 2025

shakespeare.txt.

Princeton University COS 302 / SML 305 · 2020

Language Models are Unsupervised Multitask Learners.

Radford, A., et al. · 2019

PyTorch: An Imperative Style, High-Performance Deep Learning Library.

Paszke, A., et al. · 2019 · NeurIPS 2019

Back to Topics

LearnAdvanced Training & AdaptationBuild GPT from Scratch Lab

⚡HardFine-Tuning & Training

Build GPT from Scratch Lab

Build and train a tiny GPT end to end on Shakespeare: tokenize with GPT-style subwords, remap active token IDs, run causal self-attention, track validation loss, save a checkpoint, and sample text.

15 min read

Learning path

Step 87 of 143 in the full curriculum

Pre-training Data at Scale Continued Pretraining for Domain Shift

Build GPT from Scratch Lab

Why this lab feels different

Small GPT runs still have one non-negotiable contract:

tokenize raw text into IDs
let position t attend only to positions <= t
predict token t + 1
repeat across many packed windows
judge checkpoint with both validation loss and sampled text

Five moving parts

Piece	What it does	Failure mode
Corpus	supplies raw language distribution	using too much text for CPU lab or too little text for recognizable output
Tokenizer	turns text into GPT-style subword IDs	forgetting that token IDs depend on exact tokenizer
Active-vocab remap	shrinks sparse GPT-2 token IDs into small local range	keeping 50k-way output head in toy CPU lab
Block packer	slices token stream into fixed contexts and shifted labels	off-by-one windows or broken train/val split
Decoder loop	runs masked self-attention, loss, checkpoint, and generation	missing causal mask or trusting loss without sampling

Runnable lab

This lab is more useful if first result is still half-baked. So notebook flow is staged:

train short warmup run
sample rough checkpoint
keep training same model longer
sample again and compare

First cell builds whole pipeline in plain PyTorch, then only trains for 80 steps. That is enough to learn something, but not enough to look good.^[5]

build_gpt_from_scratch.py

from pathlib import Path
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import tiktoken

# Fix random seed so article output is reproducible.
torch.manual_seed(7)

# Load raw text exactly as learner will download it.
all_text = Path("assets/shakespeare.txt").read_text(encoding="utf-8")
corpus = all_text

# Use GPT-2 BPE so tokenization matches real GPT-style subword models.
encoding_name = "gpt2"
encoder = tiktoken.get_encoding(encoding_name)
corpus_token_ids = encoder.encode(corpus)

# GPT-2 token ids are sparse across a 50k vocabulary. This run only needs
# tokens that actually appear inside bundled Shakespeare corpus, so remap them to 0..N-1.
# That keeps embedding table and output head much smaller without changing
# which subword pieces the tokenizer produced.
active_token_ids = sorted(set(corpus_token_ids))
token_to_local = {token_id: idx for idx, token_id in enumerate(active_token_ids)}
local_to_token = {idx: token_id for token_id, idx in token_to_local.items()}
ids = [token_to_local[token_id] for token_id in corpus_token_ids]

# Each training example needs block_size input tokens plus 1 next-token label.
block_size = 48
chunk_size = block_size + 1
train_flat = []
val_flat = []

# Walk through token stream in fixed windows. Every 10th chunk goes to
# validation so train and val still come from same overall distribution.
for chunk_index, start in enumerate(range(0, len(ids) - chunk_size, chunk_size)):
    chunk = ids[start:start + chunk_size]
    target = val_flat if chunk_index % 10 == 0 else train_flat
    target.extend(chunk)

# Convert Python lists into tensors once so batch sampling is cheap.
train_ids = torch.tensor(train_flat, dtype=torch.long)
val_ids = torch.tensor(val_flat, dtype=torch.long)
batch_size = 12

print(
    f"tokens={len(ids)} active_vocab={len(active_token_ids)} "
    f"train={len(train_ids)} val={len(val_ids)}"
)

def sample_batch(source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    # Pick random starting positions from long token stream.
    starts = torch.randint(0, len(source) - block_size - 1, (batch_size,))

    # x is current context window.
    x = torch.stack([source[s:s + block_size] for s in starts])

    # y is same window shifted one token to left. This is whole learning target.
    y = torch.stack([source[s + 1:s + block_size + 1] for s in starts])
    return x, y

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int = 96, n_heads: int = 4):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        # One linear layer projects each position into query, key, and value vectors.
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)

        # Lower-triangular mask blocks attention to future positions.
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(block_size, block_size, dtype=torch.bool)),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seqlen, width = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)

        def split_heads(tensor: torch.Tensor) -> torch.Tensor:
            # Turn [batch, time, width] into [batch, heads, time, head_dim].
            return tensor.view(batch_size, seqlen, self.n_heads, self.head_dim).transpose(1, 2)

        q = split_heads(q)
        k = split_heads(k)
        v = split_heads(v)

        # Attention score = query-key similarity, scaled to keep softmax stable.
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # Any position above diagonal is future token, so hide it from model.
        attn = attn.masked_fill(~self.mask[:seqlen, :seqlen], float("-inf"))
        attn = attn.softmax(dim=-1)

        # Weighted sum of value vectors produces contextualized representation.
        out = attn @ v
        out = out.transpose(1, 2).contiguous().view(batch_size, seqlen, width)
        return self.proj(out)

class Block(nn.Module):
    def __init__(self, d_model: int = 96, n_heads: int = 4):
        super().__init__()

        # Pre-LN transformer block: normalize, attend, add residual, then MLP.
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

class TinyGPT(nn.Module):
    def __init__(self, vocab_size: int, d_model: int = 96, n_heads: int = 4, n_layers: int = 2):
        super().__init__()

        # Token embeddings say "which subword is this?".
        self.token_emb = nn.Embedding(vocab_size, d_model)

        # Position embeddings say "where is this token inside current window?".
        self.pos_emb = nn.Embedding(block_size, d_model)
        self.blocks = nn.ModuleList([Block(d_model, n_heads) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d_model)

        # Final linear layer turns hidden state back into next-token logits.
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, seqlen = x.shape
        positions = torch.arange(seqlen)

        # GPT adds token meaning and position meaning before any attention happens.
        h = self.token_emb(x) + self.pos_emb(positions)[None, :, :]
        for block in self.blocks:
            h = block(h)
        h = self.ln_f(h)
        return self.head(h)

# Build model and optimizer.
model = TinyGPT(vocab_size=len(active_token_ids))
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3)

def evaluate(source: torch.Tensor) -> tuple[float, float]:
    # Average across a few validation batches so accuracy is less noisy.
    model.eval()
    losses = []
    accuracies = []
    with torch.no_grad():
        for _ in range(8):
            x, y = sample_batch(source)
            logits = model(x)
            losses.append(F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item())
            accuracies.append((logits.argmax(dim=-1) == y).float().mean().item())
    model.train()
    return sum(losses) / len(losses), sum(accuracies) / len(accuracies)

def sample_completion(prompt: str, steps: int = 80) -> str:
    # Reset sampling seed so article output stays reproducible.
    torch.manual_seed(17)

    prompt_token_ids = encoder.encode(prompt)
    prompt_local_ids = [token_to_local[token_id] for token_id in prompt_token_ids]
    context = torch.tensor([prompt_local_ids], dtype=torch.long)

    model.eval()
    with torch.no_grad():
        for _ in range(steps):
            # If sample gets longer than block size, GPT only sees most recent window.
            x = context[:, -block_size:]
            logits = model(x)

            # Only final position matters for next-token sampling.
            next_logits = logits[:, -1, :]

            # Restrict to top candidates so toy model doesn't wander too wildly.
            top_values, top_indices = torch.topk(next_logits, k=8, dim=-1)
            probs = torch.softmax(top_values / 0.9, dim=-1)
            sampled_index = torch.multinomial(probs, num_samples=1)
            next_local_id = top_indices.gather(-1, sampled_index)

            # Append sampled token and continue autoregressive loop.
            context = torch.cat([context, next_local_id], dim=1)

    # Convert local ids back to original GPT-2 token ids, then decode to text.
    return encoder.decode([local_to_token[int(idx)] for idx in context[0]])

for step in range(81):
    # 1. Draw random training batch.
    x, y = sample_batch(train_ids)

    # 2. Predict next-token logits for every position.
    logits = model(x)

    # 3. Compare logits against shifted targets.
    loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))

    # 4. Backpropagate and update weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Print occasional train/val snapshots so learner can watch first useful learning happen.
    if step % 40 == 0:
        val_loss, val_acc = evaluate(val_ids)
        print(
            f"step={step:03d} train={loss.item():.3f} "
            f"val={val_loss:.3f} val_acc={val_acc:.3f}"
        )

# Save first-checkpoint metrics so later cell can compare improvement directly.
early_val_loss = val_loss
early_val_acc = val_acc

# Save enough state to keep sampling compatible with trained weights.
checkpoint = {
    "encoding_name": encoding_name,
    "active_token_ids": active_token_ids,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
}

Warmup training log

tokens=1255253 active_vocab=17485 train=1129695 val=125538
step=000 train=9.887 val=9.799 val_acc=0.000
step=040 train=6.522 val=6.415 val_acc=0.145
step=080 train=6.259 val=6.140 val_acc=0.154

Second cell samples that early checkpoint. Prompt is chosen to be not present verbatim in training corpus, so sample is less about copied heading and more about genuine recombination.

build_gpt_from_scratch_warmup_sample.py

# Prompt is intentionally not copied from training corpus verbatim.
prompt = "Good sir,\nSpeak plain.\n"
assert prompt not in corpus

# This first sample should still look rough and undertrained.
sample = sample_completion(prompt)
print(f"prompt_seen_verbatim={prompt in corpus}")
print("sample:")
print("\n".join(line.rstrip() for line in sample.splitlines()))

Warmup sample

prompt_seen_verbatim=False
sample:
Good sir,
Speak plain.

I , a to , and , a . and

And .

And of ,

And

That .
But .

That .
The the , and ,
And .

To ,
And I , and the the , and , and ,

And ,
O

That checkpoint is still half-baked. It knows Shakespeare-ish punctuation and function words, but it doesn't yet hold a stable thought.

Third cell keeps exact same model and optimizer state, trains longer, and prints whether held-out metrics improved.

build_gpt_from_scratch_continue_training.py

# Continue from exact same checkpoint instead of restarting from scratch.
for step in range(81, 321):
    x, y = sample_batch(train_ids)
    logits = model(x)
    loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 80 == 0:
        late_val_loss, late_val_acc = evaluate(val_ids)
        print(
            f"step={step:03d} train={loss.item():.3f} "
            f"val={late_val_loss:.3f} val_acc={late_val_acc:.3f}"
        )

print(f"val_loss_improved_by={early_val_loss - late_val_loss:.3f}")
print(f"val_acc_improved_by={late_val_acc - early_val_acc:.3f}")

checkpoint = {
    "encoding_name": encoding_name,
    "active_token_ids": active_token_ids,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
}

Longer training log

step=160 train=5.871 val=5.883 val_acc=0.166
step=240 train=6.057 val=5.718 val_acc=0.175
step=320 train=5.614 val=5.534 val_acc=0.179
val_loss_improved_by=0.606
val_acc_improved_by=0.025

Fourth cell samples again from same prompt so you can compare rough checkpoint against longer-trained checkpoint.

build_gpt_from_scratch_late_sample.py

# Same prompt, same sampling settings. Only model weights changed.
sample = sample_completion(prompt)
print(f"prompt_seen_verbatim={prompt in corpus}")
print("sample:")
print("\n".join(line.rstrip() for line in sample.splitlines()))

Longer-training sample

prompt_seen_verbatim=False
sample:
Good sir,
Speak plain.
What , he , my love , and the other man ?

He in my heart ,
And to her , and not in my life

And to you are his own time ,
That ,
What !

A father .

I am my lord ?

I have a good heart ;
I am the king .

He ?

That later checkpoint is still not good writing, but it is clearly more accurate than warmup checkpoint:

validation loss drops from 6.140 to 5.534
validation next-token accuracy rises from 0.154 to 0.179
sample shifts from punctuation soup toward dialogue-like clauses and role words

Why subword IDs are closer to real GPT

Character tokens are useful for first contact because vocabulary is tiny and code is easy. But modern GPT-family models usually train on subword tokenizers, not raw characters.^[3]

That changes three practical things:

one token can span part of word, whole word, punctuation, or whitespace prefix
active vocabulary is much bigger than a character alphabet
generation quality becomes more word-shaped much earlier

This is why lab uses GPT-2 BPE plus local remap:

GPT-2 BPE keeps tokenizer behavior close to real GPTs
local remap keeps tiny model small enough for laptop CPU

If you skip remap and keep full 50k-way vocabulary head, toy model becomes needlessly heavy without teaching anything extra.

Why one-token shift still matters

Nothing about subwords changes core causal objective. Loss still compares logits at position t against ground-truth token at t + 1.

If you forget shift, model stops learning continuation and starts learning reconstruction. Loss may still fall, but checkpoint is optimizing wrong problem.

For packed token block:

text

x: [15496, 995, 11, 262]
y: [  995, 11, 262, 995]

Those are still "predict next token" pairs, even though each integer now stands for BPE token instead of character.

Packing and splitting data

Real GPT pretraining treats tokenized text like one long stream. Fixed windows are cut from stream, not from tidy sentence rows.

This lab keeps same idea:

encode text once
flatten token stream
slice fixed block_size windows
shift targets by one token

Train/val split is chunk-based instead of one big tail split. That keeps validation text closer in distribution to training text for tiny corpus, while still holding out separate windows.

Validation loss versus sample quality

Notice progression:

short warmup checkpoint is mostly punctuation soup
longer run lowers validation loss and raises validation accuracy
later sample has more dialogue-like clauses and role words

That is point of staged notebook. You can see model move from "barely shaped" to "somewhat useful" instead of pretending first checkpoint was already good.

In bigger runs you still don't choose between metrics and generations. You need both.

Common mistakes

Treating pieces as separate chores

Changing tokenizer after checkpoint exists

Tokenizer defines meaning of every token ID. If you change tokenizer, embedding rows and output logits stop matching intended text units.

Keeping sparse token IDs in tiny lab

GPT-2 BPE token 50256 is valid, but tiny lab doesn't need to allocate all rows between 0 and 50256 if corpus only touched 17,485 distinct tokens.

Trusting train loss alone

Train loss can keep improving while generations become repetitive or validation drifts up. That is overfitting, not success.

Choosing context window by convenience

block_size is a learning constraint, not only a speed knob. If window is too short for patterns you want model to learn, sample can stay locally plausible while losing longer dialogue structure.

What you should be able to defend

Skill	What strong answer includes
Coherent GPT loop	connects corpus, tokenizer, packed blocks, causal mask, shifted loss, checkpoint, and generation as one contract
Runnable small run	explains GPT-2 BPE tokenization, active-token remap, shifted targets, and held-out validation checks
Failure diagnosis	catches unshifted labels, train-only evaluation, tokenizer/checkpoint mismatch, tiny-context limits, and bad sample behavior

Practice answers

Prompt	Answer sketch
Why is this chapter different from earlier language-modeling and Transformer chapters?	Earlier chapters taught ingredients in isolation. This lab forces them to cooperate in a real GPT-style loop where tokenizer IDs, packed blocks, causal attention, validation, checkpointing, and generation all have to match.
What is most common silent bug in causal language-model training code?	Forgetting one-token shift between logits and labels. If logits at position `t` train against token `t` instead of token `t + 1`, loss can move while objective is wrong.
Why sample text after checkpoint instead of trusting loss alone?	Lower loss doesn't guarantee healthy generations. Sampling exposes repetition, collapse, context-length bugs, and tokenizer mistakes that scalar loss can hide.

What to remember

Real GPTs are token-ID machines, not string machines.
Subword BPE is closer to production GPT tokenization than character-level alphabets.
Tiny labs can still use real subword tokenization if they remap active tokens into compact local IDs.
GPT training loop is still corpus -> token IDs -> packed windows -> masked self-attention -> shifted loss -> checkpoint -> generation.
Validation loss and sampled text should both move in right direction before you trust checkpoint.

Next Step

Continue to Continued Pretraining and Domain Adaptation

PreviousPre-training Data at Scale

Share this article

X Facebook LinkedIn Bluesky Reddit Hacker News Email

References

CS336: Language Modeling from Scratch.

Stanford University · 2026

nanoGPT.

Karpathy, A. · 2025

shakespeare.txt.

Princeton University COS 302 / SML 305 · 2020

Language Models are Unsupervised Multitask Learners.

Radford, A., et al. · 2019

PyTorch: An Imperative Style, High-Performance Deep Learning Library.

Paszke, A., et al. · 2019 · NeurIPS 2019

Build GPT from Scratch Lab

Build GPT from Scratch Lab

Why this lab feels different

Five moving parts

Runnable lab

Why subword IDs are closer to real GPT

Why one-token shift still matters

Packing and splitting data

Validation loss versus sample quality

Common mistakes

Treating pieces as separate chores

Changing tokenizer after checkpoint exists

Keeping sparse token IDs in tiny lab

Trusting train loss alone

Choosing context window by convenience

What you should be able to defend

Practice answers

Why does this lab remap GPT-2 token IDs into a compact local vocabulary instead of training directly on the original token integers?

If this tutorial already uses subword tokens, why do we still need shifted labels and a causal mask?

Warmup sample is mostly punctuation soup, but later sample starts producing role words and dialogue-like clauses. What changed?

What to remember

Build GPT from Scratch Lab

Build GPT from Scratch Lab

Why this lab feels different

Five moving parts

Runnable lab

Why subword IDs are closer to real GPT

Why one-token shift still matters

Packing and splitting data

Validation loss versus sample quality

Common mistakes

Treating pieces as separate chores

Changing tokenizer after checkpoint exists

Keeping sparse token IDs in tiny lab

Trusting train loss alone

Choosing context window by convenience

What you should be able to defend

Practice answers

Why does this lab remap GPT-2 token IDs into a compact local vocabulary instead of training directly on the original token integers?

If this tutorial already uses subword tokens, why do we still need shifted labels and a causal mask?

Warmup sample is mostly punctuation soup, but later sample starts producing role words and dialogue-like clauses. What changed?

What to remember