Stanford CS336 Language Modeling from Scratch | Spring 2025 | Lecture 8: Parallelism 2

This lecture focuses on practical strategies to scale and speed up training of large language models when a single GPU or naïve setup becomes a bottleneck. It builds on two foundational ideas: data parallelism (splitting the data across devices) and pipeline parallelism (splitting the model across devices). While both approaches increase throughput and enable training larger models, each introduces its own bottleneck: data parallelism requires frequent network-wide gradient aggregation (an all-reduce), and pipeline parallelism suffers from pipeline bubbles, where some devices wait idle while others compute. The goal here is to introduce techniques that directly address these pain points without changing model quality.

The first big tool is gradient accumulation and the closely related idea of virtual batch size. Virtual batch size is the effective batch you want the optimizer to “see,” even if your GPU cannot hold that many samples at once. Gradient accumulation achieves this by processing several smaller mini-batches sequentially, summing their gradients, and then applying a single optimizer update. This lets you enjoy many benefits of large batches—more stable gradients, potentially faster convergence or better final quality—without needing the memory for all samples simultaneously. The price you pay is extra compute time per update, because you must do multiple forward/backward passes before stepping the weights.

Next, the lecture revisits data parallelism in light of gradient accumulation: if all-reduce synchronization is slowing you down, you can accumulate gradients locally for a few mini-batches on each device and then synchronize less often. This approach reduces the number of all-reduces and can significantly speed training on clusters where network bandwidth is limited. Importantly, per-device batch sizes and accumulation steps should be the same across devices to keep the workload balanced and avoid stragglers.

The lecture then returns to pipeline parallelism and its central challenge: pipeline bubbles that happen while a stage waits to receive activations or gradients. One technique to cut bubble time is interleaved pipeline stages. Instead of giving each device one contiguous block of layers (e.g., layers 1–3 on device 1, 4–6 on device 2), you interleave (e.g., device 1 runs layers 1, 5, 9; device 2 runs 2, 6, 10; etc.). This alternation reduces how long devices wait between tasks, improving average utilization. However, it increases communication because activations hop more frequently between devices.

Finally, gradient accumulation is combined with pipeline parallelism. By splitting a large batch into multiple microbatches and streaming them through the pipe, you keep all stages busy while deferring the weight update until after several microbatches. This approach reduces idle time across the pipe but requires memory to store and combine gradients from multiple microbatches. As with all distributed strategies, the best setup depends on your model size, dataset, number and type of devices, network speed, and memory constraints.

This lecture is aimed at learners who already understand basic deep learning, training loops (forward pass, loss, backward pass, optimizer step), and the concepts of gradients and batches. It is appropriate for intermediate students who have seen data and pipeline parallelism at a high level and want to master the practical knobs to squeeze more efficiency out of their hardware. You should be comfortable with the idea that distributed training has both compute and communication components, and that performance is often limited by whichever is slower.

By the end, you will be able to: define and use virtual batch size; implement gradient accumulation to mimic large-batch training within small memory; combine accumulation with data parallelism to reduce all-reduce frequency; explain pipeline bubbles and apply interleaving to reduce idle time; and orchestrate gradient accumulation with pipeline parallelism to keep stages busy while controlling memory. You’ll also be able to choose among these tools based on constraints, and articulate the relationship between model parallelism and pipeline parallelism (pipeline is a specific kind of model parallelism).

The lecture is structured as follows. It starts with a quick review of data parallelism and pipeline parallelism with their bottlenecks: all-reduce overhead and pipeline bubbles, respectively. It then introduces virtual batch size and gradient accumulation, including the key identity B_v = B × A (effective batch equals actual per-device batch times accumulation steps). After that, it shows how accumulation reduces synchronization cost in data parallelism and how interleaving reduces idle time in pipeline parallelism. It closes by combining accumulation with pipeline parallelism to keep the pipe fuller, and ends with practical guidance on choosing the right mix of techniques for your setup and a clarification that pipeline parallelism is a form of model parallelism.