Communication: AllReduce, AllGather, Ring

Collective Operations

Distributed training relies on a small set of collective communication primitives. Understanding each operation's data movement pattern is essential for analyzing distributed training performance.

AllReduce Algorithms

There are several algorithms to implement AllReduce, each with different trade-offs in latency, bandwidth utilization, and scalability.

Ring AllReduce Deep Dive

Ring AllReduce is the most widely used algorithm in deep learning because it achieves optimal bandwidth utilization regardless of the number of workers. It proceeds in two phases.

Step-by-step for 4 GPUs: Each GPU divides its data into 4 chunks. In the scatter-reduce phase (3 steps), each GPU sends one chunk to its successor and receives one from its predecessor, accumulating (summing) received chunks with its local data. After 3 steps, each GPU holds the fully reduced version of one unique chunk. The allgather phase (3 more steps) circulates these reduced chunks until every GPU has all 4 reduced chunks.

Bandwidth Analysis

Communication performance is modeled using the alpha-beta model: the time for a message of size m bytes is T = α + m/β, where α is latency (seconds per message) and β is bandwidth (bytes/second).

Ring AllReduce time:

T_ring = 2(p−1) × α + 2(p−1)/p × N/β

For large messages (N ≫ α × β), the bandwidth term dominates: T ≈ 2N/β. This is independent of p — ring AllReduce scales perfectly with the number of GPUs when bandwidth-bound.

For small messages (many small tensors), the latency term 2(p−1) × α dominates. This is why PyTorch DDP uses gradient bucketing: small gradient tensors are concatenated into large buckets (default 25 MB) before AllReduce, amortizing the per-message latency cost.

# Bandwidth utilization calculation
def allreduce_time(num_gpus, data_bytes, latency_s, bandwidth_bps):
    # Alpha-beta model for ring AllReduce
    alpha_term = 2 * (num_gpus - 1) * latency_s
    beta_term = 2 * (num_gpus - 1) / num_gpus * data_bytes / bandwidth_bps
    return alpha_term + beta_term

# Example: 8 GPUs, 1GB gradients, NVLink
# NVLink: ~600 GB/s bidirectional, ~1μs latency
t_nvlink = allreduce_time(8, 1e9, 1e-6, 600e9)  # ≈ 0.0047s

# Same setup over InfiniBand 200 Gb/s (~25 GB/s)
t_ib = allreduce_time(8, 1e9, 2e-6, 25e9)      # ≈ 0.14s

# 30× slower over IB vs NVLink — this is why
# intra-node topology matters enormously

Network Topology

Modern GPU clusters have a hierarchical network topology with vastly different bandwidths at each level. Understanding this hierarchy is critical for placement decisions and parallelism strategy.

Overlap Strategies

The key to high GPU utilization in distributed training is overlapping communication with computation. If we can transfer gradients while computing the next layer's gradients, communication cost can be hidden.

Gradient bucketing (DDP): PyTorch DDP fires AllReduce for a bucket as soon as all gradients in that bucket are ready. Since backward computes gradients in reverse layer order, the first bucket (last layers' gradients) starts AllReduce while earlier layers are still computing. With proper bucket sizing, up to 60–80% of communication can be overlapped.

# DDP with tuned bucket size for overlap
model = torch.nn.parallel.DistributedDataParallel(
    model,
    device_ids=[local_rank],
    bucket_cap_mb=25,             # Default 25MB — tune for your model
    gradient_as_bucket_view=True, # Memory optimization
    find_unused_parameters=False, # Faster if all params used every step
)

# Smaller buckets → more overlap but more latency overhead
# Larger buckets → less overlap but better bandwidth utilization
# Rule of thumb: start with 25MB, increase if network is slow

Prefetch in FSDP: FSDP performs AllGather for the next layer's parameters while the current layer is computing. This prefetching strategy hides AllGather latency when the compute time exceeds the communication time for one layer's worth of parameters.

Debugging Communication

Distributed communication bugs are among the hardest to debug in ML infrastructure. Here are the essential tools and techniques.

# Enable NCCL debug logging
# Set before launching training
import os
os.environ["NCCL_DEBUG"] = "INFO"        # Shows topology, algorithm selection
os.environ["NCCL_DEBUG_SUBSYS"] = "ALL"  # All subsystems

# Use monitored_barrier for hang detection
import torch.distributed as dist

# This will timeout and report WHICH rank is stuck
dist.monitored_barrier(timeout=timedelta(minutes=5))

# Verify all ranks can communicate
tensor = torch.ones(1).cuda()
dist.all_reduce(tensor)
assert tensor.item() == dist.get_world_size(), "AllReduce failed!"