Assignment 1 · Basics

What you build

Assignment 1 is the spine of the course: everything later assignments optimize, parallelize, scale, and align is the thing you build here. The handout breaks it into four implementation arcs plus an experimental campaign:

1. A byte-level BPE tokenizer — training (vocabulary initialization, regex pre-tokenization, iterative merges with special-token handling) and a Tokenizer class that encodes text to integer IDs and decodes back, including lazy streaming encoding for files too large for memory.
2. A Transformer LM — built module by module: Linear and Embedding (with the handout’s truncated-normal initializations), RMSNorm, a SwiGLU feed-forward network, RoPE, softmax, scaled dot-product attention with masking, causal multi-head self-attention, the pre-norm block, and the assembled LM.
3. Training machinery — numerically stable cross-entropy, AdamW as a torch.optim.Optimizer subclass, cosine learning-rate schedule with warmup, gradient clipping, a memmap-backed data loader, checkpointing, and a configurable training script.
4. Decoding — sampling from the trained model with temperature scaling and top-p (nucleus) sampling.

You then run it all: train tokenizers on TinyStories and OpenWebText, tokenize both datasets, train LMs, run ablations (RMSNorm removal, post-norm vs pre-norm, NoPE vs RoPE, SwiGLU vs SiLU), and submit an OpenWebText run to the leaderboard.

The from-scratch rule is strict: no definitions from torch.nn, torch.nn.functional, or torch.optim except nn.Parameter, container classes, and the Optimizer base class. Everything else you may use. The repo gives you adapters.py (glue code hooking your implementations to the staff tests) and test_*.py; the tests are your ground truth, run with uv run pytest. Submission is a typeset writeup.pdf plus code.zip to Gradescope.

Which lectures feed it

The handout is self-contained, but the lectures are where the design decisions get argued for rather than stated:

• Lecture 1 covers tokenization end to end — Unicode, UTF-8 bytes, why byte-level BPE wins over word- and character-level — and is the direct source for Section 2. Rewatch it if the merge procedure or pre-tokenization’s purpose feels arbitrary.
• Lecture 2 is the workhorse: einsum notation and einops (which the handout strongly recommends for every tensor operation you’ll write), plus the FLOPs and memory accounting you need for the transformer_accounting and adamw_accounting written problems. The 2mnp matmul rule and the forward:backward 1:2 ratio come straight from here.
• Lecture 3 justifies the architecture you’re told to build — pre-norm blocks, RMSNorm, SwiGLU, RoPE, no biases — and gives the hyperparameter intuitions behind the ablations in Section 7.
• Lecture 4 rounds out the architecture picture; the core of this assignment leans most heavily on the first three.

Order of attack

Follow the handout’s order — it is dependency order — but budget unevenly:

Tokenizer first, and expect it to fight back. BPE training is the one component where naive code is unusably slow. Pre-tokenization is the bottleneck; the handout tells you to parallelize it with multiprocessing (chunking at special-token boundaries, with linked starter code) and to speed up merging by incrementally updating pair counts instead of rescanning. Profile with cProfile or py-spy before optimizing. Two traps worth reading twice: special tokens must act as hard segmentation boundaries that never contribute to merges, and frequency ties break by preferring the lexicographically greater pair.

Model modules in sequence, tests as you go. Each module is small once the shapes are right; the compounding difficulty is batch-like leading dimensions (attention must accept arbitrary batch dims) and getting RoPE’s precomputed sin/cos buffers sliced by token positions. Use einsum notation and the shape errors mostly disappear.

Numerical stability is a recurring theme, not a footnote. Softmax subtracts the max; cross-entropy cancels log and exp; RMSNorm upcasts to float32 and casts back. The tests check correctness, but instability only shows up later, during training — get the habits right here.

Training loop and experiments last. Memory-map your tokenized data (np.memmap), checkpoint model + optimizer + iteration, and make hyperparameters command-line configurable — Section 7 has you running many small experiments, and the handout asks for an experiment log with loss curves against both steps and wall-clock time.

Budgets, targets, and the leaderboard

The handout gives concrete resource envelopes — treat blowing past them as a bug signal:

• TinyStories BPE (10K vocab): at most 30 minutes and 30 GB RAM, no GPU; under 2 minutes is achievable with multiprocessing.
• OpenWebText BPE (32K vocab): at most 12 hours and 100 GB RAM.
• TinyStories training (~17M params, 327,680,000 tokens): roughly 20–30 minutes on one B200; target validation loss at most 1.45.
• Low-resource path: on CPU or Apple Silicon, cut to 40M tokens and relax the target to 2.00 — the staff solution reaches 1.80 in 36 minutes on MPS. The blue “low-resource tip” boxes (MPS device strings, torch.compile backends, avoiding TF32 on MPS) are written for exactly this route.
• Ablations: about 0.5 B200-hours each; the learning-rate study gets 2.

The finale is open-ended: 10 B200-hours of experimentation, then one leaderboard run of at most 45 minutes on a B200, OpenWebText only, minimizing validation loss — anything beating the naive 5.0 baseline qualifies, submitted as a PR to the leaderboard repo. The handout points at NanoGPT-speedrun modifications and modern LLM design docs for ideas.