CS336: Language Models From Scratch – Andy Timm

If you wish to make an apple pie LLM from scratch, you must first invent the universe systems considerations.

— Carl Sagan, probably

In December, I finished working through Stanford NLP’s CS336: Language Models From Scratch. This post is a short review of the class, and a bunch of practical tips/considerations if you’re considering auditing it as well.

The basic premise of the course is that:

Researchers/Engineers are becoming disconnected from the underlying details of LLMs.

Moving up these levels of abstraction can allow you to move faster.

BUT:

There are elements of deeper understanding, research taste, and implementation prowess that are increasingly easy to never learn.

Thus: build Modern LLMs from scratch to build deeper understanding, and improve capacity to do more fundamental research.

The course delivered on this premise- I feel better equipped to understand/iterate on recent research. I’m considerably more capable at the systems skills1 increasingly required for meaningful LLM engineering. Full courses are rarely better than fully self guided learning/projects for me, so this is a strong endorsement.

Course Contents

CS336 has great lectures that survey current best practices, but the primary meat of the course is in the assignments, which ask you to build. That is, building a full tokenizer, and basic LLM from scratch are both in the first (2 week!) assignment. The second (2 week!!) assignment has you implement both Flash Attention in Triton, and handroll a version of efficient-ish DDP2.

The expectations here are (wonderfully) high- each assignment asks for a relatively low level implementation of a conceptual building block of modern LLMs, each of which would be a meaty side project in their own right.

Fortunately, these assignments are nicely scaffolded and provide opinionated input on key design decisions, which reduces the system design burden. So you’re asked to draw the rest of the fucking owl3, but they do give you guidance on which circles to draw first.

Assignments

The course is structured around 5 assignments. The topics in each are mostly shown above, but here’s a short teaser for each:

Assignment 1

Build a Tokenizer: Tokenization is a primary source of spooky confusing bullshitTM with LLMs4. Building your own tokenizer gives you a finer grained model of how tokenizers work, what tradeoffs exist in tokenization, and (personally) really helped me feel more confident fixing tokenizer issues.

Build a basic Transformer LLM from Scratch: Implementing everything from the MLPs on up, by hand, quickly identifies and requires you to fix any basic holes in your understanding.

Ablations—why are best practices best practices? Best practices in hyperparameter choices like depth/width ratios and so on can feel like the field blindly hill climbed their way to “what worked best”. While there’s truth to this, ablating many of the core points of architecture agreement will give you stronger intuition for why certain choices are points of convergence across labs.

Assignment 2

Benchmark & Implement Flash Attention: LLM’s success is heavily the story of scaling. Implementations that scale require understanding low level details of how GPUs work, how their design influences algorithmic choices, and which resources constrain you most. Flash Attention is a marvel of all these elements coming together, so implementing (much of) it from scratch forces you to understand the details, and feel the binding resource constraints.

Implement DDP: Going beyond training/finetuning/RLing relatively small models requires intelligently using multiple GPUs. DDP definitely isn’t the whole parallelism story, but implementing it from scratch makes you really feel and develop intuition for the communications challenges inherent in this type of scale up.

Assignment 3

Feeling the Scaling Laws: Scaling Laws tell us, with a surprising degree of accuracy, what performance to expect from models as we scale them up. If you actually had to make decisions about a larger training run, what small experiments would you prioritize to make decisions for the final run? You’ll feel the scaling curves you’ve seen in your bones a bit more after this assignment.

Assignment 4

Cleaning Common Crawl: There’s SO MUCH web data, and LLMs, like all ML models, are pretty garbage in, garbage out. There are so, so many issues with web scrapes—you can’t fix all of them, and you’re always going to be in triage mode. What cleaning has emerged as best practice?

Tinkering with Data Decisions: By ablating various cleaning decisions, you can develop a bit more feel for WHY certain cleaning steps get prioritized, and where remaining alpha lies.

Assignment 5

Implementing SFT/Simple RLVR/GRPO RLVR/(Optionally) RLHF via DPO: There’s so, so much beyond training a good base model that goes into making a capable LLM that is nice to work with. What are all of these steps, how has implementing...