Skip to main content
Back to path

Welcome to a Grasp lesson

Create a free account to keep studying and create your own study path.

Lesson 82

Collective Communication in Multi-GPU Training

Recommended resources

Introduction

In our previous lessons, we've delved deep into optimizing performance on a single GPU, culminating in a comparison between the high-level productivity of Triton and the granular control of raw CUDA. This equipped you to wring maximum performance from a single device. However, the relentless growth of AI models, now containing hundreds of billions or even trillions of parameters, means that even the most powerful single GPU is often insufficient for training. The next frontier in high-performance computing for AI lies in efficiently harnessing the power of multiple GPUs working in concert.

This lesson marks our transition from single-device optimization to distributed systems. We will explore how to scale training across multiple GPUs, a process that introduces a new and critical bottleneck: communication. You will learn about the role of specialized libraries designed to manage this communication, focusing on the industry standard, NVIDIA's Collective Communication Library (NCCL). We will dissect the most important operation for distributed training, the All-Reduce, and understand why it is the linchpin of gradient synchronization.

The Challenge of Scaling: Distributed Data Parallelism

When a model can fit on a single GPU but the dataset is massive, the most common strategy to accelerate training is Distributed Data Parallelism (DDP). The core idea is intuitive: if one worker can process a batch of data, multiple workers can process multiple batches of data simultaneously.

The process for DDP unfolds in these steps:

  1. Replication: The model is copied onto each of the N GPUs involved in the training.
  2. Sharding: The global batch of training data is split into N smaller "micro-batches." Each GPU receives its own micro-batch.
  3. Local Computation: Each GPU independently performs a forward pass with its micro-batch to compute predictions and loss. Then, it performs a backward pass to compute the gradients for its local model parameters.

At this point, we arrive at the central challenge of DDP. Each GPU has computed gradients based on a different slice of data. If each GPU were to update its model using only its local gradients, the models would immediately diverge, and the parallel effort would be wasted. For the training to be correct, all model replicas must remain identical at the end of each step. This requires a gradient synchronization step, where all GPUs communicate to ensure they apply the exact same update.

The following video explains this entire workflow, showing how PyTorch's DDP orchestrates this process and where the need for communication arises.

Lecture 17: NCCL

This section from the GPU MODE channel explains the use case for NCCL within the context of Distributed Data Parallel (DDP) training.

Please watch the segment from the DDP workflow. Focus on understanding why, after each GPU computes its local gradients, a communication step is necessary to average them and maintain model consistency across all GPUs.

NCCL: The GPU Communication Backbone

The process of aggregating gradients from all GPUs and distributing the result back is a type of collective operation—an operation that involves a group of communicating processes. While one could orchestrate this using basic cudaMemcpy operations from each GPU to the CPU, then performing the averaging on the CPU, and finally copying the result back to each GPU, this approach would be incredibly inefficient. The multiple transfers across the slow PCIe bus would create a massive communication bottleneck, negating the speedup gained from parallel computation.

This is where the NVIDIA Collective Communication Library (NCCL) becomes essential. NCCL is a highly optimized library that implements collective communication routines specifically for NVIDIA GPUs. It is designed to be "topology-aware," meaning it automatically detects the fastest hardware paths between GPUs and uses them to maximize communication bandwidth.

Frameworks like PyTorch and TensorFlow use NCCL as a backend. When you initialize a distributed process group in PyTorch with backend='nccl', you are instructing it to use NCCL for all its collective communication needs, primarily for the gradient synchronization step we just discussed.

The key advantage of NCCL is that it performs these communications directly between GPU memories whenever possible, avoiding the CPU as a bottleneck. This is especially powerful in systems where GPUs are connected by high-speed interconnects.

A comparison of communication backends. (a) A CPU-centric backend like GLOO uses host CPU memory as a staging area, requiring data to move from GPU to CPU and back. (b) NCCL, when possible, performs communications directly between GPU memories, dramatically reducing latency and maximizing bandwidth.

To achieve these high speeds, NCCL leverages the best available hardware interconnects. While GPUs can communicate over the standard PCIe bus, modern multi-GPU servers often feature NVLink, a proprietary high-bandwidth, low-latency interconnect developed by NVIDIA specifically for GPU-to-GPU communication.

Understanding Multi GPU Communication and Nvidia NCCL for ...

This article from Tensorfuse provides a clear overview of NCCL and the hardware it runs on.

First, read the introduction, "What is NCCL and Why Is It Important?", to solidify your understanding of its role. Then, read the section "Meet NVLink: A Faster Way for GPUs to Talk" to appreciate the hardware-level advantage that NCCL is built to exploit. Note the stark difference in bandwidth between PCIe and NVLink.

The All-Reduce Operation

The specific collective operation used to synchronize gradients in DDP is called All-Reduce. Its function is to take an input buffer from each participating process (GPU), apply a reduction operation (like sum, average, min, or max) to all corresponding elements across the buffers, and ensure the final, reduced result is available on all processes.

A conceptual diagram of an All-Reduce operation with four GPUs. Each GPU starts with a data chunk (A, B, C, or D). The All-Reduce operation performs a summation, and at the end, every GPU holds the final result: A+B+C+D.

In the context of DDP, the process is:

  1. Each GPU provides its locally computed gradient tensor as input.
  2. The All-Reduce operation is called with the SUM or AVERAGE operator.
  3. NCCL executes the operation, summing the gradients from all GPUs.
  4. At the end, each GPU's gradient tensor is overwritten with the averaged result.
  5. The optimizer on each GPU can now perform an identical weight update.

The following video provides a clear, animated explanation of All-Reduce and then places it within a simple, practical PyTorch DDP example.

Lecture 17: NCCL

This video introduces the All-Reduce collective and demonstrates its use in a toy PyTorch example.

First, watch the explanation of the All-Reduce collective. Then, watch the walkthrough of the toy DDP model. You don't need to follow every line of code; focus on the main idea: two GPUs compute local gradients of 0 and 7, respectively. After the backward pass (which triggers the All-Reduce), both GPUs have the correct, averaged gradient of 3.5.

Under the Hood: The Ring All-Reduce Algorithm

A naive implementation of All-Reduce might involve a "parameter server" approach: all GPUs send their data to one root GPU (GPU0), which performs the reduction and then broadcasts the result back to all other GPUs. This is simple but inefficient, as the root GPU and its network links become a bottleneck.

NCCL employs far more sophisticated algorithms. One of the most classic and elegant is the Ring All-Reduce. In this algorithm, GPUs are arranged in a logical ring, and each GPU only communicates with its two immediate neighbors. This avoids a central bottleneck and ensures that the communication load is balanced across all participants. The algorithm proceeds in two main phases: a Scatter-Reduce followed by an All-Gather.

Distributed Training Overview: Scaling PyTorch Across Multiple GPUs

This article by Mandeep Singh provides an excellent, illustrated explanation of the Ring All-Reduce algorithm.

Please read the following three sections to understand the complete process: Start with "Initial State", which sets up the problem. Next, read "Scatter-Reduce Phase". In this phase, chunks of the gradient tensor are passed around the ring. Each GPU adds its local data to the chunk it receives before passing it on. After N-1 steps, each GPU holds one chunk that is the final sum for that portion of the tensor. Finally, read "All-Gather Phase". Now that each GPU has one piece of the final answer, this phase simply circulates those final chunks around the ring so that every GPU ends up with a complete copy of the fully reduced gradient tensor.

This two-phase process cleverly ensures that both the computation (the additions) and the communication are distributed evenly among all GPUs. To see this in motion, watch the following short, animated visualization.

Demystifying NCCL An In depth Analysis of GPU Communication Protocols and Algorithms - Zhiyi Hu

This video from the HOTI Symposium provides a dynamic visualization of the Ring All-Reduce algorithm.

Watch the segment from the animated example. It vividly illustrates the "receive-reduce-send" primitive of the Scatter-Reduce phase, followed by the "receive-copy-send" of the All-Gather phase.

By understanding not just what All-Reduce does, but how it can be implemented efficiently, you are connecting the high-level needs of an AI training loop to the low-level constraints of network topology and bandwidth—a core skill for a GPU kernel engineer working on large-scale systems.

Conclusion

In this lesson, we made the critical leap from single-GPU to multi-GPU systems. We saw that scaling training via Distributed Data Parallelism introduces a communication bottleneck that must be managed. You learned how specialized libraries like NCCL address this challenge by providing highly optimized collective operations that leverage the fastest available hardware interconnects. We dissected the All-Reduce operation, the cornerstone of gradient synchronization in DDP, and even explored the elegant Ring All-Reduce algorithm that powers its efficient execution.

Key Takeaways:

  • Distributed Data Parallelism (DDP) is a standard technique for accelerating training by splitting data across multiple GPUs.
  • DDP necessitates a gradient synchronization step to keep model replicas consistent. This is a collective communication problem.
  • NCCL is the industry-standard library for performing high-performance collective operations on NVIDIA GPUs, using fast interconnects like NVLink to bypass CPU bottlenecks.
  • All-Reduce is the key collective operation for DDP. It performs a reduction (e.g., sum) across all GPUs and distributes the final result to everyone.
  • The Ring All-Reduce algorithm is an efficient, decentralized method for implementing All-Reduce, balancing communication and computation load across all participating GPUs.

You have now seen how the principles of high-performance computing extend from optimizing computation within a single kernel to optimizing communication between a fleet of devices. In our next and final lesson of this module, we will zoom in on a fundamental mechanism that enables this direct GPU-to-GPU communication: Peer-to-Peer (P2P) memory access.