Training at Scale

Spreading one model across thousands of GPUs: data, tensor, and pipeline parallelism, FSDP/ZeRO sharding, activation checkpointing, and the interconnect that decides which of them you can afford.

Explainers

• The Post-Training Pipeline (SFT → RLHF → DPO) — A base model only predicts the next token — it has never been told to be helpful. Post-training is how it learns: supervised fine-tuning on demonstrations, a reward model distilled from pairwise preferences, then RLHF where four models fight in memory on a KL leash — or DPO, which folds the reward model into a single classification loss. We draw the reward curve, animate the four-model dance per batch, watch a policy hack a length-biased reward, and lay SFT/DPO/PPO/GRPO flat on a trade-off table.
• Sharding in JAX — A 70-billion-parameter model does not fit on one accelerator, so JAX spreads each array across many — without you rewriting the math. Name your devices into a mesh, say which array dimension rides which mesh axis with a PartitionSpec, and the compiler places the bytes and inserts the communication. We build it up from one device to a full mesh, watch what jit does with a sharding — propagation, the gaps the compiler fills, the collectives it drops in — weigh shard_map against auto-sharding honestly, and face the Shardy cutover. Drawn, computed, and animated.
• Distributed Training, End to End — A 70-billion-parameter model needs 1.1 TB just to hold its weights, gradients, and Adam states — fourteen times what fits on one GPU. Compute the bytes, then install every fix in order: data parallelism with ring all-reduce, FSDP/ZeRO sharding, tensor and pipeline parallelism with their bubbles, activation checkpointing, and the interconnect that decides which of them you can afford — drawn, computed, and animated.
• Anatomy of a Production JAX LLM Trainer — One training loop, five subsystems that all have to agree: a device mesh, a table of sharding rules, a jitted train step, an asynchronous checkpointer, and a data pipeline that never hands two devices the same example. We build the mesh from four integers, watch a logical axis resolve onto physical hardware, dial a live mesh and shard a toy tensor, then follow one step through jit, one checkpoint through an async write, and one host through a crash — the whole trainer, end to end.