Train Your GPT

A trained powerlifter does not walk into the gym, load the bar with their one-rep max, and grind until failure. They follow a schedule. Warm-up sets to prime the nervous system. Working sets at prescribed weight and reps. Accessory work. A cool-down. Rest days between sessions. A mesocycle that ramps intensity for weeks and then deloads. Skip the warm-up and you pull a hamstring at rep one. Skip the cool-down and tomorrow's session is garbage. Skip the rest days and you overtrain — the body stops adapting and starts breaking down.

Training a GPT is that program. The architecture you built in the last few lessons — embeddings, attention, MLP, residual stream, LayerNorm — is the athlete. The training loop is a single workout. The recipe — optimizer, learning-rate schedule, weight decay, gradient clipping, batch size, precision — is the weeks-long program that turns a twitchy noob into a model that can actually lift. This isn't “run gradient descent for a while.” It's a scheduled progression, and if any block is wrong the whole cycle collapses.

This is the training recipe: the set of choices that determines whether your loss curves are smooth and monotone or whether the run diverges in the first thousand steps. Most of it comes from the GPT-2 paper, Chinchilla, and Karpathy's nanoGPT. None of it is obvious from first principles. All of it matters. We'll walk through every block of the program — warm-up, peak, cool-down, rest, recovery — in rough order of how much damage a bad choice does.

The optimizer that trains every large transformer is AdamW — Adam with decoupled weight decay. Adam keeps a per-parameter running estimate of the gradient's first moment (mean) and second moment (uncentered variance), and uses them to compute a parameter-specific step size. Think of it as a smart spotter: it remembers how each muscle group has been responding across recent reps and scales the load accordingly. Decoupled weight decay is a separate pull toward zero, applied on top of the adaptive step instead of folded into the gradient — this matters because it interacts correctly with the per-parameter learning rate.

Two knobs set the speed: the base learning rate η and the schedule that modulates it over the run. The schedule is what separates a textbook optimizer from one that actually trains a 1B-parameter model. Start too hot and early gradients are enormous — the model hasn't settled into any sensible region of parameter space yet, it's walking in cold — and the loss diverges. Hold peak LR too long and the model orbits the minimum instead of descending into it; the equivalent of grinding working sets for three straight hours without tapering. The standard GPT program: linear warm-up for the first 1–5% of steps, then cosine decay down to 10% of the peak LR for the cool-down.

Four schedules plotted against step. Constant is the naive baseline — walk into the gym, slap 400 pounds on the bar, get under it. Simple, brittle, and your first few reps will be the last. Linear warm-up alone fixes the cold start but leaves you grinding at peak intensity for the whole workout. Cosine decay alone gives a clean cool-down but skips the ramp-up and still blows up at step zero. Warm-up plus cosine is the real program: ramp the bar up so the optimizer's moment estimates can calibrate, hold intensity through the hypertrophy block, then smoothly anneal to let the model refine its form. Every major GPT-style training run since 2020 uses something close to this shape.

Now zoom out from one workout to the whole mesocycle. You have a fixed compute budget — say 10²² FLOPs, which is what a GPT-3-era run burned. How big should the model be (the lifter)? How many tokens should it see (the reps)? This is the scaling law question, and Hoffmann et al. 2022 (the Chinchilla paper) answered it with a loss curve you can fit.

The chart shows loss as a function of compute for several model sizes. The envelope — the lower boundary across all model sizes — is the Chinchilla-optimal curve. For any fixed compute, there is a sweet spot: too small a model and you're compute-bound (a lightweight doing endless reps with no progressive overload), too big and you're data-bound (a heavyweight who walked into the gym once and called it a career). Kaplan et al. 2020 originally suggested you should go bigger than you'd think; Chinchilla 2022 corrected that with a better fit: 1:20 parameters to tokens. GPT-3 (175B params, 300B tokens) was radically under-trained by this standard. Llama 1 (7B params, 1T tokens) was closer. Llama 3 (8B params, 15T tokens) is deliberately over-trained — because tokens are cheap at inference time and parameters are not.

Enough theory. Here is the code — three layers, as usual. First, the AdamW update rule in pure Python so you can see what's happening under the hood — one rep on one parameter. Then a full PyTorch training loop with the whole schedule: warm-up ramp, cosine cool-down, gradient clipping, mixed precision — the workout as actually programmed. Then a taste of distributed training with torch.distributed for multi-GPU runs, where you are effectively running the same program across a whole team of athletes in sync.

One nastier problem you will meet: the loss spike. You are 20k steps in, the schedule has been executing cleanly, loss has been smoothly dropping, and then in the space of 50 steps it doubles and starts climbing. No code change, no data change, no obvious culprit. This is a famous, legitimately-not-fully-understood phenomenon in large transformer training — the equivalent of a lifter who hit every rep for a month suddenly failing a warm-up weight. The standard fixes, in order of effort: (1) lower the learning rate globally by 2×, (2) restart from the pre-spike checkpoint with a lower LR — the ML version of deloading and ramping back up, (3) add gradient clipping if you somehow didn't already, (4) check for a pathological batch of data (very long repeats, or rare tokens), (5) switch to a more stable optimizer variant like LION or Schedule-Free AdamW. In practice (2) solves most spikes — which is why you must checkpoint optimizer state, not just model weights. Without the moment estimates, the restart is cold and the schedule starts over.

One block we haven't talked about yet: rest days. In the gym they're literal — no lifting, recovery only. In training, rest days are eval steps. Every few hundred updates you pause the workout, switch the model to eval(), and run it on a held-out validation set without updating weights. You're checking the progress without adding load. Train loss alone will happily keep dropping while the model starts memorizing its repsrather than learning the movement — that's overtraining, and it's real. The tell is the gap between train loss and validation loss: when train keeps falling and val flattens or rises, you're overfitting, and the program's supposed to end. Cosine decay gives you a natural finish line; eval steps tell you whether to end early. Knowing when to stop is a feature of the schedule, not a bug in the athlete.

What you now have. A full, reproducible training program that works at scale: AdamW with weight decay split across parameter groups, cosine LR with linear warm-up and a real cool-down, gradient clipping, gradient accumulation, bfloat16 mixed precision, eval rest days, and a checkpoint format that can resume a run without regression. Plus the Chinchilla rule of thumb telling you what size model is worth training for your compute budget. This is, no exaggeration, the same schedule the major labs use — the 2024 Llama 3 tech report describes a training loop that differs from the one above mostly in scale.

Next up — code-gpt. We've talked about the program in the abstract, and the optimizer is wired. But there's one piece still off-screen: the model itself, assembled end-to-end. Embedding layer, a stack of transformer blocks, the language-model head, the tied weights, the forward pass that takes a batch of token IDs and returns a loss scalar. That's the athlete the schedule is training — and the next lesson builds it, module by module, from the primitives you already have. After that the full loop becomes runnable, end-to-end, on a GPU you own.