Collective communication lets multiple processes exchange or synchronize tensor data as a group. In distributed training, these operations help devices coordinate model updates, especially when data-parallel replicas must share gradients. section 16.12 section 16.10
Two useful operations illustrate the pattern:
- Broadcast sends data from one source process to every process in the group. For example, rank 0 can distribute a tensor containing 123.0 to all ranks. section 16.2
- Allreduce combines each process’s data with a reduction such as sum or max, then gives the combined result back to every process. Summing rank-based tensors makes the same total available everywhere. section 16.2
- PyTorch provides these operations through torch.distributed. Gloo is CPU-oriented, NCCL is GPU-oriented, and MPI supports distributed training across multiple machines. section 16.2
Broadcast and allreduce
import torch.distributed as dist if rank == 0: tensor = torch.tensor([123.0]) else: tensor = torch.zeros(1) dist.broadcast(tensor, src=0) tensor = torch.tensor([float(rank)]) dist.all_reduce(tensor, op=dist.ReduceOp.SUM) section 16.2
Figure 16.3 Illustration of the broadcast (left) and allreduce (right) collective communication operations
Use the figure to compare broadcast’s one-to-many transfer with allreduce’s combine-and-share pattern. section 16.2
Listing 16.6 PyTorch’s broadcast and all_reduce operations
import torch.distributed as dist
...
def perform_broadcast(rank):
if rank == 0:
tensor = torch.tensor([123.0])
else:
tensor = torch.zeros(1)
dist.broadcast(tensor, src=0) #1
def perform_all_reduce(rank):
tensor = torch.tensor([float(rank)])
dist.all_reduce(tensor, op=dist.ReduceOp.SUM) #2
...
# Out[]
Process 1 performing broadcast
Process 2 performing broadcast
Process 0 performing broadcast
Process 2 has tensor tensor([0.])
Process 1 has tensor tensor([0.])
Process 0 has tensor tensor([123.])
Process 0 received tensor: tensor([123.])
Process 1 received tensor: tensor([123.])
Process 2 received tensor: tensor([123.])
Process 0 performing all_reduce
Process 2 performing all_reduce
Process 1 performing all_reduce
Process 0 has tensor tensor([0.])
Process 1 has tensor tensor([1.])
Process 2 has tensor tensor([2.])
Process 0 received tensor: tensor([3.])
Process 2 received tensor: tensor([3.])
Process 1 received tensor: tensor([3.])
Use the listing to inspect the corresponding torch.distributed calls and observe the tensors held by each process. section 16.2
