Project Week: train your very own small language model

1 From toy to real — scaling up with nanoGPT

In Session 19 you built a mini-transformer by hand — embeddings, self-attention, a feed-forward block, and a next-token loss. That was the "toy car." It taught you that the whole machine is understandable. The problem with hand-rolled code is that it's slow and missing the unglamorous bits (efficient data loading, mixed precision, checkpointing, learning-rate schedules). nanoGPT is a famous, deliberately small, readable codebase by Andrej Karpathy that has all those bits while staying simple enough to read in an afternoon.

The four stages of any GPT project

nanoGPT organises the work into four scripts, and this is a great mental model for every language-model project:

Let's name the parts. The training pipeline is this whole assembly line. prepare.py turns raw text into a stream of token IDs saved to disk. train.py reads those tokens, builds the GPT, and runs gradient descent. sample.py loads a saved checkpoint (the model's learned weights on disk) and generates new text.

Getting nanoGPT and seeing the model

# Clone Karpathy's nanoGPT and install the essentials
git clone https://github.com/karpathy/nanoGPT
cd nanoGPT
pip install torch numpy tiktoken datasets tqdm

class Block(nn.Module):           # one transformer layer
    def __init__(self, config):
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)   # masked self-attention
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp  = MLP(config)                    # the feed-forward part

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # residual + attention
        x = x + self.mlp(self.ln_2(x))    # residual + MLP
        return x

GPU and compute realities

Here's the honest part nobody tells beginners. Training language models is bottlenecked by compute. The key piece of hardware is a GPU (Graphics Processing Unit) — a chip that does the huge matrix multiplies of a neural net thousands of times faster than a CPU. A few realities:

2 Pretraining a small GPT — the full run

Pretraining means training a language model from scratch (random weights) on a large amount of raw text, with one job: predict the next token (recall Sessions 16 & 19 — this is the core language-modeling objective). The model isn't told facts or rules; it just learns the statistics of language by guessing the next token billions of times.

Step 1 — Dataset preparation & tokenization (BPE)

The model can't read letters; it reads tokens — integer IDs for sub-word chunks. We turn text into tokens with a tokenizer. The standard method is Byte-Pair Encoding (BPE), which we introduced back in Session 16 (and which the GenAI track covers in depth). BPE keeps frequent chunks whole (like "the") and breaks rare words into smaller known pieces, so the model can represent any text with a fixed vocabulary.

import numpy as np, tiktoken, os

# 1. Read raw text (point this at your own corpus, e.g. input.txt)
with open('input.txt', 'r') as f:
    data = f.read()
n = len(data)
train_data = data[:int(n*0.9)]   # 90% train
val_data   = data[int(n*0.9):]   # 10% validation

# 2. BPE-encode with GPT-2's tokenizer (cross-ref Session 16)
enc = tiktoken.get_encoding("gpt2")
train_ids = enc.encode_ordinary(train_data)
val_ids   = enc.encode_ordinary(val_data)
print(f"train has {len(train_ids):,} tokens")

# 3. Save as a flat array of uint16 token IDs (vocab < 65536)
np.array(train_ids, dtype=np.uint16).tofile('train.bin')
np.array(val_ids,   dtype=np.uint16).tofile('val.bin')

Step 2 — The config (choosing the model size)

The config is just the set of dials that define your model and training run. For a small GPT trainable on one GPU, sensible starting values are:

That gives roughly a 10M-parameter model — small enough for a free Colab GPU, big enough to write coherent text after training.

Step 3 — The training run

The training loop is the same recipe from Sessions 3 and 19: get a batch, forward pass to get the loss, backward pass for gradients, optimizer step. Here it is, complete and runnable, using nanoGPT's GPT model:

import torch, numpy as np
from model import GPT, GPTConfig   # from nanoGPT

device = 'cuda' if torch.cuda.is_available() else 'cpu'
block_size, batch_size = 256, 32

# memory-map the token stream we built in prepare.py
train_data = np.memmap('train.bin', dtype=np.uint16, mode='r')

def get_batch():
    # pick batch_size random start positions
    ix = torch.randint(len(train_data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy(train_data[i:i+block_size].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(train_data[i+1:i+1+block_size].astype(np.int64)) for i in ix])
    return x.to(device), y.to(device)   # y is x shifted by one token

# build the model from our config
cfg = GPTConfig(n_layer=6, n_head=6, n_embd=384,
                block_size=block_size, vocab_size=50257, dropout=0.2)
model = GPT(cfg).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for step in range(5000):
    x, y = get_batch()
    logits, loss = model(x, y)        # forward: predict next token, get cross-entropy loss
    optimizer.zero_grad()
    loss.backward()                   # backprop (Session 4) computes gradients
    optimizer.step()                  # nudge weights downhill (Session 3)
    if step % 500 == 0:
        print(f"step {step}: loss {loss.item():.3f}")

torch.save(model.state_dict(), 'ckpt.pt')   # save the checkpoint

Step 4 — Sampling (generating text)

Once trained, generation is the autoregressive loop you know: feed in a prompt, predict the next token's probabilities, sample one, append it, repeat.

import torch, tiktoken
from model import GPT, GPTConfig

enc = tiktoken.get_encoding("gpt2")
model = GPT(GPTConfig(...))           # same config as training
model.load_state_dict(torch.load('ckpt.pt'))
model.eval()                          # turn off dropout for inference

prompt = enc.encode("Once upon a time")
x = torch.tensor(prompt, dtype=torch.long)[None, ...]   # add batch dim

with torch.no_grad():
    out = model.generate(x, max_new_tokens=200, temperature=0.8, top_k=200)
print(enc.decode(out[0].tolist()))

3 Parameter-efficient fine-tuning with LoRA

After pretraining, you usually want to fine-tune — keep training a model on a smaller, specific dataset so it specialises (e.g. answering in your company's style, or on medical text). The obvious way is full fine-tuning: update all the model's weights. But that's expensive.

Why full fine-tuning hurts

• Memory: you must store gradients and optimizer state for every one of the billions of parameters — often 3–4× the model's own size in GPU memory.
• Storage: every fine-tuned copy is a full-size model. Ten tasks = ten giant files.
• Data hunger & overfitting: nudging all the weights with a small dataset risks forgetting what pretraining taught (called catastrophic forgetting).

The low-rank trick, made concrete

Inside the transformer, much of the work is a big matrix multiply y = W·x, where W might be 768 × 768 (≈590,000 numbers). LoRA doesn't touch W. Instead it learns a small correction ΔW and represents it as the product of two skinny matrices:

y = W·x + (B·A)·x,   where A is r × 768 and B is 768 × r.

Here r is the rank — a small number like 4, 8, or 16. The insight (from the LoRA paper) is that the update a model needs for a new task has low "intrinsic rank," so two skinny matrices capture it just fine.

Code sketch — LoRA in PyTorch

import torch, torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r=8, alpha=16):
        super().__init__()
        self.base = base                 # the frozen pretrained weight
        for p in self.base.parameters():
            p.requires_grad = False   # FREEZE the big matrix
        d_in, d_out = base.in_features, base.out_features
        self.A = nn.Parameter(torch.randn(r, d_in) * 0.01)  # skinny, random
        self.B = nn.Parameter(torch.zeros(d_out, r))            # skinny, zero-init
        self.scale = alpha / r           # LoRA scaling factor

    def forward(self, x):
        # frozen path + tiny trainable low-rank correction
        return self.base(x) + self.scale * (x @ self.A.T @ self.B.T)

from peft import LoraConfig, get_peft_model

config = LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"],
                    lora_dropout=0.05)
model = get_peft_model(base_model, config)
model.print_trainable_parameters()   # e.g. "trainable: 0.24% of all params"
# ...then train normally; save just the tiny adapter, not the whole model.

4 Wrapping up the course — where this all leads

You've reached the end of Neural Networks & Computer Vision. Step back and see how the pieces connect into the systems behind modern AI.

How everything connects to modern LLMs

A frontier model like GPT-4, Claude, or Gemini is not magic beyond what you now understand — it's this course's ideas scaled enormously:

Where to go next

• Scaling laws & bigger pretraining — how model size, data, and compute trade off.
• Alignment — instruction tuning, RLHF, and DPO to make models helpful and safe.
• Retrieval (RAG) & agents — giving models tools and fresh knowledge (covered in the GenAI track).
• Efficiency — quantization, distillation, and faster attention for cheaper inference.
• Multimodality — combining the vision skills (CNNs) you learned with language.

Final-project guidance

Your capstone is to train and fine-tune a small language model end-to-end. A solid submission:

Keep it small, get it running end to end, and document your loss curves and sample outputs. See the class handout and the final-project guide (linked in References) for the full rubric.

★ Putting it all together

This session is where the whole course pays off. Here's the one-paragraph story tying it together:

Quick self-check

📚 References & Further Reading

Class material

• SST Deep Learning handout (Session 20) — your course handout, including the course roadmap and the final-project guide.

Papers, docs & deep dives

• SST Deep Learning — course site (roadmap & final-project guide) — the online companion with the roadmap and capstone rubric for this session.
• Karpathy — nanoGPT (GitHub) — the small, readable codebase we use today; clone it and read model.py.
• LoRA: Low-Rank Adaptation of Large Language Models (Hu et al., 2021) — the original paper behind the low-rank adapter trick.
• HuggingFace PEFT documentation — the practical library for LoRA and other parameter-efficient fine-tuning methods.
• Attention Is All You Need (Vaswani et al., 2017) — the transformer architecture every GPT is built on.
• OpenAI tiktoken — the fast BPE tokenizer used in our data-prep step.