Building a Transformer LM

Lectures behind it: L01 Overview and tokenization, L02 PyTorch, resource accounting, L03 Architectures, hyperparameters · forward pointer for sampling: L10 Inference

This assignment exists to make the rest of the course legible. Everything after it — kernels, parallelism, scaling laws, data pipelines, alignment — is a modification of, or an optimization on, code you wrote here. That is why the handout forbids torch.nn.Linear, nn.Embedding, F.softmax, F.scaled_dot_product_attention and torch.optim.AdamW: if you call the library, you never learn that a Transformer's parameter count and FLOP count are things you can derive on paper, and the resource-accounting problems in §3.6 and §4.3 become unanswerable. What you are allowed is the interesting boundary — torch.nn.Parameter, the container classes (Module, ModuleList, Sequential), and the torch.optim.Optimizer base class — so you get autograd and the training-loop plumbing for free while writing every mathematical operation yourself. The handout also asks you to turn off AI autocomplete, on the grounds that it makes it much harder to engage deeply with the content.

Map of the assignment

Thirty-eight graded problems. Point values sum to 107 from the values the handout prints; two problems (cross_entropy and learning_rate_schedule) carry no printed point value in v1.0.6, so the true total is 107 plus whatever those two are worth. The Graded by column names the test file that actually runs your adapter; problems with no test file are graded from writeup.pdf alone.

Setup: environment, data, tests

The repo is deliberately almost empty. cs336_basics/ contains only pretokenization_example.py — a helper for finding chunk boundaries at special-token positions, which you are explicitly allowed to use verbatim. Everything else in that package is yours to design. The environment is managed with uv: uv run <file> solves and activates the environment on demand, and pyproject.toml already pins torch 2.6, einops, einx, jaxtyping, regex, tiktoken (used only by the tests, as the reference to match), psutil, submitit and wandb. Read the uv projects guide once; you will be adding dependencies.

The grading contract is the adapter pattern. tests/adapters.py is 592 lines of stubs that all raise NotImplementedError; each one has a fully typed signature and a docstring specifying tensor shapes. Your job is to make each stub call into your code — the handout is explicit that your changes to adapters.py should not contain any substantive logic; this is glue code. The tests import only from adapters, never from cs336_basics, which is why you are free to structure your package however you like. uv run pytest runs everything; uv run pytest -k test_rmsnorm runs one. Do not edit the test files.

Most of the model tests are snapshot tests. conftest.py loads a reference state dict from tests/fixtures/ts_tests/model.pt, feeds it through your adapter, and compares the output against a stored .npz in tests/_snapshots/. That means "passes" is numerical agreement with the staff's implementation to a tight tolerance — atol=1e-6 for RoPE-attention, atol=1e-4, rtol=1e-2 for the full LM — not merely "runs without error". The practical consequence: build bottom-up and get each module green before assembling, because a failing test_transformer_lm gives you no information about which of the twelve submodules is wrong. The fixtures also pin the shapes the tests use, which is useful when a shape error is ambiguous.

Data is two large plaintext files. On the course cluster they sit in /data on any non-head node; at home you download them with the commands in the README: TinyStories from Hugging Face (TinyStoriesV2-GPT4-train.txt, valid), and a course-hosted OpenWebText subsample (owt_train.txt.gz, owt_valid.txt.gz) which you gunzip in place. Note that only the OWT files the course provides are legal for the leaderboard.

If you have no GPU, the handout's low-resource boxes are not a consolation prize — the staff report training a fluent TinyStories model on an M3 Max in under 5 minutes on MPS and about 30 minutes on CPU. The adjustments are: drop total tokens from 327.68M to 40M, relax the loss target from 1.45 to 2.00, end the cosine decay exactly at your final step, and do not set torch.set_float32_matmul_precision('high') on MPS — the staff found silently broken TF32 kernels there in torch 2.6. torch.compile(model) helps on CPU; on MPS use backend="aot_eager" because Inductor is unsupported.

2 · Byte-pair encoding tokenizer

The mental model: you are building a compression scheme whose vocabulary is learned from the corpus. Start from the 256 byte values so nothing is ever out-of-vocabulary, then repeatedly find the most frequent adjacent pair of tokens and mint a new token for it. Vocabulary size ends up as 256 + number of merges + number of special tokens. Two structural choices make this tractable and are where most of the work lives: pre-tokenization, which chops text into word-ish units so that merges never cross word boundaries and counts can be accumulated per unique pre-token rather than per corpus position; and incremental pair counting, so that after each merge you only update the pair counts that actually changed. Section 2 is 42 of the assignment's points — nearly 40% — and almost all the difficulty is performance, not correctness.

Problem (unicode1): Understanding Unicode — 1 point

Deliverable: three one-sentence answers about chr(0). This is a warm-up that makes a real point: chr(0) is the null character, its __repr__() shows an escape sequence while printing it emits an invisible control character, and concatenated into a string it silently disappears from terminal output while still occupying a position and a byte. The lesson you need for the rest of §2 is that Python string length, printed length and byte length are three different numbers, and your tokenizer has to be correct about the third one. Purely a write-up problem — no adapter, no test.

Problem (unicode2): Unicode Encodings — 3 points

Deliverable: three short written answers on why UTF-8, why a naive per-byte decoder is wrong, and an example of an undecodable two-byte sequence. The substance is that UTF-8 is variable-width, so a multi-byte character's continuation bytes are meaningless on their own — decoding byte-by-byte with bytes([b]).decode("utf-8") raises or mangles as soon as you hit anything non-ASCII. This matters directly for your decode implementation later: you must concatenate all the token bytes first and decode once, with errors='replace' so malformed output becomes U+FFFD rather than an exception. UTF-16 and UTF-32 are rejected because they inflate ASCII by 2× and 4× respectively and introduce byte-order concerns; UTF-8 keeps a 256-symbol alphabet with no BOM problem. Write-up only.

Problem (train_bpe): BPE Tokenizer Training — 15 points

Deliverable: a function that takes an input text path, a vocab_size and a list of special_tokens, and returns (vocab: dict[int, bytes], merges: list[tuple[bytes, bytes]]) with merges in creation order. This is the single largest problem in the assignment and the one people lose the most time on. The correctness half is short: strip and split on special tokens first (so no merge ever crosses a document boundary), pre-tokenize each fragment with the GPT-2 regex using re.finditer from the regex package, represent each pre-token as a tuple of single-byte bytes objects, count pairs weighted by pre-token frequency, and merge the argmax. The performance half is where the 15 points actually live: the provided test_train_bpe_speed trains a 500-token vocabulary on a small fixture and fails if it takes more than 1.5 seconds — the staff reference does it in 0.38s, and a naive re-scan-every-merge implementation takes about 3s. You get there by maintaining an index from each pair to the pre-tokens containing it and updating only the affected counts after each merge, plus multiprocessing over chunk boundaries for pre-tokenization.

The pre-tokenizer regex must be exactly this one, taken from tiktoken PR #234 — it is what GPT-2 used, and the reference merges in the fixtures were produced with it:

PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""Handout §2.4, "Pre-tokenization"

The other detail that silently fails the reference comparison is tie-breaking:

When computing merges, deterministically break ties in pair frequency by preferring the lexicographically greater pair.Handout §2.4, "Compute BPE merges"

In Python that is just max(counts.items(), key=lambda kv: (kv[1], kv[0])) over bytes pairs — but if you compare decoded strings instead of raw bytes, or use min, your merge list diverges from the reference at the first tie and every subsequent merge is wrong. The handout notes you may implement the hot loop in C++ (via cppyy) or Rust (via PyO3) as an optional stretch, but warns that the GPT-2 regex is poorly supported and slow in most non-Python engines; Python's regex package is, if anything, faster than Oniguruma here.

• tests/adapters.py · run_train_bpe — (input_path: str | os.PathLike, vocab_size: int, special_tokens: list[str], **kwargs) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]
• tests/test_train_bpe.py — test_train_bpe_speed (1.5s wall-clock ceiling on corpus.en, vocab 500), test_train_bpe (exact match against train-bpe-reference-merges.txt and the reference vocab), test_train_bpe_special_tokens (pickle snapshot; catches merging across <|endoftext|>)
• cs336_basics/pretokenization_example.py — chunk-boundary finder you may copy verbatim for the multiprocessing split

Problem (train_bpe_tinystories): BPE Training on TinyStories — 2 points

Deliverable: two short written answers — wall time and peak memory for training a 10,000-token vocabulary on TinyStories with <|endoftext|> as a special token, plus the longest token in the vocabulary and whether it makes sense; then a profile identifying the bottleneck. Serialize the vocab and merges to disk, because everything downstream reuses them. The budget is ≤ 30 minutes with no GPU and ≤ 30 GB RAM, but the handout's hint is that under 2 minutes is achievable once you parallelize pre-tokenization on <|endoftext|>-aligned chunks and special-case the delimiter before merges. Profile with cProfile or scalene; the answer is almost always pre-tokenization, not the merge loop, which is the point of the exercise. Develop against the TinyStories validation split first — 22K documents instead of 2.12M, same bottlenecks, minutes instead of hours.

Problem (train_bpe_expts_owt): BPE Training on OpenWebText — 2 points

Deliverable: the same longest-token question for a 32,000-token vocabulary trained on OpenWebText, plus a compare-and-contrast of the two tokenizers. The interesting finding is that the two vocabularies encode their corpora's character: TinyStories is a small, clean, synthetic children's-story distribution, so its long tokens are ordinary English words and phrases; OpenWebText is scraped web text, so the longest tokens tend to be boilerplate, markup residue and site-specific junk that appears thousands of times. This is the first place in the course where you see a data-quality problem show up as an artifact in a model artifact — the theme L13 and L14 and Assignment 4 take up in full. Budget 12 hours and up to 100 GB of RAM; this is the run to start before you go to sleep.

Problem (tokenizer): Implementing the tokenizer — 15 points

Deliverable: a Tokenizer class constructed from a vocab, a merge list and optional special tokens, exposing from_files, encode, encode_iterable and decode. Encoding mirrors training: pre-tokenize, then apply the learned merges in creation order to each pre-token independently. The two traps are special tokens and streaming. Special tokens must be split out before the regex ever sees them and mapped straight to their IDs; test_overlapping_special_tokens specifically checks that when one special token is a prefix of another you match the longest, and there are dedicated tests for special tokens followed by trailing newlines and by non-whitespace. Streaming is what encode_iterable is for: it takes an iterable of strings (typically a file handle) and must yield IDs lazily, so that tokenizing a 5 MB file does not materialize the whole thing — the test enforces this with an rlimit, and it is deliberately contrasted with test_encode_memory_usage, which is marked xfail because plain encode is expected to blow the same 1 MB budget. The remaining ~25 tests all compare your output byte-for-byte against tiktoken loaded with the GPT-2 vocab, across ASCII, Unicode, German and a TinyStories sample, in both round-trip and match-tiktoken form.

• tests/adapters.py · get_tokenizer — (vocab: dict[int, bytes], merges: list[tuple[bytes, bytes]], special_tokens: list[str] | None = None) -> Any; returns your tokenizer object, so the adapter is one constructor call
• tests/test_tokenizer.py — round-trip and tiktoken-equivalence over empty strings, single characters, Unicode, an address fixture, German text and TinyStories; plus test_encode_iterable_memory_usage (Linux-only, rlimit-gated) and the xfail test_encode_memory_usage
• tests/common.py — gpt2_bytes_to_unicode(), the byte↔printable-character mapping the fixtures are stored in; you need it to read the reference vocab
• U+FFFD replacement character — what bytes.decode(errors='replace') substitutes for malformed sequences, which decode must do rather than raise

Problem (tokenizer_experiments): Experiments with tokenizers — 4 points

Deliverable: four short answers plus, critically, the tokenized datasets you will train on. Sample 10 documents from each corpus and report compression ratio in bytes per token for each tokenizer on its own corpus; then cross-apply the TinyStories tokenizer to OpenWebText and watch the ratio collapse, because a vocabulary learned on simple children's prose has no merges for web boilerplate and falls back toward byte-level. Estimate throughput in bytes/second and extrapolate to the Pile's 825 GB — the number is usually sobering and is the reason tokenization gets parallelized in practice. Part (d) is the one that gates the rest of the assignment: encode both corpora's train and validation splits to uint16 NumPy arrays on disk. uint16 is the right dtype because it holds 0–65,535, comfortably above both the 10K and 32K vocabularies while halving the file size versus uint32 — and file size is what your memmap loader pages through every step. Write-up plus artifacts; no test.

3 · Transformer LM architecture

The model is a modern decoder-only Transformer, which means four specific departures from Vaswani et al. 2017, and L03 exists to justify each: pre-norm instead of post-norm (a clean residual stream from embeddings to output, plus a final norm before the LM head), RMSNorm instead of LayerNorm (no mean subtraction, no bias), SwiGLU instead of a ReLU MLP (three weight matrices, d_ff ≈ 8/3 · d_model rounded to a multiple of 64), and RoPE instead of learned absolute position embeddings (applied to Q and K inside every attention layer, never to V and never at the embedding). No biases anywhere. Build it bottom-up in the order below; each module has its own snapshot test, and skipping ahead means debugging blind.

Two remarks from §3.3 pay for themselves. First, write every operation to tolerate arbitrary leading batch dimensions — "... d_in, d_out d_in -> ... d_out" — because the tests feed 3-D and 4-D tensors to the same adapters. Second, the handout strongly recommends einops (einsum and rearrange) over chains of view/transpose; multi-head attention's head-splitting is where this stops being a style preference. Third, and easy to miss: the handout uses column-vector math notation (y = W x) while PyTorch is row-major, so a literal transcription of the equations into @ needs a transpose. Store the parameter as W, not W transposed, for memory-ordering reasons — the tests load the staff's state dict into your module, so your shapes must match theirs exactly.

Initialization, given once in §3.4.1 and reused everywhere: linear weights are truncated normal with mean 0 and variance 2/(d_in + d_out), truncated at ±3σ; embeddings are truncated normal with variance 1, truncated at ±3; RMSNorm gains start at 1. Use torch.nn.init.trunc_normal_. The tests overwrite your weights with the reference state dict, so initialization does not affect them — it affects whether your training runs converge, which is where the points are.

Problem (linear): Implementing the linear module — 1 point

Deliverable: a Linear(in_features, out_features, device=None, dtype=None) subclass of nn.Module with no bias. Trivial in substance, load-bearing in convention: it establishes that you subclass nn.Module, call super().__init__(), wrap the weight in nn.Parameter, and store it with shape (d_out, d_in). The adapter is expected to load the given weights into your module (Module.load_state_dict is the intended route), so if your parameter is named or shaped differently from the reference this is the first place you find out. Where people get stuck is transposing: y = W x with row-major storage means x @ W.T, or in einsum, "... d_in, d_out d_in -> ... d_out".

• tests/adapters.py · run_linear — (d_in: int, d_out: int, weights: Float[Tensor, " d_out d_in"], in_features: Float[Tensor, " ... d_in"]) -> Float[Tensor, " ... d_out"]
• tests/test_model.py · test_linear — feeds layers.0.ffn.w1.weight from the reference checkpoint and snapshot-compares the output

Problem (embedding): Implement the embedding module — 1 point

Deliverable: an Embedding(num_embeddings, embedding_dim, device=None, dtype=None) module that indexes a (vocab_size, d_model) parameter with a LongTensor of token IDs. The forward pass is one advanced-indexing operation — self.weight[token_ids] — and it already broadcasts over arbitrary leading dimensions, which is why the type hint says Int[Tensor, " ..."]. Store d_model as the final dimension. The reason this gets its own problem rather than being folded into the LM is that the embedding table is usually the largest single parameter tensor in a small model, and §3.6's accounting asks you to notice that.

• tests/adapters.py · run_embedding — (vocab_size: int, d_model: int, weights: Float[Tensor, " vocab_size d_model"], token_ids: Int[Tensor, " ..."]) -> Float[Tensor, " ... d_model"]
• tests/test_model.py · test_embedding — indexes the reference token_embeddings.weight with the fixture's token IDs

Problem (rmsnorm): Root Mean Square Layer Normalization — 1 point

Deliverable: RMSNorm(d_model, eps=1e-5, device=None, dtype=None) as an nn.Module. The operation is out_i = a_i / RMS(a) · g_i where RMS(a) = sqrt(mean(a²) + eps) and g is a learnable d_model-length gain. Compared to LayerNorm it drops the mean subtraction and the bias, which is both cheaper and — per Zhang and Sennrich 2019 — empirically just as good. The one non-obvious requirement is the dtype dance: upcast the input to float32 before squaring, because in bf16 or fp16 the sum of squares overflows on realistic activations, then downcast the result back to the input dtype before returning. The handout gives the exact skeleton for this. A subtle failure people hit: putting eps outside the square root instead of inside it — the handout's formula has it under the radical, added to the mean.

• tests/adapters.py · run_rmsnorm — (d_model: int, eps: float, weights: Float[Tensor, " d_model"], in_features: Float[Tensor, " ... d_model"]) -> Float[Tensor, " ... d_model"]
• tests/test_model.py · test_rmsnorm — snapshot against tests/_snapshots/test_rmsnorm.npz
• L03 Architectures, hyperparameters — the lecture flags 11:23 as the RMSNorm explanation to re-watch if this test fails

Problem (positionwise_feedforward): Implement the position-wise feed-forward network — 2 points

Deliverable: a SwiGLU feed-forward network, FFN(x) = W2 (SiLU(W1 x) ⊙ W3 x), with W1, W3 of shape (d_ff, d_model) and W2 of shape (d_model, d_ff). Two things combine here: SiLU (also called Swish), x · sigmoid(x), from Hendrycks and Gimpel and Elfwing et al.; and the gated linear unit of Dauphin et al. 2017, an elementwise product of two projections. Shazeer 2020 put them together and reported it beats ReLU and ungated SiLU, while famously declining to explain why. Sizing: d_ff ≈ 8/3 · d_model, rounded to a multiple of 64 for tensor-core alignment — the 8/3 exists so that three matrices of that width cost about the same parameters as two matrices at the classic 4× width. You are explicitly permitted torch.sigmoid here for numerical stability. Note that there is also a bare run_silu adapter with its own test but no problem number of its own; implement SiLU as a standalone function and both are satisfied.

• tests/adapters.py · run_swiglu — (d_model: int, d_ff: int, w1_weight, w2_weight, w3_weight, in_features: Float[Tensor, " ... d_model"]) -> Float[Tensor, " ... d_model"], with w1_weight and w3_weight shaped (d_ff, d_model) and w2_weight shaped (d_model, d_ff)
• tests/adapters.py · run_silu — (in_features: Float[Tensor, " ..."]) -> Float[Tensor, " ..."]; no dedicated handout problem, graded by test_silu_matches_pytorch
• tests/test_model.py · test_swiglu — snapshot at atol=1e-5 using layers.0.ffn.{w1,w2,w3}.weight

Problem (rope): Implement RoPE — 2 points

Deliverable: a RotaryPositionalEmbedding(theta, d_k, max_seq_len, device=None) module whose forward(x, token_positions) takes (..., seq_len, d_k) and returns the same shape. RoPE from Su et al. 2021 injects position by rotating each adjacent pair of embedding dimensions by an angle that depends on the token's absolute position and the pair index — θ(i,k) = i / Θ^((2k-2)/d) — so that the dot product between a query at position i and a key at position j depends only on i - j. Do not build the full d × d block-diagonal matrix; precompute cos and sin tables of shape (max_seq_len, d_k/2) in __init__ and register them with register_buffer(persistent=False) — buffers, not parameters, because there is nothing to learn. Then index the tables by token_positions, which is why the signature takes positions rather than assuming 0..seq_len-1: the truncated-input test and any KV-cache you write later need that flexibility. The canonical way to fail this test, called out in L03 (37:49), is applying the rotation once at the embedding layer rather than to Q and K inside each attention layer.

• tests/adapters.py · run_rope — (d_k: int, theta: float, max_seq_len: int, in_query_or_key: Float[Tensor, " ... sequence_length d_k"], token_positions: Int[Tensor, " ... sequence_length"]) -> Float[Tensor, " ... sequence_length d_k"]
• tests/test_model.py · test_rope — snapshot with explicit pos_ids; test_multihead_self_attention_with_rope checks it in context at atol=1e-6

Problem (softmax): Implement softmax — 1 point

Deliverable: a function taking a tensor and a dimension, returning a normalized distribution along that dimension. The whole point is the stability trick: subtract the maximum along the target dimension before exponentiating. Softmax is invariant to adding a constant to all inputs, so this changes nothing mathematically, but it caps the largest exponent at exp(0) = 1 and prevents inf/inf = NaN on large logits. You will reuse this in attention, in cross-entropy (where the same trick reappears in log-space) and in your decoder. Note the test name is test_softmax_matches_pytorch, not test_softmax — pytest -k test_softmax works, but the handout spells the full name.

• tests/adapters.py · run_softmax — (in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, " ..."]
• tests/test_nn_utils.py · test_softmax_matches_pytorch — compares directly against torch.nn.functional.softmax

Problem (scaled_dot_product_attention): Implement scaled dot-product attention — 5 points

Deliverable: Attention(Q, K, V, mask) = softmax(Q Kᵀ / sqrt(d_k)) V, handling any number of leading batch-like dimensions and an optional boolean mask. Five points for four lines of math, because the shape discipline is the lesson: queries are (..., n, d_k), keys (..., m, d_k), values (..., m, d_v), and the scores matrix is (..., n, m). The 1/sqrt(d_k) scale keeps the pre-softmax variance from growing with d_k and pushing softmax into its saturated regime. The mask convention is the one that catches people, and it is worth reading twice:

Canonically (and slightly confusingly), a value of True at position (i, j) indicates that the query i does attend to the key j, and a value of False indicates that the query does not attend to the key.Handout §3.5.4, "Masking"

So you add -inf where the mask is False — the opposite polarity from PyTorch's attn_mask convention in some APIs, and the reason a correct-looking implementation produces a snapshot mismatch that looks like a transposed causal mask. Use masked_fill(~mask, -inf) or torch.where on the pre-softmax scores, never on the probabilities.

• tests/adapters.py · run_scaled_dot_product_attention — (Q: Float[Tensor, " ... queries d_k"], K: Float[Tensor, " ... keys d_k"], V: Float[Tensor, " ... values d_v"], mask: Bool[Tensor, " ... queries keys"] | None = None) -> Float[Tensor, " ... queries d_v"]
• tests/test_model.py — test_scaled_dot_product_attention (third-order tensors) and test_4d_scaled_dot_product_attention (fourth-order); the second is there specifically to catch hard-coded rank assumptions

Problem (multihead_self_attention): Implement causal multi-head self-attention — 5 points

Deliverable: causal multi-head self-attention as an nn.Module taking d_model and num_heads, with d_k = d_v = d_model / num_heads. Three moves: project the input once each to Q, K and V with full-width (d_model, d_model) matrices (the handout is explicit that you should end up with exactly three matmuls, not 3 × num_heads of them), rearrange the output dimension into (head, d_k) and move head in front of the sequence dimension so it batches, then apply attention with a causal mask and project back through W_O. The causal mask is torch.triu or a broadcast index comparison j <= i — and because your attention already supports masks, that is all the causality you need. RoPE goes on Q and K only, after the head split, so each head gets the same rotation applied independently; it never touches V. The reference weight layout matters: attn.q_proj.weight is the concatenation of the per-head matrices along dim 0, so head h occupies rows [h·d_k : (h+1)·d_k]. Splitting along the wrong axis is the most common cause of a near-miss snapshot.

• tests/adapters.py · run_multihead_self_attention — (d_model: int, num_heads: int, q_proj_weight, k_proj_weight, v_proj_weight, o_proj_weight, in_features: Float[Tensor, " ... sequence_length d_in"]) -> Float[Tensor, " ... sequence_length d_out"]
• tests/adapters.py · run_multihead_self_attention_with_rope — same plus max_seq_len: int, theta: float and token_positions: Int[Tensor, " ... sequence_length"] | None = None; no separate handout problem, but it has its own test
• tests/test_model.py — test_multihead_self_attention and test_multihead_self_attention_with_rope (atol=1e-6)

Problem (transformer_block): Implement the Transformer block — 3 points

Deliverable: the pre-norm block, two sublayers, each of the form y = x + Sublayer(RMSNorm(x)) — attention first, feed-forward second. Structurally this is assembly, but the adapter tells you exactly what the staff's module hierarchy is named, and matching it saves you from writing key-remapping glue: attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight, attn.output_proj.weight, ln1.weight, ffn.w1.weight, ffn.w2.weight, ffn.w3.weight, ln2.weight. Name your submodules that way and the adapter is block.load_state_dict(weights). Note the ordering trap: ln1 is the norm before attention and ln2 the norm before the FFN — in a post-norm implementation the same names would sit after the sublayers, which is exactly the confusion the post-norm ablation later makes you confront on purpose. The adapter also takes max_seq_len and theta so it can construct or pass in a RoPE module; the docstring says either wiring is fine.

• tests/adapters.py · run_transformer_block — (d_model: int, num_heads: int, d_ff: int, max_seq_len: int, theta: float, weights: dict[str, Tensor], in_features: Float[Tensor, " batch sequence_length d_model"]) -> Float[Tensor, " batch sequence_length d_model"]; the docstring enumerates all nine weight keys and their shapes
• tests/test_model.py · test_transformer_block — strips the layers.0. prefix from the reference state dict and snapshot-compares

Problem (transformer_lm): Implementing the Transformer LM — 3 points

Deliverable: the full model — token embedding, num_layers blocks, a final RMSNorm, and a linear LM head to vocab_size logits. Return logits, not probabilities: the softmax lives in your loss function and in your decoder, and applying it here would double-normalize. Reference key names are token_embeddings.weight, layers.N.* (the block keys above), ln_final.weight and lm_head.weight — note the input and output embeddings are separate matrices, i.e. untied, which is one of the first things the leaderboard problem suggests you try changing. context_length is a constructor argument even though RoPE needs no learned position table, because it sizes the precomputed cos/sin buffers. There are two tests: the ordinary one, and test_transformer_lm_truncated_input, which feeds a half-length sequence to catch implementations that hard-code the sequence length when building the causal mask or slicing RoPE tables.

• tests/adapters.py · run_transformer_lm — (vocab_size: int, context_length: int, d_model: int, num_layers: int, num_heads: int, d_ff: int, rope_theta: float, weights: dict[str, Tensor], in_indices: Int[Tensor, " batch_size sequence_length"]) -> Float[Tensor, " batch_size sequence_length vocab_size"]
• tests/test_model.py — test_transformer_lm (atol=1e-4, rtol=1e-2) and test_transformer_lm_truncated_input (atol=1e-4)

Problem (transformer_accounting): Transformer LM resource accounting — 5 points

Deliverable: five written answers, no code. Count the trainable parameters of a GPT-2 XL-shaped model and the memory to load it in fp32; enumerate every matrix multiply in a forward pass and total the FLOPs; say which components dominate; repeat for GPT-2 small, medium and large and describe how the proportions shift with scale; then push GPT-2 XL's context to 16,384 and describe what changes. The only rule you need is that multiplying A ∈ R^(m×n) by B ∈ R^(n×p) costs 2mnp FLOPs — each output entry is an n-length dot product, n multiplies and n adds. The finding the problem is engineered to produce: at ordinary context lengths the feed-forward and projection matmuls dominate and attention's quadratic QKᵀ and softmax·V terms are a minority, and as d_model grows they become an even smaller fraction — but part (e) reverses this by scaling the context 16×, at which point the terms quadratic in sequence length take over. That single crossover is the motivation for FlashAttention in L06 and Assignment 2. L02 is a rehearsal for exactly this arithmetic.

4 · Loss and optimizer

Three pieces stand between a model that runs and a model that trains: a numerically careful cross-entropy, an AdamW you wrote yourself, and the two stabilizers that every real pretraining run uses — a cosine learning-rate schedule with linear warmup, and global gradient-norm clipping. The section also asks you to feel the learning rate rather than read about it, with a toy sweep before the real one.

Problem (cross_entropy): Implement Cross entropy — no point value printed in v1.0.6

Deliverable: a function from logits and integer targets to a scalar average loss. Mathematically ℓ_i = -log softmax(o_i)[x_{i+1}], but implementing it that way overflows: the handout's requirements are to subtract the maximum logit for stability and to cancel the log and exp wherever possible, which means computing logsumexp(o) - o[target] directly instead of building probabilities and taking their log. Handle arbitrary leading batch dimensions with the vocabulary dimension last, and average over the batch. Perplexity, defined in the same section as exp(mean(ℓ)), is not separately graded but is what you report in every experiment from §7 onward. Worth knowing: this problem's point value is genuinely absent from the v1.0.6 handout — it is not an OCR artifact.

• tests/adapters.py · run_cross_entropy — (inputs: Float[Tensor, " batch_size vocab_size"], targets: Int[Tensor, " batch_size"]) -> Float[Tensor, ""]
• tests/test_nn_utils.py · test_cross_entropy — checks the value and, separately, behaviour on large logits where a naive implementation returns NaN

Problem (learning_rate_tuning): Tuning the learning rate — 1 point

Deliverable: one or two sentences describing what happens to the toy SGD example at learning rates 1e1, 1e2 and 1e3 over 10 iterations. The handout hands you a complete SGD(torch.optim.Optimizer) implementation with a 1/sqrt(t+1) decay and a ten-line training loop on (weights**2).mean(); you just change one number. The expected observation is a monotone story — faster decay, then much faster decay, then divergence — and the reason it is worth a point is that it plants the "edge of stability" intuition you are asked to investigate properly in Problem (learning_rate). Read the provided code carefully anyway: it is also your template for the AdamW problem, showing how self.state[p], param_groups and the optional closure argument fit together.

Problem (adamw): Implement AdamW — 2 points

Deliverable: an AdamW subclass of torch.optim.Optimizer taking lr, betas, eps and weight_decay, following algorithm 2 of Loshchilov and Hutter 2019. Per parameter you keep first and second moment estimates m and v in self.state[p], update them with β1 and β2, compute a bias-corrected step size α_t = α · sqrt(1 - β2^t) / (1 - β1^t), apply θ ← θ - α_t · m / (sqrt(v) + ε), and then — decoupled, as a separate term — θ ← θ - α·λ·θ. The decoupling is the entire point of AdamW versus Adam with L2 regularization: weight decay must not flow through the moment estimates. Two off-by-one traps: t starts at 1, not 0, so initialize your step counter accordingly or the first bias correction is wrong; and the handout's algorithm applies weight decay with the base α, not the bias-corrected α_t. (β1, β2) = (0.9, 0.999) is the classic default, but LLaMA and GPT-3 use (0.9, 0.95) — a difference the tuning problem gives you room to explore.

• tests/adapters.py · get_adamw_cls — () -> Any; returns the class, not an instance, so the test can construct it with its own hyperparameters
• tests/test_optimizer.py · test_adamw — runs a fixed number of steps on a toy problem and snapshot-compares the resulting weights against tests/_snapshots/test_adamw.npz

Problem (adamwAccounting): Resource accounting for training with AdamW — 2 points

Deliverable: four written answers. (a) An algebraic expression for peak fp32 memory decomposed into parameters, activations, gradients and optimizer state, in terms of batch_size and the model hyperparameters, assuming d_ff = 4 · d_model; the handout enumerates exactly which activations to count, so you are not guessing at scope. (b) Instantiate for GPT-2 XL to get a · batch_size + b and solve for the largest batch that fits in 80 GB. (c) The FLOPs of one AdamW step. (d) The days to train GPT-2 XL for 400K steps at batch 1024 on one A100, at 50% MFU against a 19.5 teraFLOP/s fp32 peak, assuming the backward pass costs 2× the forward. The structural insight is that AdamW costs 3× the parameter memory beyond the weights themselves — gradients plus two moment buffers — which is why optimizer-state sharding is the first thing Assignment 2 asks you to implement, and why activation memory is what actually limits batch size. Part (d) is your first end-to-end MFU calculation; the number of days it produces is the argument for everything in the systems half of the course.

Problem (learning_rate_schedule): Implement cosine learning rate schedule with warmup — no point value printed in v1.0.6

Deliverable: a pure function of (t, α_max, α_min, T_w, T_c) returning the learning rate at step t, in three regimes: linear warmup α_t = (t / T_w) · α_max for t < T_w; cosine annealing α_t = α_min + 0.5 · (1 + cos(π · (t - T_w)/(T_c - T_w))) · (α_max - α_min) for T_w ≤ t ≤ T_c; and a constant α_min thereafter. This is LLaMA's schedule. Two things to watch: the adapter is named run_get_lr_cosine_schedule, while the handout's prose refers to it as adapters.get_lr_cosine_schedule — trust the file; and the boundary conditions are tested, so make sure t = T_w gives exactly α_max and t = T_c gives exactly α_min. When you later run fewer steps than the default (as the CPU/MPS path suggests), set T_c to your actual final step so the decay finishes exactly where training stops. Like cross_entropy, this problem carries no printed point value in v1.0.6.

• tests/adapters.py · run_get_lr_cosine_schedule — (it: int, max_learning_rate: float, min_learning_rate: float, warmup_iters: int, cosine_cycle_iters: int)
• tests/test_optimizer.py · test_get_lr_cosine_schedule — checks the value at every step across all three regimes

Problem (gradient_clipping): Implement gradient clipping — 1 point

Deliverable: a function taking an iterable of parameters and a maximum ℓ2 norm, modifying p.grad in place and returning nothing. Compute the norm over all gradients jointly — one global norm, not per-tensor — and if it exceeds M, scale every gradient by M / (norm + ε) with ε = 1e-6, PyTorch's default. Two places to slip: skipping parameters whose .grad is None, and computing per-parameter norms instead of the global one, which changes the update direction rather than just its length. The ε in the denominator means the clipped norm lands just under M, which the handout calls out so you do not chase the discrepancy; the test compares against torch.nn.utils.clip_grad_norm_.

• tests/adapters.py · run_gradient_clipping — (parameters: Iterable[torch.nn.Parameter], max_l2_norm: float) -> None; in-place, returns nothing
• tests/test_nn_utils.py · test_gradient_clipping — compares your in-place result against PyTorch's clip_grad_norm_

5 · Training loop

Everything now exists except the plumbing that keeps a multi-hour run alive: a batch sampler that does not load the dataset into RAM, checkpoints you can resume from, and a configurable script. These three problems are worth 7 points but determine whether the 23 H100 hours of experiments in §7 are pleasant or miserable, so treat the script as infrastructure rather than as a one-off.

Problem (data_loading): Implement data loading — 2 points

Deliverable: a function taking a NumPy token array, a batch size, a context length and a device string, returning (inputs, targets) both of shape (batch_size, context_length) on that device, where targets is inputs shifted one position right. Sampling is trivially easy because the corpus is one long concatenated stream: pick batch_size random start offsets in [0, n - context_length - 1] and slice. No padding, no document boundaries to respect, uniform sequence lengths, perfect hardware utilization. The real requirement is in the surrounding prose rather than the function signature — load the on-disk array with np.memmap or np.load(..., mmap_mode='r'), with a dtype matching what you saved (uint16), so the OS pages in only what you touch. Get the dtype wrong and you silently read garbage token IDs; the handout suggests explicitly checking that no value exceeds your vocabulary size, which is a two-line sanity check worth writing. Cast to torch.long when you build the tensors, since uint16 is not a valid index dtype.

• tests/adapters.py · run_get_batch — (dataset: npt.NDArray, batch_size: int, context_length: int, device: str) -> tuple[torch.Tensor, torch.Tensor]
• tests/test_data.py · test_get_batch — checks shapes, device placement, the input/target offset relationship, and that sampled starts cover the valid range

Problem (checkpointing): Implement model checkpointing — 1 point

Deliverable: save_checkpoint(model, optimizer, iteration, out) and load_checkpoint(src, model, optimizer), the latter returning the saved iteration number. A checkpoint needs three things to make a run resumable: model state_dict, optimizer state_dict (AdamW's moment estimates are not recoverable otherwise, and restarting without them causes a visible loss spike), and the step counter, without which your cosine schedule restarts from warmup. Both functions must accept a path or a file-like object — the test uses tmp_path, but the signature's BinaryIO | IO[bytes] is there deliberately, and hard-coding open() on the argument fails. One point, ten minutes, and it is the difference between losing a four-hour run to a preemption and not.

• tests/adapters.py · run_save_checkpoint — (model: torch.nn.Module, optimizer: torch.optim.Optimizer, iteration: int, out: str | os.PathLike | BinaryIO | IO[bytes])
• tests/adapters.py · run_load_checkpoint — (src: str | os.PathLike | BinaryIO | IO[bytes], model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> int
• tests/test_serialization.py · test_checkpointing — trains a toy model, saves, perturbs, reloads, and checks both the weights and the returned iteration

Problem (training_together): Put it together — 4 points

Deliverable: a training script. The handout asks for four capabilities: all model and optimizer hyperparameters configurable (command-line arguments are the suggested route), memmap-backed loading of train and validation sets, checkpoint serialization to a user-provided path, and periodic logging of train and validation metrics to the console and/or Weights & Biases. There is no test — this is graded from your code and, implicitly, from whether the §7 experiments you report were runnable. Invest here: you are about to launch on the order of a dozen runs varying one thing at a time, and every hyperparameter you hard-code now becomes a code edit and a lost comparison later. Log wallclock time alongside step count from the start, because learning_rate, batch_size_experiment and leaderboard all require wallclock-axis curves and you cannot reconstruct them after the fact.

6 · Generating text

The model emits a distribution over the next token at every position; sampling turns that into text. Take the logits at the final position, temperature-scale, optionally truncate to a nucleus, sample, append, repeat until <|endoftext|> or a token budget. This is the one place in A1 where the course's inference material (L10) is relevant, and it arrives seven lectures early — you are writing the naive quadratic-recompute version, with no KV cache, which is exactly the baseline L10 improves on.

Problem (decoding): Decoding — 3 points

Deliverable: a decode function supporting four features: completion from a user prompt until <|endoftext|>, a maximum-new-tokens cap, temperature scaling softmax(v/τ), and top-p (nucleus) sampling per Holtzman et al. 2020. Temperature interpolates between greedy (τ → 0, one-hot on the argmax) and the raw distribution (τ = 1); top-p keeps the smallest set of tokens whose probabilities sum to at least p and renormalizes over them, which prunes the long tail of individually-unlikely-but-collectively-common garbage that makes small models incoherent. Implementation: sort probabilities descending, cumulative-sum, cut at the first index where the cumsum reaches p, zero the rest, renormalize, then torch.multinomial. Apply temperature before the nucleus cut, since it changes which tokens are in the nucleus. Two practical points: remember to slice the logits to the last position only, and remember to truncate the context to context_length as the generation grows past it. No adapter and no test — this is graded through Problem (generate)'s text dump.

7 · Experiments

Ten problems, 22 points, and roughly 23 H100 hours of compute. The framing in §7.1 is the important part: run small, run systematically, and keep records — you are asked to submit an experiment log and learning curves against both gradient steps and wallclock time, so the logging you build in experiment_log is a prerequisite for every problem after it. The base configuration is fixed for you so ablations are comparable; you tune only the learning rate, warmup, AdamW betas/epsilon and weight decay.

Problem (experiment_log): Experiment logging — 3 points

Deliverable: experiment-tracking infrastructure plus a written log of everything you tried for the problems below. Track validation loss against both gradient steps and wallclock time — the leaderboard is scored on a wallclock budget, so a step-axis-only curve cannot answer the question the leaderboard asks. Log the full config with each run so that six runs later you can still tell which curve was which. The write-up half matters as much as the code: an experiment log that records the runs that diverged and what you concluded is the artifact, not a gallery of your best curves. Three points, and it is the cheapest three points in §7 if you build it before you start burning GPU hours rather than reconstructing it afterwards.

Problem (learning_rate): Tune the learning rate — 3 points · 4 H100 hrs

Deliverable: (a) learning curves for a sweep over learning rates with your search strategy explained, and a model hitting the target; (b) curves including at least one divergent run, with an analysis relating the divergence point to the best rate. The bar is concrete:

A model with validation loss (per-token) on TinyStories of at most 1.45Handout §7.2, Problem (learning_rate), deliverable (a) — relaxed to 2.00 on the CPU/MPS path

Sweep in multiplicative steps (a factor of 2 or 3 per point, not a linear grid) until you find divergence, then work back. The folk wisdom part (b) probes — that the best learning rate sits just below the one that blows up — is worth taking seriously because it is the practical procedure people actually use at scale, and because seeing your own loss curve go vertical is more instructive than reading that it can. If nothing diverges, your gradient clipping is masking the instability; try turning it off for the diagnostic runs. If you are on CPU or MPS, remember to move the cosine decay's endpoint to your reduced step count.

Problem (batch_size_experiment): Batch size variations — 1 point · 2 H100 hrs

Deliverable: learning curves across batch sizes from 1 up to the GPU memory limit, including typical values like 64 and 128, with the learning rate re-tuned where necessary, plus a few sentences of findings. The instructive result is that the relationship is not monotone in the way you might expect: very small batches are slow per token and noisy, very large batches waste compute on redundant gradient information, and the useful regime is wide and flat. Because the total token budget is fixed, doubling the batch halves the number of optimizer steps — so this is really a measurement of how much of training's progress comes from steps versus from tokens. Keeping the learning rate fixed while changing the batch size confounds the comparison, which is why the deliverable explicitly permits re-tuning.

Problem (generate): Generate text — 1 point · uses your trained checkpoint

Deliverable: at least 256 tokens of generated text (or up to the first <|endoftext|>), plus a comment on its fluency naming at least two factors that affect quality. The handout prints reference samples for both the full and the low-resource configurations so you can calibrate — the full-budget sample is coherent if repetitive children's-story prose, the 40M-token sample is recognizably English but loses the thread mid-paragraph. Factors worth discussing: total tokens processed, the decoding parameters themselves (temperature and top-p can rescue or ruin the same checkpoint), model capacity, and the narrowness of the TinyStories distribution, which flatters a small model enormously — a point that main_experiment makes by contrast.

Problem (layer_norm_ablation): Remove RMSNorm and train — 1 point · 1 H100 hr

Deliverable: two learning curves — one at your previously optimal learning rate, one at the best rate you can find — plus a few sentences on RMSNorm's impact. Removing normalization from a pre-norm Transformer typically diverges immediately at the tuned learning rate, and recovers only at a substantially smaller one, and then to a worse final loss. The reason is that without normalization the residual stream's scale grows layer over layer, so the same learning rate that was appropriate at layer 1 is far too large at layer 4. This is the cheapest, most legible demonstration in the assignment that "stability trick" is not a euphemism for "minor detail".

Problem (pre_norm_ablation): Implement post-norm and train — 1 point · 1 H100 hr

Deliverable: a learning curve for a post-norm Transformer compared against your pre-norm baseline. Post-norm is the original 2017 arrangement — z = RMSNorm(x + Attn(x)), y = RMSNorm(z + FFN(z)) — which puts a normalization inside the residual path rather than beside it. The expected result is trainable but worse and more sensitive, matching Nguyen and Salazar 2019 and Xiong et al. 2020. At four layers the gap is modest; the literature's point is that it widens sharply with depth, and that post-norm models need a warmup schedule to train at all. If you named your norms ln1 and ln2 as the reference does, this ablation is a few lines — move the calls from the sublayer inputs to the sublayer outputs.

Problem (no_pos_emb): Implement NoPE — 1 point · 1 H100 hr

Deliverable: a learning curve comparing RoPE against no position information at all. The surprise is that NoPE works — often nearly as well at this scale — because a causal decoder can infer position from the mask structure itself: the token at position i attends to exactly i+1 positions, and that count is recoverable information. Tsai et al. 2019 and Kazemnejad et al. 2023 are the references; the latter argues NoPE can even generalize better to longer sequences than explicit schemes. Implementation is a flag that skips the RoPE call on Q and K — do not remove the module, since you need both configurations for the comparison.

Problem (swiglu_ablation): SwiGLU vs. SiLU — 1 point · 1 H100 hr

Deliverable: a learning curve comparing SwiGLU against an ungated FFN_SiLU(x) = W2 SiLU(W1 x), at approximately matched parameter counts, plus a discussion. The parameter matching is the whole design of the experiment and the handout gives you the recipe: the ungated version uses d_ff = 4 · d_model with two matrices, against SwiGLU's 8/3 · d_model with three, so the two configurations have about the same parameter count and the comparison isolates gating rather than capacity. Expect a small but real win for SwiGLU — which is exactly Shazeer 2020's finding, and small enough that running the ablation yourself is the only way to believe it.

Problem (main_experiment): Experiment on OWT — 2 points · 3 H100 hrs

Deliverable: a learning curve on OpenWebText with the same architecture and iteration count as TinyStories, an interpretation of how the losses differ, and a generated text sample with an explanation of why the quality is worse at identical compute. The losses are not comparable across the two runs, and saying why is most of the point: different tokenizers means different vocabularies (10K versus 32K) and different bytes-per-token, so per-token cross-entropy is measured in different units. Beyond the units issue, OpenWebText is simply a far higher-entropy distribution — arbitrary web prose, many domains, code, markup, names — where TinyStories is a synthetic corpus with a small vocabulary and formulaic structure. A 17M-parameter model can nearly memorize the latter's regularities and cannot begin to cover the former's. The handout warns you may need to retune the learning rate or batch size.

Problem (leaderboard): Leaderboard — 6 points · 10 H100 hrs

Deliverable: the final validation loss, a learning curve with a wallclock x-axis under the budget, a description of what you did, and a pull request to the leaderboard repo. The Spring 2025 rules are short:

Runtime Your submission can run for at most 1.5 hours on an H100. […] Data You may only use the OpenWebText training dataset that we provide. Otherwise, you are free to do whatever your heart desires.Handout §7.5, "Rules for the leaderboard" — the target is to beat a naive baseline of 5.0 loss

This is a fixed-compute quality competition, which is a different optimization problem from the one you have been solving: anything that improves loss-per-second wins, including choices that would be wrong at larger scale. Ideas the handout points at: weight tying between the input embedding and the LM head (see Vaswani et al. §3.4 and PaLM §2 — and reduce the embedding init standard deviation if you do it), architectural borrowings from Llama 3 or Qwen 2.5, and the accumulated tricks in modded-nanogpt. Test candidate modifications on TinyStories or an OWT subset before spending a 1.5-hour slot. The handout's own caveat is worth heeding: some of what wins here will not survive scaling, which is precisely what L09 and L11 and Assignment 3 are about.

What you hand in

Two files to Gradescope. writeup.pdf answers every written question — typeset, not photographed — which means all of §2's Unicode and tokenizer-experiment answers, both accounting problems, the learning-rate observations, every learning curve from §7, the generated text dumps, and your experiment log. code.zip contains everything you wrote. The repo ships make_submission.sh to build it: it runs the full test suite to test_results.xml (with || true, so a failing test does not abort the packaging) and zips the tree into cs336-spring2025-assignment-1-submission.zip, excluding caches, virtualenvs, and every data-shaped extension — .txt, .json, .pkl, .npy, .npz, .pt, and the test fixtures and snapshots. Check the resulting zip before you upload: if you stored a serialized vocabulary as .json or a merge list as .txt inside the repo, the script strips it.

The leaderboard entry is separate: a pull request to assignment1-basics-leaderboard adding a row to the results table, sorted by increasing loss, with your final validation loss, a link to a publicly viewable learning curve (a W&B report link — open it in a private window to confirm it is public — or an image committed to the repo's images/ folder), and a description of what you did. Compute discipline for the whole assignment: roughly 23 H100 hours across §7 if you follow the printed budgets, of which the leaderboard alone is 10, plus up to 12 non-GPU hours for the OpenWebText BPE training. Nothing except the leaderboard run has a hard wall-clock limit, but the tokenizer's test_train_bpe_speed does, and it is the only place in A1 where slow code is a failing grade rather than an inconvenience.

Materials, verified

• Assignment 1 handout (PDF), v1.0.6 — the source this page walks through, pinned at the spring2025 tag
• tests/adapters.py — the 21 stubs that define the grading contract; every signature quoted on this page was read from here
• README.md · pyproject.toml — setup, the data download commands, and the pinned dependency set
• Neural Machine Translation of Rare Words with Subword Units — Sennrich, Haddow, Birch (2016) · the BPE-for-NLP paper the tokenizer section is built on, including the stylized low/lower/widest/newest example
• Neural Machine Translation with Byte-Level Subwords — Wang, Cho, Gu (2019) · why the base alphabet is 256 bytes rather than Unicode code points
• Language Models are Unsupervised Multitask Learners (GPT-2) — Radford et al. (2019) · source of the pre-tokenization regex; the exact pattern is in tiktoken PR #234
• Improving Language Understanding by Generative Pre-Training (GPT-1) — Radford et al. (2018) · cited for the output-embedding / LM-head formulation in §3.2
• Attention Is All You Need — Vaswani et al. (2017) · §3.2.1 scaled dot-product attention, §3.2.2 multi-head, §3.3 the FFN you replace, §3.4 weight tying, §5.4 residuals
• Layer Normalization — Ba, Kiros, Hinton (2016) · the baseline RMSNorm simplifies
• Root Mean Square Layer Normalization — Zhang, Sennrich (2019) · equation 4 is the one you implement
• Transformers without Tears — Nguyen, Salazar (2019) · and On Layer Normalization in the Transformer Architecture — Xiong et al. (2020) · the pre-norm case, and the reference for the post-norm ablation
• Bridging Nonlinearities and Stochastic Regularizers (GELU/SiLU) — Hendrycks, Gimpel (2016) · and Sigmoid-Weighted Linear Units — Elfwing, Uchibe, Doya (2017) · the two independent introductions of SiLU
• Language Modeling with Gated Convolutional Networks — Dauphin et al. (2017) · the GLU · and GLU Variants Improve Transformer — Shazeer (2020) · SwiGLU, and the source of the "divine benevolence" line
• RoFormer: Enhanced Transformer with Rotary Position Embedding — Su et al. (2021) · RoPE
• LLaMA: Open and Efficient Foundation Language Models — Touvron et al. (2023) · the source of the RMSNorm choice, the no-bias convention and the cosine schedule you implement
• The Llama 3 Herd of Models — Grattafiori et al. (2024) · and Qwen2.5 Technical Report — Yang et al. (2024) · current SwiGLU-and-RoPE production architectures; suggested reading for leaderboard ideas
• PaLM: Scaling Language Modeling with Pathways — Chowdhery et al. (2022) · the no-bias convention, weight tying, and the definition of MFU used in adamwAccounting (d)
• Adam: A Method for Stochastic Optimization — Kingma, Ba (2015) · also the handout's pointer for typical hyperparameter values · and Decoupled Weight Decay Regularization — Loshchilov, Hutter (2019) · algorithm 2 is the AdamW you implement
• Language Models are Few-Shot Learners (GPT-3) — Brown et al. (2020) · cited for the (0.9, 0.95) beta choice at scale
• Scaling Laws for Neural Language Models — Kaplan et al. (2020) · and Training Compute-Optimal Language Models (Chinchilla) — Hoffmann et al. (2022) · the source of the "backward is 2× forward" convention used in the accounting problems
• The Curious Case of Neural Text Degeneration — Holtzman et al. (2020) · top-p / nucleus sampling
• Transformer Dissection — Tsai et al. (2019) · and The Impact of Positional Encoding on Length Generalization in Transformers — Kazemnejad et al. (2023) · why a causal decoder can work without position embeddings at all
• TinyStories: How Small Can Language Models Be and Still Speak Coherent English? — Eldan, Li (2023) · the training corpus · files: HF dataset, train.txt, valid.txt
• OpenWebText Corpus — Gokaslan, Cohen, Pavlick, Tellex (2019) · the course subsample: stanford-cs336/owt-sample, owt_train.txt.gz, owt_valid.txt.gz
• modded-nanogpt — Keller Jordan et al. · the NanoGPT speedrun repo the handout points at for leaderboard ideas
• einops · einx · jaxtyping · uv — the four libraries the handout asks you to learn along the way
• torch.nn containers · torch.optim.Optimizer · nn.init.trunc_normal_ · np.memmap — the exact APIs you are allowed and asked to use
• PyTorch on Apple Metal · PyTorch MPS notes — for the low-resource path
• U+FFFD, the Unicode replacement character — what your decode must emit for malformed byte sequences
• Field map extra: karpathy/minbpe and Let's build the GPT Tokenizer — a clean reference implementation and a two-hour walkthrough of exactly §2; read them after your own attempt, not before
• Field map extra: The Annotated Transformer — the 2017 architecture line by line, useful for seeing precisely which four things §3 changes
• Field map extra: karpathy/nanoGPT — the closest public analogue to what §5 asks you to build, including the memmap loader and checkpoint-resume patterns