Lecture 12: Efficient Training of Large Models

What this lecture covers

This lecture steps outside natural language processing to answer a practical question every student hits during a final project: why does training run out of GPU memory, and what can you do about it? Shikhar Murty covers how numbers are represented on a GPU, how mixed precision training cuts memory use without hurting accuracy, how multiple GPUs coordinate through techniques like DDP and ZeRO/FSDP, and how parameter-efficient fine-tuning methods like LoRA let you adapt a large pretrained model by training only a small number of extra parameters.

It comes after the course's lectures on transformers and pretraining, and is meant to be directly usable for the course's final projects, where students fine-tune large models with limited compute. After watching, you should be able to explain what mixed precision training does and why it needs loss scaling, describe what DDP, ZeRO stages 1 through 3, and FSDP each shard across GPUs, and know when to reach for LoRA instead of full fine-tuning.

Key ideas

• Floating-point formats trade off range and precision: fp32 uses 4 bytes per number with wide range and high precision; fp16 halves memory but has a much smaller range, which can round small gradients to zero and large ones to NaN.
• Mixed precision training: keep fp32 "master weights," run the forward and backward pass in fp16, scale the loss up before the backward pass to keep small gradients from underflowing, then unscale and update the master weights.
• bfloat16 avoids loss scaling: it keeps fp32's exponent range (so no small-number underflow) but sacrifices more precision, and needs no gradient scaler, though it requires newer GPUs like the A100 or H100.
• DDP (distributed data parallel): each GPU holds a full copy of the model, gradients, and optimizer state, processes a different data slice, and synchronizes gradients with an all-reduce operation after each backward pass.
• ZeRO shards optimizer state, gradients, and parameters: stage 1 shards optimizer state, stage 2 also shards gradients, and stage 3 (also called FSDP) also shards the model parameters themselves; stages 1 and 2 save memory at no extra communication cost, while stage 3 adds real communication overhead.
• MPI communication primitives: all-reduce, reduce-scatter, and all-gather are the building blocks; an all-reduce is mathematically equivalent to a reduce-scatter followed by an all-gather, which is why ZeRO stage 1 and 2 save memory "for free."
• Model activations also consume GPU memory, scaling linearly with batch size; none of the sharding techniques covered address this, which is why batch size limits still exist.
• Parameter-efficient fine-tuning (LoRA): instead of updating all of a pretrained model's weights, LoRA freezes them and adds a trainable low-rank update (B * A, with rank much smaller than the weight matrix's dimensions) to specific weight matrices, drastically cutting the number of trainable parameters and storage per task.

Walkthrough

Floating-point representation and why memory runs out (2:07)

The lecture starts from first principles: fp32 numbers use 32 bits, split between sign, exponent (range), and mantissa (precision), and every neural network parameter in fp32 costs 4 bytes of GPU memory. Naively switching everything to fp16 to save memory introduces two problems shown concretely: small gradients underflow to zero because fp16 has a much smaller exponent range, and values like 1.1 round to 1 because of reduced precision. Both problems matter directly for training, since a diagram of real training gradients shows more than half of them collapsing to zero under plain fp16.

Mixed precision training and bfloat16 (6:11)

The fix is to keep an fp32 "master" copy of the model, run the forward and backward pass in fp16 for speed and memory savings, and scale the loss by a large constant (for example, 1,000) before backpropagating so that small gradient values don't underflow to zero. After computing gradients in fp16, they're cast back to fp32, unscaled, and used to update the master weights, which are then copied back into the fp16 model. This works but requires manually tuning the scaling factor. bfloat16 avoids that complexity by keeping fp32's 8-bit exponent (so its range matches fp32 and nothing underflows) while using fewer precision bits, meaning no gradient scaler is needed at all, at the cost of requiring newer GPU architectures such as the A100 or H100.

Multi-GPU basics: distributed data parallel (14:15)

With multiple GPUs, the simplest approach (DDP) splits the dataset across GPUs while keeping a full, synchronized copy of the model, gradients, and Adam-style optimizer state (momentum and variance, which need fp32 storage) on every GPU. Each GPU runs its own forward and backward pass on its data slice, producing different gradients, which are then merged and redistributed to every GPU using an all-reduce operation costing 2 bytes of communication per parameter (since gradients are fp16). The problem with DDP is that this optimizer state, duplicated on every GPU, scales memory poorly.

ZeRO stages 1 and 2: sharding optimizer state and gradients (18:19)

ZeRO (Zero Redundancy Optimizer, from Microsoft's DeepSpeed project) shards state across GPUs instead of duplicating it. In stage 1, only the optimizer state is sharded: each GPU computes the full gradient for its data, a reduce-scatter sends each GPU only the gradient chunk it's responsible for, each GPU updates its own parameter shard, and an all-gather resynchronizes the full parameter set. Because an all-reduce equals a reduce-scatter plus an all-gather, this saves memory with no added communication cost compared to DDP. Stage 2 goes further and shards gradients too, computing and immediately sending each layer's gradient to the responsible worker rather than storing the full gradient vector, again without extra communication overhead.

ZeRO stage 3 (FSDP): sharding the model itself (28:35)

Stage 3, also known as fully sharded data parallel (FSDP), shards the model parameters themselves across GPUs, which is necessary when even the model alone doesn't fit on one GPU. Unlike stages 1 and 2, this isn't free: for each layer, GPUs must all-gather the full parameters before a forward pass, discard them afterward, all-gather them again for the backward pass, compute gradients, and reduce-scatter them to the responsible GPU. The lecture also explains how FSDP overlaps these all-gather operations with computation by prefetching the next layer's parameters while the current layer runs, and notes that how a model is divided into "FSDP units" is architecture-specific, which is why off-the-shelf sharding policies work well for transformers but may need custom tuning for novel architectures.

Activations, and a decision flowchart for fine-tuning (35:40)

The lecture corrects an earlier simplification: GPU memory also holds activations from the forward pass, needed for the backward pass, and this scales linearly with batch size, which none of the sharding techniques address. This motivates a practical checklist for final projects: always use mixed precision (bfloat16 on newer GPUs); if batch size 1 doesn't fit, try ZeRO stage 2 (free) or stage 3 (has overhead but shards the model); if full fine-tuning still doesn't fit even at batch size 1, move to parameter-efficient fine-tuning.

Why parameter-efficient fine-tuning, and LoRA (38:41) / (45:51)

Beyond memory limits, the lecture motivates parameter-efficient fine-tuning with the growing gap between compute demand for state-of-the-art models and global compute capacity, the environmental cost of large training runs, and the risk of concentrating model development among a few well-resourced organizations. Full fine-tuning updates every parameter of a pretrained model and requires storing a full parameter set per task, which becomes impractical for models the size of GPT-3. LoRA instead freezes the pretrained weights and adds a trainable low-rank update, the product of two small matrices B and A, to selected weight matrices (typically the query and value projections in self-attention), scaled by a factor alpha that trades off retaining pretrained knowledge against learning new task-specific behavior. As the rank increases, LoRA approaches full fine-tuning; in practice, a small rank (the lecture suggests starting around 8) with alpha set to 1 already performs close to full fine-tuning while requiring far fewer trainable and storable parameters per task.

Before you watch

• Review the earlier lecture on transformers and self-attention, since LoRA is applied to specific weight matrices inside the attention mechanism.
• Familiarity with how automatic differentiation computes gradients layer by layer (covered in an earlier lecture) helps with the ZeRO stage 2 and FSDP walkthroughs.
• Basic familiarity with the Adam optimizer (momentum and variance terms) is assumed when the lecture discusses optimizer state memory.

Check your understanding

1. Why does naive fp16 training risk gradients underflowing to zero, and how does loss scaling in mixed precision training address this?
2. Why does bfloat16 avoid the need for gradient scaling, and what does it sacrifice compared to fp16?
3. What is sharded in ZeRO stage 1, stage 2, and stage 3 (FSDP), and why do stages 1 and 2 add no communication overhead compared to DDP while stage 3 does?
4. What kind of GPU memory does none of the techniques covered (mixed precision, DDP, ZeRO, FSDP) reduce, and why does it still limit batch size?
5. In LoRA, what does the rank of the update matrices control, and why does a smaller rank still often perform close to full fine-tuning?