Stanford CS336 Language Modeling from Scratch | Spring 2025 | Parallelism 1

Key Concepts

• Data Parallelism: Copying model parameters across GPUs, splitting the batch across GPUs.
• Model Parallelism: Splitting the model parameters across GPUs.
• Activation Parallelism: Splitting the activation memory across GPUs.
• All-Reduce: Collective communication operation where each rank has a piece of data, the data is reduced (e.g., summed), and the result is copied to every rank.
• Broadcast: Collective communication operation where a single rank's data is copied to all other ranks.
• Reduce-Scatter: Collective communication operation where each rank has a piece of data, the data is reduced, and the result is scattered such that each rank receives a different part of the reduced result.
• All-Gather: Collective communication operation where each rank's data is gathered to all other ranks.
• Optimizer State Sharding: Distributing the optimizer state across GPUs to reduce memory footprint.
• Pipeline Parallelism: Splitting the model into stages (layers) and assigning each stage to a GPU, passing activations between GPUs.
• Tensor Parallelism: Splitting the matrix multiplications within a layer across GPUs.
• Sequence Parallelism: Splitting pointwise operations across the sequence dimension.
• Activation Recomputation: Recomputing activations during the backward pass to reduce memory footprint (e.g., FlashAttention).
• Zero Redundancy Optimizer (ZeRO): A family of data parallelism techniques that shard optimizer states, gradients, and parameters across GPUs.
• Fully Sharded Data Parallelism (FSDP): PyTorch's implementation of ZeRO stage 3, where parameters, gradients, and optimizer states are sharded.
• Context Parallelism/Ring Attention: Splitting the computation and activation cost of large attention mechanisms by passing keys and values around different machines in a ring-like fashion.
• Expert Parallelism: Splitting a large MLP into smaller expert MLPs and scattering them across different machines, with sparse activation of experts.
• 3D/4D/5D Parallelism: Combining data, model, and pipeline parallelism (and potentially other forms) to maximize scaling.

1. Introduction: The Need for Multi-Machine Optimization

• The lecture focuses on parallelism across machines to train large models that don't fit on a single GPU.
• Addresses both compute and memory concerns.
• Communication across GPUs is heterogeneous, leading to different parallelization paradigms.
• The lecture covers popular parallelization strategies, how to combine them, and examples of their use in large-scale distributed training.
• The lecture is structured into three parts: networking basics, mapping networking hardware to parallelization strategies, and case studies.
• While GPU scaling is impressive, multi-machine parallelism is necessary for rapidly scaling compute and memory.
• The world's fastest supercomputers rely on exaflops of compute for training large language models.
• Memory is a key constraint, as models with billions of parameters don't fit on a single GPU.

2. Networking Basics and Collective Communication

• GPUs in a cluster are typically in machines with multiple GPUs connected via fast interconnects (e.g., NVLink).
• Communication between machines is slower (e.g., HDR Infiniband).
• This hardware hierarchy impacts model parallelization strategies.
• Collective Communication Operations:
  • All-Reduce: Sums data from all ranks and distributes the result to all ranks. Cost: roughly 2x the total data size.
  • Broadcast: Copies data from one rank to all other ranks. Cost: roughly 1x the total data size.
  • Reduce: Sums data from all ranks and sends the result to one rank.
  • All-Gather: Each rank copies its data to all other ranks.
  • Reduce-Scatter: Data is reduced and scattered, with each rank receiving a different part of the result.
• Equivalence: All-Reduce can be replaced with Reduce-Scatter followed by All-Gather in the bandwidth-limited regime.
• GPU vs. TPU Networking:
  • GPUs: Nodes with multiple GPUs connected by switches, all-to-all connectivity up to 256 GPUs.
  • TPUs: Toroidal mesh, each TPU chip talks to its neighbors.
  • TPU networking is optimized for collective communications.

3. Parallelization Strategies: Data, Model, and Activation Parallelism

• Goal: Linear memory and compute scaling with the number of GPUs.
• Algorithms implemented using collective communication primitives.
• Three main types of parallelism:
  • Data Parallelism: Copy parameters across GPUs, split the batch across GPUs.
  • Model Parallelism: Split the model parameters across GPUs.
  • Activation Parallelism: Split the activation memory across GPUs.

3.1 Data Parallelism

• Naive data parallelism: Split the batch, compute gradients on each GPU, synchronize gradients (all-reduce), update parameters.
• Compute Scaling: Good, each GPU gets B/M examples (where B is batch size and M is the number of GPUs).
• Communication Overhead: Transmit 2x the number of parameters every batch (due to all-reduce).
• Memory Scaling: Poor, every GPU replicates the parameters and optimizer state.
• Memory usage is dominated by optimizer states (e.g., Adam's first and second moment estimates).
• Optimizer State Sharding (ZeRO Stage 1): Split optimizer states across GPUs.
  • Each GPU computes gradients for everything but updates parameters only for the shard it owns.
  • Process:
    1. Each GPU gets a different data point and computes a full gradient.
    2. Reduce-scatter the gradients so each GPU has the gradient information for its parameters.
    3. Each GPU updates its parameters.
    4. All-gather the updated parameters back to all ranks.
  • Communication cost is the same as all-reduce (2x the number of parameters).
  • Zero stage one is "free" in the bandwidth-limited regime.
• Gradient Sharding (ZeRO Stage 2): Shard gradients across machines.
  • As gradients are computed backwards, they are immediately reduced to the GPU that owns the corresponding parameters.
  • Requires synchronizing layer by layer.
  • Communication remains the same.
• Parameter Sharding (ZeRO Stage 3/FSDP): Shard everything (parameters, gradients, optimizer states).
  • Send and request parameters on demand during forward and backward passes.
  • Total communication cost: 3x the number of parameters.
  • Overlapping communication and computation is key to efficiency.
  • Example: Computation graph w1 + w2 * w0 * x = y. The CPU dispatches commands to fetch parameters and perform matrix multiplies. The GPU performs all-gather, forward steps, and backward steps, overlapping communication and computation.
• Data Parallelism Summary:
  • Stage 1: Free, same communication as naive data parallel, shards optimizer state.
  • Stage 2: More overhead, incremental freeing of gradients, same bandwidth consumption.
  • Stage 3: Most involved, 3x communication cost, but can be efficient with clever masking.
  • Data parallelism doesn't care about architecture.
  • FSDP allows fitting larger models by saving memory.

3.2 Data Parallelism Limitations

• Batch size is a critical resource.
• Cannot parallelize greater than the batch size.
• Diminishing returns to batch sizes beyond a certain point.
• Data parallelism alone cannot achieve arbitrarily large parallelism.
• Zero stages 1 and 2 don't scale memory.
• Zero stage 3 can be slow and doesn't reduce activation memory.

3.3 Model Parallelism

• Goal: Scale up memory without changing the batch size.
• Split parameters across GPUs.
• Two types:
  • Pipeline Parallelism
  • Tensor Parallelism

3.3.1 Pipeline Parallelism

• Split the model into layers, assign each layer to a GPU, pass activations between GPUs.
• Naive implementation leads to poor GPU utilization (most GPUs are idle).
• Pipeline: Have a sequence of things that need to be processed by each GPU.
• Micro-batching reduces the bubble size.
• Batch size is a resource: can be used to make the pipeline bubble smaller or for data parallelism.
• Pipelines help save memory and can have good communication properties.
• Pipeline parallelism is often used on slower network links (inter-node or across different racks).
• Batch sizes are key to hiding the size of the bubble.
• Zero Bubble Pipelining/Dualpipe: Split the backward pass into two components: backpropagating activations and computing gradients for the weights. Reschedule the computation of the weights to fill in the bubbles.
• Pipeline parallelism is infrastructurally very complicated.

3.3.2 Tensor Parallelism

• Split matrix multiplications across GPUs.
• Decompose matrices into submatrices and do partial sums.
• Example: MLP with matrix multiply X * A = Y. Split A into A1 and A2. Each GPU gets the same input X and operates on it with A1 or A2. All-reduce to sum the results.
• Requires high-speed interconnects.
• Tensor parallelism is typically applied within a single node (up to 8 GPUs).
• Throughput decreases with increasing tensor parallelism.
• No bubble, relatively low complexity.
• Larger communication overhead than pipeline parallelism.
• Tensor parallelism is used whenever there are low latency, high bandwidth interconnects.

3.4 Activation Parallelism

• Activations are a significant part of memory usage.
• Tensor and pipeline parallelism can linearly reduce most things, but not all activation memory.
• Activation memory per layer: SBH * 34 + 5 * A * S^2 / H (where S is sequence length, B is batch size, H is hidden size, and A is attention heads).
• Tensor parallelism reduces most of the activation memory, but some terms remain (layer norm, dropouts, inputs to attention and MLP).
• Sequence Parallelism: Split pointwise operations across the sequence dimension.
  • Layer norms across different positions in the sequence do not interact.
  • All-gather and reduce-scatter operations are used to synchronize the parallel computations.
• Combining tensor and sequence parallelism allows dividing all activation memory by the number of devices.
• Activation recomputation (e.g., FlashAttention) can further reduce memory usage.

3.5 Other Parallelism Strategies

• Context Parallel/Ring Attention: Split computation and activation cost of large attention mechanisms by passing keys and values around different machines in a ring-like fashion.
• Expert Parallelism: Split a large MLP into smaller expert MLPs and scatter them across different machines, with sparse activation of experts.

4. Comparison of Parallelism Strategies

| Strategy | Overhead per Batch | Memory Scaling | Bandwidth Properties | Consumes Batch Size | | -------------------- | ------------------ | -------------- | -------------------- | ------------------- | | DDP/ZeRO Stage 1 | Yes | No | Reasonable | Yes | | FSDP (ZeRO Stage 3) | No | Yes | Higher | No | | Pipeline Parallel | No | Yes | Depends | Yes | | Tensor Parallel | No | No | Very High | No |

• Memory, bandwidth, compute, and batch size are limited resources.
• The optimal parallelism strategy depends on the ratio of batch size to the number of GPUs.
• If the batch size is too small, the system is communication-bound.
• Mixing FSDP and tensor parallelism can lead to a compute-bound regime.
• If the batch size is large enough, pure FSDP is sufficient.

5. 3D/4D Parallelism: Putting It All Together

• First, fit the model and activations in memory.
• Use tensor parallelism up to the number of GPUs per machine.
• Use FSDP or pipeline parallelism across machines until the model fits in memory.
• Scale the rest of the way with data parallelism.
• If the batch size is small, use gradient accumulation.

6. Examples and Case Studies

• Megatron-LM paper shows linear gains in aggregate flops with careful 3D parallelism.
• Tensor parallel size of 8 is often optimal.
• Activation recomputation enables larger batch sizes.
• Recent language models (OMO, DeepSeek, e, Llama 3, Gemma 2) use various combinations of FSDP, tensor parallelism, pipeline parallelism, expert parallelism, and context parallelism.
• Llama 3's report provides details about their networking and parallelism strategy.
• GPU failures and data corruption are common challenges when training large models.

7. Conclusion

• Scaling beyond a certain point requires multi-GPU, multi-node parallelism.
• There is no single solution; combine different approaches to leverage strengths.
• Simple and interpretable rules of thumb can guide the execution of parallelism in practice.