Communication Overlap — infra.bea.ai

Communication overlap is the technique of running network operations concurrently with GPU compute, hiding latency behind useful work. Done well, communication overhead approaches zero.

1. The Overlap Opportunity

A training step has sequential dependencies:

Forward must complete before backward starts
Backward for layer $L$ must complete before layer $L-1$

But DP gradient sync doesn’t depend on the full backward. As soon as layer $L$ ‘s gradients are ready, AllReduce can start — while backward continues for layers $L-1, L-2, \ldots$ .

2. Gradient Bucketing

Instead of AllReducing each parameter tensor individually (too many small messages), PyTorch groups parameters into buckets (~25 MB default):

Backward computes:  grad(layer 80) → grad(layer 79) → ... → grad(layer 1)

Bucket filling:     [layer 80, 79, 78]  → Bucket A (25 MB) → AllReduce starts
                    [layer 77, 76, 75]  → Bucket B (25 MB) → AllReduce starts
                    ...                                      (while backward continues)

Communication is pipelined with backward compute. By the time backward finishes for layer 1, most AllReduce operations are already complete.

PyTorch DDP Configuration

model = DistributedDataParallel(
    model,
    device_ids=[local_rank],
    bucket_cap_mb=25,              # bucket size (tune for your network)
    gradient_as_bucket_view=True,  # avoid extra memory copy
    find_unused_parameters=False,  # set True only if needed (slower)
)

Bucket Size Tuning

Bucket size	Effect
Too small (1 MB)	Too many AllReduce calls, high latency overhead
Too large (500 MB)	AllReduce starts late, less overlap with backward
Sweet spot (10-50 MB)	Good overlap, reasonable number of operations

3. FSDP Communication Overlap

FSDP has more communication to overlap (AllGather + ReduceScatter per layer):

Forward:
  Layer 1: AllGather params → Compute → Free non-local params
  Layer 2: [AllGather prefetch during Layer 1 compute] → Compute → Free
  ...

Backward:
  Layer L: AllGather params → Compute grads → ReduceScatter grads
  Layer L-1: [AllGather prefetch during Layer L backward] → Compute → ReduceScatter
  ...

Prefetching: Start gathering the next layer’s parameters while the current layer is computing. This requires extra memory (one layer’s params in advance) but hides AllGather latency.

# FSDP2 prefetching
fully_shard(layer, mp_policy=mp_policy)  # prefetch is enabled by default in FSDP2

4. TP + DP Overlap

With TP=8 and DP=32, both communication patterns must be managed:

Communication	When	Can overlap with
TP AllReduce (per layer)	After attention + FFN	Nothing — on critical path
DP AllReduce (per step)	After each bucket fills	Backward compute of earlier layers
PP P2P (per micro-batch)	Between stages	Other micro-batch compute

TP communication cannot be overlapped — the next layer’s forward depends on the AllReduce result. This is why TP overhead is “exposed” and must use NVLink.

DP communication overlaps well — gradient bucketing pipelines AllReduce with backward. Typical overlap: 80-95% of DP AllReduce hidden behind backward.

5. Measuring Overlap Effectiveness

# PyTorch profiler shows overlapped vs exposed communication
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    with_stack=True,
) as prof:
    train_step()

# In the trace: look for NCCL operations running in parallel with compute kernels
# Exposed communication = NCCL operations with no overlapping compute

Target: Exposed communication < 5% of step time.

For communication cost formulas, see Communication Overhead Estimation.

Self-Check

Explain gradient bucketing and why bucket size matters
Describe how backward compute overlaps with DP AllReduce
Explain FSDP’s AllGather prefetching mechanism
Identify which communication (TP vs DP) can be overlapped and which cannot
Use PyTorch profiler to measure exposed vs overlapped communication