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

Key Concepts

• Parallelism: Single GPU, Multi-GPU, Multi-Node
• Collective Operations: Broadcast, Scatter, Gather, Reduce, All Gather, Reduce Scatter
• Hardware: PCIe, NVLink, NV Switch, HBM (High Bandwidth Memory)
• Software: Nickel, PyTorch Distributed
• Distributed Training: Data Parallelism (DDP), Tensor Parallelism, Pipeline Parallelism
• Arithmetic Intensity: Maximizing computation relative to data transfer.
• Synchronization: Ensuring all processes are at the same point in execution.

1. Introduction to Multi-GPU and Multi-Node Training

• The lecture focuses on leveraging hardware for faster model training, specifically parallelism across multiple GPUs.
• The goal is to structure computation to avoid data transfer bottlenecks and maintain high arithmetic intensity.
• Last week's lecture covered parallelism within a single GPU, including techniques like fusion and tiling to minimize HBM access.
• This week extends to multi-GPU and multi-node training, requiring model replication and sharding.

2. Memory Hierarchy and Communication

• Memory Hierarchy (Fastest to Slowest):
  • L1 Cache (Single GPU)
  • HBM (Single GPU)
  • NVLink (Between GPUs on the same node)
  • NV Switch (Across nodes)
• Minimizing data transfer is crucial, but the mechanics differ across the hierarchy.
• NVLink and NV Switch optimize communication between GPUs, bypassing the CPU and Ethernet.
• H100 GPUs have 18 NVLink generation 4 connections, providing 900 GB/s bandwidth.
• The cost of accessing HBM is still lower than using NVLink, but new Blackwell architectures are improving NVLink speeds.

3. Collective Operations: Building Blocks for Distributed Programming

• Collective operations are primitives for distributed programming, providing a higher-level abstraction than point-to-point communication.
• Terminology:
  • World Size: Number of devices (e.g., 4 GPUs).
  • Rank: Device ID (0, 1, 2, 3 for four devices).
• Collective Operation Definitions:
  • Broadcast: Copy data from one rank to all others.
  • Scatter: Distribute different data from one rank to different ranks.
  • Gather: Collect data from all ranks onto one rank.
  • Reduce: Combine data from all ranks using an associative and commutative operation (e.g., sum) onto one rank.
  • All Gather: Gather data from all ranks onto all ranks.
  • Reduce Scatter: Reduce data and scatter the result across ranks.
• Key Equivalency: All Reduce = Reduce + All Gather
• Mnemonic:
  • "Reduce" implies an associative and commutative operation.
  • "Scatter" is the inverse of "Gather."
  • "All" means all destinations (devices).

4. Hardware Implementation and Software Libraries

• Traditional Setup: GPUs connected to CPUs via PCIe bus, nodes connected via Ethernet. This introduces overhead.
• Modern Setup: GPUs directly connected via NVLink and NV Switch, bypassing CPU and Ethernet.
• Nickel: Nvidia's Collective Communication Library, translates collective operations into low-level packets for GPU communication. It handles topology discovery and optimizes communication paths.
• PyTorch Distributed: Provides a high-level interface for collective operations in PyTorch. Supports multiple backends (Nickel for GPU, Gloo for CPU). Includes features like FSDP (not used in this class).

5. Examples of Collective Operations in PyTorch

• A utility function is used to run a function in multiple processes using Python's multiprocessing library.
• Each process initializes itself and connects to a single host for coordination.
• torch.distributed.barrier() is used for synchronization.
• All Reduce Example:
  • Each rank initializes a tensor with different values.
  • torch.distributed.all_reduce() sums the tensors across all ranks.
  • The result is that each rank now has the sum of all the original tensors.
• Reduce Scatter Example:
  • An input tensor is created with a dimension equal to the world size.
  • torch.distributed.reduce_scatter() performs a reduction (sum) and scatters the result across ranks.
  • Each rank receives a different portion of the reduced tensor.
• All Gather Example:
  • The output of the reduce scatter is used as input.
  • torch.distributed.all_gather() gathers the tensors from all ranks onto all ranks.
  • Each rank now has a concatenation of all the original tensors.
• The examples demonstrate that Reduce Scatter + All Gather = All Reduce.

6. Benchmarking Collective Operations

• Benchmarking is performed on a single node with four GPUs.
• All Reduce Benchmarking:
  • A tensor of 100 million elements is used.
  • A warm-up run is performed before timing.
  • torch.cuda.synchronize() and torch.distributed.barrier() are used for accurate timing.
  • Bandwidth is calculated based on the number of bytes transferred and the duration.
  • The number of bytes transferred is calculated as element_size * 4 bytes * (world_size - 1) * 2 (factor of 2 for send and receive).
  • Observed bandwidth is around 277 GB/s.
• Reduce Scatter Benchmarking:
  • Similar process to all reduce benchmarking.
  • The number of bytes transferred is calculated as element_size * 4 bytes * (world_size - 1) (no factor of 2).
  • Observed bandwidth is around 70 GB/s.
• The difference in bandwidth between all reduce and reduce scatter is speculated to be due to hardware optimizations for all reduce.

7. Distributed Training Strategies: Data, Tensor, and Pipeline Parallelism

• The lecture provides barebones implementations of data, tensor, and pipeline parallelism using deep MLPs.
• MLPs are used because they are compute-bound, even though they are simple architectures.
• Data Parallelism (DDP):
  • The data is split along the batch dimension.
  • Each rank gets a different slice of the data.
  • Each rank has a copy of the entire model.
  • After the backward pass, torch.distributed.all_reduce() is used to average the gradients across workers.
  • The optimizer updates the parameters as usual.
  • Losses are different across ranks, but gradients are reduced to be the same.
• Tensor Parallelism:
  • The model is split along the hidden dimension.
  • Each rank gets every layer, but only a part of each layer.
  • All data and activations are transferred around.
  • torch.distributed.all_gather() is used to gather activations from all ranks.
  • Requires high interconnect bandwidth.
• Pipeline Parallelism:
  • The model is split by layers.
  • Each rank gets all the data, but different layers.
  • The batch is divided into microbatches to mitigate pipeline bubbles.
  • torch.distributed.send() and torch.distributed.recv() are used for point-to-point communication between ranks.
  • The implementation is naive and does not overlap communication and computation.

8. Jax and TPUs

• Jax allows defining the model and sharding strategy, and the Jax compiler handles the rest.
• Lavanter is a toolkit developed based on Jax.
• Jax provides a higher-level abstraction than PyTorch for distributed training.

9. Conclusion

• The lecture covered various ways to parallelize training, including data, tensor, and pipeline parallelism.
• Each method involves splitting either the model or the data along some dimension.
• Recomputation is a recurring theme, with tradeoffs between recomputing, storing in memory, or communicating.
• Hardware is constantly improving, but models will continue to grow, requiring these parallelization techniques.
• The hierarchical structure of computer systems will always be relevant.