Fundamental Concepts of Distributed Computing

Training contemporary deep learning models presents significant computational hurdles. Models with billions of parameters might not fit onto a single accelerator's memory, and iterating through terabytes of data on one machine can stretch training times from days to weeks or even months. Distributed computing offers a path forward, enabling the pooling of resources from multiple devices (like GPUs) across one or several machines to tackle these large scale problems. Before implementing specific PyTorch strategies like DistributedDataParallel (DDP) or Fully Sharded Data Parallelism (FSDP), it's essential to grasp the fundamental vocabulary and communication patterns used in these distributed setups.

Core Terminology in Distributed Training

When discussing distributed training, several terms appear frequently. Understanding their precise meaning is important for configuring and debugging distributed jobs.

• Node: Refers to a single computational machine in your setup. This could be a physical server in a rack or a virtual machine instance in the cloud. A node typically contains one or more processing units (CPUs, GPUs).
• Process / Worker: An independent instance of your Python training script running on a node. In typical GPU based training, you often launch one process per GPU to maximize hardware utilization. These processes execute concurrently and need to coordinate.
• Rank: A unique integer identifier assigned to each process participating in the distributed computation. Ranks typically range from 0 to $N-1$, where $N$ is the total number of processes involved. By convention, rank 0 often has special responsibilities, like logging or saving checkpoints, although this isn't a strict requirement.
• Size: The total number, $N$, of processes cooperating in the distributed training job. If you are training across 4 nodes, each with 8 GPUs, and running one process per GPU, the size is $4 \times 8 = 32$.
• Process Group: A defined subset of all processes (the group). By default, all processes belong to a single group. However, PyTorch allows creating subgroups, which is useful for more complex parallelism schemes like hybrid data and model parallelism, where different types of communication might happen among different sets of workers.
• Backend: The underlying communication library that facilitates message passing between processes. PyTorch's torch.distributed package supports several backends:
  • NCCL (NVIDIA Collective Communications Library): The preferred backend for GPU based training on NVIDIA hardware. It's highly optimized for inter GPU communication, both within a node (using NVLink) and across nodes (using network interfaces like InfiniBand or Ethernet).
  • Gloo: A platform agnostic backend that works for CPU based communication and communication between GPUs across different node types or network setups where NCCL might not be optimal or available. It also supports GPUs but is generally slower than NCCL for GPU collectives.
  • MPI (Message Passing Interface): A standard for high performance computing communication. Can be used if your cluster environment is already configured for MPI, but NCCL or Gloo are more common within the PyTorch ecosystem.

Collective Communication Operations

Distributed training relies heavily on collective communication operations, where multiple processes synchronize and exchange data simultaneously. These are the primitives upon which higher level strategies like DDP are built. torch.distributed provides functions for these operations:

• Broadcast (torch.distributed.broadcast): Sends a tensor from a single designated process (the src rank) to all other processes in the group. This is commonly used at the start of training to ensure all workers begin with the exact same initial model parameters.

• Reduce (torch.distributed.reduce): Collects tensors from all processes in the group, applies a specified reduction operation (like SUM, AVG, MAX, MIN), and stores the result on a single destination process (the dst rank).

• All Reduce (torch.distributed.all_reduce): Similar to Reduce, but the final result of the reduction operation is distributed back to all processes in the group. This is the foundation of DDP, where gradients computed independently on each worker are averaged across all workers, ensuring consistent model updates everywhere.

• Scatter (torch.distributed.scatter): Takes a list of tensors on a single source process (src) and distributes one tensor from the list to each process in the group (including itself). The $i$ th tensor in the list goes to the process with rank $i$.

• Gather (torch.distributed.gather): The inverse of Scatter. Each process sends its tensor to a designated destination process (dst), which collects them into a list of tensors ordered by rank.

• All Gather (torch.distributed.all_gather): Similar to Gather, but every process in the group receives the concatenated list of tensors from all other processes. Useful when every worker needs the complete set of results from all other workers, for example, gathering embeddings computed in parallel.

• Reduce Scatter (torch.distributed.reduce_scatter): Performs an element wise reduction (like All Reduce) on a list of input tensors across all processes, and then scatters the reduced results, so each process receives a portion of the final reduced tensor. This can be more efficient than separate Reduce and Scatter operations in some scenarios.

Visualizing these operations can aid understanding. Consider a simple All Reduce operation for summing gradients in a 4 process setup:

Each process computes its local gradient ($G_i$). During the All-Reduce step, these gradients are communicated and summed across all processes. The final summed gradient is then made available back on every process, ready for the optimizer step.

Connecting Concepts to Training Strategies

These fundamental concepts directly map to the distributed training techniques discussed later in this chapter:

• Data Parallelism (DDP): Relies heavily on All-Reduce to average gradients calculated on different data batches across workers. Broadcast is used initially to synchronize model weights.
• Model/Pipeline Parallelism: Involves more targeted point to point communication (using send/recv primitives, not detailed here but part of torch.distributed) or specific collective operations like Scatter and Gather between specific ranks holding different parts of the model or processing different micro batches.
• Fully Sharded Data Parallelism (FSDP): Uses a combination of All-Gather to reconstruct full parameters for the forward/backward pass within a layer, and Reduce-Scatter to average gradients and shard them back efficiently across workers.

Understanding these building blocks, nodes, processes, ranks, backends, and collective communication patterns, provides a solid foundation for effectively implementing and troubleshooting distributed training workflows in PyTorch. With this vocabulary established, we can proceed to examine how PyTorch orchestrates these elements for different parallelization strategies.