Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 8: Parallelism — Annotated Lecture Notes & Key Takeaways

Course Information

Institution: Stanford

Original Course: Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 8: Parallelism

Instructor Bio: This lecture is delivered by **Percy Liang**, Associate Professor of Computer Science at Stanford University and core faculty of the Stanford Institute for Human-Centered Artificial Intelligence (HAI). Percy Liang leads the Stanford Natural Language Processing Group, with research spanning the theoretical foundations and practical systems of language modeling, machine learning alignment, compositional semantics, and NLP evaluation. He received his PhD in Computer Science from the University of California, Berkeley and his BA in Mathematics from Harvard University. His work has been recognized with the NSF CAREER Award, Google Faculty Research Award, and multiple best paper awards at top-tier machine learning and NLP conferences. He is also widely known for creating influential benchmarks and open-source frameworks for language model research and assessment.

Course Description: This second lecture on parallelism covers advanced, hybrid parallelism strategies for training state-of-the-art large language models. It explores combined 3D parallelism (data + tensor + pipeline), sequence parallelism for long context workloads, and expert parallelism for Mixture-of-Experts (MoE) architectures. The session also discusses practical distributed training frameworks, communication optimization techniques, and real-world cluster deployment considerations for training models at scale.

If you missed this one, here's the setup. Percy taught the previous session on the mechanics of collective communication; this lecture takes those primitives and asks what they're for. The instructor's stated goal is close to a knowledge dump — the messy, practical details of how large language models actually get trained across big clusters — and he is upfront that the endpoint is not one clever trick but a stack of them running at once. He announces early that he'll eventually put "4D parallelism" on a slide, then quietly admits it's really more than four. By the time you reach the last ten minutes, you understand why: every dimension you add solves one resource problem and spends another.
 

He also frames the assignment before he teaches any of it. You are expected to look at a network topology and a model and work out the right parallelization strategy, and to implement an FSDP-style wrapper yourself. That framing matters, because most of the lecture is really a set of rules for reasoning about which resource you are spending.
 

Two Bottlenecks, Two Kinds of Wire, and the Data Center as the Unit

 

The justification for all of this is two separate walls. One is raw compute: no single accelerator gets you to the exaflop scale that top machines reach, so you link many of them. The other is memory: models don't fit on one device, so weights, gradients, and optimizer bookkeeping have to be split somehow.
 

The distinction he returns to again and again is intra-node versus inter-node. Inside a box, links are fast enough that you can afford communication-hungry strategies. Across boxes they are not, and the strategies you deploy there have to respect the channel. Everything is discussed at the level of collectives rather than packets, and the one identity he drills on is that an all-reduce costs about the same as a reduce-scatter followed by an all-gather. A student asked why that particular decomposition deserves attention; the answer was refreshingly honest — there are other decompositions, this one just happens to be the one the algorithms below need.
 

Then comes a hardware detour that he labels as optional but clearly enjoys. TPUs are wired as a toroidal mesh: every chip talks to neighbors, neighbors wrap around, and the number of neighbors stays constant no matter how big the system gets. GPUs are wired the opposite way, closer to a fat tree with all-to-all ambition — fast at the bottom, pods above, spine switches above that, and cost and complexity that grow with node count. The consequence, as he tells it, is that TPUs are excellent for predictable neighbor-to-neighbor traffic and dense models with predictable partitions, while GPUs handle the messy, stochastic routing patterns of things like mixture-of-experts.
 

He then notes that Google announced new TPU hardware the very morning of the lecture, and that the training side of it looked much more like a switched all-to-all fabric — he named the higher-level networking layer Virgo — with the explanation that today's models are MoEs and MoE inference routes tokens everywhere. The larger point is the one worth writing down: workloads are now dictating network design, not the other way around.
 

Two more examples follow. Groq is the answer to "if SRAM is so good, why not all SRAM?" — and it works. Huawei's Ascend is the answer to "if all-to-all is good, why not connect everything with fiber?" — and it also works, in the sense that a rack of roughly 384 weaker chips can be stitched together, at the cost of about four times the power of an equivalent Nvidia system. I liked this section for the intuition, but it's the loosest part of the lecture: the power multiple and the rack size are asserted from memory, not sourced, and the argument is essentially "brute force gets you somewhere, elegance gets you somewhere else." Fine as intuition, weak as evidence.
 

The section ends with a line that reframes the whole course: the new unit of compute is not the GPU, it's the data center.
 

Data Parallelism and the ZeRO Ladder

 

Plain data parallelism is the easy case. Take a batch of size B, cut it across M machines, have each compute gradients on its slice, then sum them. He uses plain SGD here rather than Adam, deliberately, to keep the accounting clean. Compute scales fine as long as each device still gets enough examples. Communication costs about two times the parameter count per step. Memory savings: exactly none, since every device holds a full copy of everything.
 

And memory is where things get grim. His rule of thumb is roughly sixteen bytes per parameter, or about five copies of the weights, depending on precision. There are the parameters themselves, a place to accumulate gradients, possibly a higher-precision accumulator, and then — the real culprit — Adam's first and second moments, which may need to be kept in high precision for stability. That optimizer state is the biggest single slice of memory in a naive setup. He shows a bar chart where going from full replication to sharding everything drops per-device memory from 120 to 1.9, and preloads the punchline: you would expect to pay for that in communication, and remarkably, mostly you don't.
 

ZeRO stage one shards only optimizer state. Everyone keeps full parameters and gradients, computes a full gradient locally, then reduce-scatters it so each worker receives just the slice of gradient space it is responsible for updating. After updating, workers all-gather the fresh parameters back. Because reduce-scatter plus all-gather is equivalent to an all-reduce, this has the same communication cost as naive DDP (16:50). The memory saving is genuinely free, in the literal sense that the byte count on the wire is unchanged.
 

Stage two shards gradients as well, which looks impossible at first because you can no longer materialize a full gradient vector. The fix is a systems trick: sweep backwards through the graph, and as soon as a layer's gradient is computed, reduce it to the owning worker and free it. Incremental or all at once, the totals work out the same.
 

Stage three — FSDP, the one you'll implement — shards parameters too. Each device holds only a slice of weights, gradients, and optimizer state at any moment, and materializes what it needs on demand: all-gather the weights for layer i, run the forward, free them; all-gather them again for the backward, reduce-scatter the gradients, free. That's two all-gathers and one reduce-scatter, so about three times the parameter traffic instead of two. Which should be bad, and looks insane when you consider it happens at every layer.
 

The reason it isn't is prefetching. He walks through a figure from the PyTorch FSDP write-up showing several streams — CPU issue, GPU compute, GPU communication — where the all-gather for layer i+1 is issued while layer i is still computing. There are bubbles, but if your network is fast and your compute per layer is large, the communication hides underneath the computation. His practical claim is that FSDP on an A100-class device takes you from not fitting a 7B model at all to fitting something on the order of 50B, and that measured GPU utilization lands close to single-device performance.
 

A good chunk of the question time is worth your attention, because the confusion is predictable. One student asked whether this means taking gradients from the next GPU to compute gradients for the previous one; the answer is that this is pipelining, and not what's happening here. Every device runs the whole model from start to finish. The only difference is that no device holds the whole set of weights at once — it requests and frees them around each computation. Another student asked why per-layer communication isn't multiplied by layer count into something ruinous; the answer is that yes, there are more operations, but each one is tiny compared to all-reducing an entire network.
 

Where I'd push back slightly: "free" is doing a lot of work in this section. Stage one and two really are free by the wire-count argument. Stage three is not — it costs an extra all-gather, and calling it free is a bet that overlap will hold. He does say the condition out loud (enough compute, fast enough network), but the lecture moves on quickly, and if your per-layer compute is small or your interconnect is congested, that bet is exactly the one that fails.
 

Model Parallelism: When Activations, Not Weights, Start Flying Around

 

Before leaving data parallelism he raises the reason it can't be the whole story. Data parallelism spends batch size as its resource, and batch size is not infinite — past a critical batch size, adding another example buys you less progress than spending that example on another optimizer step. There's a real trade-off between idling your hardware with small batches and taking an optimization hit with large ones. Data parallelism also does nothing for activation memory.
 

The conceptual switch with model parallelism is what moves between devices. Under FSDP, weights fly around and the computation is otherwise unchanged. Under model parallelism, layers or matrices live in different places and the intermediate activations are what get shipped. Pipeline parallel cuts along depth, tensor parallel cuts along width, expert parallel shards experts.
 

Pipeline parallel and its bubble

 

The naive picture is depressing and he draws it that way: four devices each holding a quarter of the layers, one active at a time, the rest idle, and the same in reverse on the backward pass. That idle region is the bubble. The fix is microbatching — feed elements through continuously so the stages stay busy — and the utilization ends up governed roughly by the ratio of pipeline stages to microbatches. Bubble shrinks like one over the number of microbatches, so you need a large batch to make pipelining pay. This is the same batch-size budget being spent a second way.
 

Why bother at all? Pipeline traffic is point-to-point and proportional to batch × sequence × hidden, which is usually far smaller than shipping whole parameter matrices, so pipelines are the right tool for the slowest links in your system — across pods, or across data centers. He cites the Megatron paper's parameter sweeps for how utilization behaves, mentions a DeepSeek-style schedule that interleaves forward and backward chunks to shrink the bubble further, and then gets to the idea he clearly finds most elegant.
 

Zero-bubble pipelining comes from noticing that the backward pass does two separable things at every node: propagate partial derivatives further down the graph, and compute the gradient with respect to that node's weights. Only the first is on the critical path — the next stage can't start until the signal arrives. The weight gradient is a leaf; you can do it whenever. So you run the B parts as fast as possible and defer the W parts into whatever gaps appear, which nearly fills the pipeline. It is, by his own admission, much more complicated than anyone would like.
 

Tensor parallel and the width axis

 

Tensor parallel is the same tiling idea that shows up everywhere in GPU kernels: split a matmul into smaller matmuls and combine partial sums. In an MLP block, the up-projection and the attention projections get cut column-wise, the down-projection and attention output get cut row-wise, and the cheap stuff — layer norm, nonlinearities, MoE routers — is simply replicated, because splitting it costs more than it saves.
 

The detail worth internalizing is the forward/backward duality. On the forward pass the entry operation is an identity (just copy the input to each shard) and the exit operation is an all-reduce. On the backward pass that flips. Write tensor parallel without getting this straight and you'll get it subtly wrong.
 

The cost is that every matmul now carries an all-reduce of activation-sized tensors, all the time. He quantifies it as roughly eight times batch × sequence × hidden per matmul, all-reduce rather than point-to-point. So tensor parallel belongs on the fastest links you have, which for GPUs means inside one box — up to the eight devices on NVLink — and performance falls off a cliff once you cross nodes. This is also where the TPU point from the opening pays off: on a mesh with no sharp intra-box boundary, you can push tensor parallel far wider than eight, and TPU people will say so.
 

His comparison table is simple: tensor parallel has no bubble and is conceptually easy, but it's communication-hungry; pipeline parallel is communication-cheap but bubble-prone and batch-hungry. Use tensor parallel where the wires are fast, pipeline everywhere else.
 

Activation memory and sequence parallel

 

The section I found most useful is the one where he stops treating memory as "just parameters." A profiler trace shows the real picture: optimizer state, weights, then a big dynamic hump of activations, with peak usage occurring not at peak activation count but a bit later, once you've started sweeping backwards and still need to hold activations. For larger models at moderate sequence lengths, activations dwarf parameter memory. Any memory strategy that ignores them is incomplete.
 

The accounting he gives for storing everything is about 34 times sequence × batch × hidden, plus a second term involving attention heads and the sequence length squared over hidden size (the transcript renders it awkwardly; it's the quadratic attention-score-and-dropout term, and flash attention or recomputation makes it go away). Tensor parallel divides most of that — the MLP and attention portions — by the tensor parallel degree. The problem is the residue: layer norms, dropout masks, and the residual inputs to attention and MLP blocks don't get divided, because tensor parallel doesn't split those operators. So even at absurd tensor parallel widths you still eat roughly ten times sequence × batch × hidden.
 

Sequence parallel is the patch, and he warns that the name is misleading — context parallel would be the natural name for it. The leftover terms are cheap in compute, so you shard them along the sequence axis instead of the hidden axis, materializing them with all-gathers and reduce-scatters exactly when needed. It's FSDP's pattern applied to activations, with the same forward/backward reversal of gather and scatter. Combined with recomputation to kill the quadratic term, the floor lands at about 34 × sequence × batch × hidden divided by tensor parallel width — a lower bound he suggests memorizing for back-of-envelope fit calculations. A student asked why you don't also recompute the MLP part; the answer is that you can, but running the MLP again in the backward pass is expensive, whereas attention recomputation is cheaper because it's tiled and it's the quadratic cost you most want to avoid.
 

Expert Parallelism and Context Parallelism: Where the Clean Story Frayed

 

MoE models are now standard, so expert parallelism gets its own treatment, and the lecturer presents it as tensor parallel's cousin: you're splitting the FFN across devices, you pay communication, and you get activation savings. For MoE layers, though, the guidance from Megatron is to prefer expert parallel over tensor parallel, for reasons that double as a list of tensor parallel's weaknesses. Cutting matrices too finely leaves you with small matmuls and poor utilization. Routing sparse token activations is easier than moving dense tensor-parallel activations around. And if you have experts anyway, you might as well spread them over the devices you have.
 

He then spends real time on why this is hard in practice rather than on a slide. MoE traffic is a stream of all-to-all dispatches, issued at every MLP, and it's latency-critical because the compute is blocked waiting for tokens to arrive. He points at DeepSeek's DPP library and Nvidia's Hybrid EP as evidence of how deep this goes, and adds a detail I found genuinely striking: the DeepSeek team reportedly dug out undocumented GPU machine-code instructions to shave their communication kernels. The lesson being that frontier efficiency work lives at that level of detail.
 

The genuinely confusing part — and he flags it as messy — is composition. Most strategies combine like building blocks; expert parallel doesn't. In older libraries, the data parallel and expert parallel replicas are the same split, which caps how far expert parallelism can go and constrains how it interacts with tensor parallel. Worse, MoE changes only the MLPs and leaves attention alone, so expert parallel applies to the model unevenly. You want a high tensor parallel degree to cut up attention, but a low one so your expert matmuls stay large. The resolution in recent systems is to decouple the two: attention gets one tensor parallel configuration, MoE layers get another.
 

Context parallel, or ring attention, gets about ninety seconds. It splits very long sequences across accelerators and passes pieces around in a ring, which suits a mesh topology; it's used in long-context extension and in serving. He skips it on the grounds that it overlaps conceptually with what's already been covered.
 

Choosing a Mix: Rooflines, Prescriptions, and What Real Runs Do

 

The synthesis section is where the lecture earns its length. He puts up a table of every strategy with their drawbacks highlighted, and the reason for the highlighting is the thesis: no single strategy dominates. FSDP is elegant but doesn't touch activations and burns global batch size. Tensor parallel cuts activations without touching batch size but demands fast networking. Pipeline parallel is the thing you reach for across slow links. Large systems use many of them at once because each one is the right answer to a different constraint.
 

Then a roofline-flavored argument. You can compute, per layer and per strategy, how much compute you do and how much communication you must move, and plot the ratio against batch size. While compute time exceeds communication time, you can hide communication and stay fully utilized; below that line you're waiting on the network. His worked example: with a per-chip batch of about 2000, FSDP alone sits comfortably in the compute-bound region. As batch size drops, FSDP alone crosses into communication-bound territory, and you add tensor parallel to push the curve back out. Add strategies until you're compute-bound again. That, he says, is all "3D" or "4D" parallelism really means — not a taxonomy, just the practice of stacking strategies until the hardware stays busy.
 

The prescription he distills is short enough to memorize, and he presents it as the practical inversion of everything abstract he'd said before: cut the model up by whatever means necessary until it fits; use tensor or expert parallel of about eight across the fast interconnect inside a box; use pipeline parallel or ZeRO-3 for the rest of the fitting; once it fits, spend every remaining device on data parallel; and if the batch ends up too small, use gradient accumulation. He notes this matches Megatron's own guidance, which reads in reverse order: minimize model parallelism, maximize data parallelism, keep expert and tensor parallel within the NVLink domain, use pipeline parallel to go multi-node, prefer expert parallel for MoEs, and add context parallel for long sequences.
 

He then cites an older Nvidia paper with Stanford involvement (he names Matei and Deepak) whose large-scale sweeps show exactly this pattern: data parallel maxed out, tensor parallel climbing until it hits eight and stopping there, pipeline parallel growing from that point on, and at extreme scale data parallel actually dropping to around six because you need so much tensor and pipeline capacity just to fit the model. Utilization stays remarkably flat even at enormous GPU counts, which is his explanation for why giant data centers — even cross-data-center training — are viable at all. One observation from that paper deserves its own note: doing more work via activation recomputation can improve utilization, because the freed memory converts into batch size, and batch size converts into utilization. Counterintuitive, and he says so.
 

The closing tour of real runs is the most quotable part of the lecture. A 7B OLMo-scale model was trained purely with FSDP, which is his evidence that many models that size need nothing fancier. DeepSeek V1 used data parallel at ZeRO stage one plus tensor, sequence, and pipeline parallel. DeepSeek V3, being an MoE, swapped tensor parallel for expert parallel at 64-way width, grouping eight machines per expert domain. Yi used the classic data-plus-tensor-plus-pipeline combination. Gemma 2 used only FSDP plus tensor and sequence parallel — no pipeline at all — which he offers as the TPU bet in practice: a big enough mesh means you may never need pipelines. Mixtral-class models in Nvidia's Megatron Bridge configs show expert parallel of eight, pipeline parallel of four, and a separate tensor parallel of four applied to attention. Qwen 3 follows the DeepSeek recipe with expert parallel of 32, pipeline parallel of eight, and tensor parallel of two. The pattern he extracts: everyone maximizes data parallel, tensor parallel almost always stays at or below eight, and expert parallel has been allowed to get large — largely, he thinks, because DeepSeek V3 built the infrastructure to make it work.
 

One aside that stuck with me, because it's the kind of thing lectures usually omit: hardware fails. He mentions that during Llama 3 405B training, GPUs failed something like 148 times, so redundancy and checkpointing are as much a part of large-scale parallelism as any sharding scheme.
 

Where I found the argument thinner

 

A few things I'd want a second opinion on. The Gemma 2 example is presented as evidence for the TPU approach, but the lecturer himself says it's unclear to him whether skipping pipelines can scale indefinitely — that's an honest caveat, and it means the example supports "this works at Gemma scale," not "this is the better philosophy." DeepSeek's 64-way expert parallel is described as reusing pipelining tricks to avoid idle periods, and then dropped; we're told it's complicated, not how it works. The Huawei and Groq material in the opening is impressionistic and unsourced. And the summary table's drawbacks are explicitly his own subjective coloring — useful, but not a measurement.
 

My other reservation is structural: the lecture assumes you already have Lecture 7's collective primitives in working memory. The all-reduce equals reduce-scatter plus all-gather identity is reviewed in about thirty seconds and then load-bearing for the entire ZeRO argument. If you're reading these notes without that background, go back and get it first; the "free lunch" result won't make sense otherwise.
 

Final Takeaway

 

The through-line is less a list of techniques than a way of budgeting. Parallelism here means trading four currencies — memory, bandwidth, batch size, and latency — against each other under a specific physical topology, and the strategies are just the exchange rates. Weights or activations, fast links or slow links, bubble or fragmentation: every choice in this lecture is one side of such a trade. What makes the lecture worth the eighty minutes is that it ends in a genuinely simple rule, arrived at after showing you why the space is not simple: fit the model with tensor or expert parallel inside the box, stretch with pipeline parallel or ZeRO-3, then spend everything else on data parallel.
 

Content Disclaimer:

This article is for general reference only and does not constitute professional R&D guidance, production process advice or quality certification. All material performance data has specific test premises; readers should verify parameters against actual equipment and working conditions.