Chapter 16: Training models on multiple GPUs

Introduction

This introduction explains why growing models and datasets can make training on a single GPU infeasible.

• Limits of single-GPU training: Explain why training on one GPU can become infeasible as models and datasets grow, using the LLaMA 3.1 memory example to illustrate the scale of large models.
  Why it matters: This explains the problem that motivates training across multiple GPUs.

Key terms: single GPU, large language models, parameters, inference

16.1 Introduction to parallel programming

Introduces parallel programming as a way to speed up training, illustrates the benefit with concurrent work, and distinguishes threads from processes.

• Parallel tasks and speedup: Explain that parallel programming splits a problem into smaller tasks that run simultaneously, and describe how the section’s toy example contrasts sequential and concurrent execution.
  Why it matters: This is the basic reason to run training work concurrently instead of sequentially.
• Threads versus processes: Distinguish threads from processes by their memory spaces and creation and switching costs, and state why the section says processes are typically used for distributed training.
  Why it matters: The distinction explains why the threaded timing example does not imply that distributed GPU training typically uses threads.

Key terms: parallel programming, sequential, concurrently, threads, processes, memory space, distributed training

16.1.1 Distributed computing terminology

Defines the process, machine, and communication terms used to describe distributed training across GPUs.

• World Size and Rank: Explain that world size counts the processes in a distributed system and that each process has a unique rank from 0 to world size minus 1. Relate process counts to GPU counts under the section’s usual one-GPU-per-process arrangement.
  Why it matters: These terms let you identify how many processes are training and which process is which.
• Nodes and Rank Scopes: Distinguish a node from a process and explain how global rank and local rank identify a process across the whole system versus within one node. State why the chapter can simply say “rank” for its single-machine focus.
  Why it matters: Rank scope becomes important when training extends beyond one machine.
• Process Group Communication: Describe a process group as processes able to communicate with one another, and identify the PyTorch call that creates one.
  Why it matters: It names the set of processes whose communication PyTorch manages.
• SPMD Execution Model: Explain SPMD as multiple processes executing the same program with their own data, and relate that pattern to running the same model on different parts of training data.
  Why it matters: It explains how multiple processes can share a training script while working on different data.

Key terms: world, world size, rank, node, global rank, local rank, process group, torch.distributed.init_process_group, SPMD, Single Program Multiple Data

16.1.2 Hardware requirements

This subsection explains what hardware and operating system are practical for running the chapter’s distributed-training examples.

• GPU and CPU options: Explain that distributed training typically uses multiple GPUs, but the chapter’s examples can also run on a single machine using only CPUs. Distinguish using CPUs to experiment with the concepts from using multiple GPUs to observe a training speedup.
  Why it matters: You can try the distributed-training concepts even if you do not have multiple GPUs.
• Operating system support: State that PyTorch’s distributed APIs have limited support on Windows, so the section recommends Linux or macOS, with WSL as an option for Windows users.
  Why it matters: Choosing a supported environment can help you run the distributed examples.

Key terms: distributed training, multiple GPUs, CPUs, speedup, PyTorch’s distributed APIs, Windows, Linux, macOS, Windows Subsystem for Linux (WSL)

16.1.3 Initializing a distributed program

This subsection explains how to launch processes and initialize communication for a distributed PyTorch program, either manually or with torchrun.

• Spawning distributed processes: Describe how torch.multiprocessing.spawn creates multiple processes that run the same function, and how each process receives its index and the supplied arguments.
  Why it matters: Launching the processes is the first step toward using multiple GPUs.
• Bootstrapping process communication: Explain how the example bootstraps communication with a TCPStore and initializes a process group, including the role of MASTER_ADDR and MASTER_PORT and the optional store argument.
  Why it matters: Starting processes alone is not enough; they also need communication to be initialized.
• Torchrun launch alternative: Describe how torchrun simplifies launching a distributed job compared with manually setting bootstrap environment variables and calling mp.spawn.
  Why it matters: torchrun removes much of the manual launch and bootstrap setup.

Key terms: torch.multiprocessing.spawn, process index, num_processes, torch.distributed, TCPStore, MASTER_ADDR, MASTER_PORT, init_process_group, torchrun, mp.spawn

16.2 Collective communication

This section explains how distributed processes communicate tensor data, which backends carry that communication, and what broadcast and allreduce do.

• Distributed process group communication: Explain that torch.distributed.init_process_group() sets up communication between processes so they can exchange tensor data, and why communication performance matters during training.
  Why it matters: Communication overhead can limit the benefit of distributed training.
• Communication backend choices: Distinguish the section's common communication backends by their stated uses, and explain the performance cost of using gloo for supported GPU-tensor collectives.
  Why it matters: The backend affects where communication happens and can add costly data copies.
• Broadcast and allreduce results: Describe what each process contributes and receives in broadcast versus allreduce, using the section's rank-0 broadcast and summed allreduce examples.
  Why it matters: These collectives provide different ways to share tensor data among distributed processes.

Key terms: distributed process group, torch.distributed.init_process_group(), tensor data, gloo, nccl, mpi, backend, collective communication, broadcast, allreduce

16.3 Introduction to parallelisms

This section explains why GPUs can process separate data shards independently during training but must synchronize gradients before updating copies of the same model.

• Independent passes over data shards: Explain how copies of a model on different GPUs can perform forward and backward passes on separate portions of a batch, and why their losses and gradients differ.
  Why it matters: Independent work lets multiple GPUs process a batch, while the differing gradients explain why another step is needed.
• Gradient synchronization before updates: Explain why separately calling optimizer.step() would make model copies diverge, and describe the gradient-reduction step that keeps them consistent.
  Why it matters: Synchronizing gradients turns independent GPU computations into updates to one consistent model.

Key terms: forward pass, backward pass, shard, gradients, optimizer.step(), gradient reduction, allreduce, synchronized gradients

16.4 Data parallelism

This section explains how data-parallel training synchronizes gradients, how PyTorch’s DistributedDataParallel handles that work, and why each device needs a distinct share of the data.

• Gradient synchronization across devices: Trace a data-parallel training step in which devices process separate portions of data, compute gradients, synchronize those gradients, and then update model parameters.
  Why it matters: The order of gradient synchronization and parameter updates is the central coordination step in data-parallel training.
• DistributedDataParallel model wrapping: Explain why the section uses PyTorch’s DistributedDataParallel instead of relying on its basic manual implementation.
  Why it matters: DDP provides the practical alternative to manually synchronizing gradients in the example.
• Distinct data shards per device: Explain why a data-parallel data loader uses DistributedSampler.
  Why it matters: Synchronizing gradients is not enough if devices repeatedly load the same examples instead of dividing the dataset.

Key terms: data parallelism, loss.backward(), dist.all_reduce, opt.step(), gradients, DistributedDataParallel, DDP, gradient synchronization, DistributedSampler, shard the dataset

16.5 Model parallelism

Explains why models too large for one GPU need model parallelism and how partitions exchange activations and gradients during training.

• Model parallelism motivation: Explain why data parallelism alone cannot train a model that exceeds one GPU’s memory, and how model parallelism addresses that limit.
  Why it matters: This distinction explains why adding GPUs through data parallelism does not, by itself, make an oversized model fit.
• Forward activation handoff: Trace how a batch moves through model partitions and describe the communication between devices in the example.
  Why it matters: Passing activations between partitions lets the divided model produce the same output as the original model.
• Backward gradient communication: Describe what the distributed model must do during the backward pass as well as the forward pass.
  Why it matters: Splitting the forward computation is not enough: training needs gradient information across the model partitions too.

Key terms: data parallel, model parallelism, replicating the model, forward pass, activations, peer-to-peer (P2P) operations, send(), recv(), backward pass, gradients

16.5.1 Pipeline parallelism

Explains how pipeline parallelism overlaps work across model stages and microbatches, its tradeoffs, and the PyTorch setup shown in the example.

• Stage and microbatch overlap: Explain how pipeline parallelism splits a model into device-assigned stages and input into microbatches so different devices can work concurrently, unlike basic model parallelism where devices may sit idle.
  Why it matters: This is the reason pipelining can make better use of multiple devices.
• Pipeline suitability and tradeoffs: Describe why pipeline parallelism suits large language models with repeated transformer layers, and what limits its benefits.
  Why it matters: More concurrent work does not automatically mean better performance if stages communicate heavily or progress unevenly.
• PyTorch pipeline setup: Trace the example’s setup: specify model split points, build a stage for each rank using the pipeline API, and define a schedule for processing microbatches. State the forward-then-backward ordering of the example’s GPipe schedule.
  Why it matters: Split points determine where model stages are formed, while the schedule determines how microbatches move through the pipeline.

Key terms: pipeline parallelism, stages, microbatches, repeated blocks, transformer layers, communication overhead, bottlenecks, stragglers, torch.distributed.pipelining, split_spec

16.5.2 Tensor parallelism

This section explains how tensor parallelism splits parameters within neural-network modules across devices and uses communication to preserve the model’s output.

• Within Module Parameter Splitting: Explain how tensor parallelism divides a model’s parameters across devices, and distinguish its within-module split from pipeline parallelism’s module-level stages.
  Why it matters: The granularity of the split is the key distinction between these two forms of model parallelism.
• Distributed Linear Layer Computation: Trace the section’s linear-layer example: weights and biases are distributed across two devices, each device performs matrix multiplication with the input tensor, and devices exchange intermediate results so the model produces the correct output shape.
  Why it matters: Splitting the work does not eliminate the need for devices to share results.
• Row Wise and Column Wise Splits: State that model weights can be split row-wise along dimension 0 or column-wise along dimension 1, and give the stated aim of minimizing communication time and accelerating training.
  Why it matters: The choice of split is presented as a way to limit the communication cost of tensor parallelism.

Key terms: tensor parallelism, model parallelism, pipeline parallelism, parameters, tensors, linear layer, weights and biases, input tensor, matrix multiplication, intermediate results

16.5.3 Deciding between pipeline and tensor parallelism

Explains how model shape and per-device memory limits guide the choice between pipeline and tensor parallelism.

• Parallelism Choice by Model Shape: Explain why deep, relatively narrow models often suit pipeline parallelism, while shallow models with large layers tend to suit tensor parallelism. State why tensor parallelism is necessary when a layer cannot fit on one device.
  Why it matters: The model’s architecture and memory needs point to different ways of dividing work across GPUs.

Key terms: pipeline parallelism, tensor parallelism, skinnier and deeper, wider and shallower, memory constraints

16.6 n-dimensional parallelism

The section explains why large-scale training combines forms of parallelism and how DeviceMesh organizes communication for a combination of model partitioning and replication.

• Layered parallelism at scale: Explain why n-dimensional parallelism combines strategies rather than relying on one, and how communication costs influence where different strategies are used.
  Why it matters: The mix of strategies affects how efficiently a large model uses many GPUs.
• DeviceMesh communication groups: Describe how DeviceMesh represents a mesh of devices with process groups, using the section’s two-replica, four-partition example to explain which groups partition each replica and which synchronize corresponding weights.
  Why it matters: The mesh lets you specify communication among the right subsets of GPUs instead of treating every GPU as one group.

Key terms: n-dimensional parallelism, tensor parallelism, pipeline parallelism, allreduce, DeviceMesh, device mesh, process groups, ranks, replicate

16.7 Fully sharded data parallelism

This section explains how Fully Sharded Data Parallel (FSDP) shards model state across GPUs, gathers it temporarily for computation, and helps train models too large for ordinary data parallelism.

• Full parameter sharding: Explain what FSDP shards across devices and why this reduces replication when training large models.
  Why it matters: Sharding model state lets multiple GPUs share the memory burden of a large model.
• Temporary weight gathering: Trace how FSDP makes sharded weights available for computation and returns them to a sharded state.
  Why it matters: Temporary gathering allows computations to use needed weights without keeping every weight replicated throughout training.
• FSDP sharding setup: Describe what must be specified when applying FSDP and where this section applies fully_shard in its example.
  Why it matters: Applying sharding to the intended modules and devices is how the example puts FSDP into use.
• FSDP selection tradeoff: Justify when the section recommends FSDP and identify the limitation that may call for other parallelism.
  Why it matters: The choice to shard a model depends on both GPU memory limits and the cost of communication.

Key terms: Fully Sharded Data Parallel, FSDP, full parameter sharding, all_gather, reduce_scatter, forward pass, backward pass, torch.distributed.fsdp, fully_shard, device mesh

16.8 Large language model–specific parallelisms

This section introduces context and expert parallelism as GPU strategies relevant to large language models, without explaining how they work.

16.8.1 Context parallelism

Explains how context parallelism splits long sequences across ranks to reduce training memory while preserving causal attention at a communication cost.

• Context parallelism tradeoff: Explain how context parallelism distributes a long sequence across ranks, preserves autoregressive constraints, and trades synchronization for lower per-rank memory.
  Why it matters: This explains when splitting the context can make long-sequence training practical despite added communication.

Key terms: context parallelism, ranks, causal masking, activation memory, intra-node interconnects

16.8.2 Expert Parallelism

Explains how expert parallelism distributes a mixture-of-experts model across devices and uses sparse routing to scale it.

• Expert networks across devices: Describe how a mixture-of-experts model replaces a feedforward network with specialized expert networks and how expert parallelism places subsets of those experts on different devices.
  Why it matters: This identifies what is being split across devices in expert parallelism.
• Sparse routing and communication: Trace how a router selects a few experts per token, how selected tokens and outputs move between devices, and why this permits larger models without proportional per-device computation or memory growth.
  Why it matters: Sparse selection and token exchange explain both the communication requirement and the scaling benefit.

Key terms: Mixture of Experts (MoE), expert parallelism, expert networks, sparse computation, router, top-k, auxiliary loss, all-to-all communication

16.9 Tying all parallelisms together

The section explains why combining distributed-training strategies requires additional engineering and presents TorchTitan as a practical PyTorch reference.

• Combined Parallelism Challenges: Explain why combining parallelism strategies in a real training environment is harder than using each strategy alone, citing checkpointing, coordination across devices, or resource utilization.
  Why it matters: Combining strategies adds system-level problems beyond implementing each one separately.
• TorchTitan Practical Reference: Describe TorchTitan as a native-PyTorch reference for distributed large-language-model training and state what it provides to help launch a training job.
  Why it matters: TorchTitan gives you a concrete reference for putting distributed-training pieces together.

Key terms: parallelism strategies, checkpointing, resource utilization, communication across devices, TorchTitan, native PyTorch APIs, ready-to-use configurations, tokenizers

16.10 Conclusion

The conclusion connects distributed training’s use of multiple GPUs with gradient synchronization and scaling to larger models and datasets.

• Gradient synchronization and scaling: Explain that collective communication synchronizes gradients across devices in distributed training, which helps scale training to larger models and datasets.
  Why it matters: It connects coordination across GPUs to the reason for using distributed training.

Key terms: collective communication, synchronizing gradients, distributed training

16.11 Exercises

16.12 Summary