Preface

I wrote this book because I got tired of treating the pieces of a modern reasoning-model pipeline as separate black boxes. I could serve a model. I could point an eval harness at an endpoint and read a number off the screen. I could, on a good day, run a fine-tune. What I could not do was hold the whole loop in my head at once and say, honestly, why each number came out the way it did. The eval score was a number the harness handed me. The reward was a function I copied from somebody's repo. The gradient step was whatever the trainer did between two log lines. I wanted to close the loop and see the seams.

So this is a book about one loop, run end to end, on one GPU, with the math left in.

The loop is the spine

Everything here hangs off a single cycle: serve a model, evaluate it on a task, turn each response into a score, use those scores to train, then evaluate again and see whether the number moved. That is the whole argument of the book compressed into five words, and I come back to it constantly.

flowchart LR
    A[Serve<br/>open-weight model] --> B[Evaluate<br/>run the task]
    B --> C[Score<br/>reward per response]
    C --> D[Train<br/>update the weights]
    D --> E[Re-evaluate<br/>did the number move?]
    E -->|next iteration| A

That is Figure 0.0, and it is the mental model I want you to carry into every chapter. When we are deep in the weeds of a KL penalty or a vLLM flag or a bootstrap confidence interval, the question is always the same: which arrow of this diagram am I standing on right now, and what does this piece feed into next? "Evals as rewards," the title, is really a claim about the middle of this diagram: that the thing you measure with (the eval) and the thing you train against (the reward) are two readings of the same underlying signal, and that if you understand one you are most of the way to understanding the other. A well-posed eval is a reward function you have not plugged into an optimizer yet. A reward function is an eval you have decided to differentiate through. Keeping those two ideas in the same frame is the point.

The reason the loop matters more than any single stage is that reasoning models are trained by closing it. You do not get a model that reasons better by staring at a static benchmark. You get it by scoring responses, pushing the weights toward the responses that scored well, and checking that the benchmark actually moved and did not just wobble inside its own noise. If you only ever run the top arc (serve and evaluate) you are doing measurement. If you only ever run the bottom arc (score and train) you are optimizing something you cannot see. The whole point is to run the circle.

Why one 16GB GPU

The obvious way to write a book like this is to assume a cluster, or at least a rented 80GB card, and wave at the parts that do not fit. I deliberately did the opposite. Every lab in this book runs on a single NVIDIA RTX 5080 with 16GB of VRAM, in the machine I describe in the hardware baseline chapter, and that constraint is not an apology. It is a design input.

Sixteen gigabytes is enough to be real and small enough to hurt, which is exactly the regime where you are forced to understand what is actually consuming memory. When the whole model plus its KV cache plus optimizer state plus activations has to fit inside a budget you can count on your fingers, you cannot treat memory as infinite and you cannot hand-wave the arithmetic. You end up knowing, to within a few hundred megabytes, where every gigabyte went, because you had to. That knowledge is transferable in a way that "it fit on the A100" never is. Throughout the book I keep a running vram-budget admonition precisely so the memory accounting stays in the foreground rather than being the thing that mysteriously blows up at 2am.

The constraint also keeps the book honest about tradeoffs. Quantization, LoRA and its low-rank cousins, gradient checkpointing, paged KV caches, small-batch RL: these stop being optional cleverness and become the load-bearing structure of every lab. You learn them because you have no choice, which is the best reason to learn anything.

Three readers, one book

I am writing for three audiences at once, and I owe each of them a different thing.

The first reader is future me, six months from now, trying to re-run a lab on a machine whose driver has moved on. That reader needs reproducibility above charm: exact commands, pinned versions, artifacts written to disk, and numbers stamped with the date and driver they were measured under. Wherever this book feels almost pedantically specific about how something was run, that is me writing to future me, who will not remember and will not forgive vagueness.

The second reader is the Substack reader, arriving one chapter at a time. The book is built so that chapters extract cleanly into posts, which is why each one opens with intuition before it earns the right to show you an equation, and why I seed shareable framings in substack-seed admonitions as I go. If you are that reader, you should be able to drop into a chapter, get the idea, and leave with something you could explain to a colleague over coffee.

The third reader is a thesis committee, and this one sets the floor for rigor. Every derivation has to survive being read by someone whose job is to find the step I skipped. Every statistic (every confidence interval, every claim that a training run "improved" something) has to be defensible, not vibes. When I compute an uncertainty I show the estimator; when I claim a difference is real I say how I know it is not noise. The derivation admonitions are where I pay that debt in full, working things from definitions rather than asserting them.

Serving three readers at once is a discipline, not a compromise. Intuition first keeps the Substack reader; math in full keeps the committee; exact commands and dated numbers keep future me. A chapter is done when all three are satisfied.

The main line and the Thesis Thread

The book runs on two tracks. The main line is generic: it works with open-weight reasoning models and standard tasks, and nothing in it depends on my particular research problem. You can follow the main line, learn the loop, and apply it to whatever you care about.

Braided through the main line is a recurring worked example I call the Thesis Thread, flagged wherever it appears with a thesis-thread admonition. The thread follows one concrete problem in verifiable reasoning (an SDA-flavored task where a response can be checked, not just liked) from first eval to trained model. It exists to show the whole loop applied to a single problem with real stakes, so the generic machinery has somewhere to land. If you only want the general skills, you can skip the thread and lose nothing structural. If you want to see how the pieces actually compose on a problem someone cared enough to write a thesis about, follow it. Keeping the two tracks visibly separate is deliberate: it lets the general book stay general while still getting the payoff of a sustained example.

What this book is not

To keep the contract clean, here is what I am leaving out on purpose.

This is an open-source-stack book. Every tool in it is something you can read the source of, run locally, and pin to a version. I lean on hosted APIs only as a point of comparison, never as a dependency, because a book about closing the loop on your own machine cannot have its core steps happen on someone else's. If a lab needs a component, that component is open weights or open source, full stop.

It is also not a business book. There is nothing here about pricing a product, sizing a market, or deciding whether any of this is worth doing commercially. Those are real questions and other people write about them well. My scope is the technical loop and the understanding required to run it, on the assumption that you already decided you want to. Where a real deployment would raise a cost or product question, I will name it and move on rather than pretend to answer it.

That is the deal. One loop, one GPU, all the math, three readers, two tracks, open source only. The next chapter teaches you how the book is put together so the conventions never get in your way. After that we start unboxing the machine.

How to read this book

Before you hit chapter one, spend a few minutes here. This book has a fixed rhythm, a small vocabulary of labeled boxes, a shorthand for the references I lean on, and a couple of non-negotiable rules about how the code runs. None of it is complicated, but knowing it up front means the conventions get out of your way instead of tripping you up later. Think of this chapter as the legend on a map.

The Theory, Tooling, Lab rhythm

Most chapters move through three phases in the same order, and the order is the pedagogy.

Theory comes first. I try to give you the intuition for what we are about to do and why it should work before I write a single equation, because an equation you were handed cold is just notation, whereas an equation you were expecting is an insight. This is also the part that extracts most cleanly into a blog post, so it is written to stand more or less on its own.

Tooling comes second. Here the intuition earns its math, and the math earns its implementation. This is where I derive the thing properly (the loss, the estimator, the update rule) and then connect each symbol to the actual library, flag, or function that computes it. The goal of this phase is that you never have to take the tooling on faith: you can point at a line of output and say which term of which equation produced it.

Lab comes last, and every lab ends with something written to disk. You run the commands, you get an artifact (a metrics JSON, a set of scored responses, a checkpoint, a plot), and that artifact is the proof the loop closed. The reason theory precedes tooling precedes lab is simple: I want you to understand the thing, then see how it is built, then build it yourself, in that order, so that when the lab misbehaves you already have the model in your head to debug it. Reversing the order (code first, understanding later) is how you end up with a working script you cannot reason about, which is exactly the situation this whole book exists to get you out of.

Not every chapter uses all three phases with equal weight. A reference chapter may be almost all tooling; the preface and this chapter are all theory. But when a chapter is teaching a piece of the loop, expect the three-beat structure.

The admonition vocabulary

Throughout the book I use a small set of labeled callout boxes (mdBook calls them admonitions). Each type means one specific thing, so you can decide at a glance whether a box is essential to your path or safely skippable. Here is the whole vocabulary.

If you read nothing but the prose and the gotcha boxes you will still be able to run every lab. If you read the derivation boxes too you will be able to defend every result. The rest are there to enrich, extend, or share.

Reference short keys

I lean on seven books repeatedly, and rather than spell out full citations every time, I use short keys in the text and in read-along boxes. Here is the key.

When I write "see [S&B] on the policy gradient theorem" or drop a read-along box pointing at [BRM], those keys are what I mean. The intent is not to make you buy seven books. It is to be honest about where an idea came from and to give you a well-worn path to more depth than a single chapter can hold. Where two references cover the same ground differently (and [S&B] versus [RLHF] on reward often will), I will say which one I am following and why.

uv everywhere

Every piece of Python in this book runs under uv. There are no bare pip install lines and no hand-rolled python -m venv incantations anywhere in the labs, and this is a rule, not a preference.

The reason is reproducibility, which is the promise I made to future me in the preface. uv gives every lab a pinned, lockfile-backed environment that resolves the same way tomorrow as it did today, and it does it fast enough that spinning up a fresh environment per lab is painless rather than a chore. So when a lab needs dependencies you will see uv managing them, when it runs a script you will see uv run, and when it pins versions you will see them in a lockfile that is itself an artifact. If you are used to pip and venv, everything you know transfers; uv is just the tool that makes the pinning automatic instead of aspirational. If you have never used it, the chapter-zero setup walks you through installing it once, and after that it is invisible.

Every lab ends in an artifact

I said it under the rhythm and I will say it again on its own because it is the load-bearing rule of the whole book: a lab is not done when the script exits zero. A lab is done when it has written a durable artifact to disk. A metrics file, a table of scored responses, a checkpoint directory, a saved figure. Something you can come back to, diff against a later run, and point at in an argument.

This is what makes the serve, evaluate, score, train, re-evaluate loop auditable rather than anecdotal. "The model got better" is a claim. Two dated metrics JSONs, one before training and one after, with the estimator and its uncertainty recorded in each, are evidence. The artifacts are also what let future me re-run a lab and check whether the new hardware gives the old answer. So expect every lab to tell you exactly what it wrote and where, and expect me to treat that file, not the terminal output, as the real result.

Running the book on the same hardware

This book is written to be run, not just read, and it is written against one specific machine, described in full in the next chapter. If you have that machine or something close to it, you can execute every lab as written and compare your numbers to mine directly. If your hardware differs, the labs still run; what changes is the measured numbers, and because every measured number in this book is stamped with the date and driver it came from, you will always be able to tell which figures are hardware-relative and recompute them on your own box.

Practically, that means the ideal way to read this book is with the machine on and a terminal open. Read the theory, follow the tooling, then actually run the lab and watch your artifact land on disk next to mine. The book will make sense from the armchair, but it was built for the workbench. The next chapter pins down exactly what that workbench is.

The hardware baseline

Every measured number in this book is relative to one machine. If I tell you a serving configuration hits some throughput, or a training step takes some seconds, or a KV cache eats some fraction of VRAM, that number only means something once you know exactly what it ran on. So before any measurement appears, here is the machine it was measured on, pinned down to the parts that matter.

The baseline machine is an MSI Aegis R2, model A2NVV9-2218US, a prebuilt desktop I chose specifically because it is an ordinary, buyable configuration rather than a bespoke workstation. The whole premise of the book is that the loop runs on hardware a normal person can own, so the baseline had to be hardware a normal person can buy.

The full specification

Whenever I refer to "the baseline machine" anywhere in the book, this is the row-by-row thing I mean.

Why each part matters

The GPU is the whole story, and two of its numbers set the terms for everything else. It is a Blackwell-generation card, which matters because the software stack (the CUDA version, the driver branch, the kernels that vLLM and the training libraries dispatch to) has to actually support Blackwell, and at the time I built this that support was new enough to be worth stating explicitly rather than assuming. And it has 16GB of VRAM, which is the constraint the entire book is organized around. Every vram-budget admonition is accounting against this 16GB ceiling. The roughly 960 GB/s of memory bandwidth is the number that quietly governs inference speed for the memory-bound regime we usually live in: for autoregressive decoding, throughput is often set more by how fast you can stream weights and KV cache through memory than by raw compute, so the bandwidth figure is one I will keep returning to when a measurement looks surprising.

The CPU is an Intel Core Ultra 9 285, an Arrow Lake part. It is not doing the heavy lifting, but it is not idle either. Data loading, tokenization, the orchestration around a serving engine, and the CPU-side of any offloading strategy all land here, and when a lab is unexpectedly slow it is worth knowing whether the GPU is waiting on the CPU. Arrow Lake also brings its own quirks (its power and scheduling behavior differs from older Intel parts), which is exactly the kind of thing that is invisible until it explains a weird measurement.

System RAM is 32GB of DDR5. That is comfortable but not lavish, and it becomes relevant the moment a technique spills off the GPU: CPU offload of optimizer state or KV cache, memory-mapped weight loading, or just holding a dataset in memory. 32GB means offloading is possible but bounded, and I will flag the cases where it is the thing that fits or does not.

Storage is deliberately split into two tiers, and the split is a design choice, not an accident of what was in the box.

Working storage versus archive

The 1TB NVMe SSD is working storage. It holds the things a running lab touches constantly: the model weights currently being served or trained, the active dataset, checkpoints being written mid-run, the scratch space for a training job. NVMe is here because weight loading and checkpoint I/O are fast enough to not be the bottleneck, which matters when a lab loads a multi-gigabyte model or writes a checkpoint every few hundred steps.

The 5TB NAS is the archive tier. It holds the things I want to keep but am not actively hammering: past checkpoints, completed run artifacts, datasets I am not currently using, the accumulated sediment of every lab I have ever run. It is network-attached and therefore slower than the NVMe, which is fine, because nothing in the hot path should be reaching across the network for it. The discipline is to keep the working set on NVMe and everything else on the NAS, so the fast disk stays fast and the archive grows without pressure.

That is the high-level shape. The actual mechanics (how the mount is set up, how artifacts get promoted from working to archive, how I keep the two from drifting) are a chapter 0.5 problem, and I will lay out the whole storage workflow there. For now the thing to hold onto is just the two-tier split: NVMe for what is hot, NAS for what is kept.

Dating and driver-stamping every number

Here is the discipline that makes this baseline useful rather than decorative. Every genuinely measured number in this book is recorded with the date it was measured and the driver it was measured under. Not "about 40 tokens per second" but "40 tokens per second, measured on the baseline machine on such-and-such date under driver 570-open." Where a number in the text is a real measurement, you will see it written as (measured on the baseline machine — record value, date, driver) until it is filled in from an actual run, precisely so that an unfilled measurement is impossible to mistake for a real one.

The reason is that hardware and drivers move, and a number without a stamp is a number you cannot trust later. A driver update can change throughput. A new CUDA release can change what fits in memory. A kernel improvement can make last month's careful workaround unnecessary. If every measurement carries its date and driver, then when I upgrade the driver next year I can re-run the labs and diff the new numbers against the old ones and actually see what the upgrade changed. If the numbers were undated, that diff would be meaningless, because I would not know whether a difference was the upgrade or just drift I never pinned down.

This is the same promise I made to future me in the preface, made concrete. A dated, driver-stamped measurement is reproducible in the only sense that matters: I can tell you the conditions precisely enough that you, or a later me, can reproduce them and check. An undated number is a rumor. So the rule for the whole book is that measured quantities are never bare. They come stamped, or they come marked as not yet measured, and there is no third category.

With the machine pinned and the stamping discipline stated, every measurement from here on has a fixed frame of reference. The next chapter starts the actual work: unboxing this machine and getting it ready to serve its first model.

Why this machine, this stack

Goal. Orient the reader: what the loop is, and why one number, 16 GB of VRAM, drives every downstream choice in this book.

Covers. The loop diagram (Figure 0.1); the 16 GB constraint as the central design input; working (NVMe) versus archive (NAS) storage tiers; the open-data boundary and why any business stack stays separate.

Theory

The loop is the thesis

Everything in this book is one loop. I serve an open-weight reasoning model, evaluate it on tasks with checkable answers, score those answers into a scalar reward, use that reward to train the model, and then re-evaluate to see whether the number moved. That is the whole machine. The rest of the book is me refusing to treat any box in that loop as opaque.

flowchart LR
    S["Serve<br/>(vLLM, OpenAI-compatible)"] --> E["Evaluate<br/>(Inspect tasks)"]
    E --> C["Score<br/>(scorers → scalar reward)"]
    C --> T["Train<br/>(Unsloth GRPO / LoRA)"]
    T --> R["Re-evaluate<br/>(same tasks, held-out split)"]
    R -.->|"reasoning delta"| S
    R --> L["Log<br/>(MLflow: metrics, artifacts, SHAs)"]
    L -.-> S

Figure 0.1 - the serve → evaluate → score → train → re-evaluate loop. The dotted return edge is the only thing I actually care about: did the model get better at reasoning, measured the same way twice?

The reason I draw it as a cycle and not a pipeline is that the interesting quantity lives on the return edge. A pipeline ends with a deliverable. This loop ends with a comparison: the same evaluation, run before and after training, on data the model never trained on. I call that difference the reasoning delta, and Part 7 is entirely about measuring it honestly. If I cannot trust the evaluation, the delta is noise dressed up as progress, so a surprising amount of this book is spent making the "evaluate" and "score" boxes boring and defensible before I ever let the "train" box touch a weight.

One GPU is a feature, not an apology

The subtitle says "on one GPU," and I mean it as a design constraint I chose, not a limitation I am making excuses for. The baseline machine is an MSI Aegis R2 (A2NVV9-2218US): an Intel Core Ultra 9 285, an NVIDIA RTX 5080 with 16 GB of VRAM (Blackwell, roughly 960 GB/s of memory bandwidth per NVIDIA's RTX 5080 product specification), 32 GB of DDR5, a 1 TB NVMe system drive, and a 5 TB NAS for archive. Ubuntu 24.04 LTS, NVIDIA 570-open driver. Every measured number in this book is relative to exactly that machine, which is why I date and driver-stamp all of them.

The single most consequential number in that list is 16 GB. Not the core count, not the 960 GB/s, not the 32 GB of system RAM. Sixteen gigabytes of VRAM is the constraint that, once you internalize it, predicts almost every architectural decision I make later. So it is worth sitting with the number before we get anywhere near a command line.

Let me make the weights arithmetic explicit, because it is the seed of everything. A parameter count $N$ stored at $b$ bytes per parameter needs

$$\text{weight bytes}=N\times b.$$

At bf16, $b=2$. At 4-bit quantization, $b=0.5$, plus a small overhead for scales and zero-points. So the same 7B model is about 14 GB at bf16 and roughly 4 GB at 4-bit. On a 16 GB card, the difference between those two numbers is the difference between "will not load" and "loads with 12 GB left over for cache and activations." I will derive the KV-cache and activation terms carefully in Part 1 and Part 2; here I only want the reader to feel the shape of the constraint. VRAM is a fixed pie. Every design choice is really a question of which slice something is allowed to take.

Two storage tiers, because weights are heavy and cheap to lose

The second design input is mundane and I ignored it at my peril the first time: model weights and datasets are large, and the two things I want from them, fast access and durable retention, pull in opposite directions. Fast access wants the local NVMe. Durable retention wants somewhere I will not accidentally rm -rf during a cache cleanup. So I split storage into two tiers from day one.

The working tier is the 1 TB NVMe. It holds the Hugging Face cache, the currently active model weights, in-flight checkpoints, and scratch. It is fast (an NVMe Gen4/Gen5 SSD delivers multiple GB/s of sequential read; the exact figure is a machine-log entry, measured on the baseline machine; record value, date, driver) and it is treated as disposable. Anything on the working tier should be reconstructible from a revision hash and a lockfile. If the NVMe died tomorrow, I should lose time, not results.

The archive tier is the 5 TB NAS. It holds things I want to keep: final checkpoints I have decided are worth keeping, acceptance reports, MLflow artifact stores I have exported, and datasets I have curated. The NAS is slower over the network than local NVMe, and that is fine, because I touch it deliberately and rarely. Chapter 0.5 makes this split concrete with HF_HOME, hf_transfer, and a written retention policy; here the point is only that "where does this file live" is a decision I make on purpose, tier by tier, and not something I let a tool decide for me by defaulting into ~/.cache.

The open-data boundary

There is one more line I draw before any code, and it is as much about discipline as architecture. Everything in this book runs on open weights and open data. Open-weight models (Qwen, Llama, and friends, pinned to specific revisions), open evaluation datasets, open benchmarks. Nothing in the loop depends on a private dataset, a proprietary API, or anything I could not hand to a thesis committee and say "here, reproduce it."

That boundary exists for two reasons. The obvious one is reproducibility: a thesis result that leans on data nobody else can see is not a result, it is an anecdote. The less obvious one is contamination of purpose. I also do consulting and product work, and that work has its own private data, its own models, its own infrastructure. If I let that stack bleed into this one, two bad things happen. First, I can no longer publish cleanly, because I cannot separate the open result from the private input. Second, I create a compliance and confidentiality hazard where private client data could end up in a public checkpoint or a Substack post. So the business stack stays on separate machines, separate accounts, separate storage, with no shared caches and no shared credentials. This book never touches it, and I will not gesture at it again except to say: if you are reading this to build your own loop, keep your paid work and your published work in different houses.

Why reasoning models, and why open weights

One more design input deserves a paragraph, because it shapes what the loop is even for. The subject of this book is reasoning models: models trained or coaxed to produce an explicit chain of intermediate steps before committing to an answer. That choice is not incidental to the 16 GB constraint, it is synergistic with it. Reasoning happens at inference time, in tokens, which means a modestly sized model that thinks for longer can close a surprising amount of the gap to a much larger model that answers immediately. On a card that cannot hold a 70B model, "spend more tokens, not more parameters" is exactly the trade the hardware wants me to make, and it is the trade the loop is built to measure and then reinforce.

Open weights are the other half. I need models whose weights I can download, quantize, serve, and fine-tune locally, with no API gatekeeper between me and the tensors. Only open-weight families (Qwen, Llama, and their kin) let me do the full serve-and-train loop on my own GPU, and only open weights let me pin a revision hash and hand a reviewer something they can reproduce byte for byte (Chapter 0.5). A closed model behind an API can be evaluated but never trained here, and never fully inspected, which would make half the loop a black box and violate the book's founding promise. So the loop runs on open reasoning models specifically because that is the intersection where I can both measure a model honestly and change it.

Tooling

This chapter has no command surface of its own, but it is worth naming the tools each box in Figure 0.1 will become, so the later chapters land in a frame that already exists in the reader's head.

The serve box is vLLM, run as an OpenAI-compatible server on localhost. I chose it because it gives me paged KV-cache memory management and continuous batching, which are exactly the two things that let a 16 GB card serve a real model at usable throughput. Part 2 opens the box.

The evaluate box is Inspect (the UK AI Safety Institute's evaluation framework). Evaluations are defined as tasks; the model under test is just an endpoint. Because the model is an endpoint, the same eval runs against the served base model and the trained model with no change, which is what makes the before/after comparison fair. Part 3 opens this box.

The score box is Inspect scorers plus, later, my own reward functions. A scorer turns a model's answer into a number: exact-match, a parsed numeric answer, a judged rubric. In Part 7 those same scorers become the reward signal for training, which is the pun in the book's title: evals as rewards. The eval and the reward are the same function, used twice.

The train box is Unsloth running GRPO (and LoRA/QLoRA for the parameter-efficient plumbing). Unsloth is what makes reinforcement-style training of a reasoning model fit on 16 GB at all. Parts 5, 6, and 7 open this box, math first.

The log box is MLflow on localhost, and it is the least glamorous and most important tool in the stack. Every run in this book logs its git SHA, its uv lockfile hash, the driver version, the VRAM peak, and its seeds, so that a number in a figure can always be traced back to the exact state of the world that produced it. Chapter 0.6 stands it up.

A note on why the boxes are tools and not one monolithic script. It would be possible to write a single Python program that loads a model, evaluates it, computes rewards, trains, and re-evaluates, all in one process. I deliberately do not, and the reason connects back to the 16 GB constraint and to the endpoint seam. A monolith would have to hold serving state and training state in the same process and the same environment, which the memory budget forbids and the dependency split (Chapter 0.4) makes painful. More importantly, a monolith hides its own seams: when a number moves, I want to know which box moved it, and a pipeline of separable tools with a logged interface between each stage lets me answer that. So each box is its own tool with its own environment, talking to the next through a defined interface (an HTTP endpoint, a scored dataset, a checkpoint on disk), and the whole thing is coordinated by small scripts rather than fused into one process. The loop is modular because modularity is what lets me attribute a change to a cause.

Wrapping all of it is uv for Python environments and Docker for the few services (like MLflow) that are cleaner as containers. The doctrine that serve and train live in two separate uv projects (Chapter 0.4) is a direct consequence of the 16 GB constraint: two environments means I can lock two mutually incompatible dependency trees (vLLM's and Unsloth's) without either poisoning the other, and it means I never accidentally hold both heavyweights resident at once.

Lab

There is no runnable lab in this chapter. The deliverable is the annotated stack diagram below, which I keep as the reference picture for the whole book. Everything I build later is an implementation of one of these layers, and when I get lost in a later chapter, this is the diagram I come back to.

flowchart TB
    subgraph HW["Hardware - MSI Aegis R2 (A2NVV9-2218US)"]
        GPU["RTX 5080 · 16 GB VRAM · Blackwell · ~960 GB/s"]
        CPU["Core Ultra 9 285 · 32 GB DDR5"]
        NVME["1 TB NVMe - working tier"]
        NAS["5 TB NAS - archive tier"]
    end

    subgraph OS["OS & drivers - Ubuntu 24.04 LTS"]
        KERN["HWE kernel"]
        DRV["NVIDIA 570-open"]
        CT["Docker + nvidia-container-toolkit"]
    end

    subgraph ENV["Environments - uv, two projects"]
        SERVEENV["serve/ - vLLM (locked)"]
        TRAINENV["train/ - Unsloth (locked)"]
    end

    subgraph LOOP["The loop"]
        SERVE["Serve - vLLM OpenAI endpoint"]
        EVAL["Evaluate - Inspect tasks"]
        SCORE["Score - scorers → reward"]
        TRAIN["Train - GRPO / LoRA on 16 GB"]
        REEVAL["Re-evaluate - reasoning delta"]
    end

    subgraph SPINE["Tracking spine"]
        MLF["MLflow - runs, metrics, artifacts"]
    end

    HW --> OS --> ENV
    SERVEENV --> SERVE
    TRAINENV --> TRAIN
    SERVE --> EVAL --> SCORE --> TRAIN --> REEVAL
    REEVAL -.->|reasoning delta| SERVE
    LOOP --> SPINE
    NVME -. "cache, checkpoints, scratch" .- ENV
    NAS -. "final checkpoints, reports" .- SPINE

    classDef boundary fill:#f6f6f6,stroke:#999,stroke-dasharray:4 3;

Figure 0.2 - the annotated stack. Read it top to bottom: hardware constrains the OS and drivers, which host two locked uv environments, which run the loop, which logs into MLflow. The 16 GB card at the top is the reason every layer below it looks the way it does.

What you should see. Not on a screen (there is nothing to run yet), but in your head: a single card at the top with a hard 16 GB ceiling, and a chain of decisions falling out of it. Weights must be quantized or the model will not load. Serving and training must be separated in time, because they cannot both hold their heavyweight state at once. Storage must be tiered, because weights are big and results are precious. And a boundary must be drawn around open data, because a result you cannot reproduce is not a result. If those four consequences feel inevitable rather than arbitrary by the end of this diagram, the chapter did its job, and the rest of the book is just me making each of them concrete, one command at a time.

Firmware first: BIOS and the Windows send-off

Goal. Get firmware current and validate the hardware while the shipped Windows install is still around to help, then wipe with confidence.

Covers. The MSI BIOS update path; why a prebuilt machine ships with stale firmware; the license-in-firmware quirk that makes the Windows send-off painless; and the preflight hardware validation I run before I erase the only OS the vendor tested.

Theory

Why firmware is the first thing, not an afterthought

It is tempting to treat the BIOS as something you only touch when something is broken. On a brand-new prebuilt machine that instinct is exactly backwards, and the reason is a supply-chain timing gap. A system like the MSI Aegis R2 (A2NVV9-2218US) is assembled, imaged, boxed, warehoused, and shipped. The firmware flashed onto the board happened at the earliest point in that chain, often months before the box reaches me. Firmware, meanwhile, is where a lot of the platform's hardware-enablement work lives: memory training for the DDR5 kit, PCIe link stability, CPU microcode, fan and thermal curves, and compatibility fixes for exactly the kind of new silicon this machine has.

The baseline machine is a genuinely new platform on two fronts at once. The CPU is an Intel Core Ultra 9 285 (Arrow Lake), and the GPU is an RTX 5080 (Blackwell). New CPU platforms in particular tend to accumulate a rapid series of early BIOS releases that fix memory compatibility, boot reliability, and microcode issues discovered after launch. The firmware that shipped on my board predates most of those fixes by construction. So updating the BIOS before I do anything else is not superstition; it is me pulling in months of platform fixes that the factory image could not have included.

License-in-firmware: why the Windows send-off is clean

There is a small, delightful fact about modern OEM Windows that makes wiping the shipped install stress-free: the Windows license is not a sticker anymore, and it is not a file on the disk. Since the OEM Activation 3.0 scheme, the Windows product key for a prebuilt machine is embedded in the firmware, in an ACPI table called MSDM (Microsoft Digital Marker). The key lives in the BIOS. That means I can erase the entire drive, install Linux, and the Windows entitlement is still sitting safely in firmware. If I ever need Windows back on this hardware (for a firmware tool that only ships as a Windows binary, say), a fresh Windows installer reads the MSDM key from the board and activates automatically, no key entry, no phone call.

You can see the key before you wipe, which is a nice way to confirm the machine is genuinely licensed and that the send-off will be clean. From the shipped Windows, an elevated PowerShell prompt reports it:

(Get-WmiObject -Query 'SELECT * FROM SoftwareLicensingService').OA3xOriginalProductKey

If that returns a key, the MSDM table is populated and the entitlement will survive the wipe. Later, from Linux, you can confirm the same table is present without any Windows at all:

# The MSDM ACPI table is what holds the OEM Windows key in firmware.
sudo test -f /sys/firmware/acpi/tables/MSDM && echo "MSDM present (Windows entitlement lives in firmware)" || echo "no MSDM table"

The practical upshot: I do not need to hoard a product key, and I do not need to feel precious about the shipped OS. The license is in the board. I keep the shipped Windows just long enough to do the firmware update and the burn-in, then I erase it without a second thought.

Why validate hardware before wiping

Here is the reasoning that gives this chapter its shape. The shipped Windows install is the only software configuration the vendor actually tested against this exact hardware. Its drivers are matched, its diagnostic tools run natively, and if a component is dead on arrival, the vendor's own utilities will say so in language their support desk understands. The moment I wipe it, I take on the burden of proving that every hardware fault is not my fault. So I do my destructive-testing and my burn-in while the vendor's known-good environment is still on the disk. If the memory is bad, the GPU is unstable, or a fan does not spin, I want to find out now, with a clean RMA story ("your shipped OS, your diagnostic tool, your fault"), not three weeks into a Linux setup where support will reasonably ask whether I broke it myself.

There is also a firmware-update safety argument for keeping Windows briefly. MSI's most convenient BIOS-flashing tools historically run as Windows applications, and while the board also supports flashing from a USB stick via the BIOS itself (M-Flash), doing the first update from the vendor's own environment removes a variable. So the order is deliberate: update firmware from the tested OS, burn in on the tested OS, capture the preflight results, and only then erase.

What a burn-in is actually looking for

It helps to be precise about what sustained-load testing catches, because "stress test" is a vague phrase that hides three distinct failure modes. The first is infant mortality: components that are marginal from the factory tend to fail early, under their first real thermal cycling, rather than randomly across their lifetime. A half-hour of full load is a compressed version of that early stress, so if a solder joint, a fan bearing, or a VRM is going to die young, I would much rather provoke it now, inside the warranty window and on the vendor's OS, than three weeks into a training run. The second is thermal design error: a heatsink that was not seated flat, a fan curve that ramps too late, or case airflow that recirculates hot air will not show up in a five-second check but will show up as a temperature that keeps climbing toward the throttle point under sustained load. The third is power delivery instability: a PSU that sags under combined CPU-plus-GPU load, which manifests as a crash or reset only when both are hammered at once, which is why I run a brief combined-load pass and not just each component alone.

None of these three is a correctness problem, which is the important distinction from the memory test. MemTest86 is looking for silent wrongness (a bit that flips and corrupts data with no crash). The burn-in is looking for instability and heat (a machine that crashes, throttles, or degrades under load). A thorough preflight needs both, because they fail differently and a machine can pass one while failing the other. A GPU with perfect memory can still thermal-throttle a long training run into uselessness, and a perfectly cool machine can still have one bad DRAM cell that quietly poisons a checkpoint.

Tooling

Three tools carry this chapter: MSI's firmware update path, a memory tester, and a GPU/CPU stress harness. I will name the actual surfaces here and drive them in the Lab.

MSI BIOS updates come from the product's own support page. Every MSI board reports its exact model and current BIOS version, and the support page lists BIOS releases with dates and per-release changelogs. There are two ways to apply one. M-Flash reboots into the BIOS and flashes from a FAT32 USB stick holding the BIOS file, and it is the method I trust for the actual flash because it does not depend on a running OS. MSI Center (Windows) can also fetch and apply updates, which is convenient for discovering that an update exists. My rule: discover the version any way, flash with M-Flash.

You can read the current firmware version from either OS. From Linux:

sudo dmidecode -s bios-version
sudo dmidecode -s bios-release-date
sudo dmidecode -s baseboard-product-name
sudo dmidecode -s system-product-name   # should report the Aegis R2 SKU

MemTest86 is the memory tester. Bad or mistrained DDR5 is the single most insidious fault for this kind of work, because it produces silent corruption, not clean crashes: a flipped bit in a weight tensor or a training batch just makes your numbers subtly wrong. MemTest86 boots from its own USB stick, outside any OS, and walks the full address space with a battery of patterns. It is the one test I never skip.

Stress and burn-in is about heat and sustained load, not correctness. A prebuilt machine can pass a five-second check and then throttle or crash under twenty minutes of real load because a heatsink is poorly seated or a fan curve is wrong. On the shipped Windows I use FurMark for the GPU and Prime95 (or MSI's built-in stress tooling) for the CPU, watching temperatures and clocks. On Linux later I will re-run an equivalent sustained-load test as part of acceptance (Chapter 0.7); here the goal is only to shake out infant mortality while the warranty environment is intact.

Lab

The artifact for this chapter is a filled-in preflight checklist committed alongside the machine log. It is the document I would show a warranty desk, and later the first page of the machine's provenance. I run the whole sequence from the shipped Windows (for the flash and the stress tests) plus a couple of confirmations from a Linux live USB, and I record every result with a value, a date, and the firmware version in effect.

Step 1 - record the as-shipped state

Before changing anything, capture what arrived. Boot a Linux live USB (or use the shipped Windows) and read the identifiers:

#!/usr/bin/env bash
# Capture the as-shipped hardware and firmware identity.
# Run from a Linux live USB or the shipped OS. Output is the top of the preflight log.
set -euo pipefail

OUT="preflight-$(date +%Y%m%d).log"
{
  echo "# Preflight record - $(date --iso-8601=seconds)"
  echo
  echo "## Firmware"
  echo "BIOS version : $(sudo dmidecode -s bios-version)"
  echo "BIOS date    : $(sudo dmidecode -s bios-release-date)"
  echo "Board        : $(sudo dmidecode -s baseboard-product-name)"
  echo "System       : $(sudo dmidecode -s system-product-name)"
  echo
  echo "## CPU"
  lscpu | grep -E 'Model name|Socket|Core\(s\)|Thread'
  echo
  echo "## Memory"
  sudo dmidecode -t memory | grep -E 'Size|Speed|Manufacturer|Part Number' | grep -v 'No Module'
  echo
  echo "## GPU (PCI view; driver not required)"
  lspci | grep -Ei 'vga|3d|display'
  echo
  echo "## Windows entitlement"
  if sudo test -f /sys/firmware/acpi/tables/MSDM; then
    echo "MSDM present - Windows license lives in firmware, safe to wipe"
  else
    echo "MSDM absent - investigate before wiping"
  fi
} | tee "$OUT"
echo "Wrote $OUT"

Step 2 - update the BIOS

Find the current BIOS version (Step 1) and compare it to the latest release on the MSI product support page for this exact SKU. If the board is behind, download the latest BIOS, unzip it to a freshly FAT32-formatted USB stick, and flash with M-Flash:

1. Copy the unzipped BIOS file to the root of a FAT32 USB stick.
2. Reboot; press Del to enter BIOS setup.
3. Select M-Flash; confirm the reboot into the flash utility.
4. Choose the BIOS file on the USB stick; confirm.
5. Wait. The board writes flash and reboots itself, possibly twice.
   Do NOT power off, do NOT press keys until it returns to POST.
6. Re-enter BIOS; confirm the new version string; load optimized defaults.

Record the before and after versions in the log. This is the single most important line in the whole preflight, because it timestamps the platform state that every later measurement in the book assumes.

Step 3 - memory test

Boot MemTest86 from its own USB stick and run at least one full pass (all tests). A full pass on 32 GB of DDR5 takes a while (the exact wall-clock is a machine-log entry; measured on the baseline machine; record value, date, firmware). Zero errors is the only acceptable result. A single error means RMA the memory now, before it silently corrupts a checkpoint later.

Step 4 - sustained-load burn-in

From the shipped Windows, run the GPU and CPU under sustained load and watch temperatures and clocks. The purpose is to catch throttling and infant mortality, not to benchmark:

GPU  : FurMark, stress preset, ~20-30 min. Watch for artifacts,
       driver resets, or thermal throttling. Log peak GPU temp (°C)
       and whether clocks held.
CPU  : Prime95 (Small FFTs) or MSI's stress tool, ~20-30 min.
       Log peak package temp (°C) and whether any core threw errors.
Both : run together briefly to check the PSU and case airflow under
       combined load.

All temperatures and clocks here are machine-log entries; measured on the baseline machine; record value, date, firmware. The pass condition is qualitative and strict: no crashes, no artifacts, no driver resets, and temperatures that plateau below throttling rather than climbing without bound.

Step 5 - fill and commit the checklist

The artifact is this filled-in checklist, saved as preflight-<date>.md next to the log from Step 1 and committed to the machine-log repository (the same repo Chapter 0.3 uses for the nvidia-smi capture):

# Preflight checklist - MSI Aegis R2 (A2NVV9-2218US)

- Date: ____________________    Operator: ____________________
- BIOS before: __________  →  BIOS after: __________  (release date: ______)
- CPU reported: Intel Core Ultra 9 285              [ ] confirmed
- GPU reported: NVIDIA RTX 5080 16GB                [ ] confirmed
- Memory: 32 GB DDR5, ____ MT/s, all slots seen     [ ] confirmed
- NVMe: 1 TB present, SMART healthy                 [ ] confirmed
- MSDM present (Windows license in firmware)        [ ] confirmed
- MemTest86: ____ full pass(es), errors: ____        [ ] PASS (0 errors)
- GPU burn-in: ____ min, peak ____ °C, clocks held?  [ ] PASS
- CPU burn-in: ____ min, peak ____ °C, errors: ____  [ ] PASS
- Combined load: stable? ____                        [ ] PASS
- Decision: [ ] proceed to wipe   [ ] RMA - reason: ______________

What you should see. By the end of this sequence you should be holding a single committed document that says, in effect: the firmware is current as of a named release and date, the memory passed a full MemTest86 sweep with zero errors, and the GPU and CPU held stable temperatures under half an hour of sustained load each. Every blank is filled with a real value and a real date. The Windows entitlement is confirmed to live in firmware, so the shipped OS is now expendable. That document is your license to wipe: from here on, every problem you hit is a software problem on a hardware base you have personally proven, and the next chapter can erase Windows without a shred of anxiety.

Ubuntu 24.04 on Blackwell

Goal. Install Ubuntu 24.04 with a kernel and driver that actually drive Blackwell and Arrow Lake, and prove the GPU link is what the spec promises.

Covers. Point-release and HWE kernels, and why a recent kernel matters for Arrow Lake and the Wi-Fi module; Secure Boot and MOK enrollment; nomodeset recovery; the NVIDIA 570-open driver and why Blackwell requires the open kernel modules; and verifying the PCIe Gen5 x16 link.

Theory

Why the kernel version is the first Linux decision

On older, well-worn hardware you can install more or less any Ubuntu and it just works, because the kernel that ships in the installer already contains drivers for silicon that has been in the wild for years. The baseline machine is the opposite case. Both the CPU (Core Ultra 9 285, Arrow Lake) and the platform around it (the Wi-Fi module, the chipset, the integrated graphics on the CPU) are new enough that hardware enablement for them landed in the Linux kernel only recently. Install too old a kernel and the symptoms are exactly the annoying ones: no Wi-Fi, flaky suspend, missing sensors, an integrated GPU the kernel does not recognize.

Ubuntu solves this with two mechanisms, and it is worth understanding both because they interact. The first is point releases: 24.04 gets periodic refreshes (24.04.1, 24.04.2, and so on), and each point release rolls in a newer kernel and updated installer media. Installing from the latest point-release ISO, rather than the original 24.04.0 media, is the single cheapest way to start on a kernel that already knows about your hardware. The second is the HWE (Hardware Enablement) stack: an LTS release can track a newer "hardware enablement" kernel backported from a later Ubuntu, so that a 24.04 LTS machine can run a kernel several versions ahead of the one it launched with while keeping the LTS userland. For new silicon, HWE is the difference between "supported eventually" and "supported now."

The reasoning that follows from this: I install from the newest point-release media I can get, I make sure the HWE kernel meta-package is installed so that future kernel updates track hardware enablement, and I keep a wired fallback for the first boot. A recent kernel is not a nice-to-have on this machine; it is the precondition for the machine being fully usable at all.

Secure Boot, MOK, and signed kernel modules

Modern Ubuntu ships with Secure Boot enabled, and Secure Boot's whole job is to refuse to load kernel code that is not signed by a key the firmware trusts. That is a good default, but it collides with the fact that the NVIDIA driver is an out-of-tree kernel module that has to be compiled and inserted on your machine. A freshly built NVIDIA module is not signed by anything the firmware knows about, so with Secure Boot on, the kernel will refuse to load it, and you get a system that boots to a black screen or falls back to software rendering with no NVIDIA driver in sight.

Ubuntu's answer is MOK, the Machine Owner Key. When the driver is installed via DKMS, Ubuntu generates a local signing key, signs the module with it, and asks you to enroll that key into the firmware's trust store on the next reboot, through a blue MOK-management screen that appears before the OS loads. You confirm enrollment with a password you set during install, and from then on the firmware trusts modules signed by your local key. It is a one-time dance per key, and skipping it (or clicking through it wrong) is the most common reason a correctly installed NVIDIA driver still will not load.

nomodeset: the recovery hatch

There is a failure mode worth pre-loading into your head before it happens, because it happens at the worst time: you finish the install, reboot, and get a black screen or an unusable low-resolution desktop, because the graphics stack tried to mode-set on a GPU whose driver is not ready. The escape hatch is a kernel boot parameter, nomodeset, which tells the kernel not to hand graphics mode-setting to the GPU driver at boot, falling back to a generic framebuffer. That gets you a working (if ugly) desktop or console from which you can fix the real driver situation.

You reach it from the GRUB menu at boot: highlight the Ubuntu entry, press e, find the linux line, and append nomodeset before booting. It is temporary (it does not survive a reboot unless you edit the GRUB config), which is exactly what you want: use it to get in, fix the driver, and reboot back to a normal graphical boot. On a Blackwell card especially, where the correct driver is non-negotiable (see below), nomodeset is the difference between "recoverable" and "reinstall from USB."

Why Blackwell requires the open kernel modules

Here is the piece that is specific to this generation and trips people up. NVIDIA ships its Linux driver in two kernel-module flavors: the long-standing proprietary modules, and the newer open modules (often written "open" or "-open", e.g. nvidia-driver-570-open). "Open" here refers to the kernel-module source being open; the userspace CUDA stack is unchanged. For older GPUs you could choose either. For Blackwell you cannot: the RTX 5080 and its GB20x-generation siblings are supported only by the open kernel modules. The proprietary modules do not support these GPUs at all. So on the baseline machine, "NVIDIA 570-open" is not a preference, it is the only thing that will drive the card.

The reason is architectural. NVIDIA moved the driver's hardware-management logic into a GPU System Processor (GSP) firmware blob that runs on the card itself, and the open kernel modules are built around that GSP-firmware model. Newer architectures (Turing onward, and mandatorily on the latest generations) are designed for this split, which is why the open modules are the supported path for them and the sole path for Blackwell. Practically: I install nvidia-driver-570-open (the 570 branch is the one that carries Blackwell support), and if I ever see advice to install the plain nvidia-driver-570 proprietary variant on this card, I ignore it, because it will not load.

Verifying the Gen5 x16 link

The last piece of theory is a habit: never assume the GPU negotiated the link you paid for. The RTX 5080 sits in a PCIe Gen5 x16 slot, and a Gen5 x16 link is a lot of bandwidth. But a reseated card, a slot that shares lanes with an NVMe drive, a BIOS setting, or a marginal riser can silently negotiate down to x8, or to Gen4, and nothing will crash. The only symptom is throughput that is quietly lower than it should be, which is precisely the kind of bug that poisons benchmarks. So verifying the link speed and width after install is a standing acceptance step, and I capture the result into the machine log the same way I capture nvidia-smi.

Why apt and DKMS instead of NVIDIA's .run installer

There are two ways to put the NVIDIA driver on a Linux box, and the choice matters more on this stack than it looks. NVIDIA ships a self-contained .run installer that compiles and installs the driver directly, and Ubuntu ships the same driver as apt packages that build the kernel module through DKMS (Dynamic Kernel Module Support). I use the apt/DKMS path, and the reason is the interaction between three moving parts: the HWE kernel (which updates), Secure Boot (which demands signed modules), and the driver (which is an out-of-tree module that must be rebuilt for each kernel).

DKMS automates exactly that interaction. When the HWE kernel updates (and on this machine it will, because hardware enablement is still landing), DKMS automatically rebuilds the NVIDIA module against the new kernel and re-signs it with the enrolled MOK key, so the driver keeps working across kernel upgrades with no manual intervention. The .run installer does none of that: it builds the module once against the current kernel, and the next kernel update leaves you with a driver that will not load, at which point you are re-running the installer from a text console. On a Secure Boot machine the .run path is worse still, because you have to arrange module signing yourself. So the packaged path is not just more convenient, it is the one that composes correctly with the two things this chapter already committed to: a kernel that moves and a Secure Boot chain that must stay signed.

Tooling

The tools here are the Ubuntu installer, apt plus ubuntu-drivers, mokutil, and nvidia-smi and lspci for verification.

The installer is Ubuntu Desktop 24.04, latest point release. I use the normal graphical installer, choose to erase the disk (the shipped Windows is already validated and expendable, per Chapter 0.2), and set a Secure Boot / MOK password when prompted so third-party drivers can be signed and enrolled.

Driver installation goes through Ubuntu's packaging rather than NVIDIA's .run installer, because apt + DKMS handle the MOK signing and the kernel-rebuild-on-update automatically. ubuntu-drivers devices lists the recommended driver; I install the -open package explicitly rather than trusting autoinstall to pick the flavor.

mokutil inspects and manages Secure Boot and MOK state. nvidia-smi is the ground truth that the driver is actually driving the card. lspci -vv reports the negotiated PCIe link speed and width.

Lab

The full sequence is: install Ubuntu from current media, get onto the HWE kernel, install the 570-open driver, enroll MOK, reboot, then capture nvidia-smi and the PCIe link speed and commit both to a machine-log repository. The artifact is that committed capture.

Step 1 - install and reach the HWE kernel

Install Ubuntu Desktop 24.04 (latest point release), erasing the disk. On first boot, get online (wired or a USB dongle if built-in Wi-Fi is not yet happy), then update and ensure the HWE kernel stack is installed:

#!/usr/bin/env bash
# Get fully updated and onto the HWE (hardware-enablement) kernel line.
set -euo pipefail

sudo apt update && sudo apt full-upgrade -y

# HWE kernel meta-package for 24.04 (the -generic-hwe metapackage already
# pulls the matching linux-image, so I do not name it separately).
sudo apt install -y \
  linux-generic-hwe-24.04 \
  build-essential dkms mokutil pciutils

echo "Reboot, then confirm the running kernel:"
echo "  uname -r    # expect the HWE kernel, not the base GA kernel"

Reboot and confirm with uname -r that you are on the HWE kernel line.

Step 2 - install the 570-open driver

Check what Ubuntu recommends, then install the open flavor explicitly:

#!/usr/bin/env bash
# Install the NVIDIA 570 OPEN kernel modules - mandatory for Blackwell (RTX 5080).
set -euo pipefail

# See what Ubuntu detects/recommends (informational).
ubuntu-drivers devices || true

# Install the OPEN 570 driver explicitly. Do NOT install the non-open variant on Blackwell.
# The 570 branch ships in 24.04.2+; on older point releases it may be absent from
# the archive, in which case add the graphics-drivers PPA first:
#   sudo add-apt-repository ppa:graphics-drivers/ppa && sudo apt update
sudo apt update
sudo apt install -y nvidia-driver-570-open

echo
echo "During install, DKMS signs the module with a local MOK key."
echo "On the NEXT reboot you MUST complete MOK enrollment:"
echo "  - a blue 'Perform MOK management' screen appears BEFORE the OS loads"
echo "  - choose 'Enroll MOK' -> Continue -> Yes -> enter the password you set"
echo "  - the machine reboots and the module can now load under Secure Boot"

Step 3 - enroll MOK and reboot

Reboot. Before the OS loads, the blue MOK-management screen appears. Choose Enroll MOK, Continue, Yes, and enter the password. The machine reboots into a normal graphical session with the NVIDIA driver loaded. Confirm Secure Boot and enrollment state:

mokutil --sb-state          # expect: SecureBoot enabled
mokutil --list-enrolled | grep -i -A2 'Signature key' || true

If the driver still will not load, this is where nomodeset earns its keep: reboot, edit the GRUB linux line to append nomodeset, boot to a working desktop, and re-check the driver install and MOK state from there.

Step 4 - capture the proof and commit it

Now produce the artifact. This script captures the driver state and the negotiated PCIe link and writes a dated log for the machine-log repository:

#!/usr/bin/env bash
# Capture nvidia-smi + PCIe link speed/width as the machine-log artifact for this install.
set -euo pipefail

STAMP="$(date --iso-8601=seconds)"
OUT="gpu-state-$(date +%Y%m%d).log"

# Find the NVIDIA GPU's PCI address for the link query (select by NVIDIA vendor id 10de).
GPU_ADDR="$(lspci -d 10de: | grep -Ei 'vga|3d' | head -n1 | awk '{print $1}')"

{
  echo "# GPU state capture - $STAMP"
  echo "kernel        : $(uname -r)"
  echo "secure boot   : $(mokutil --sb-state 2>/dev/null || echo 'unknown')"
  echo
  echo "## nvidia-smi"
  nvidia-smi
  echo
  echo "## Driver / CUDA versions"
  nvidia-smi --query-gpu=driver_version,name,memory.total --format=csv
  echo
  echo "## PCIe link (negotiated) for $GPU_ADDR"
  # LnkSta is the negotiated state; expect Speed 32GT/s (Gen5), Width x16.
  sudo lspci -vv -s "$GPU_ADDR" | grep -E 'LnkCap|LnkSta'
  echo
  echo "## PCIe link (as nvidia-smi sees it)"
  nvidia-smi --query-gpu=pcie.link.gen.current,pcie.link.width.current --format=csv
} | tee "$OUT"

echo "Wrote $OUT - commit this to the machine-log repo."

Commit the result to a small, dedicated machine-log git repository (the same one that holds the Chapter 0.2 preflight checklist). That repo becomes the machine's permanent provenance record: firmware version, install date, driver version, and link speed, all timestamped.

cd ~/machine-log
cp ~/gpu-state-*.log ./
git add gpu-state-*.log preflight-*.md
git commit -m "Install capture: 570-open driver + PCIe link on Aegis R2"

What you should see. The nvidia-smi table should print with the card named as an RTX 5080, 16 GB (16384 MiB) total memory, and a driver version on the 570 branch. Secure Boot should report enabled, and the driver should still load, which together prove MOK enrollment worked. The PCIe capture should show the negotiated link at Gen5 (32 GT/s) and width x16 in both the lspci LnkSta line and the nvidia-smi query; if either reads x8 or Gen4, stop and investigate the slot, the BIOS, and the card seating before trusting any throughput number later. The concrete values (exact driver build, temperatures, clocks) are machine-log entries; measured on the baseline machine; record value, date, driver. The committed gpu-state-*.log is the artifact: from here on, any performance claim in the book can be traced to a machine whose driver and link speed are on the record.

uv and the two-environment doctrine

Goal. Establish one canonical way to manage Python, and separate the serving stack from the training stack so their dependencies never fight.

Covers. uv as a toolchain (Python pinning, virtual environments, lockfiles, uv run); why serve and train stay in separate projects; and a repository layout of two uv projects, serve/ and train/, each with a committed uv.lock.

Theory

One tool for the whole Python surface

Python environment management has historically been a pile of tools that each do one slice: pyenv to pick an interpreter, venv to make an environment, pip to install, pip-tools or poetry to lock. Every seam between them is a place for "works on my machine" to hide. uv collapses that whole pile into one fast tool. It installs and pins Python interpreters, creates virtual environments, resolves and installs dependencies, writes a cross-platform lockfile, and runs commands inside the environment, all from one binary. For a book whose entire premise is that nothing should be a black box, having one tool with one mental model for "what Python, which packages, exactly which versions" is worth a lot.

The reason I standardize on uv, and forbid myself bare pip and python -m venv for the rest of the book, is reproducibility with teeth. A requirements.txt full of loose version ranges tells you what you asked for, not what you got. uv's lockfile records the exact resolved version and hash of every package in the dependency graph, transitively, so that uv sync on another day rebuilds the identical environment or fails loudly if it cannot. When a reasoning-delta number in Part 7 looks surprising, the first question is always "did the environment change," and a committed uv.lock answers it definitively instead of with a shrug.

Why serve and train are two environments, not one

Here is the doctrine, and it falls directly out of the 16 GB constraint and out of dependency reality. The serving stack (vLLM) and the training stack (Unsloth, and the PyTorch/transformers/TRL/bitsandbytes constellation around it) are two large, fast-moving, deeply-pinned dependency trees. They both pin specific versions of PyTorch, CUDA-linked libraries, transformers, and low-level kernels, and those pins do not agree. Trying to satisfy both in one environment forces the resolver into a compromise that makes at least one of them unhappy, and the failure is often not a clean resolver error but a subtle runtime breakage: a kernel compiled against the wrong CUDA minor, an attention backend that silently falls back to a slow path, a transformers version that one side needs and the other rejects.

So I keep them apart, and the separation buys three distinct things. First, clean resolution: each project locks its own coherent dependency graph with no cross-contamination, so vLLM gets exactly the stack it wants and Unsloth gets exactly the stack it wants. Second, temporal separation on the GPU: the loop already serves and trains at different times (Chapter 0.1), never both at once on 16 GB, so there is no reason to co-resident their software either; two environments make "one heavy thing at a time" the path of least resistance. Third, independent evolution: I can bump vLLM to chase a serving improvement without touching the trainer, and re-lock only serve/, so a serving upgrade cannot silently perturb a training result. Each environment moves on its own clock, and each uv.lock timestamps that clock.

Why not conda, poetry, or bare pip

It is fair to ask why uv specifically, given that conda, poetry, pipenv, and plain pip-plus-venv all exist and all "work." The honest answer is a combination of speed, correctness, and scope. Conda solves the interpreter-and-native-library problem well but is slow to resolve and pulls packages from its own channels, which drifts from the PyPI reality that vLLM and Unsloth actually publish to. Poetry locks nicely but does not manage the Python interpreter itself and has historically been slow on large graphs. Bare pip does not lock at all in any reproducible way without bolting on pip-tools, and python -m venv leaves you managing interpreters by hand. uv does the whole span (interpreter, environment, resolution, locking, running) in one fast tool, with a lockfile that pins exact versions and hashes, which is precisely the reproducibility property this book needs. So the choice is not fashion; it is that uv is the one tool whose scope matches "pin everything and rebuild it identically," which is the only workflow a measurement-driven thesis can trust.

The corollary is the rule I enforce on myself for the rest of the book: no bare pip install, ever, and no python -m venv, ever. The moment I pip install something into an environment imperatively, that environment stops being a pure function of the committed lockfile and becomes a snowflake that only exists on this machine, at this moment, in a state no one can reproduce. Every dependency enters through uv add so it lands in pyproject.toml and uv.lock, and every command runs through uv run so it uses the locked environment. It is a small discipline that pays off the first time I need to reproduce a three-month-old result and find that uv sync on the old commit just works.

The evaluator is a thin third thing, and that is deliberate

A fair question: if serve and train each get an environment, where does the evaluation harness (Inspect) live? The answer is that the evaluator talks to the model over the OpenAI-compatible HTTP endpoint, so it only needs a light client stack, not the heavy serving or training dependencies. In practice Inspect can live in its own small environment or ride along in whichever project is convenient, because it never imports vLLM or Unsloth; it makes HTTP calls. I will pin down its exact home when Part 3 builds the eval harness. The point for now is that the two heavyweight environments are the ones the doctrine is about, and the evaluator's independence from both is a feature of the endpoint seam, not an afterthought.

The lockfile is a contract, and CI can enforce it

A lockfile is only worth committing if it is treated as a contract rather than a suggestion, and the discipline that makes it a contract is small: I never edit uv.lock by hand, I only ever regenerate it through uv, and I treat an unexpected change to it in a diff as a signal that something moved that I should understand before I commit. The payoff is that the pair (pyproject.toml, uv.lock) becomes a checkable claim. If I check out an old commit to reproduce an old result, uv sync rebuilds the exact environment that produced it, or it fails loudly because a wheel is no longer available, in which case I know immediately that the old result cannot be reproduced with today's package index rather than discovering it silently through a subtly different number.

This is also where the two-environment doctrine pays a second dividend. Because each environment is locked independently, I can verify each one in isolation, and I can bisect a problem to one side. If a serving benchmark regresses, I check whether serve/uv.lock changed since the last good run; the training stack is provably irrelevant because it is a different lockfile that was not touched. Without the split, every dependency change is a change to one giant shared environment, and I lose the ability to say "the serving stack is identical, so the regression is not in the serving dependencies." The separation turns "the environment changed" from an undifferentiated worry into a precise, attributable fact.

```admonish under-the-hood title="Why uv sync is safe to run anywhere, anytime" uv sync is idempotent and total: it makes the .venv exactly match uv.lock, adding what is missing, removing what should not be there, and downgrading what drifted, without consulting pyproject.toml's loose ranges at all. That means running it is never a gamble. It cannot silently upgrade you, because upgrades only happen when you deliberately re-resolve with uv add or uv lock. This is the property that makes the lockfile a contract rather than a starting point: the environment is a pure function of the committed lock, and uv sync is how you evaluate that function. It is why every chapter can open with uv sync and trust exactly what it gets.

## Tooling

The uv command surface I actually use is small and worth memorizing, because it recurs in every later chapter.

- `uv init <dir>` creates a new project: a `pyproject.toml`, a pinned Python, and the scaffolding for a `.venv`.
- `uv python pin <version>` records which Python this project uses, writing `.python-version`.
- `uv add <pkg>` adds a dependency, resolves it, updates `pyproject.toml`, and updates `uv.lock`.
- `uv sync` materializes the `.venv` to exactly match `uv.lock` (installing, removing, or downgrading as needed).
- `uv lock` re-resolves and rewrites `uv.lock` from `pyproject.toml` without changing the running environment.
- `uv run <cmd>` runs a command inside the project's environment, syncing first if needed, so you never manually `source .venv/bin/activate`.

The habit that makes this reproducible: I never install into an environment imperatively and hope. I declare dependencies in `pyproject.toml` (via `uv add`), commit the resulting `uv.lock`, and reconstruct anywhere with `uv sync`. `uv run` is how every script in this book is invoked, so the environment is always the locked one, never whatever happens to be on `PATH`.

```admonish read-along title="Why `uv run`, always"
Activating a virtual environment by hand is a stateful, forgettable act - the classic bug is running a script against the wrong `.venv` because you forgot which one was active. `uv run python train.py` binds the command to *this project's* locked environment every time, with no activation step to forget. Throughout the book, whenever you see a Python invocation, it is `uv run ...` for exactly this reason: the environment is a property of the project directory, not of your shell's mutable state.

Two extra flags matter for this hardware. Some GPU packages publish CUDA-specific wheels on their own package indexes rather than only on PyPI, so a project may declare an extra index (for the CUDA-matched PyTorch build, for instance) in pyproject.toml under a [tool.uv] section. And uv can pin the Python version per project, which matters because vLLM and Unsloth may support different Python version windows; each project pins its own.

Lab

I create both projects, lock vLLM in serve/ and Unsloth in train/, and commit both pyproject.toml + uv.lock pairs. The artifact is those two pairs, printed in full. The exact resolved versions are recorded from the lock on the machine; the pyproject.toml contents below are realistic declarations, and the resulting uv.lock files (large, hash-bearing, machine-generated) are committed but not reprinted here in full.

Step 1 - repository layout

thesis-loop/
├── serve/            # vLLM serving environment
│   ├── pyproject.toml
│   ├── uv.lock       # committed
│   └── .python-version
├── train/            # Unsloth training environment
│   ├── pyproject.toml
│   ├── uv.lock       # committed
│   └── .python-version
└── README.md

Step 2 - create and lock the serve environment

#!/usr/bin/env bash
# Create the serving environment: vLLM, locked.
set -euo pipefail

uv init serve
cd serve

# Pin a Python the serving stack supports (record the exact value on the machine).
uv python pin 3.12

# Add vLLM. uv resolves the full graph (torch, CUDA-linked libs, etc.) and writes uv.lock.
uv add vllm

# Add the light client bits used to poke the endpoint during bring-up.
uv add httpx

# Materialize the environment from the lock and smoke-test the import.
uv sync
uv run python -c "import vllm; print('vllm', vllm.__version__)"

[project]
name = "serve"
version = "0.1.0"
description = "vLLM serving environment for the evals-as-rewards loop"
requires-python = ">=3.12,<3.13"
dependencies = [
    "vllm>=0.8",     # Blackwell-capable serving engine; exact pin lands in uv.lock
    "httpx>=0.27",   # for poking the OpenAI-compatible endpoint during bring-up
]

[tool.uv]
# vLLM pulls a CUDA-matched torch build; if a dedicated index is needed for the
# Blackwell (sm_120) wheels, declare it here so resolution is reproducible.
# index-url / extra-index-url are recorded on the machine once the exact build is chosen.

3.12

The generated serve/uv.lock is a large machine-written file that pins every transitive dependency with an exact version and a content hash (vLLM, torch, the CUDA-linked libraries, and everything they pull). It is committed verbatim. Its exact contents, including the resolved vLLM and torch versions, are recorded from the lock on the machine; record the resolved versions, date, and driver alongside it.

Step 3 - create and lock the train environment

#!/usr/bin/env bash
# Create the training environment: Unsloth (+ the GRPO/LoRA stack), locked.
set -euo pipefail

uv init train
cd train

# Pin a Python the training stack supports (record the exact value on the machine).
uv python pin 3.12

# Add Unsloth and the post-training constellation. uv resolves torch/transformers/TRL/peft/bnb.
uv add unsloth
uv add trl peft bitsandbytes datasets

# Materialize and smoke-test.
uv sync
uv run python -c "import unsloth, trl, peft; print('unsloth stack imported OK')"

[project]
name = "train"
version = "0.1.0"
description = "Unsloth training environment (GRPO/LoRA) for the evals-as-rewards loop"
requires-python = ">=3.12,<3.13"
dependencies = [
    "unsloth>=2025.1",   # 16GB-friendly training; exact pin lands in uv.lock
    "trl>=0.12",         # GRPO trainer lives here
    "peft>=0.13",        # LoRA/QLoRA adapters
    "bitsandbytes>=0.44",# 4-bit quantization for QLoRA
    "datasets>=3.0",     # dataset loading
]

[tool.uv]
# Unsloth pins a CUDA-matched torch; if a dedicated index is needed for the
# Blackwell (sm_120) wheels, declare it here. Recorded on the machine.

3.12

As with serve/, the generated train/uv.lock pins the entire graph (Unsloth, torch, transformers, TRL, peft, bitsandbytes, datasets, and all transitive dependencies) with exact versions and hashes, and is committed verbatim. The resolved versions are recorded from the lock on the machine; record the resolved versions, date, and driver.

Step 4 - commit both pairs

cd thesis-loop
git add serve/pyproject.toml serve/uv.lock serve/.python-version
git add train/pyproject.toml train/uv.lock train/.python-version
git commit -m "Two locked uv projects: serve (vLLM) and train (Unsloth)"

# Record the resolved headline versions into the machine log for quick reference.
echo "serve: $(cd serve && uv run python -c 'import vllm; print(vllm.__version__)')" >> ~/machine-log/env-versions.txt
echo "train: $(cd train && uv run python -c 'import unsloth,trl,peft; print(\"unsloth+trl+peft OK\")')" >> ~/machine-log/env-versions.txt

What you should see. Two sibling directories, each with a pyproject.toml you wrote, a .python-version pinning the interpreter, and a uv.lock that uv generated and that you did not hand-edit. In serve/, uv run python -c "import vllm; print(vllm.__version__)" prints a version and exits cleanly. In train/, the Unsloth stack imports without error. Neither environment mentions the other's heavyweight packages: serve/uv.lock has no Unsloth, train/uv.lock has no vLLM, which is the whole point. The committed lockfiles are the artifact. From now on, any machine that runs uv sync in serve/ gets byte-identical serving dependencies, and the same for train/, so a reasoning-delta result can never be quietly undermined by a dependency that drifted between the serve step and the train step. The exact resolved versions live in the lockfiles and in ~/machine-log/env-versions.txt; record them with the date and driver in effect.

Storage tiers and cache discipline

Goal. Split working storage from archive storage on purpose, and make "which model" mean a revision hash rather than a name.

Covers. The NVMe working tier versus the NAS archive tier; HF_HOME and hf_transfer; a checkpoint retention policy; and what "reproducible" actually means for weights (pin revisions and commit hashes, not names).

Theory

Two tiers, two jobs

Storage on this machine has to do two jobs that pull in opposite directions, and the mistake is trying to make one device do both. The first job is fast working access: when I serve a model or resume a training run, the weights need to stream off disk quickly, and scratch files need to be written and deleted without ceremony. The second job is durable retention: the handful of things I actually want to keep (a final checkpoint, an acceptance report, a curated dataset) need to survive a cache purge, a reinstall, or a dead system drive. Fast working access wants the local NVMe. Durable retention wants somewhere I will not casually rm -rf. So I run two tiers.

The working tier is the 1 TB NVMe. It holds the Hugging Face cache, the currently active model weights, in-flight checkpoints, and scratch. Its defining property is that it is disposable: everything on it should be reconstructible from a revision hash plus a lockfile. If the NVMe died tomorrow I should lose time re-downloading, not lose results. That disposability is a discipline, not just a description: I never put the only copy of something I care about on the working tier.

The archive tier is the 5 TB NAS. It holds things whose loss would actually hurt: final checkpoints I have decided are worth keeping, acceptance reports, exported MLflow artifact stores, and curated datasets. It is reached over the network, so it is slower than local NVMe (the exact throughput is a machine-log entry; measured on the baseline machine; record value, date, driver), and that is fine, because I touch it deliberately and rarely. The NAS is where a result goes to become permanent.

What "reproducible" means for weights

Here is the idea that changes how you think about models: a model name is not an identity. "Qwen2.5-7B-Instruct" names a moving target. The repository behind that name can be updated: a tokenizer fix, a re-uploaded shard, a config tweak. If my thesis says "I evaluated Qwen2.5-7B-Instruct" and someone downloads that name six months later, they may get subtly different weights and get different numbers, and neither of us will know why. The name is a pointer, and pointers move.

The fix is to pin the revision: every Hugging Face model repository is a git repository, so every state of it has a commit hash. When I say "the model," I mean a repository at a specific commit, e.g. Qwen/Qwen2.5-7B-Instruct at revision <40-hex-commit>. That hash is immutable. Downloading that revision gives byte-identical files today, next year, and on someone else's machine. So the reproducibility rule for this book is: pin revisions, not names, everywhere a model or dataset is loaded, and record the pinned hash in the run log alongside the git SHA and the lock hash.

The same discipline for datasets

Everything I just argued about model weights applies with equal force to datasets, and it is easy to forget because a dataset feels more static than a model. It is not. A Hugging Face dataset repository is git-backed exactly like a model repository, so "the GSM8K split I evaluated on" is only well-defined if I pin its revision too. A dataset owner can fix a label, add examples, or re-shuffle a split, and if my eval loaded the dataset by bare name it would silently pull the changed version on the next run, which would move my metric for a reason that has nothing to do with the model. So the pin-the-revision rule covers both nouns: every model load and every dataset load in this book carries an explicit revision=<hash>, and both hashes are recorded in the run log. The reasoning delta is "same model change, same eval data, measured twice," and "same eval data" is only true if the data is pinned as hard as the weights.

There is a subtler contamination angle here too, which Part 3 develops in full but which the storage layer enables. Pinning dataset revisions is also part of how I reason about test-set contamination: if I know the exact commit of the eval set I used, I can check whether that exact data could have leaked into a model's training corpus, and I can hold the eval data fixed while I vary the model. A dataset that drifts under its name makes contamination analysis impossible, because I can no longer say precisely which examples were in play. So pinning is not just reproducibility hygiene; it is a precondition for the contamination arguments the thesis will need to make.

The cache is a shared resource, so control where it lives

Hugging Face libraries default to caching under ~/.cache/huggingface, which on this machine would land on the system drive by default and mixes cached weights in with your home directory. I move the whole cache onto the working tier explicitly via HF_HOME, for three reasons. First, placement: the cache belongs on the fast disk I designated as the working tier, not wherever the default happens to fall. Second, visibility: one environment variable means both the serve and train environments share one cache, so I download a model once and both use it, and I always know where the big files are when the disk fills. Third, portability of the decision: HF_HOME is a single knob I set in a dotfile, so the storage layout is a documented choice rather than an accident of defaults.

hf_transfer is the other half of cache discipline. Model repos are large (tens of GB), and the default download path is not always the fastest. hf_transfer is a Rust-based accelerated download backend for the Hugging Face hub; enabling it (an env var plus the package) parallelizes and speeds up large downloads so the first pull of a model does not dominate my afternoon. The actual speedup is a machine-log entry; measured on the baseline machine; record value, date, driver.

A retention policy, because checkpoints breed

Training produces checkpoints, and checkpoints breed faster than anything else on the working tier. A single training run might save the model state every few hundred steps, and each save of a 7B model is gigabytes. Left ungoverned, a week of experiments buries the NVMe under dozens of intermediate checkpoints, ninety percent of which I will never load again. So the storage policy needs a rule for what to keep, and the rule has to be simple enough that I actually follow it at 11pm when a run finishes.

My rule has three tiers of fate. Intermediate checkpoints live on the working tier only while their run is active, and exist so I can resume if the run crashes. When a run finishes, exactly two checkpoints earn keeping: the one that scored best on the held-out eval (the actual product of the run) and the final one (for continuing training later if I want). Everything else is deleted from the working tier the moment those two are safely copied to the NAS. This keeps the fast disk lean and pushes the small number of genuinely valuable artifacts onto durable storage, which is precisely the division of labor the two tiers exist for. The "best-by-eval" choice is deliberate: it ties the retention decision back to the loop's own measurement, so the checkpoint I keep is the one the evaluation actually preferred, not just the last one that happened to be written.

```admonish gotcha title="Never hand-rm a Hugging Face snapshot directory" The Hugging Face cache stores each model as a tree of content-addressed blobs with per-revision snapshot directories full of symlinks into those blobs. If you rm -rf a snapshot directory to reclaim space, you can leave dangling references or orphan blobs that other revisions still point at, and in the worst case corrupt the cache so a later download half-fails in confusing ways. Reclaim space with huggingface-cli delete-cache (or scan-cache to see what is there first), which understands the blob-sharing and removes revisions safely. The working tier is disposable, but disposing of it wrong still costs you an afternoon.

## Tooling

The tools here are environment variables (`HF_HOME`, `HF_HUB_ENABLE_HF_TRANSFER`), the `huggingface_hub` download API with a pinned `revision`, a NAS mount managed via `/etc/fstab`, and a written retention policy that I actually follow.

**`HF_HOME`** relocates the entire Hugging Face cache (models, datasets, tokens) to a directory I choose on the working tier. **`HF_HUB_ENABLE_HF_TRANSFER=1`** plus the `hf_transfer` package switches on the accelerated download backend. Both are set once in a dotfile so every shell and every `uv run` inherits them.

**The NAS mount** is a network filesystem (NFS or SMB) mounted at a stable path like `/mnt/nas`, declared in `/etc/fstab` so it comes back after a reboot. I mount it read-write for the archive workflow but treat it as append-mostly: things go there to be kept, not churned.

**Pinned downloads** use `huggingface_hub`'s `snapshot_download` (or `hf_hub_download`) with an explicit `revision=<hash>`. That single argument is the whole reproducibility story for weights.

```admonish under-the-hood title="What HF_HOME actually moves"
`HF_HOME` sets the root under which the hub libraries put everything: the `hub` cache of downloaded repos (each stored as a content-addressed blob tree with symlinked snapshots per revision), the datasets cache, and the stored auth token. Because snapshots are symlinked into a per-revision directory pointing at deduplicated blobs, keeping several revisions of the same model costs only the bytes that differ between them, not a full copy each. This is also why you should delete revisions through the cache tooling rather than by hand: naive `rm` on a snapshot can orphan or break the shared blob store.

Lab

I set the env vars, mount the NAS, do a first pinned model download onto the working tier, and write the storage-policy doc. The artifacts are a storage-policy document plus the dotfile that encodes the layout.

Step 1 - the dotfile that encodes the layout

# Storage + cache discipline for the evals-as-rewards loop.
# Sourced from ~/.bashrc so every shell and every `uv run` inherits it.

# Working tier: the fast NVMe. All Hugging Face caches live here.
export HF_HOME="/data/hf"                 # NVMe working tier
export HF_HUB_ENABLE_HF_TRANSFER=1        # accelerated large-file downloads

# Archive tier: the NAS mount point (see /etc/fstab).
export NAS_ROOT="/mnt/nas"
export ARCHIVE_ROOT="${NAS_ROOT}/thesis-loop"   # where kept things go

# Fail loudly if the working tier is missing, rather than silently
# scattering a 14 GB download into the wrong place.
[ -d "$HF_HOME" ] || echo "WARN: HF_HOME $HF_HOME does not exist yet" >&2

echo 'source ~/.config/thesis-loop/storage.env' >> ~/.bashrc

# /data is the NVMe working root. On a fresh machine it is owned by root, so I
# create it with sudo and then hand ownership to my user; otherwise the very
# first HF download fails writing into an unwritable /data.
sudo mkdir -p /data/hf                     # NVMe working tier
sudo chown -R "$USER:$USER" /data          # own the whole working root
source ~/.config/thesis-loop/storage.env

Step 2 - mount the NAS archive tier

sudo apt install -y nfs-common    # NFS client; without it the fstab NFS mount fails
sudo mkdir -p /mnt/nas

# Example NFS entry; adjust host/export/protocol to your NAS.
# Appended to /etc/fstab so the mount survives reboot.
echo 'nas.local:/volume1/thesis  /mnt/nas  nfs  defaults,noatime,_netdev  0  0' | sudo tee -a /etc/fstab

sudo mount -a
mkdir -p /mnt/nas/thesis-loop/{checkpoints,reports,datasets,mlflow-exports}
df -h /mnt/nas    # confirm it mounted and shows the NAS capacity

Step 3 - first pinned model download

This script downloads a model by revision onto the working tier, using hf_transfer, and records the exact hash it pinned. Run it from the serve/ project so it uses that locked environment:

"""Download a model by pinned revision onto the working tier (HF_HOME).

Reproducibility rule: pin the commit hash, not the name. Record what we pinned.
Run with: uv run python scripts/fetch_model.py
"""
import os
import json
import datetime
from huggingface_hub import snapshot_download, HfApi

REPO = "Qwen/Qwen2.5-7B-Instruct"
# Pin an explicit commit. Resolve the current main -> hash once, then hard-code it.
# Replace with the exact 40-hex commit you intend to freeze on.
REVISION = "REPLACE_WITH_40_HEX_COMMIT"

def resolve_latest(repo: str) -> str:
    """Helper: print the current main commit so you can copy it into REVISION."""
    info = HfApi().repo_info(repo)
    return info.sha

if __name__ == "__main__":
    if REVISION.startswith("REPLACE"):
        print("Current main commit for", REPO, "->", resolve_latest(REPO))
        print("Copy that hash into REVISION and re-run to freeze the download.")
        raise SystemExit(0)

    path = snapshot_download(
        repo_id=REPO,
        revision=REVISION,               # the whole reproducibility story
        # lands under $HF_HOME/hub; hf_transfer accelerates it via the env var
    )

    record = {
        "repo": REPO,
        "revision": REVISION,
        "local_path": path,
        "hf_home": os.environ.get("HF_HOME"),
        "fetched_at": datetime.datetime.now().astimezone().isoformat(),
    }
    print(json.dumps(record, indent=2))
    # Append to a pins log so every frozen model is on the record.
    with open(os.path.expanduser("~/machine-log/model-pins.jsonl"), "a") as f:
        f.write(json.dumps(record) + "\n")

cd serve
uv add hf_transfer                       # accelerated backend, into the locked env
# First run prints the current main commit; paste it into REVISION, then:
uv run python scripts/fetch_model.py

Step 4 - the storage-policy document (artifact)

# Storage policy - evals-as-rewards loop

## Tiers
- **Working tier - 1 TB NVMe (`/data`, `HF_HOME=/data/hf`).**
  Fast, disposable. Holds the HF cache, active weights, in-flight checkpoints,
  scratch. Rule: nothing here is the only copy of anything I care about.
- **Archive tier - 5 TB NAS (`/mnt/nas/thesis-loop`).**
  Durable, slow, append-mostly. Holds kept checkpoints, acceptance reports,
  MLflow exports, curated datasets.

## Reproducibility rule for weights
- Always load models/datasets by **revision commit hash**, never by bare name.
- Every frozen model is appended to `~/machine-log/model-pins.jsonl`
  (repo, revision, path, HF_HOME, timestamp).

## Checkpoint retention policy
- **In-flight checkpoints** live on the working tier during a run.
- **Keep** only: the best-by-eval checkpoint of a run, plus the final one.
- **Promote** kept checkpoints to `${ARCHIVE_ROOT}/checkpoints/<run-id>/` on the NAS.
- **Delete** all other intermediate checkpoints from the working tier once the
  run's best/final are safely on the NAS.
- **Cache hygiene**: prune unused HF revisions with `huggingface-cli delete-cache`
  (never hand-`rm` snapshot dirs - it can orphan the shared blob store).

## Env (see ~/.config/thesis-loop/storage.env)
- `HF_HOME=/data/hf`
- `HF_HUB_ENABLE_HF_TRANSFER=1`
- `NAS_ROOT=/mnt/nas`, `ARCHIVE_ROOT=/mnt/nas/thesis-loop`

What you should see. After sourcing the dotfile, echo $HF_HOME prints /data/hf and df -h /mnt/nas shows the NAS capacity, confirming both tiers exist. The first run of fetch_model.py prints the current main commit hash; once you paste that into REVISION and re-run, the model downloads into /data/hf/hub (visibly on the NVMe, not in your home directory), and a line lands in ~/machine-log/model-pins.jsonl recording exactly which commit you froze. The download speed and the on-disk size are machine-log entries; measured on the baseline machine; record value, date, driver. The two artifacts, docs/storage-policy.md and ~/.config/thesis-loop/storage.env, are committed. From here on, "the model" in this book always means a repo at a pinned commit sitting on the working tier, and anything worth keeping has a defined path on the NAS.

Containers and the tracking spine

Goal. Stand up the tracking spine that every later chapter logs into, and fix the schema for what a run is.

Covers. Docker plus the nvidia-container-toolkit; an MLflow server on localhost; the run / experiment / artifact schema the whole book reuses; and the five things that get logged on every single run (git SHA, uv lock hash, driver version, VRAM peak, seeds).

Theory

Why a tracking spine at all

Without a tracking spine, an experiment-heavy project degenerates into a graveyard of terminal scrollback and half-remembered numbers. You run an eval, see 61.4%, tweak something, see 63.1%, and three weeks later you cannot answer the only question that matters: what exactly was different between those two runs? Which model revision, which code, which dependencies, which seed? A reasoning-delta thesis lives or dies on that question, because the entire claim is "this number moved because of this change, and nothing else." The tracking spine is the infrastructure that makes "nothing else" checkable rather than hopeful.

MLflow is the spine I use. It gives me three nouns that map cleanly onto how I actually work. An experiment is a named bucket for related runs (say, "grpo-gsm8k-qwen7b"). A run is a single execution with its parameters, metrics, and artifacts. An artifact is any file a run produces and wants to keep: a config, a plot, a checkpoint pointer, an eval transcript. Every number in Part 9's figures will be pulled from MLflow runs, so the discipline of logging into it from day one is what lets the book assemble itself from data instead of from memory.

Docker, and why the GPU needs a toolkit to enter a container

A couple of services in this book (MLflow here, and later some serving and burst workflows) are cleaner to run as containers than as host installs, because a container pins their entire environment and keeps them from polluting the host. Docker provides that. But there is a specific wrinkle for GPU work: a container is isolated from the host's devices by default, so a process inside it cannot see the NVIDIA GPU unless something bridges the driver across the isolation boundary. That something is the nvidia-container-toolkit. It hooks into the container runtime so that --gpus all (or a Compose deploy.resources reservation) injects the host's GPU devices and the matching driver libraries into the container, letting CUDA code inside the container talk to the real card.

The MLflow server itself does not need the GPU (it is a metadata and file service), so the very first container I stand up is a clean test of Docker and Compose without the GPU complication. But I install and verify the nvidia-container-toolkit in the same chapter, because later chapters will run GPU workloads in containers, and the canonical proof that the toolkit works is running nvidia-smi inside a CUDA container and seeing the same card the host sees.

```admonish gotcha title="--gpus all without the toolkit = 'could not select device driver'" If you try to run a GPU container before installing and configuring the nvidia-container-toolkit, Docker fails with an error like could not select device driver "" with capabilities: [[gpu]]. Installing the toolkit is not enough on its own; you must also register it with the Docker runtime (nvidia-ctk runtime configure) and restart the Docker daemon so it knows about the NVIDIA runtime. The verification step - nvidia-smi inside a container - is the only proof that matters; a green apt install is not.

### The five things every run logs, no exceptions

Here is the heart of the chapter, and it is a policy, not a tool. Every run in this book, whether it is a serve benchmark, an eval, or a training step, logs the same five provenance fields, always, so that any number can be reconstructed. If a run cannot produce these five, it does not get to contribute a number to a figure.

1. **Git SHA of the code.** The exact commit of the loop repository that produced the run, including a dirty flag if the working tree had uncommitted changes. This answers "which code."
2. **uv lock hash.** A hash of the relevant `uv.lock` (serve's or train's, per Chapter 0.4), so the run is tied to the exact resolved dependency graph. This answers "which environment."
3. **Driver version.** The NVIDIA driver version from `nvidia-smi` (Chapter 0.3), so a run is tied to the exact GPU software substrate. This answers "which driver."
4. **VRAM peak.** The peak GPU memory the run reached, so I can see how close to the 16 GB ceiling it ran and catch a regression that would OOM. This answers "did it fit, and how tightly."
5. **Seeds.** Every RNG seed the run set (Python, NumPy, Torch, and any sampler seed), so stochastic results are reproducible. This answers "was it luck."

Together with the pinned model revision from Chapter 0.5, these five turn a run from an anecdote into a reproducible object. The git SHA and lock hash pin the software; the driver pins the substrate; the VRAM peak pins the resource envelope; the seeds pin the randomness. A reviewer with those six facts (five plus the model revision) can, in principle, reproduce the run exactly. That is the whole point of the spine.

```admonish derivation title="Why these five and not others"
Think of a run's output as a function of inputs: $\text{result} = f(\text{code}, \text{deps}, \text{driver}, \text{model}, \text{data}, \text{seed}, \text{hardware})$. Hardware is fixed by the baseline machine and the preflight (Chapters 0.2–0.3). Model and data are pinned by revision (Chapter 0.5). That leaves code, deps, driver, and seed as the free inputs a run must record to be reproducible - which is fields 1, 2, 3, and 5. The VRAM peak (field 4) is not an input but a witness: it is the cheapest early-warning that a change pushed the run's resource envelope, which on a 16 GB card is the difference between "runs" and "OOMs." So four fields pin the inputs, and one field guards the constraint.

Params, metrics, tags, artifacts: using the four slots correctly

MLflow gives a run four kinds of slot, and using them consistently is what makes the store queryable later instead of a junk drawer. Parameters are the inputs I set: the model revision, the learning rate, the batch size, the number of GRPO groups. They are logged once and never change during a run, and they are what I filter and group by in Part 9. Metrics are the outputs I measure, and they are time-series: a training loss logged every step, an eval accuracy logged at each checkpoint, the VRAM peak logged at the end. Because metrics carry a step index, a training curve is a real series I can plot, not a single final number. Tags are metadata about the run itself rather than its inputs or outputs, which is exactly where the five provenance fields live: git SHA, lock hash, driver, and seeds are tags, so they annotate the run without pretending to be experimental variables. Artifacts are files: configs, plots, eval transcripts, and pointers to the checkpoint on the NAS.

The discipline is to always put a thing in the same slot. A learning rate is always a param, never a tag; an eval score is always a metric, never a param; provenance is always a tag. When that convention holds across every run in the book, Part 9 can do things like "pull every run in experiment grpo-gsm8k-qwen7b that shares this git SHA and model revision, and plot their eval-accuracy metric against training step," and get a clean answer, because the fields it filters on are reliably in the slots it expects. A store where the learning rate is sometimes a param and sometimes buried in a run name is a store you cannot query, only browse.

Tooling

The tools are Docker + Docker Compose, the nvidia-container-toolkit, the MLflow server, and a small Python helper that stamps the five provenance fields onto every run.

Docker Compose describes the MLflow service declaratively in a docker-compose.yml, so standing up the spine is one command and its configuration is version-controlled. The nvidia-container-toolkit is installed on the host and registered with Docker so later GPU containers work; I verify it here even though MLflow itself does not use the GPU. The MLflow server runs on localhost with a SQLite backend store and a filesystem artifact store on the working tier. The provenance helper is a tiny module every run imports, which reads the git SHA, hashes the lockfile, reads the driver version, tracks the VRAM peak, and logs the seeds, so the five-field policy is enforced by code rather than by my memory.

Lab

I install and verify the nvidia-container-toolkit, stand up MLflow via Docker Compose, write the provenance helper, and log a dummy run that carries all five fields. The artifacts are the docker-compose.yml and an MLflow-conventions document.

Step 1 - install and verify the nvidia-container-toolkit

#!/usr/bin/env bash
# Install Docker Engine + NVIDIA container toolkit, then prove the GPU crosses
# the container boundary.
set -euo pipefail

# --- Docker Engine, from Docker's official apt repo ---
sudo apt update
sudo apt install -y ca-certificates curl gnupg
sudo install -m 0755 -d /etc/apt/keyrings
sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
sudo chmod a+r /etc/apt/keyrings/docker.asc
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] \
https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME") stable" \
  | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
sudo apt update
sudo apt install -y docker-ce docker-ce-cli containerd.io \
  docker-buildx-plugin docker-compose-plugin

# --- NVIDIA container toolkit, from NVIDIA's libnvidia-container apt repo ---
# Stock Ubuntu has no nvidia-container-toolkit package; add the repo first.
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
  | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
  | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
  | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list > /dev/null
sudo apt update
sudo apt install -y nvidia-container-toolkit

# Register the toolkit with the Docker runtime, restart the daemon.
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# The only proof that matters: run nvidia-smi INSIDE a container.
sudo docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu24.04 nvidia-smi

That final nvidia-smi inside the container should print the same RTX 5080 the host sees. If it errors with a device-driver message, the toolkit is not registered with the runtime; re-run the nvidia-ctk step and restart Docker.

Step 2 - the MLflow service

# MLflow tracking spine - localhost, single user.
# Tracking store: SQLite. Artifact store: a dir on the working tier.
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:latest
    container_name: thesis-mlflow
    command: >
      mlflow server
      --host 0.0.0.0
      --port 5000
      --backend-store-uri sqlite:////mlflow/store/mlflow.db
      --artifacts-destination /mlflow/artifacts
    ports:
      - "127.0.0.1:5000:5000"          # bind to localhost only
    volumes:
      - /data/mlflow/store:/mlflow/store          # SQLite DB on the NVMe working tier
      - /data/mlflow/artifacts:/mlflow/artifacts  # artifacts on the NVMe working tier
    restart: unless-stopped

sudo mkdir -p /data/mlflow/{store,artifacts}   # on the NVMe working tier
sudo chown -R "$USER:$USER" /data/mlflow       # own it so the container can write
cd tracking
docker compose up -d
# UI now at http://127.0.0.1:5000
curl -s http://127.0.0.1:5000/health && echo "  <- MLflow healthy"

Step 3 - the provenance helper

"""Stamp the five mandatory provenance fields onto every MLflow run.

Every run in the book imports start_run() from here. No run contributes a
number to a figure unless it carries: git SHA, uv lock hash, driver version,
VRAM peak, and seeds.
"""
import os
import sys
import json
import hashlib
import random
import subprocess
import contextlib

import mlflow

MLFLOW_URI = "http://127.0.0.1:5000"


def _git_sha() -> str:
    sha = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()
    # Dirty if there are unstaged OR staged-but-uncommitted changes.
    dirty = (
        subprocess.call(["git", "diff", "--quiet"]) != 0
        or subprocess.call(["git", "diff", "--cached", "--quiet"]) != 0
    )
    return sha + ("-dirty" if dirty else "")


def _lock_hash(lock_path: str) -> str:
    with open(lock_path, "rb") as f:
        return hashlib.sha256(f.read()).hexdigest()[:16]


def _driver_version() -> str:
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"]
        )
        return out.decode().strip()
    except Exception:
        return "unknown"


def set_all_seeds(seed: int) -> dict:
    random.seed(seed)
    seeds = {"python": seed, "seed_arg": seed}
    try:
        import numpy as np
        np.random.seed(seed)
        seeds["numpy"] = seed
    except ImportError:
        pass
    try:
        import torch
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
        seeds["torch"] = seed
    except ImportError:
        pass
    return seeds


@contextlib.contextmanager
def start_run(experiment: str, lock_path: str, seed: int, run_name: str | None = None):
    """Open an MLflow run pre-stamped with the five mandatory provenance fields."""
    mlflow.set_tracking_uri(MLFLOW_URI)
    mlflow.set_experiment(experiment)
    seeds = set_all_seeds(seed)
    with mlflow.start_run(run_name=run_name) as run:
        # Fields 1,2,3,5 up front (VRAM peak, field 4, is logged at the end).
        mlflow.set_tags({
            "git_sha": _git_sha(),                 # 1
            "uv_lock_hash": _lock_hash(lock_path), # 2
            "driver_version": _driver_version(),   # 3
            "seeds": json.dumps(seeds),            # 5
        })
        try:
            yield run
        finally:
            # Field 4: VRAM peak, if torch/CUDA is present in this env.
            try:
                import torch
                if torch.cuda.is_available():
                    peak = torch.cuda.max_memory_allocated() / (1024**3)
                    mlflow.log_metric("vram_peak_gb", round(peak, 3))
            except Exception:
                pass

Step 4 - log a dummy run

"""Log a dummy run to prove the spine and the five-field policy work end to end.

Run with: uv run python smoke_run.py
"""
import mlflow
from provenance import start_run

# lock_path points at whichever env's uv.lock this run belongs to.
LOCK = "../serve/uv.lock"

with start_run(experiment="smoke", lock_path=LOCK, seed=1234, run_name="hello-spine"):
    mlflow.log_param("dummy_param", "hello")
    for step in range(5):
        mlflow.log_metric("dummy_metric", 0.5 + 0.1 * step, step=step)
    with open("note.txt", "w") as f:
        f.write("first artifact on the spine\n")
    mlflow.log_artifact("note.txt")
    print("Logged dummy run. Open http://127.0.0.1:5000 to see it.")

cd tracking
uv run --with mlflow python smoke_run.py

Step 5 - the conventions document (artifact)

# MLflow conventions - evals-as-rewards loop

## Server
- Localhost only: tracking URI http://127.0.0.1:5000 (see tracking/docker-compose.yml).
- Backend store: SQLite on /data/mlflow/store (NVMe working tier).
- Artifact store: /data/mlflow/artifacts (NVMe), exported to NAS to make permanent.

## Nouns
- **Experiment**: named bucket of related runs, e.g. `grpo-gsm8k-qwen7b`.
- **Run**: one execution; carries params, metrics (time-series), tags, artifacts.
- **Artifact**: any file a run keeps (config, plot, transcript, checkpoint pointer).

## The five mandatory fields (enforced by tracking/provenance.py)
Every run logs, no exceptions:
1. `git_sha`      - code commit (+ `-dirty` if the tree was modified)
2. `uv_lock_hash` - sha256[:16] of the relevant uv.lock (serve or train)
3. `driver_version` - from nvidia-smi
4. `vram_peak_gb` - torch.cuda.max_memory_allocated at run end
5. `seeds`        - JSON of every RNG seed set (python/numpy/torch/sampler)
Plus the pinned model **revision** (Chapter 0.5) logged as a param on any run
that loads a model.

## Naming
- Experiment names: `<method>-<task>-<model>`, lowercase, hyphenated.
- Run names: short and human, e.g. `baseline`, `grpo-step-200`, `ablation-no-kl`.

## Retention
- Artifact store lives on the working tier; export to
  `${ARCHIVE_ROOT}/mlflow-exports/` on the NAS to make a run permanent.

What you should see. nvidia-smi inside the CUDA container prints the same RTX 5080 as the host, proving the toolkit bridges the GPU. curl http://127.0.0.1:5000/health returns healthy, and the MLflow UI loads in a browser. After the smoke run, the UI shows an experiment named smoke with one run named hello-spine, and that run carries four tags (git_sha, uv_lock_hash, driver_version, seeds), a vram_peak_gb metric (which will read near zero for this non-GPU dummy; real GPU runs will show a real peak, measured on the baseline machine; record value, date, driver), a five-point dummy_metric curve, and one artifact note.txt. The two artifacts, tracking/docker-compose.yml and docs/mlflow-conventions.md, are committed. From here on, every chapter that produces a number logs it through start_run, so the five fields are never optional and never forgotten.

Acceptance: proving the machine

Goal. Prove the machine is healthy with a scripted, repeatable acceptance suite before trusting any number it produces.

Covers. What a healthy card looks like at idle; memory-bandwidth and GEMM sanity tests; thermals under sustained load; and first tokens from PyTorch. The suite ends in an acceptance report archived to the NAS.

Theory

Why acceptance is a gate, not a vibe

Everything downstream in this book is a measurement, and a measurement is only as trustworthy as the instrument that produced it. The preflight in Chapter 0.2 proved the hardware was sound under the vendor's OS; the install in Chapter 0.3 proved the driver and link were correct. Acceptance is the third gate: it proves that the fully assembled Linux stack, driver plus CUDA plus PyTorch, produces sane, stable, in-spec numbers before I let it produce research numbers. It is deliberately a gate and not a feeling. Either the card hits a plausible fraction of its rated bandwidth, sustains a GEMM without errors, holds temperature under load, and emits correct tokens, or it does not, and if it does not, I fix the machine before I trust a single reasoning-delta figure.

The reason to script it, rather than eyeball nvidia-smi once and move on, is that acceptance is something I re-run. After a driver update, after a kernel bump, after moving the machine, after any "huh, that's slower than I remember" moment, I re-run the same suite and diff the report against the last known-good one. A scripted suite turns "the machine feels fine" into "the machine matches acceptance report v1 within tolerance," which is a claim I can actually check.

What a healthy card looks like at idle

Before any load, a healthy RTX 5080 at idle has a recognizable signature in nvidia-smi: low power draw, low temperature, near-zero memory used, near-zero utilization, and the correct static facts (16 GB total memory, the 570-branch driver, PCIe Gen5 x16). The idle numbers are a baseline, not a benchmark; their job is to catch gross problems like a card stuck at high power draw, a runaway process already holding VRAM, or a link that renegotiated down since install. All the specific idle values are machine-log entries; measured on the baseline machine; record value, date, driver.

Bandwidth and GEMM: the two numbers that predict everything

Two microbenchmarks predict almost all of the performance behavior this book cares about, and they map onto the two regimes I will formalize with the roofline model in Part 2.

The first is memory bandwidth. The RTX 5080 has roughly 960 GB/s of memory bandwidth (per NVIDIA's RTX 5080 product specification). A bandwidth test streams large buffers through the memory system and reports the achieved rate. It will never hit the theoretical peak (no real workload does), but it should reach a healthy fraction of it. This matters because the token-generation phase of LLM inference is memory-bound: each token's latency is dominated by reading the model weights and KV cache out of VRAM, so achieved bandwidth is a direct predictor of decode throughput. If the bandwidth test comes back far below the rated figure, every tokens-per-second number later will be quietly disappointing, and I want to know that is a hardware/driver issue, not a vLLM configuration issue.

The second is GEMM throughput. A general matrix multiply (GEMM) is the compute primitive at the heart of every transformer layer, and a large GEMM in the card's native low precision (bf16/fp16, and on Blackwell the newer low-precision tensor-core paths) exercises the compute-bound regime: prompt prefill and training are dominated by these matmuls. A GEMM sanity test does two jobs at once. It checks correctness (the result matches a reference within floating-point tolerance, which catches a miscompiled kernel or unstable memory), and it estimates throughput in FLOP/s, which I compare against the card's rated compute to confirm the tensor cores are actually engaged and not silently falling back to a slow path.

Thermals: sustained load is a different test than peak load

A card can pass a two-second GEMM and still fail in practice, because a two-second burst does not heat the silicon the way a twenty-minute training run does. Thermal acceptance is about sustained load: I run the GPU flat out long enough for temperatures to reach steady state, and I watch two things. Does the temperature plateau below the throttling threshold, and do the clocks hold, or do they sag as the card thermal-throttles to protect itself? A machine that throttles under sustained load will produce results that silently degrade the longer a run goes, which is poison for a training loop that runs for an hour. All temperatures, clocks, and throttle observations here are machine-log entries; measured on the baseline machine; record value, date, driver.

First tokens: the whole stack, end to end

The final acceptance check is the most holistic: load a small model in PyTorch and generate a few tokens. This exercises the entire software stack, PyTorch, CUDA, the driver, the GPU, and the tokenizer, in the same shape the real workloads use, and confirms it produces coherent output. It is not a benchmark; it is a smoke test that says "the thing that will run the loop can, in fact, turn a prompt into tokens on this GPU." If bandwidth, GEMM, and thermals all pass but first-tokens fails, the problem is in the Python/CUDA stack, not the hardware, and I have narrowed it before I even start.

Tolerances, and what "healthy fraction" actually means

A gate needs a threshold, and a threshold needs a number, so it is worth being honest about where the acceptance tolerances come from. I do not have the card's internal design targets, so I anchor the tolerances to two things: the published spec (roughly 960 GB/s of memory bandwidth, and the card's rated low-precision compute from NVIDIA's RTX 5080 specification) and the first clean run of the suite itself, which becomes the reference. A microbenchmark never reaches theoretical peak, because the peak assumes perfect overlap and no overhead, so I expect the bandwidth test to land at some fraction of 960 GB/s and the GEMM test to land at some fraction of rated compute. The exact fractions I treat as healthy are recorded from the first good run; measured on the baseline machine; record value, date, driver. The point of writing them down is that "healthy fraction" stops being a vibe and becomes a stored number I can diff against.

Once that reference exists, the tolerance is relative, not absolute, which is the more useful form. I do not actually care whether the bandwidth test hits a specific GB/s figure to three digits; I care whether a later run drops meaningfully below the reference the machine established when it was known-good. A ten-or-fifteen-percent drop from the reference bandwidth is a red flag worth chasing (a thermal issue, a link renegotiation, a driver regression), even if the absolute number still looks superficially fine. So acceptance v1 is not just a pass/fail certificate, it is the baseline against which every future acceptance run is a relative comparison. That is why the report records the actual measured numbers and not just checkmarks: the numbers are the reference the machine will be held to for the rest of its working life.

Tooling

The suite is a bash driver that calls small, focused checks: nvidia-smi for idle and for load monitoring, a CUDA bandwidth measurement, a PyTorch GEMM correctness-and-throughput check, a sustained-load thermal watch, and a PyTorch generate check. Everything runs from the train/ uv environment (it has PyTorch), and the whole run logs into MLflow through the Chapter 0.6 provenance helper so the acceptance report itself carries the five mandatory fields. The report is written to disk and archived to the NAS.

Lab

The suite is one bash script that orchestrates a Python check module, prints a human-readable report, and archives it to the NAS. The artifact is acceptance report v1 on the NAS. Every measured value in the report is a placeholder to be filled on the machine.

The Python checks

"""GPU acceptance checks: GEMM correctness+throughput, bandwidth, first tokens.

Run individual checks via the CLI, e.g.:  uv run python checks.py gemm
"""
import sys
import time
import json

import torch


def idle_summary() -> dict:
    """Static facts the report should confirm."""
    p = torch.cuda.get_device_properties(0)
    return {
        "name": p.name,
        "total_mem_gb": round(p.total_memory / (1024**3), 2),
        "capability": f"{p.major}.{p.minor}",   # Blackwell reports sm_120
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
    }


def gemm_check(n: int = 8192, iters: int = 50, dtype=torch.bfloat16) -> dict:
    """Large bf16 GEMM: correctness vs fp32 reference + throughput estimate."""
    torch.manual_seed(0)
    a = torch.randn(n, n, device="cuda", dtype=dtype)
    b = torch.randn(n, n, device="cuda", dtype=dtype)

    # Correctness on a small tile against an fp32 reference.
    ref = (a[:512, :512].float() @ b[:512, :512].float())
    got = (a[:512, :512] @ b[:512, :512]).float()
    max_rel_err = ((got - ref).abs() / (ref.abs() + 1e-3)).max().item()

    # Throughput: 2*n^3 FLOP per matmul.
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        c = a @ b
    torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / iters
    tflops = (2 * n**3) / dt / 1e12

    return {
        "n": n, "dtype": str(dtype),
        "max_rel_err": max_rel_err,
        "sec_per_gemm": dt,
        "tflops_est": round(tflops, 1),   # compare to rated compute
        "correct": max_rel_err < 5e-2,
    }


def bandwidth_check(bytes_gb: float = 2.0, iters: int = 50) -> dict:
    """Streaming copy bandwidth: read+write a large buffer, report GB/s."""
    n = int(bytes_gb * (1024**3) / 4)     # fp32 elements
    x = torch.randn(n, device="cuda")
    y = torch.empty_like(x)

    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        y.copy_(x)                        # reads n*4 bytes, writes n*4 bytes
    torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / iters
    moved_bytes = 2 * n * 4               # read + write
    gbps = moved_bytes / dt / 1e9         # decimal GB/s, same units as the rated figure
    return {"gb_per_s_est": round(gbps, 1), "sec_per_iter": dt}   # compare to ~960 GB/s rated


def first_tokens(model_id: str, revision: str, n_new: int = 20) -> dict:
    """Whole-stack smoke test: load a small model by pinned revision, generate."""
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(model_id, revision=revision)
    model = AutoModelForCausalLM.from_pretrained(
        model_id, revision=revision, dtype=torch.bfloat16, device_map="cuda"
    )
    prompt = "The capital of France is"
    ids = tok(prompt, return_tensors="pt").to("cuda")
    out = model.generate(**ids, max_new_tokens=n_new, do_sample=False)
    text = tok.decode(out[0], skip_special_tokens=True)
    return {"prompt": prompt, "completion": text, "generated_ok": len(text) > len(prompt)}


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    result = {}

    # Guard each check so one failure (e.g. first_tokens with a placeholder
    # revision) records an error instead of aborting the whole report.
    def guard(name, fn):
        try:
            result[name] = fn()
        except Exception as e:
            result[name] = {"error": repr(e)}

    if which in ("idle", "all"):
        guard("idle", idle_summary)
    if which in ("gemm", "all"):
        guard("gemm", gemm_check)
    if which in ("bandwidth", "all"):
        guard("bandwidth", bandwidth_check)
    if which in ("first_tokens", "all"):
        # Use a small model pinned by revision; replace the hash on the machine.
        guard("first_tokens", lambda: first_tokens(
            "Qwen/Qwen2.5-0.5B-Instruct", revision="REPLACE_WITH_40_HEX_COMMIT"
        ))
    print(json.dumps(result, indent=2))

The thermal watch

"""Sustained-load thermal check: hammer the GPU, sample temp/clock/power, report.

Run: uv run python thermal_watch.py 300   # seconds of sustained load
"""
import sys
import time
import subprocess

import torch


def sample() -> dict:
    q = "temperature.gpu,clocks.sm,power.draw,utilization.gpu"
    out = subprocess.check_output(
        ["nvidia-smi", f"--query-gpu={q}", "--format=csv,noheader,nounits"]
    ).decode().strip()
    t, clk, pw, util = [s.strip() for s in out.split(",")]
    return {"temp_c": float(t), "sm_mhz": float(clk), "power_w": float(pw), "util_pct": float(util)}


def main(seconds: int):
    n = 8192
    a = torch.randn(n, n, device="cuda", dtype=torch.bfloat16)
    b = torch.randn(n, n, device="cuda", dtype=torch.bfloat16)
    samples = []
    t_end = time.time() + seconds
    while time.time() < t_end:
        for _ in range(50):
            c = a @ b                       # keep the tensor cores busy
        torch.cuda.synchronize()
        samples.append(sample())
        time.sleep(1)
    peak_t = max(s["temp_c"] for s in samples)
    min_clk = min(s["sm_mhz"] for s in samples)
    max_clk = max(s["sm_mhz"] for s in samples)
    # Throttling shows up as clocks sagging from their early peak under sustained heat.
    print({"peak_temp_c": peak_t, "clock_min_mhz": min_clk, "clock_max_mhz": max_clk,
           "clock_sag_mhz": round(max_clk - min_clk, 1), "n_samples": len(samples)})


if __name__ == "__main__":
    main(int(sys.argv[1]) if len(sys.argv) > 1 else 300)

The suite driver

#!/usr/bin/env bash
# Scripted acceptance suite. Produces acceptance-report-vN and archives it to the NAS.
# Run from the train/ uv project (it has torch+transformers): bash run_acceptance.sh
set -euo pipefail

source ~/.config/thesis-loop/storage.env    # HF_HOME, ARCHIVE_ROOT
STAMP="$(date --iso-8601=seconds)"
VER="${1:-1}"
REPORT="acceptance-report-v${VER}-$(date +%Y%m%d).md"

{
  echo "# Acceptance report v${VER} - MSI Aegis R2 (A2NVV9-2218US)"
  echo "generated: ${STAMP}"
  echo
  echo "## Idle / static facts"
  echo '```'
  nvidia-smi --query-gpu=name,memory.total,driver_version,pcie.link.gen.current,pcie.link.width.current,temperature.gpu,power.draw --format=csv
  echo '```'
  echo
  echo "## GEMM, bandwidth, first-tokens (JSON)"
  echo '```json'
  uv run python checks.py all
  echo '```'
  echo
  echo "## Thermals under sustained load (300 s)"
  echo '```'
  uv run python thermal_watch.py 300
  echo '```'
  echo
  echo "## Pass/fail (fill in against tolerances)"
  echo "- [ ] card named RTX 5080, 16 GB total, driver on 570 branch, link Gen5 x16"
  echo "- [ ] GEMM correct (max_rel_err < 5e-2) and tflops_est within tolerance of rated"
  echo "- [ ] bandwidth gb_per_s_est a healthy fraction of ~960 GB/s rated"
  echo "- [ ] thermals: peak temp below throttle, clock_sag small (clocks held)"
  echo "- [ ] first_tokens: generated_ok true, completion coherent"
} | tee "$REPORT"

# Archive to the NAS to make it permanent.
mkdir -p "${ARCHIVE_ROOT}/reports"
cp "$REPORT" "${ARCHIVE_ROOT}/reports/"
echo "Archived $REPORT to ${ARCHIVE_ROOT}/reports/"

Logging acceptance as an MLflow run

"""Wrap the suite as an MLflow run so acceptance carries the five provenance fields.

Run from the train/ project (which has torch); pull in mlflow for the logging:
  uv run --with mlflow python log_acceptance.py acceptance-report-v1-YYYYMMDD.md
"""
import sys
import mlflow
sys.path.insert(0, "../tracking")
from provenance import start_run   # noqa: E402

report_path = sys.argv[1]
# The suite runs from acceptance/, which has no lockfile of its own; point at the
# train/ env's lockfile (that is the env the checks actually run under).
with start_run(experiment="acceptance", lock_path="../train/uv.lock", seed=0, run_name="v1"):
    mlflow.log_artifact(report_path)     # the report becomes an MLflow artifact too
    mlflow.set_tag("acceptance_version", "1")
    print("Logged acceptance report to MLflow.")

What you should see. The suite prints a report in which: the idle line names the card as an RTX 5080 with 16 GB total, a 570-branch driver, and a Gen5 x16 link; the GEMM check reports correct: true with a small max_rel_err and a tflops_est in a plausible fraction of the card's rated compute; the bandwidth check reports a gb_per_s_est that is a healthy fraction of the roughly 960 GB/s rated figure; the thermal watch reports a peak_temp_c below the throttling threshold with a small clock_sag_mhz (meaning clocks held under sustained load); and first_tokens reports generated_ok: true with a coherent completion of "The capital of France is." Every one of those numbers is a placeholder until you run it; measured on the baseline machine; record value, date, driver. The report is saved locally, copied to ${ARCHIVE_ROOT}/reports/ on the NAS, and logged as an MLflow run carrying the five mandatory fields. That archived acceptance-report-v1 is the artifact and the machine's clean bill of health: from here on, when a later number looks wrong, the first move is to re-run the suite and diff against v1, and the whole rest of the book is allowed to trust the instrument.

Tensors, autograd, and number formats

Everything downstream in this book is a program that moves numbers around a GPU and occasionally differentiates them. Before I can talk honestly about attention, memory budgets, or reward gradients, I need the three primitives those programs are built from: the tensor (how numbers are laid out in memory), autograd (how the machine computes derivatives without me deriving them by hand), and the number formats (how a real number gets squeezed into 32, 16, 8, or even 4 bits, and what precision I lose when it does). None of these is a black box once you look at it, and all three come back to bite you at 2am if you treat them as one. So this chapter opens Part I by making them concrete, with the two derivations that the rest of the book leans on: reverse-mode differentiation, and the representation error of floating point.

Theory

A tensor is typed memory plus a shape

The word "tensor" scares people off with its physics connotations, but in this stack a tensor is a boringly concrete thing: a flat, contiguous block of bytes in memory, plus a small amount of metadata that tells the library how to read those bytes as a multidimensional array. The metadata is what makes an array a tensor. It is a shape (how many elements along each axis), a stride (how many elements to skip in the flat buffer to move one step along each axis), a dtype (how to interpret each element's bytes as a number), and a device (which piece of hardware's memory the buffer lives in). That is essentially all of it.

Take a $2\times 3$ matrix of 32-bit floats. The library allocates $2\times 3=6$ elements times 4 bytes each, so 24 contiguous bytes. The shape is $(2,3)$. The stride, in row-major (C) order, is $(3,1)$: to move down a row I jump 3 elements in the buffer, to move across a column I jump 1. Indexing element $(i,j)$ is then pure arithmetic on the flat buffer,

$$\text{offset}(i,j)=i\cdot s_0+j\cdot s_1=3i+j\text{\hspace{1em}(1.1)}$$

and reading that element means grabbing 4 bytes at base + 4 * offset and interpreting them under the dtype. This is worth internalizing because it explains a pile of behavior that otherwise looks like magic. A transpose does not move any bytes; it just swaps the stride to $(1,3)$ and the shape to $(3,2)$, which is why x.T is instant and free but leaves you with a non-contiguous tensor whose elements are no longer in reading order. A reshape that can be expressed as a stride change is free; one that cannot forces a copy. A slice is a new shape, a new stride, and an offset into the same buffer, which is why slices alias their parent and writing through one mutates the other. When a kernel later demands a contiguous input and you have handed it a transposed view, the .contiguous() call that fixes it is a full memory copy, and on a 16GB card those copies are exactly the kind of thing that quietly doubles your footprint. The stride model is not trivia; it is the difference between a view and a copy, and copies are where VRAM goes to die.

The dtype is the other half of the metadata, and it is the whole back third of this chapter, so hold that thought. The device matters because the RTX 5080's 16GB of VRAM is a separate address space from the 32GB of system DDR5, and moving a tensor across the PCIe bus is neither free nor implicit. x.to("cuda") is a real transfer. I will come back to this constantly.

Autograd is reverse-mode differentiation, and here is why it is cheap

Training is optimization, optimization needs gradients, and the gradients of a modern network with billions of parameters are not something anyone derives by hand. Autograd derives them for you, and the specific algorithm it uses (reverse-mode automatic differentiation, also called backpropagation) is not an approximation and not finite differences. It is the exact chain rule, evaluated in a specific order that makes it astonishingly cheap for the shape of function we care about. That shape is: many inputs (the parameters), one scalar output (the loss). Reverse mode is the right tool for exactly that shape, and it is worth deriving why.

The practical upshot for how PyTorch works: as your forward code runs, the library records each operation as a node in a directed acyclic graph, storing the inputs it will need to compute that node's VJP. Calling .backward() on the scalar loss walks that graph in reverse topological order, applying equation (1.5) at each node and accumulating adjoints into .grad on the leaf tensors (your parameters). The graph is built fresh each forward pass ("define-by-run"), which is why ordinary Python control flow just works. A tensor participates in this only if it has requires_grad=True; the activations recorded for the backward pass are held alive until you call .backward(), which is why torch.no_grad() at inference time is not a nicety but a large memory saving. It tells the engine to skip recording entirely.

Floating point, bit by bit

Now the dtypes, which is where the 16GB constraint really starts to shape decisions. A real number like $1/3$ has infinitely many digits; a fixed-width binary format has to round it to one of finitely many representable values. Every format in this stack is a specific answer to "how do I spend my bits between range and precision," and the answer determines both how much memory a model eats and how much numerical error it carries. The common scaffold is the IEEE-754-style layout: one sign bit, some exponent bits, and some mantissa (fraction) bits.

With equation (1.8) in hand, the formats read off cleanly. I will give each as (sign, exponent, mantissa), its machine epsilon $2^{-p}$, its dynamic range, and where it shows up.

FP32 (1, 8, 23), bias 127. Machine epsilon $2^{-23}\approx 1.2\times 10^{-7}$ (about 7 decimal digits), max finite $\approx 3.4\times 10^{38}$, smallest normal $\approx 1.2\times 10^{-38}$. This is the "full precision" reference. Master weights, optimizer state, and loss accumulation live here because you want the small numbers to survive being added to big ones. Four bytes per element.

TF32 (1, 8, 10), NVIDIA's tensor-core format. It keeps FP32's 8 exponent bits (so full FP32 range) but truncates the mantissa to 10 bits, giving $2^{-10}\approx 10^{-3}$ precision. It is stored in a 32-bit slot but multiplied at reduced precision inside the tensor cores; it is a compute mode, not a storage format. Matmuls silently run in TF32 by default on recent hardware unless you ask for true FP32.

FP16 (1, 5, 10), bias 15. Machine epsilon $2^{-10}\approx 9.8\times 10^{-4}$, max finite $65504$, smallest normal $\approx 6.1\times 10^{-5}$. Ten mantissa bits is decent precision, but only five exponent bits means the range is narrow, and that narrowness is the classic FP16 training failure: gradients smaller than $\sim 6\times 10^{-5}$ flush to zero (underflow) and activations above $65504$ blow up to infinity (overflow). This is exactly why FP16 training needs loss scaling, multiplying the loss by a big constant to shove the gradients back up into representable range before they underflow, then dividing it back out.

BF16 (1, 8, 7), bias 127. This is the format that made mixed precision pleasant. It keeps FP32's full 8-bit exponent (identical range, $\approx 3.4\times 10^{38}$, no overflow or underflow drama) and spends only 7 bits on the mantissa, giving coarse precision $2^{-7}\approx 7.8\times 10^{-3}$. The trade is deliberate: for training, range matters more than precision, because you can tolerate a noisy gradient but you cannot tolerate a gradient that became infinity. BF16 needs no loss scaling because nothing underflows the way it does in FP16. On the RTX 5080 (Blackwell), BF16 is the default working dtype for both my inference and my training labs.

FP8 comes in two flavors, both 8 bits, differing in how they split the remaining 7 bits after the sign. E4M3 (1, 4, 3), bias 7, favors precision: max value $448$ (it reclaims the would-be infinity encoding to extend range slightly, so it has no infinities, only NaN). E5M2 (1, 5, 2), bias 15, favors range: max $57344$, and it keeps IEEE-style inf/NaN. The rule of thumb the hardware vendors use: E4M3 for the forward pass (weights and activations, where you want precision) and E5M2 for gradients (where you want range). FP8 shows up in the fastest inference paths and in FP8 training on Hopper/Blackwell tensor cores; on 16GB it is mostly an inference-time trick.

INT8 and INT4 are not floating point at all. They are plain signed integers (range $[-128,127]$ and $[-8,7]$) paired with a floating-point scale $s$ and sometimes a zero-point $z$, so the real value is reconstructed as $x\approx s\cdot (q-z)$ where $q$ is the stored integer. Because the grid is uniform (evenly spaced by $s$, unlike floating point's per-binade spacing), integer quantization is a great fit for weight distributions that are roughly bounded and bell-shaped, and it is the workhorse of weight-only quantization. The scale is chosen per-tensor or, better, per-channel/per-group so a few outlier weights do not stretch the grid and wreck everyone else's precision. INT4 weight quantization (with a group size like 128) is how a model that would need many gigabytes in BF16 fits in a couple; I spend a whole chapter on the theory in Part II.

MXFP4 is the microscaling format that makes 4-bit floating point actually usable, and it is worth understanding because gpt-oss ships in it. The element type is E2M1 (1 sign, 2 exponent, 1 mantissa), which by equation (1.6) can represent only the eight magnitudes $\{0,0.5,1,1.5,2,3,4,6\}$ and their negatives, a laughably coarse grid on its own. The trick is the "MX" part: elements are grouped into blocks of 32, and each block shares one 8-bit power-of-two scale (an E8M0 exponent, no mantissa). So the stored cost per element is 4 bits plus $8/32=0.25$ bits of shared scale, i.e. 4.25 bits, and the shared scale slides each block's tiny 3-bit grid up or down to wherever that block's numbers actually live. It is INT-style block scaling married to a float element type, and it is the reason a 4-bit model can hold a usable dynamic range. I will read a real MXFP4 checkpoint off disk in the model-zoo chapter.

The through-line: picking a dtype is picking a point on the range-precision-memory triangle that equation (1.8) governs, and on a 16GB card the memory axis is never negotiable, so the craft is in giving up precision (and sometimes range) exactly where the math can absorb it and nowhere else.

Tooling

The tool that embodies all of this is PyTorch, and specifically two of its subsystems: the tensor/storage model and the autograd engine. They are the reason I can write the labs in this book in a few dozen lines instead of a few thousand.

On the tensor side, PyTorch exposes exactly the metadata model from the theory. Every torch.Tensor is a view onto a Storage (the flat byte buffer) with its own .shape, .stride(), .dtype, .device, and a storage_offset(). You can see the view-versus-copy distinction directly: x.t() and x[::2] and x.view(...) return tensors sharing x's storage (check with x.data_ptr()), while x.contiguous() or x.clone() or an incompatible .reshape() allocate new storage. The dtypes are first-class objects (torch.float32, torch.bfloat16, torch.float16, torch.float8_e4m3fn, torch.float8_e5m2, torch.int8, and so on), each carrying an .itemsize in bytes, which is all you need to compute a tensor's footprint: numel() * element_size(). That formula is the atom of every VRAM budget in this book.

On the autograd side, PyTorch is a define-by-run reverse-mode engine, exactly the algorithm of equations (1.2)–(1.5). Operations on tensors with requires_grad=True are recorded into a graph of Function nodes; each node stashes whatever forward values its VJP needs (visible as .saved_tensors) and exposes a .grad_fn. Calling .backward() on a scalar seeds $\bar{L}=1$ and runs the reverse sweep, accumulating into leaf .grad. torch.no_grad() disables recording (inference, and any bookkeeping you do not want differentiated); tensor.detach() cuts a subgraph loose; and mixed precision is handled by torch.autocast, which runs matmul-heavy ops in BF16/FP16 while keeping a master copy and the loss accumulation in FP32, i.e. it automates precisely the range-precision split I described. torch.cuda.amp.GradScaler is the loss-scaling machinery for the FP16 path.

Lab

The goal of this lab is to make the byte-level footprint of a real model layer visible, so that when later chapters claim "the weights are $N$ bytes" you have already watched the arithmetic hold on your own machine. I will load a small open model, walk it parameter by parameter, and print each tensor's shape, dtype, element size, and total bytes, then reconcile the sum against what CUDA reports as allocated. I will also demonstrate the same weights re-cast to BF16 and FP16 so the memory halving is not a claim but a measurement.

This is a uv project. From the repo root:

uv init labs/tensors-dtypes
cd labs/tensors-dtypes
uv add torch transformers safetensors

"""Layer-by-layer dtype and memory inspector for a small model.

Writes a CSV artifact of every parameter's footprint and prints a summary
reconciled against torch.cuda.memory_allocated().
"""
import csv
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM

MODEL = "Qwen/Qwen3-0.6B"  # small enough to load anywhere; swap freely
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

def human(nbytes: int) -> str:
    for unit in ("B", "KiB", "MiB", "GiB"):
        if nbytes < 1024 or unit == "GiB":
            return f"{nbytes:8.2f} {unit}"
        nbytes /= 1024

def inspect(model, dtype_label: str, writer) -> int:
    total = 0
    for name, p in model.named_parameters():
        nbytes = p.numel() * p.element_size()
        total += nbytes
        writer.writerow({
            "dtype_label": dtype_label,
            "param": name,
            "shape": "x".join(map(str, p.shape)),
            "dtype": str(p.dtype),
            "elem_bytes": p.element_size(),
            "numel": p.numel(),
            "bytes": nbytes,
        })
    return total

def main() -> None:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    csv_path = OUT / "param_footprint.csv"
    with csv_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=[
            "dtype_label", "param", "shape", "dtype",
            "elem_bytes", "numel", "bytes",
        ])
        writer.writeheader()

        # Load once in fp32 on CPU so element_size() is unambiguous.
        model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float32)
        n_params = sum(p.numel() for p in model.parameters())

        results = {}
        for label, dtype in [("fp32", torch.float32),
                             ("bf16", torch.bfloat16),
                             ("fp16", torch.float16)]:
            model_cast = model.to(dtype)
            results[label] = inspect(model_cast, label, writer)

        # Move the bf16 copy to the GPU and reconcile with CUDA's accounting.
        allocated = None
        if device == "cuda":
            torch.cuda.reset_peak_memory_stats()
            model_gpu = model.to(torch.bfloat16).to("cuda")
            allocated = torch.cuda.memory_allocated()
            print(torch.cuda.memory_summary(abbreviated=True))

    print(f"\nModel: {MODEL}")
    print(f"Total parameters: {n_params:,}")
    for label, total in results.items():
        implied = n_params * {"fp32": 4, "bf16": 2, "fp16": 2}[label]
        print(f"  {label}: summed {human(total)}   "
              f"(n_params x itemsize = {human(implied)})")
    if allocated is not None:
        print(f"CUDA allocated for bf16 weights on GPU: {human(allocated)}")
    print(f"\nArtifact: {csv_path.resolve()}")

if __name__ == "__main__":
    main()

Run it:

uv run python inspect_dtypes.py

What you should see. The script writes artifacts/param_footprint.csv, one row per parameter per dtype, which you can open and sort to see exactly which tensors dominate (the embedding/lm_head and the MLP up/down/gate projections will be the fattest rows, foreshadowing the memory chapter). In the terminal you get a per-dtype total that matches n_params × itemsize to the byte: for a roughly 0.6B-parameter model the FP32 total lands near 2.2 GiB and the BF16 and FP16 totals near 1.1 GiB each, cleanly halved, because equation (1.8)'s bookkeeping is just numel × element_size (the human() helper divides by 1024, so these are GiB, not the 2.4/1.2 GB the same bytes make in decimal). The CUDA-allocated figure for the BF16 copy on the GPU reads a little above the 1.1 GiB summed figure (allocator rounding plus context), and memory_summary() shows that headroom. Record the exact allocated value with date and driver on your run (measured on the baseline machine, record value, date, driver), because that gap between "the arithmetic says 1.1 GiB" and "CUDA says a bit more" is the first hint of the accounting the whole memory chapter is about.

Tokenization and embeddings

A language model does not read text. It reads integers. Everything between the string you typed and the first matrix multiply inside the network is the job of the tokenizer and the embedding matrix, and this stage is where a surprising amount of eval reproducibility quietly lives or dies. If you run the same prompt through two different chat templates you can get two different integer sequences, two different loss values, and two different eval scores, without touching a single weight. So this chapter treats the input surface as a first-class object: how byte-pair encoding turns text into tokens, what the vocabulary size actually trades off, why a chat template is itself a tokenization decision and not a cosmetic string wrapper, and how the embedding matrix (and weight tying) turns those integers into the vectors the transformer actually consumes.

Theory

Why not just use characters, or words

There are two obvious ways to turn text into integers, and both are bad. Character-level tokenization gives a tiny vocabulary (a few hundred symbols) but makes sequences brutally long, and a transformer's compute scales with the square of sequence length, so long sequences are expensive exactly where you cannot afford it. Word-level tokenization gives short sequences but a vocabulary that is effectively unbounded (every proper noun, typo, and hashtag is a new word) and helpless in front of anything it never saw at training time, the dreaded out-of-vocabulary token. The whole field settled on a middle path: subword tokenization, where common words stay whole, rare words fracture into pieces, and truly novel strings fall back to bytes so that nothing is ever out-of-vocabulary. Byte-pair encoding is the dominant recipe for building that subword vocabulary, and it is what Qwen3, gpt-oss, and essentially every model I touch in this book use.

Byte-pair encoding, derived from a greedy compression rule

BPE started life as a compression algorithm, and it is easiest to understand as one. You begin with the most granular possible alphabet (raw bytes, so 256 base symbols, which guarantees every possible input is representable) and you repeatedly find the most frequent adjacent pair of symbols in your training corpus and merge it into a new single symbol. Do that $K$ times and you have learned $K$ merges and a vocabulary of $256+K$ tokens.

The consequence of equation (2.1)'s greedy, frequency-driven rule is that the resulting vocabulary is a mirror of the training corpus. Frequent English words become single tokens; a chemistry paper's jargon fractures into many pieces; a language underrepresented in the corpus tokenizes into far more tokens per sentence than English does, which is a real and measurable cost (more tokens means more compute and, in a fixed context window, less room). This is why "tokens per word" (the fertility of a tokenizer) differs across languages and domains, and why the same eval prompt can cost noticeably more tokens under one model's tokenizer than another's.

The vocabulary-size tradeoff, quantified

How big should the vocabulary be? This is a genuine tradeoff with costs on both sides, and it is worth making the two forces explicit rather than hand-waving "bigger is better up to a point."

Chat templates are tokenization, not decoration

Here is the point most tutorials skip and that matters most for evals. A modern instruction-tuned model was trained on conversations that were serialized into a specific string format using special tokens that mark the boundaries between roles. These special tokens (things like a begin-of-turn marker, role labels for system/user/assistant, and an end-of-turn marker) are real entries in the vocabulary with their own integer IDs, and the model learned that "an assistant reply begins right after this exact token sequence." The chat template is the (usually Jinja) function that takes your list of role-tagged messages and produces exactly that serialized string, special tokens and all, before BPE ever runs.

This means the chat template is part of the tokenization, and getting it wrong is not a style error, it is feeding the model a distribution it was never trained on. Concretely, a bare prompt like Solve: 2+2 and the same content wrapped as <|im_start|>user\nSolve: 2+2<|im_end|>\n<|im_start|>assistant\n are different token sequences, and the model's next-token distribution after each is different, because only the second one puts the model in the state it was trained to answer from. Two families disagree on the marker strings, on whether a system message is injected by default, on whether a trailing "generation prompt" (the dangling assistant\n that cues the model to start replying) is appended, and on whether reasoning models wrap a hidden thinking segment in their own special tokens. Every one of those choices changes the integer sequence, and therefore the logits, and therefore any eval score computed from them. When I insist later that evals pin the chat template, this is why: it is as load-bearing as the sampler settings.

From integers to vectors: the embedding matrix

Once text is a sequence of token IDs $[t_1,\dots ,t_L]$ with each $t_i\in \{0,\dots ,V-1\}$, the model's first learned operation is to look up a dense vector for each ID. That is the embedding matrix $E\in {\mathbb{R}}^{V\times d}$: row $t_i$ is the $d$-dimensional vector fed into the first transformer layer.

At the other end of the network, the model produces a hidden vector $h\in {\mathbb{R}}^d$ for each position and must turn it into a distribution over the $V$ possible next tokens. It does so with an output projection (the "unembedding" or lm_head) $U\in {\mathbb{R}}^{V\times d}$, computing logits $\ell =Uh\in {\mathbb{R}}^V$, one score per vocabulary entry, which the softmax then normalizes (that is the whole next chapter and the objective chapter). Both $E$ and $U$ are $V\times d$ tables that map between token-space and hidden-space, just in opposite directions.

Weight tying

Since $E$ maps tokens to vectors and $U$ maps vectors to token-logits, and both are $V\times d$, a natural question is whether they should be the same matrix (up to transpose). Weight tying sets $U=E$, so the output logits are $\ell =Eh$ and the input lookup is ${\mathbf{e}}_t^{\top }E$. The argument for it is partly parametric (it removes a full $V\times d$ matrix from the model, which for Qwen3-0.6B's $151,936\times 1024$ is $\approx 0.15$ billion parameters saved, a big fraction of a small model) and partly semantic (a token's input representation and the direction you must point a hidden state to predict that token are plausibly the same geometry). Small models tie almost always, because the embedding is a huge fraction of their parameters and the savings are decisive; large models sometimes untie, because at scale the extra flexibility of a separate output matrix is affordable and slightly helpful. Qwen3's smaller variants tie; whether a given checkpoint ties is a flag in its config.json (tie_word_embeddings), which I will read directly in the model-zoo chapter. Knowing whether tying is on tells you immediately whether the embedding and the lm_head are one tensor or two on disk, which is a memory fact you can check.

Tooling

The tool here is the Hugging Face tokenizers / transformers stack, and specifically the AutoTokenizer interface plus the apply_chat_template method. A loaded tokenizer bundles four things that map straight onto the theory: the merge rules and vocabulary (equation 2.1's learned artifact, usually a fast Rust-backed BPE), the special-token registry (the IDs for the role and boundary markers), the chat template (a Jinja string stored in tokenizer_config.json), and the decode direction (turning IDs back into text). The fast tokenizers are the same tokenizers library exposed to Python, which is why they are quick enough to sit in an eval inner loop.

The two calls that matter for reproducibility are tokenizer(text), which runs plain BPE and returns input_ids, and tokenizer.apply_chat_template(messages, ...), which first renders your role-tagged message list through the model's Jinja template into the exact training-time string and then tokenizes it. The important arguments are add_generation_prompt (whether to append the dangling assistant cue so the model knows to start replying, which you want for generation and not for scoring a completed conversation) and, for reasoning models, enable_thinking or a template-specific flag that controls whether the hidden reasoning segment's special tokens are emitted. Because the template lives inside the checkpoint, using AutoTokenizer.from_pretrained(model_id) is what guarantees you serialize conversations the way that specific model expects. Hand-rolling the string yourself is how you accidentally feed a Qwen model a Llama-shaped prompt and wonder why the eval regressed.

Lab

The goal is to make "chat templates are tokenization" undeniable by measuring it: I will take one fixed prompt, render and tokenize it under three different chat templates, and count exactly how the integer sequences differ. The artifact is a table on disk showing token counts and the specific special tokens each template injects, which is the evidence behind the reproducibility rule.

uv init labs/tokenization
cd labs/tokenization
uv add transformers torch

"""Tokenize one prompt under several chat templates and diff the results.

Artifact: artifacts/template_diff.json plus a printed comparison table.
"""
import json
from pathlib import Path

from transformers import AutoTokenizer

# Three instruction models with genuinely different chat templates.
MODELS = [
    "Qwen/Qwen3-0.6B",
    "meta-llama/Llama-3.2-1B-Instruct",
    "microsoft/Phi-3-mini-4k-instruct",
]

MESSAGES = [
    {"role": "system", "content": "You are a careful math tutor."},
    {"role": "user", "content": "What is 17 * 24? Think step by step."},
]

OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

def analyze(model_id: str) -> dict:
    tok = AutoTokenizer.from_pretrained(model_id)
    rendered = tok.apply_chat_template(
        MESSAGES, tokenize=False, add_generation_prompt=True
    )
    ids = tok.apply_chat_template(
        MESSAGES, tokenize=True, add_generation_prompt=True
    )
    special_ids = set(tok.all_special_ids)
    specials_used = [
        tok.convert_ids_to_tokens(i) for i in ids if i in special_ids
    ]
    return {
        "model": model_id,
        "vocab_size": tok.vocab_size,
        "n_tokens": len(ids),
        "n_special": sum(1 for i in ids if i in special_ids),
        "specials_used": specials_used,
        "rendered_prefix": rendered[:200],
        "first_20_ids": ids[:20],
    }

def main() -> None:
    results = [analyze(m) for m in MODELS]
    (OUT / "template_diff.json").write_text(json.dumps(results, indent=2))

    print(f"{'model':38} {'vocab':>8} {'tokens':>7} {'special':>8}")
    print("-" * 65)
    for r in results:
        print(f"{r['model']:38} {r['vocab_size']:>8} "
              f"{r['n_tokens']:>7} {r['n_special']:>8}")
    print()
    for r in results:
        print(f"{r['model']}")
        print(f"  special tokens injected: {r['specials_used']}")
        print(f"  rendered head: {r['rendered_prefix']!r}\n")

    counts = [r["n_tokens"] for r in results]
    print(f"Token-count spread for the SAME prompt: "
          f"min={min(counts)}, max={max(counts)}, "
          f"delta={max(counts) - min(counts)} tokens")
    print(f"Artifact: {(OUT / 'template_diff.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python compare_templates.py

What you should see. The script writes artifacts/template_diff.json and prints a table where the same two-message prompt produces visibly different token counts across the three models, typically a spread of several tokens, driven entirely by how many special/boundary tokens each template injects and whether it prepends a default system frame. You will see genuinely different special-token strings in each row (Qwen's <|im_start|>/<|im_end|> family versus Llama's <|start_header_id|>/<|eot_id|> family versus Phi's <|system|>/<|end|> family), and the printed rendered head lets you read the exact serialized string each model expects. The closing line reports the token-count delta for identical content, which is the whole point: nothing about the meaning changed, yet the integer sequence the model sees did, and that is sufficient to move a loss or an eval number. Record the exact counts from your run (they depend on the checkpoint revisions you pulled) and keep the JSON as the receipt for why the eval config pins a tokenizer and template, not just a model name.

Attention from first principles

Attention is the operation that lets a token look at other tokens and decide, per query, which ones matter. It is the reason transformers work, and it is the single most important thing to be able to derive rather than invoke, because every later chapter (the KV cache and its memory, GQA, the roofline of inference) is an argument about the cost of this specific computation. So I am going to build scaled dot-product attention from the ground up: what queries, keys, and values are, why the dot product is the right similarity, why we divide by $\sqrt{d}$ and not by $d$ or nothing, how the softmax turns scores into weights (and what its Jacobian is, because that is what backprop needs), how causal masking makes it autoregressive, and how multi-head attention buys representational capacity for almost free. Then I will implement it in raw PyTorch and check it, number for number, against a real Hugging Face layer.

Theory

The problem attention solves

A token embedding on its own is context-free: the vector for "bank" is the same whether the sentence is about a river or a loan. To build a contextual representation, each token needs to pull in information from the other tokens that are relevant to it, and "relevant" has to be learned and content-dependent, not fixed by position. Attention is a differentiable, content-addressed lookup that does exactly this. Every token emits a query (what am I looking for?), every token exposes a key (what do I offer as a match?) and a value (what information do I carry?), and each token's output is a weighted average of everyone's values, where the weights come from how well its query matches each key. It is a soft, learnable dictionary lookup: instead of retrieving the one value whose key matches exactly, you retrieve a blend, weighted by match quality.

Queries, keys, values as learned projections

Start from the input to an attention layer: a sequence of $L$ token vectors stacked into a matrix $X\in {\mathbb{R}}^{L\times d_{\text{model}}}$, one row per position. Attention does not use $X$ directly; it first projects it into three separate spaces with three learned weight matrices:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V,\text{\hspace{1em}(3.1)}$$

with $W_Q,W_K\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$ and $W_V\in {\mathbb{R}}^{d_{\text{model}}\times d_v}$, giving $Q,K\in {\mathbb{R}}^{L\times d_k}$ and $V\in {\mathbb{R}}^{L\times d_v}$. The point of three separate projections is that the same token plays three different roles: the query space is tuned for asking, the key space for being found, the value space for what actually gets copied. These are the only learned parameters in attention (plus an output projection I will add shortly); everything else is fixed arithmetic.

Scaled dot-product attention, derived

The similarity between query $i$ and key $j$ is their dot product $q_i\cdot k_j={\sum }_{m=1}^{d_k}{q}_{im}{k}_{jm}$, which is large and positive when the two vectors point the same way and near zero when they are orthogonal. Stacking all pairs gives the score matrix $Q{K}^{\top }\in {\mathbb{R}}^{L\times L}$, entry $(i,j)$ being how much token $i$'s query matches token $j$'s key. We convert each query's row of scores into a probability distribution over the $L$ keys with a softmax, then use those probabilities to average the values. The whole operation is one line:

$$\mathrm{Attn}(Q,K,V)=\underset{A\in {\mathbb{R}}^{L\times L}}{\underbrace{\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)}}V.\text{\hspace{1em}(3.2)}$$

Equation (3.2) is the whole mechanism, and the row-stochastic matrix $A$ (each row sums to 1, all entries nonnegative) is the attention matrix: $A_{ij}$ is the fraction of its output that token $i$ pulls from token $j$'s value. The output is $AV\in {\mathbb{R}}^{L\times d_v}$, each row a convex combination of value vectors. The only piece not yet justified is that $\sqrt{d_k}$ in the denominator, and it is not cosmetic.

The softmax and its Jacobian

The softmax turns a row of scores $z\in {\mathbb{R}}^L$ into a probability vector $a\in {\mathbb{R}}^L$ via

$$a_i=\frac{e^{z_i}}{{\sum }_{j=1}^L{e}^{z_j}},\text{\hspace{1em}(3.6)}$$

which is nonnegative and sums to $1$ by construction. To train through attention, backprop needs the derivative of these probabilities with respect to the scores, and the softmax Jacobian is clean and worth deriving once, because it explains both why saturated softmaxes kill gradients and how the attention weights get their gradient.

Causal masking makes it autoregressive

For a language model, token $i$ must predict token $i+1$ using only tokens $1$ through $i$; letting it attend to future tokens would leak the answer. Attention enforces this with a causal mask applied to the scores before the softmax: set $S_{ij}=-\infty$ for all $j>i$, so that after the softmax $A_{ij}=0$ there. Since $e^{-\infty }=0$, the future tokens get exactly zero weight and drop out of every convex combination, while the surviving weights still renormalize to sum to $1$ over the allowed positions:

$$S_{ij}=\frac{q_i\cdot k_j}{\sqrt{d_k}}+M_{ij},M_{ij}=\{\begin{array}{ll}0& j\le i\\ -\infty & j>i\end{array}\text{\hspace{1em}(3.10)}$$

The mask is a fixed lower-triangular pattern, not learned. In practice frameworks add a large negative number (like the dtype's most-negative finite value) rather than literal $-\infty$ to avoid NaNs, but the effect is the same: future is unreachable. This lower-triangular structure is why the KV cache works at inference (each new token only ever attends backward, so past keys and values can be stored and reused), which is the entire subject of the KV-cache chapter in Part II.

Multi-head attention

A single attention head produces one set of weights, one "way of looking." But a token often needs to attend to several things at once for different reasons: the subject it agrees with, the clause it modifies, the earlier mention it refers to. Multi-head attention runs $h$ independent attention operations in parallel, each with its own small projections, then concatenates and mixes the results. Split the model dimension into $h$ heads of size $d_k=d_{\text{model}}/h$ each, and for head $t$:

$${\mathrm{head}}_t=\mathrm{Attn}(X{W}_Q^{(t)},X{W}_K^{(t)},X{W}_V^{(t)}),t=1,\dots ,h,\text{\hspace{1em}(3.11)}$$

then concatenate the $h$ outputs along the feature axis and apply a final output projection $W_O\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{model}}}$:

$$\mathrm{MHA}(X)=[{\mathrm{head}}_1\Vert {\mathrm{head}}_2\Vert \cdots \Vert {\mathrm{head}}_h]W_O.\text{\hspace{1em}(3.12)}$$

The elegance is in the bookkeeping: because each head is size $d_{\text{model}}/h$, the total parameter count of the $h$ heads equals that of a single full-width attention, so multi-head attention buys $h$ different views at essentially the same parameter and compute cost as one wide head. It is capacity for almost free, which is why every model uses it. In implementation the per-head projections of equation (3.11) are fused into single $d_{\text{model}}\times d_{\text{model}}$ matrices $W_Q,W_K,W_V$ and then reshaped into $h$ heads, so on disk you see one big projection per role, not $h$ small ones. When I get to GQA in the next chapter, the whole trick will be to keep $h$ separate query heads but share keys and values across groups of them, purely to shrink the KV cache; everything about why that is even coherent comes from equations (3.11)–(3.12).

Tooling

The tool that embodies attention is, at the bottom, torch.nn.functional.scaled_dot_product_attention (SDPA), the fused primitive that PyTorch exposes and that Hugging Face models call under the hood. It implements exactly equation (3.2) with the causal mask of (3.10), but it does not build the $L\times L$ matrix $A$ in memory when it can avoid it: on supported hardware it dispatches to a FlashAttention-style kernel that computes the softmax-weighted sum in tiles, streaming over the sequence and never materializing the full attention matrix. That matters enormously for the 16GB budget, because the naive $A\in {\mathbb{R}}^{L\times L}$ scales quadratically with sequence length and is the first thing to blow up at long context; the fused kernel keeps attention's memory linear in $L$ even though its compute stays quadratic. I will unpack that roofline in Part II, but the tooling point for now is that SDPA is the single call that a hand-written attention and a production model both ultimately reduce to, which is why I can verify one against the other.

On the Hugging Face side, a model's attention module (for example Qwen3Attention) is a thin wrapper: it holds the fused q_proj, k_proj, v_proj, and o_proj linear layers of equations (3.1) and (3.12), reshapes into heads, applies rotary position embeddings (next chapter) to $Q$ and $K$, and then calls SDPA. Because the wrapper is thin, I can pull the exact projection weights out of a real layer, run the arithmetic myself with the same weights, and expect a numerical match to within floating-point tolerance. That is the whole design of the lab.

Lab

The goal is to implement scaled dot-product multi-head attention in raw PyTorch, straight from equations (3.1)–(3.12), and verify it numerically against a real Qwen3 attention layer by feeding both the same input and the same weights and checking the outputs agree to floating-point tolerance. The artifact is a saved report of the max absolute difference, which is the evidence that "attention" is exactly the arithmetic derived above and nothing hidden.

uv init labs/attention
cd labs/attention
uv add torch transformers

"""Raw-PyTorch scaled dot-product attention, verified against a HF layer.

Artifact: artifacts/attention_check.json with the max abs difference.
"""
import json
import math
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM

MODEL = "Qwen/Qwen3-0.6B"
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

def manual_attention(x, layer, cfg):
    """Reproduce one attention layer's output from raw ops (eqs 3.1-3.12).

    x: (batch, seq, d_model). Uses the layer's own projection weights.
    Note: we skip RoPE here and disable it on the reference too, so this
    isolates the scaled-dot-product-attention math of this chapter. We do
    apply Qwen3's per-head q_norm/k_norm, because the reference layer applies
    them to Q and K before attention and the outputs will not match without it.
    """
    b, s, _ = x.shape
    n_q = cfg.num_attention_heads
    n_kv = cfg.num_key_value_heads
    d_head = cfg.head_dim if hasattr(cfg, "head_dim") else cfg.hidden_size // n_q

    # Qwen3 applies a per-head RMSNorm (q_norm/k_norm) over head_dim to Q and K
    # before attention; replicate it so the manual output matches the reference.
    q = layer.q_norm(layer.q_proj(x).view(b, s, n_q, d_head)).transpose(1, 2)   # (b, n_q, s, d)
    k = layer.k_norm(layer.k_proj(x).view(b, s, n_kv, d_head)).transpose(1, 2)  # (b, n_kv, s, d)
    v = layer.v_proj(x).view(b, s, n_kv, d_head).transpose(1, 2)

    # Expand KV heads to match query heads (GQA; trivial when n_kv == n_q).
    rep = n_q // n_kv
    k = k.repeat_interleave(rep, dim=1)
    v = v.repeat_interleave(rep, dim=1)

    scores = (q @ k.transpose(-2, -1)) / math.sqrt(d_head)         # eq 3.2/3.5
    causal = torch.triu(torch.full((s, s), float("-inf")), diagonal=1)
    scores = scores + causal                                       # eq 3.10
    attn = torch.softmax(scores, dim=-1)                           # eq 3.6
    out = attn @ v                                                 # eq 3.2
    out = out.transpose(1, 2).contiguous().view(b, s, -1)
    return layer.o_proj(out)                                       # eq 3.12

def main() -> None:
    torch.manual_seed(0)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, torch_dtype=torch.float32, attn_implementation="eager"
    )
    model.eval()
    cfg = model.config
    layer = model.model.layers[0].self_attn

    b, s = 1, 12
    x = torch.randn(b, s, cfg.hidden_size)

    with torch.no_grad():
        mine = manual_attention(x, layer, cfg)

        # Reference: call the same layer but neutralize RoPE by passing
        # position_embeddings of (ones cos, zeros sin) so no rotation happens,
        # and pass the same causal mask the manual path applies so both are
        # autoregressive (q_norm/k_norm live inside the layer already).
        d_head = cfg.head_dim if hasattr(cfg, "head_dim") else cfg.hidden_size // cfg.num_attention_heads
        cos = torch.ones(b, s, d_head)
        sin = torch.zeros(b, s, d_head)
        mask = torch.triu(torch.full((s, s), float("-inf")), diagonal=1)
        ref = layer(x, position_embeddings=(cos, sin),
                    attention_mask=mask.view(1, 1, s, s))[0]

    max_abs = (mine - ref).abs().max().item()
    mean_abs = (mine - ref).abs().mean().item()
    report = {
        "model": MODEL,
        "seq_len": s,
        "max_abs_diff": max_abs,
        "mean_abs_diff": mean_abs,
        "match": max_abs < 1e-4,
    }
    (OUT / "attention_check.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))
    print(f"Artifact: {(OUT / 'attention_check.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python verify_attention.py

What you should see. The script writes artifacts/attention_check.json reporting a max absolute difference between my hand-rolled attention and the real Qwen3 layer that is down at the floating-point-noise floor, on the order of $10^{-5}$ to $10^{-6}$ in FP32, with "match": true. That tiny residual is not error in the derivation; it is the accumulated rounding of doing the same real-valued arithmetic in a slightly different operation order (equation 1.8 from the tensors chapter, made visible). The takeaway is concrete and load-bearing for the rest of the book: a production attention layer is exactly equations (3.1) through (3.12), same weights, same softmax, same $\sqrt{d_k}$, and once you have watched your own code reproduce it to five decimal places, attention stops being a black box and becomes a thing you can reason about the cost of. Record your exact max-diff (measured on the baseline machine, record value, date, driver), since it depends on dtype and the transformers version you pinned.

The transformer block

Attention lets a token gather information; it does not, on its own, make a network. The transformer block is the repeating unit that stacks around attention to make a working model: a normalization, the attention sublayer, a residual connection, another normalization, a feed-forward sublayer, another residual. A modern block (Qwen3, Llama, gpt-oss) differs from the 2017 original in four specific, deliberate ways, and this chapter derives all four because each is load-bearing for the memory and speed arguments later. The feed-forward became a gated MLP (SwiGLU). LayerNorm became RMSNorm. Absolute position embeddings became rotary (RoPE), which I derive as an actual rotation. And multi-head attention became grouped-query attention (GQA), whose entire justification is KV-cache pressure on cards exactly like the 16GB one. By the end I will assemble a block from these parts and match it, dimension for dimension, against a real Qwen3 config.

Theory

The residual stream is the backbone

Before the parts, the skeleton. A transformer block does not replace its input with the sublayer output; it adds the sublayer output to its input:

$$x^{\prime }=x+\mathrm{Attn}(\mathrm{Norm}(x)),x^{\prime \prime }=x^{\prime }+\mathrm{MLP}(\mathrm{Norm}(x^{\prime })).\text{\hspace{1em}(4.1)}$$

That "$x+$" is the residual connection, and reading equation (4.1) as an addition rather than a transformation is the single most useful frame for the whole architecture. The input $x$ flows down the stack essentially untouched (the "residual stream"), and each sublayer reads from it, computes something, and writes an increment back by addition. The stream is a running sum; a 36-layer model is 72 sublayers each adding their two cents. This matters for two reasons that recur all book. First, gradients: because the forward pass adds, the backward pass through $x+f(x)$ has Jacobian $I+f^{\prime }$, so the identity term $I$ guarantees a clean gradient path straight down the stream even when $f^{\prime }$ is small, this is why deep transformers train at all, and it is the residual restatement of the "identity term keeps gradients alive" idea from the autograd chapter. Second, the norm placement in equation (4.1) puts Norm inside the sublayer, before attention and before the MLP, never on the residual stream itself. This is pre-normalization, and it is what modern models use because it keeps the residual stream un-normalized and the training stable; the original transformer put the norm after the addition (post-norm) and was famously twitchy to train without warmup.

RMSNorm

Every sublayer's input gets normalized first, to keep activations at a controlled scale so the next matmul does not see wildly varying magnitudes. LayerNorm did this by subtracting the mean and dividing by the standard deviation across the feature dimension, then applying a learned scale and shift. RMSNorm drops the mean-centering and the shift entirely, keeping only the scaling by root-mean-square:

$$\mathrm{RMSNorm}(x{)}_i=\frac{x_i}{\sqrt{\frac{1}{d}{\sum }_{j=1}^d{x}_j^2+ϵ}}\cdot g_i,\text{\hspace{1em}(4.2)}$$

where $g\in {\mathbb{R}}^d$ is a learned per-feature gain and $ϵ$ (around $10^{-6}$) guards the square root. The denominator is the root-mean-square of the feature vector, hence the name. The justification for dropping the mean subtraction is partly empirical (it works as well) and partly efficiency: RMSNorm computes one reduction (a sum of squares) instead of two (mean, then variance), and skips the subtraction, which at the scale of every-token-every-layer is a real saving. There are no bias terms and no re-centering, so on disk an RMSNorm layer is a single $d$-vector of gains, which you will see when I read a config. The important property for later: like all normalizations it is applied per-token across features, so it does not mix information between positions (that is attention's job) and it does not change the sequence length or the memory shape.

The gated MLP: SwiGLU

The feed-forward sublayer is where most of a transformer's parameters live, and its job is a per-token nonlinear transformation: take each token's $d_{\text{model}}$-vector, blow it up to a wide intermediate dimension $d_{\text{ff}}$, apply a nonlinearity, and project back down. The classic version was two matrices with a ReLU or GELU between them, $\mathrm{MLP}(x)=W_{\text{down}}\sigma (W_{\text{up}}x)$. Modern models use a gated variant, SwiGLU, which splits the up-projection into two parallel paths and uses one to gate the other.

RoPE: position as rotation

Attention as I derived it in the last chapter is permutation-equivariant: shuffle the input tokens and the outputs shuffle the same way, because nothing in $Q{K}^{\top }$ knows position. A language model obviously needs to know order, so position information must be injected. Rotary Position Embedding (RoPE) does this in the most elegant way anyone has found: instead of adding a position vector to the embedding, it rotates the query and key vectors by an angle proportional to their position, so that the dot product between a query at position $m$ and a key at position $n$ depends only on their relative offset $m-n$.

GQA: sharing keys and values to shrink the KV cache

The last modern change is the one most tied to my 16GB constraint. At inference, an autoregressive model caches the keys and values of all past tokens so it does not recompute them for every new token (the KV cache, derived fully in Part II). The size of that cache is proportional to the number of key/value heads, and for long contexts it becomes the dominant memory cost, often larger than the weights. Grouped-query attention attacks this directly: keep all $h$ query heads (you want the representational capacity of many query views, from the multi-head argument in the last chapter), but let several query heads share one key head and one value head.

The through-line for these four changes: RMSNorm and SwiGLU are efficiency and capacity refinements to the sublayers, RoPE is a cleaner way to inject position that gives relative-position awareness for free, and GQA is a direct memory optimization aimed squarely at the KV cache. Three of the four are about doing more with less memory, which is the whole ethos of running this on one card.

Tooling

The tool is again Hugging Face transformers, this time the model config and the block module. A model's config.json (loaded as a PretrainedConfig) is the exact parameterization of everything above: hidden_size is $d_{\text{model}}$, intermediate_size is $d_{\text{ff}}$ (the $2/3$-adjusted SwiGLU width from equation 4.4), num_attention_heads is $h$, num_key_value_heads is $h_{kv}$ (equation 4.10's GQA lever), head_dim is $d_{\text{head}}$, num_hidden_layers is $N$, rms_norm_eps is the $ϵ$ of equation (4.2), rope_theta is the base of equation (4.8), and hidden_act names the SwiGLU nonlinearity (silu). The block module (Qwen3DecoderLayer) is equation (4.1) in code: an input RMSNorm, the self-attention with RoPE and GQA, a residual add, a post-attention RMSNorm, the SwiGLU MLP, and another residual add. Because the config is the parameterization, I can instantiate my own block from the same config and check that its parameter shapes and count match the real layer exactly, if my SwiGLU width, my head split, and my norm shapes all agree, I have almost certainly assembled the same block.

Lab

The goal is to assemble a transformer block from the four derived parts (RMSNorm, SwiGLU, RoPE, GQA attention) in raw PyTorch, parameterized entirely by a real Qwen3 config, and verify that its parameter shapes and total count match the reference Qwen3DecoderLayer dimension for dimension. The artifact is a JSON report of every submodule's shape alongside the reference, which proves the block I built is the block Qwen3 ships.

uv init labs/transformer-block
cd labs/transformer-block
uv add torch transformers

"""Assemble a transformer block from parts and match a Qwen3 config.

Artifact: artifacts/block_shapes.json comparing my block to the real layer.
"""
import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoConfig, AutoModelForCausalLM

MODEL = "Qwen/Qwen3-0.6B"
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

class RMSNorm(nn.Module):                                    # eq 4.2
    def __init__(self, d, eps):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d))
        self.eps = eps
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.g

class SwiGLU(nn.Module):                                     # eq 4.3
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.gate = nn.Linear(d_model, d_ff, bias=False)
        self.up = nn.Linear(d_model, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d_model, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class GQAAttention(nn.Module):                               # eqs 3.x + 4.9/4.10
    def __init__(self, cfg):
        super().__init__()
        self.h = cfg.num_attention_heads
        self.h_kv = cfg.num_key_value_heads
        self.d_head = cfg.head_dim
        d = cfg.hidden_size
        self.q_proj = nn.Linear(d, self.h * self.d_head, bias=False)
        self.k_proj = nn.Linear(d, self.h_kv * self.d_head, bias=False)
        self.v_proj = nn.Linear(d, self.h_kv * self.d_head, bias=False)
        self.o_proj = nn.Linear(self.h * self.d_head, d, bias=False)
        # (RoPE tables omitted for the shape-matching lab; see eq 4.5-4.8)

class Block(nn.Module):                                      # eq 4.1
    def __init__(self, cfg):
        super().__init__()
        self.input_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn = GQAAttention(cfg)
        self.post_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp = SwiGLU(cfg.hidden_size, cfg.intermediate_size)

def shape_map(module) -> dict:
    return {n: list(p.shape) for n, p in module.named_parameters()}

def main() -> None:
    cfg = AutoConfig.from_pretrained(MODEL)
    mine = Block(cfg)

    ref_model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float32)
    ref_layer = ref_model.model.layers[0]

    mine_params = sum(p.numel() for p in mine.parameters())
    ref_params = sum(p.numel() for p in ref_layer.parameters())

    report = {
        "model": MODEL,
        "config": {
            "hidden_size": cfg.hidden_size,
            "intermediate_size": cfg.intermediate_size,
            "num_attention_heads": cfg.num_attention_heads,
            "num_key_value_heads": cfg.num_key_value_heads,
            "head_dim": cfg.head_dim,
            "gqa_group_size": cfg.num_attention_heads // cfg.num_key_value_heads,
            "ffn_ratio": round(cfg.intermediate_size / cfg.hidden_size, 3),
            "rope_theta": cfg.rope_theta,
        },
        "my_block_params": mine_params,
        "ref_layer_params": ref_params,
        "my_shapes": shape_map(mine),
        "ref_shapes": shape_map(ref_layer),
    }
    (OUT / "block_shapes.json").write_text(json.dumps(report, indent=2))

    print(f"GQA group size g = {report['config']['gqa_group_size']} "
          f"(KV cache is 1/{report['config']['gqa_group_size']} of MHA)")
    print(f"FFN ratio d_ff/d_model = {report['config']['ffn_ratio']} "
          f"(SwiGLU, ~3x hidden per eq 4.4)")
    print(f"My block params:  {mine_params:,}")
    print(f"Ref layer params: {ref_params:,}")
    print(f"Artifact: {(OUT / 'block_shapes.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python assemble_block.py

What you should see. The script writes artifacts/block_shapes.json listing every parameter tensor of my hand-built block next to the real Qwen3 layer's. The projection shapes line up: q_proj maps hidden_size to num_attention_heads * head_dim while k_proj and v_proj map to the smaller num_key_value_heads * head_dim, and you can read the GQA asymmetry straight off those two shapes, the K/V projections are literally $g$ times narrower than Q, which is equation (4.10) made visible in the weights. The printed summary reports the GQA group size (2 for the 0.6B, meaning half-size KV cache), the FFN ratio near 3 (the $2/3$-adjusted SwiGLU width of equation 4.4, not a clean power of two), and the two parameter totals within a fraction of a percent of each other, the small gap explained by the q/k norms and RoPE buffers I left out. Keep the JSON: it is the receipt that "a Qwen3 layer" is exactly RMSNorm + GQA-attention-with-RoPE + RMSNorm + SwiGLU wired residually, four derivations you can now point at.

The language-modeling objective

Everything up to here has been machinery: tensors, tokens, attention, the block. This chapter is the goal those parts are optimized toward. A language model is trained to do one thing, predict the next token, and the loss that measures how well it does that is cross-entropy, which is not an arbitrary choice but the unique consequence of maximum likelihood. Get this chapter right and a lot of later confusion dissolves, because the training loss, perplexity, the eval log-probabilities, and eventually the RL reward are all readings of the same next-token distribution. So I derive cross-entropy from maximum likelihood, derive perplexity as its exponential, nail down the off-by-one bookkeeping of label shifting, explain why loss curves have the shape they do, and then compute a loss and a perplexity by hand on a real paragraph and match the library to the decimal.

Theory

Next-token prediction as a probability model

A language model defines a probability distribution over sequences of tokens. By the chain rule of probability, any joint distribution over a sequence $t_1,t_2,\dots ,t_T$ factors into a product of conditionals, each one the probability of the next token given all the previous:

$$p_{\theta }(t_1,\dots ,t_T)=\prod _{i=1}^T{p}_{\theta }(t_i\mid t_1,\dots ,t_{i-1}).\text{\hspace{1em}(5.1)}$$

This factorization is exact, not an approximation, and it is why autoregressive next-token prediction is a complete model of language and not a simplification: if you can model every conditional $p_{\theta }(t_i\mid t_{<i})$, you have modeled the whole joint. The transformer's job is to produce each of those conditionals, and it does: for each position $i$, the network outputs a hidden vector, the unembedding turns it into a logit vector ${\ell }^{(i)}\in {\mathbb{R}}^V$, and the softmax turns that into the distribution over the next token,

$$p_{\theta }(t_i=w\mid t_{<i})=\mathrm{softmax}({\ell }^{(i)}{)}_w=\frac{e^{{\ell }_w^{(i)}}}{{\sum }_{v=1}^V{e}^{{\ell }_v^{(i)}}}.\text{\hspace{1em}(5.2)}$$

The causal mask from the attention chapter is what makes equation (5.2) honest: position $i$'s logits depend only on $t_{<i}$, never on $t_i$ or later, so the model genuinely predicts rather than peeks. And because a single forward pass produces the logits at every position simultaneously (that is what the sequence dimension is), one forward pass over a length-$T$ sequence yields all $T$ conditionals at once, which is what makes training efficient.

Cross-entropy is maximum likelihood, derived

We have a model that assigns a probability to any sequence. Training means adjusting $\theta$ so the model assigns high probability to the real text in the training corpus. That principle is maximum likelihood, and cross-entropy loss is exactly what it becomes after a couple of monotone transformations.

The second form of equation (5.5), $-{\ell }_{t_i}+\log {\sum }_v{e}^{{\ell }_v}$, is worth pausing on because that second term, the log-sum-exp, is the log of the softmax denominator, i.e. the log-partition-function. It is what couples the true token's logit to all the others: to lower the loss you can either raise the true token's logit or lower everyone else's, and the log-sum-exp is what makes those two moves trade off. It is also the term that must be computed in a numerically stable way (subtract the max logit first), which is a detail the library handles and a hand computation must respect.

Perplexity

The loss in equation (5.4) is in nats (natural-log units), which is fine for optimization but unintuitive as a report. Perplexity is the standard human-facing metric, and it is just the loss exponentiated back out of log-space.

Label shifting: the off-by-one that everyone hits

There is a bookkeeping subtlety that trips up every from-scratch implementation, and it follows directly from equation (5.2): the logits at position $i$ predict the token at position $i+1$, not the token at position $i$. So to build the (prediction, target) pairs for the loss, you shift. Given an input sequence of token IDs input_ids of length $T$, the model produces logits of shape $(T,V)$, and the correct pairing is:

$$\text{prediction at position }i⟷\text{target }{t}_{i+1},i=1,\dots ,T-1.\text{\hspace{1em}(5.9)}$$

In practice this means the logits are sliced to drop the last position (it predicts a token beyond the sequence, which has no label) and the labels are sliced to drop the first token (nothing predicts it), so you compute the loss over $T-1$ pairs: logits[:-1] against labels[1:]. Getting this shift wrong by one position is the single most common silent bug in a custom training loop; it does not crash, it just trains the model to predict the current token from itself, which produces a suspiciously low loss that collapses to nonsense. The Hugging Face models do this shift internally when you pass labels, which is convenient but hides the mechanism, so the lab does it by hand to make equation (5.9) unavoidable. Padding tokens and, in instruction tuning, the prompt tokens are masked out of the loss by setting their label to a sentinel (-100) that the loss function skips, which is how you train only on the completion and not on the question.

Why loss curves look the way they do

A training loss curve for a language model has a characteristic shape, and every part of it is explainable from the objective. It starts high, near $\log V$: an untrained model with random weights outputs roughly uniform logits, so by equation (5.5) the initial per-token loss is about $-\log (1/V)=\log V$, which for a vocabulary of $151,936$ is $\log (151936)\approx 11.9$ nats. That is a number you can predict before training starts and check on step 0, and it is a great sanity test: if your fresh model's loss is not near $\log V$, something (usually the label shift or the tokenizer) is wrong. The curve then drops fast at first, because the easiest wins are cheap: learning unigram frequencies (some tokens are just common) and basic local structure knocks the loss down quickly. Then it flattens into a long slow decline, because the remaining loss is the genuinely hard, long-range, semantic prediction that improves slowly, and because equation (5.6)'s gradient shrinks as predictions get better. The curve is often plotted on a log-x axis precisely because progress is roughly linear in log-steps over the long tail. Sharp spikes usually mean a bad batch or a learning rate too high (an optimizer step that overshot); a loss that plateaus far above $\log V$'s implied floor for the data means underfitting or a bug. Reading these curves fluently is a skill the whole training half of the book leans on, and it all comes back to equation (5.4).

Tooling

The tool is torch.nn.functional.cross_entropy (and its object form nn.CrossEntropyLoss), which fuses three things the derivation kept separate: the softmax of equation (5.2), the log, and the negative-log-likelihood gather of equation (5.5), all in one numerically stable kernel. It expects raw logits, not probabilities, precisely so it can apply the stable log-sum-exp (subtracting the max logit) internally rather than taking the log of an already-softmaxed number that may have underflowed. It gathers the true-token logit via an integer label tensor (no one-hot ever built, the same gather-versus-one-hot efficiency from the embeddings chapter), supports the -100 ignore index for masking, and returns the mean over unmasked positions by default, which is exactly equation (5.4). On the Hugging Face side, passing labels=input_ids to a causal LM's forward makes it do the label shift of equation (5.9) internally and return .loss; understanding that this convenience is just the shifted cross-entropy is the difference between trusting the number and being able to reproduce it, which is the point of the lab.

Lab

The goal is to compute a language-modeling loss and perplexity fully by hand from raw logits, honoring the label shift of equation (5.9) explicitly, and match the library's fused cross_entropy and the model's own .loss to many decimal places. The artifact is a JSON report showing the hand-computed and library values side by side, which proves the loss is exactly the shifted negative log-likelihood of the derivation and nothing more.

uv init labs/lm-objective
cd labs/lm-objective
uv add torch transformers

"""Compute LM loss and perplexity by hand, matched against the library.

Artifact: artifacts/loss_check.json with hand vs library values.
"""
import json
import math
from pathlib import Path

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen3-0.6B"
PARAGRAPH = (
    "The mitochondrion is the powerhouse of the cell. It generates most of "
    "the cell's supply of adenosine triphosphate, used as chemical energy."
)
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

def main() -> None:
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float32)
    model.eval()

    ids = tok(PARAGRAPH, return_tensors="pt").input_ids  # (1, T)
    T = ids.shape[1]

    with torch.no_grad():
        logits = model(ids).logits  # (1, T, V)

    # --- By hand: shift (eq 5.9), then per-position NLL (eq 5.5), then mean. ---
    shift_logits = logits[0, :-1, :]   # positions 1..T-1 predict next token
    shift_labels = ids[0, 1:]          # tokens 2..T are the targets
    log_probs = F.log_softmax(shift_logits, dim=-1)          # eq 5.2 in log-space
    true_logp = log_probs[torch.arange(T - 1), shift_labels] # gather eq 5.5
    per_token_nll = -true_logp
    hand_loss = per_token_nll.mean().item()                 # eq 5.4
    hand_ppl = math.exp(hand_loss)                           # eq 5.7

    # --- Library: fused cross_entropy on the same shifted tensors. ---
    lib_loss = F.cross_entropy(shift_logits, shift_labels).item()

    # --- Library convenience: model computes the shift + loss itself. ---
    with torch.no_grad():
        model_loss = model(ids, labels=ids).loss.item()

    report = {
        "model": MODEL,
        "num_tokens": T,
        "num_predicted_positions": T - 1,
        "hand_loss_nats": hand_loss,
        "library_cross_entropy": lib_loss,
        "model_dot_loss": model_loss,
        "hand_perplexity": hand_ppl,
        "log_V_floor": math.log(model.config.vocab_size),
        "max_abs_diff_hand_vs_lib": abs(hand_loss - lib_loss),
    }
    (OUT / "loss_check.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))
    print(f"\nA random model would start near log(V) = "
          f"{report['log_V_floor']:.3f} nats; this trained model is far below.")
    print(f"Artifact: {(OUT / 'loss_check.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python loss_by_hand.py

What you should see. The script writes artifacts/loss_check.json where three independently computed numbers agree: my by-hand loss, the library's cross_entropy on the same shifted tensors, and the model's own .loss from passing labels, all matching to roughly six decimal places (max_abs_diff around $10^{-6}$). The hand-computed perplexity is the exponential of that loss per equation (5.7). And the report prints the $\log V$ floor (about 11.9 nats for Qwen3's vocabulary) next to the actual loss, which sits far below it, exactly as a trained model should, that gap between "random-init loss $\approx \log V$" and "trained loss" is the whole of language modeling compressed into two numbers. Keep the JSON as the receipt that the training loss is nothing but the shifted negative log-likelihood of equation (5.4), computed by hand and confirmed. Record your exact loss and perplexity values from the run, since they depend on the checkpoint revision.

Where memory goes: training vs inference

This is the chapter the whole book's constraint hangs on. Sixteen gigabytes is a hard wall, and the difference between a lab that runs and a lab that dies with CUDA out of memory is almost always a memory-accounting mistake, not a compute one. So I am going to account for every byte, in both regimes. Inference memory is short: weights plus the KV cache, and that is essentially it. Training memory is the interesting one, because reverse-mode autograd (from chapter 1) forces you to store activations, and the AdamW optimizer keeps two extra full-size copies of every parameter, so training a model can cost eight to sixteen times what running it costs. I will derive the backprop-through-a-linear-layer rule that makes the activation-storage requirement concrete, write down the memory-accounting formulas, and then do the byte arithmetic for a 4B model in both modes on the RTX 5080, which is exactly the calculation that tells you what you can and cannot fine-tune on this card.

Theory

Inference memory: weights plus KV cache

At inference, with torch.no_grad() on, autograd records nothing, so there are no stored activations to speak of beyond the single layer being computed. The memory is two things. First, the weights, which is just parameter count times bytes-per-element: a model with $P$ parameters in a dtype of $b$ bytes needs $P\cdot b$ bytes resident. Second, the KV cache, the stored keys and values for every past token so the model does not recompute them each step, whose size I derived in the transformer-block chapter (equation 4.9):

$$M_{\text{infer}}=\underset{\text{weights}}{\underbrace{P\cdot b}}+\underset{\text{KV cache}}{\underbrace{2\cdot B\cdot L\cdot N\cdot h_{kv}\cdot d_{\text{head}}\cdot b}}+(\text{small: current activations, CUDA context}).\text{\hspace{1em}(6.1)}$$

The activations at inference are negligible because you only ever hold one layer's worth at a time (the residual stream is a single $B\times L\times d$ tensor passing through), and the CUDA context and kernels are a fixed few hundred megabytes. So inference is dominated by weights, with the KV cache growing linearly in sequence length and becoming the swing factor at long context. This is why the whole of Part II is about serving: on a fixed card, weights are fixed but the KV cache is where you win or lose, and GQA (equation 4.10) was the first lever for it.

Training memory: why it explodes

Training adds three new consumers, and together they are why training costs multiples of inference. The first is gradients: every trainable parameter needs a same-shaped gradient tensor, another $P$ elements. The second is optimizer state, and for AdamW this is the killer. The third is stored activations, forced by reverse-mode autograd. Take them in order, starting with the one that requires a derivation.

Now the accounting. Consider a full fine-tune with mixed-precision AdamW, the standard recipe. Per parameter, the framework holds: an FP32 master copy of the weight (4 bytes, kept in full precision so tiny updates are not lost to BF16's coarse mantissa, per chapter 1's equation 1.8), a BF16 working copy used in the forward and backward passes (2 bytes), a gradient (2 bytes in BF16, sometimes 4), and AdamW's two optimizer moments, the first moment $m$ and the second moment $v$, each an FP32 same-shaped tensor (4 bytes each). That is the famous accounting:

$$\underset{\text{fp32 master}}{\underbrace{4}}+\underset{\text{bf16 weight}}{\underbrace{2}}+\underset{\text{grad}}{\underbrace{2}}+\underset{\text{Adam }m}{\underbrace{4}}+\underset{\text{Adam }v}{\underbrace{4}}=16\text{ bytes per parameter}.\text{\hspace{1em}(6.4)}$$

AdamW is "2 extra copies" in the sense that $m$ and $v$ are each a full parameter-sized FP32 tensor on top of the weights and gradient; those two moments alone are 8 of the 16 bytes. So the fixed, activation-independent cost of full-fine-tuning a $P$-parameter model is $16P$ bytes, versus $2P$ bytes to merely run it in BF16. That factor of eight, before activations, is most of the "8-16x more" this chapter opened with. On top of it sit the stored activations from equation (6.3), whose size scales with batch size, sequence length, and depth:

$$M_{\text{train}}=\underset{\text{weights, grads, Adam }m,v}{\underbrace{16P}}+\underset{\text{stored activations}}{\underbrace{A\cdot B\cdot L\cdot N}}+(\text{CUDA context}),\text{\hspace{1em}(6.5)}$$

where $A$ is the per-token, per-layer activation footprint (a handful of $d$-sized and $d_{\text{ff}}$-sized tensors that must survive to backward). The activation term is the one you can actually control at run time, which is where the next idea comes in.

Gradient checkpointing

Equation (6.5)'s activation term grows with depth $N$ and can dominate for long sequences or large batches, and it is the one term you can trade against compute. Gradient checkpointing (activation checkpointing) is the deal: during the forward pass, do not store the intermediate activations for every layer; store only a sparse set of "checkpoints" (say, one per layer boundary or every $\sqrt{N}$ layers). Then during the backward pass, when equation (6.3) needs the input $X$ to a layer that was not stored, recompute it by re-running the forward pass from the nearest checkpoint. You pay one extra forward pass worth of compute (roughly a 30 percent wall-clock hit) in exchange for cutting the stored-activation memory from $O(N)$ down to $O(\sqrt{N})$ or less. It is the purest expression of the memory-versus-compute tradeoff that reverse-mode autograd sets up: autograd wants to store everything (chapter 1), checkpointing declines to, and recomputes on demand. On a 16GB card it is frequently the difference between a batch size of 1 and a usable batch size, and every training lab in this book turns it on.

The bigger structural escape, which I only flag here and derive in the post-training part, is to not update most of the parameters at all. LoRA and its quantized cousin QLoRA freeze the base weights (so they need no gradient, no master copy, and no Adam moments, collapsing the $16P$ term to $2P$ for the frozen base plus $16\times (\text{tiny})$ for the small adapter) and keep the frozen base in 4-bit (halving even that $2P$). That is how a model whose full fine-tune needs many times the card's capacity still trains on it, and equation (6.4) is exactly the cost structure LoRA is dodging.

Tooling

The tool is PyTorch's CUDA caching allocator and its introspection API, plus the torch.utils.checkpoint module for gradient checkpointing. The allocator does not hand each tensor its own cudaMalloc; it grabs large blocks from the driver and sub-allocates, which is why freed tensors do not immediately return memory to the OS (they return to PyTorch's pool) and why nvidia-smi shows a higher, stickier number than your tensors' live bytes. The functions that tell the truth about your program's own usage are torch.cuda.memory_allocated() (bytes currently held by live tensors), torch.cuda.max_memory_allocated() (the peak, which is what actually has to fit), and torch.cuda.memory_summary() (a full human-readable table broken down by allocation category). The distinction between allocated and reserved is the one to internalize: reserved is what PyTorch has taken from the driver, allocated is what your tensors are using, and the gap is pool headroom, not a leak. For checkpointing, torch.utils.checkpoint.checkpoint wraps a submodule so its activations are dropped and recomputed, and most model libraries expose a gradient_checkpointing_enable() one-liner that flips it on for every block. The whole discipline of training on 16GB is: watch max_memory_allocated, keep it under the card's capacity, and reach for checkpointing, quantization, and LoRA (equations 6.4–6.5) when it climbs too high.

Lab

The goal is to measure the memory accounting of equations (6.1)–(6.5) on a real small model: load it for inference and read the footprint, then run a training step (forward, loss, backward, optimizer step) and watch the peak climb through the predicted multiple as gradients, optimizer state, and stored activations come online. The artifact is a JSON report of measured peaks in each phase, checked against the byte arithmetic.

uv init labs/memory
cd labs/memory
uv add torch transformers

"""Measure inference vs training memory and reconcile with the formulas.

Artifact: artifacts/memory_report.json with per-phase peak allocations.
Run on a CUDA machine; on CPU it reports parameter math only.
"""
import json
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen3-0.6B"
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

def gb(nbytes: int) -> float:
    return round(nbytes / 1e9, 3)

def main() -> None:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tok = AutoTokenizer.from_pretrained(MODEL)
    P = None
    report = {"model": MODEL, "device": device}

    # --- Inference footprint (eq 6.1). ---
    model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
    P = sum(p.numel() for p in model.parameters())
    report["params"] = P
    report["predicted_weights_bf16_GB"] = gb(P * 2)
    report["predicted_full_train_16B_GB"] = gb(P * 16)

    if device == "cuda":
        model = model.to("cuda")
        torch.cuda.reset_peak_memory_stats()
        ids = tok("Explain entropy in one sentence.", return_tensors="pt").input_ids.to("cuda")
        with torch.no_grad():
            model(ids)
        report["measured_inference_peak_GB"] = gb(torch.cuda.max_memory_allocated())

        # --- One training step (eq 6.4/6.5). ---
        model.train()
        opt = torch.optim.AdamW(model.parameters(), lr=1e-5)  # allocates m, v
        torch.cuda.reset_peak_memory_stats()
        out = model(ids, labels=ids)
        out.loss.backward()          # allocates grads + stored activations (eq 6.3)
        opt.step()                   # materializes Adam m, v (eq 6.4)
        report["measured_train_step_peak_GB"] = gb(torch.cuda.max_memory_allocated())
        report["train_to_infer_ratio"] = round(
            report["measured_train_step_peak_GB"]
            / report["measured_inference_peak_GB"], 2)
        print(torch.cuda.memory_summary(abbreviated=True))
    else:
        report["note"] = "No CUDA; reporting parameter arithmetic only."

    (OUT / "memory_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))
    print(f"Artifact: {(OUT / 'memory_report.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python measure_memory.py

What you should see. The script writes artifacts/memory_report.json with the parameter count, the predicted BF16 weight footprint ($2P$) and full-train footprint ($16P$), and the measured peaks. On the GPU, the inference peak lands near the predicted weight bytes plus a little context and KV, while the training-step peak is a clear multiple of it (the train_to_infer_ratio field), because backward brought gradients and stored activations (equation 6.3) online and opt.step() materialized Adam's moments (equation 6.4). Watching that ratio jump from roughly 1 to several in a single script is the entire chapter made empirical: nothing about the model changed, only whether you asked autograd to remember the forward pass and whether an optimizer allocated its state. Record the exact peaks with date and driver (measured on the baseline machine, record value, date, driver); the memory_summary() table printed alongside shows the allocated-versus-reserved split so you can see the caching allocator's headroom directly.

Sampling and decoding

The model gives you a probability distribution over the next token. It does not give you a token. The step that turns that distribution into an actual choice is sampling, and it is where a great deal of eval reproducibility is won or lost, because two people running "the same model on the same benchmark" can get different scores purely because one used greedy decoding and the other sampled at temperature 0.7. Sampling is a set of knobs that reshape the distribution before drawing from it: temperature, top-k, top-p, min-p, repetition penalties. Each one is a small, exact transformation, and I am going to derive the most important (temperature, as a Boltzmann rescaling) and specify the rest precisely, then argue hard that evals must pin every one of these plus the seed, and be honest about where even a pinned seed stops guaranteeing determinism. The lab runs one prompt through a sweep of sampler settings and shows the outputs diverge, which is the whole reproducibility argument made visible.

Theory

Greedy versus stochastic

The simplest decoder is greedy: at each step take the single highest-probability token, $t_i=\arg {\max }_v{p}_{\theta }(v\mid t_{<i})$. It is deterministic (same input, same output, always) and it is what you often want for evals where there is one right answer and you want to measure the model's best single guess. Its weakness is that locally-greedy is not globally-optimal (the highest-probability token now can lead into a low-probability continuation later, the classic garden-path problem), and it produces repetitive, flat text because it never takes a risk. Stochastic decoding instead samples from the distribution, $t_i\sim p_{\theta }(\cdot \mid t_{<i})$, which introduces diversity and, for reasoning models, is often necessary because the interesting solution path is not always the single most probable token at every step. The knobs below all operate on stochastic decoding, reshaping the distribution you sample from. Greedy is the limit of one of them (temperature to zero), which is the first thing to derive.

Temperature as Boltzmann rescaling

Temperature is the master knob, and it comes straight from statistical physics, which is not a coincidence: the softmax is the Boltzmann distribution, and temperature is the same temperature.

Truncation: top-k, top-p, min-p

Temperature reshapes the whole distribution but never zeroes anything out, so at high temperature there is always a small chance of sampling a genuinely bad token from the long tail. The truncation knobs fix that by cutting the tail off before sampling, then renormalizing what remains. They differ in how they decide where to cut.

Top-k keeps the $k$ highest-probability tokens and zeros the rest: sort the distribution, take the top $k$, renormalize over just those. It is simple but rigid, because a fixed $k$ is too permissive when the model is confident (the true answer is in the top 2, but you keep 50) and too restrictive when the model is genuinely uncertain (the mass is spread over 200 plausible tokens but you keep 50).

Top-p (nucleus sampling) fixes the rigidity by cutting at a cumulative probability mass instead of a fixed count. Sort tokens by probability descending, then keep the smallest set whose cumulative probability first exceeds $p$ (say $0.9$), and renormalize: