Big picture: what CS336 is trying to teach you
CS336 is about learning to build a modern LLM as a system, end to end — not just reading about transformers or calling an API.
From the review of the course, the core premise is: people are sliding up abstraction layers, losing contact with the details, so the course has you “build modern LLMs from scratch” to regain that depth and systems taste (source).
Concretely, you touch almost every layer mentioned in overviews like Hao Hoang’s writeup: tokenization, data pipelines, architecture design, training, scaling, inference, even alignment techniques like SFT/RLHF (source).
The end-to-end LM pipeline you’ll build
You’re heading toward a minimal GPT-style pipeline with these major components:
Let’s unpack each piece — because you’ll actually implement them, not just call them.
| Component | What it is | What you’ll actually build / do |
|---|---|---|
| Data | Source text and how you feed it to training | Collect or load a text corpus, clean it, split into train/val/test, chunk into sequences |
| Tokenizer | Map between text and token IDs | Implement a tokenizer (likely BPE or similar) and use it to turn text into integer sequences and back |
| Model | Neural net that predicts next token | Implement a transformer-style GPT: embeddings, attention, MLP blocks, output head |
| Training | How the model learns from data | Implement the training loop: batching, forward pass, loss, backward, optimizer, checkpointing |
| Eval | How you know it works | Compute metrics like loss/perplexity, run generation, compare versions and ablations |
Mechanistically, the flow is:
- Data → Tokenizer: you start with raw text (say,
tiny_shakespeare.txt), and your tokenizer learns a vocabulary + rules to turn"Hello world"into e.g.[1513, 82, 50256]. - Tokenizer → Model: those integer token IDs go into embeddings, through the transformer, out comes a probability distribution over the next token.
- Model → Training: you compare the model’s predicted distribution to the actual next token (cross-entropy loss), backpropagate gradients, and update weights with an optimizer.
- Training → Eval: periodically, you run on a held-out validation set with no gradient updates. Lower val loss = better generalization; you can also generate text and eyeball quality.
A tiny, runnable sketch (PyTorch) of part of this pipeline, using a dummy tokenizer and model:
pythonimport torch
import torch.nn as nn
# 1) "Tokenizer": here just character to index
chars = sorted(list(set("hello world")))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}
def encode(s):
return [stoi[ch] for ch in s]
def decode(ids):
return "".join(itos[i] for i in ids)
# 2) Simple language model over characters
vocab_size = len(chars)
embed_dim = 16
class TinyLM(nn.Module):
def __init__(self, vocab_size, embed_dim):
super().__init__()
self.embed = nn.Embedding(vocab_size, embed_dim)
self.linear = nn.Linear(embed_dim, vocab_size)
def forward(self, x):
# x: (batch, seq_len)
emb = self.embed(x) # (batch, seq_len, embed_dim)
h = emb.mean(dim=1) # super dumb "context" pooling
logits = self.linear(h) # (batch, vocab_size)
return logits
model = TinyLM(vocab_size, embed_dim)
# 3) Training step on "hello" -> predict last char
tokens = torch.tensor([encode("hello")[:-1]]) # input: "hell"
target = torch.tensor([encode("hello")[-1]]) # target: "o"
logits = model(tokens)
loss = nn.CrossEntropyLoss()(logits, target)
loss.backward()
This is not a transformer, but the mechanism is exactly the same at a high level: data → IDs → model → logits → loss → gradients.
When you get lost in details (kernels, Triton, config files), come back to this 5-piece mental model: data, tokenizer, model, training, eval. Every assignment plugs into one or more of these.