Maximizing DDP Efficiency: Gradient Buckets and Asynchronous Communication in PyTorch

Distributed training often feels slower than it should.

You add more GPUs, expect near-linear speedup — and instead, training stalls. The culprit is usually not computation, but communication.

In PyTorch’s DistributedDataParallel (DDP), gradient synchronization can quickly become the bottleneck. But under the hood, DDP is doing something surprisingly clever to hide this cost.

In this post, we’ll break down two key techniques that make DDP efficient:

• overlapping communication with computation
• gradient bucketing

Source: Pydorch DDP

By the end, you’ll understand not just what DDP does, but why it’s designed this way.

The Problem: Communication Overhead

In naive DDP, after the backward pass, each parameter tensor’s gradient is all-reduced individually across GPUs. For models with millions of parameters, this can mean thousands of small all-reduce operations. Each operation carries an overhead — even if the underlying GPU communication is fast, issuing many small calls serially introduces latency.

for param in model.parameters():
    if param.grad is not None:
        dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
        param.grad /= world_size

Issues:

• Small tensors lead to high overhead because each all-reduce has startup latency.
• Communication only starts after the backward pass completes, leaving GPUs idle while waiting for gradients to sync.

Batch All-Reduce: Reducing Communication Overhead

A simple and highly effective optimization is batching the gradients. Instead of all-reducing each parameter separately, we can concatenate all gradients into a single large tensor, perform a single all-reduce, and then split the tensor back into parameter-shaped pieces:

from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors

grads = [p.grad for p in model.parameters() if p.grad is not None]
flat_grads = _flatten_dense_tensors(grads)
dist.all_reduce(flat_grads, op=dist.ReduceOp.SUM)

for g, p in zip(_unflatten_dense_tensors(flat_grads, grads), grads):
    p.grad.copy_(g / world_size)

Why this helps:

• Fewer communication calls → less per-call overhead.
• Larger messages are more bandwidth-efficient.
• Keeps GPU busy instead of repeatedly launching small kernels.

Buckets and Overlapping Communication

But we can push this idea even further.

So far, we’ve treated gradient synchronization as something that happens after the backward pass. But this assumption is actually unnecessary.

In reality, the backward pass itself is incremental — gradients are produced layer by layer, flowing from the output back to the input. This opens up an important opportunity: we can start communicating gradients before the entire backward pass finishes.

Take a simple 3-layer network as an example:

Layer 3 → grad ready
Layer 2 → grad ready
Layer 1 → grad ready

At any given moment during backpropagation, some gradients are already available while others are still being computed. Waiting for the full backward pass means we are leaving communication bandwidth idle.

In PyTorch, this is made possible by registering a backward hook on each parameter. The hook is triggered the moment a gradient is computed, allowing us to immediately launch an asynchronous communication operation:

def _make_hook(param):
    def hook(grad):
        dist.all_reduce(grad, async_op=True)
    return hook

for p in model.parameters():
    if p.requires_grad:
        p.register_hook(_make_hook(p))

This enables overlap of computation and communication: while the backward pass is computing gradients for later layers, the earlier layers’ gradients are already being averaged across GPUs.

But register_hook is suboptimal:

• register_hook triggers every time a gradient is computed for a tensor, before accumulation.
• If a parameter is used multiple times in the computation graph, the hook might trigger multiple times, leading to duplicate communication. For example:

x = torch.tensor(1.0, requires_grad=True)

def hook(grad):
    print("hook:", grad)

x.register_hook(hook)

y = x * 2 + x * 3
y.backward()

# grad₁ → all-reduce
hook: 2

# grad₂ → all-reduce
hook: 3

Instead, PyTorch provides a handy API register_post_accumulate_grad_hook ensures the hook is triggered after gradients are fully accumulated, making it safe for DDP’s asynchronous communication.

x = torch.tensor(1.0, requires_grad=True)

def hook(grad):
    print("post hook:", grad)

x.register_post_accumulate_grad_hook(hook)

y = x * 2 + x * 3
y.backward()

# grad₁ + grad₂ → final grad
post hook: 5

With this approach, communication overlaps naturally with computation, reducing idle time and improving efficiency.

However, this is still not ideal. We are now launching one all-reduce per parameter, which introduces significant overhead — especially for models with many small tensors.

This leads us to the next optimization: gradient bucketing, where gradients are grouped into “buckets,” each containing multiple parameters. Once a bucket is filled, it is asynchronously all-reduced.

grad(Layer N)   ┐
grad(Layer N-1) ├──▶ Bucket 1 (filled)──▶ all-reduced
grad(Layer N-2) ┘

Time →

Backward compute:   ███████████████████████
Comm (Bucket 1):        ███████
Comm (Bucket 2):              ███████
Comm (Bucket 3):                    ███████

This improves bandwidth utilization while still overlapping computation and communication:

• Each bucket is a tensor of up to bucket_size_mb megabytes.
• Gradients are concatenated into buckets in the reverse order of model parameters. Because backward pass computes gradients from output to input, so the last parameters in the forward order become ready first.
• Once a bucket is full, an asynchronous all-reduce is launched.

import torch
import torch.distributed as dist

class BucketedDDP:
    def __init__(self, module: torch.nn.Module, bucket_size_mb: float):
        self.module = module
        self.bucket_size_bytes = int(bucket_size_mb * 1024 * 1024)
        self.handles = []

        # 1️⃣ broadcast 
        for p in self.module.parameters():
            dist.broadcast(p.data, src=0)

        # 2️⃣ create buckets
        self.buckets = []
        self._build_buckets()

        # 3️⃣ register hook
        for bucket in self.buckets:
            bucket["ready_count"] = 0
            for p in bucket["params"]:
                if p.requires_grad:
                    p.register_post_accumulate_grad_hook(
                        self._make_hook(bucket)
                    )

    def _build_buckets(self):
        current_bucket = []
        current_size = 0

        # reverse （important!）
        params = list(self.module.parameters())[::-1]

        for p in params:
            size = p.numel() * p.element_size()

            if current_size + size > self.bucket_size_bytes and current_bucket:
                self.buckets.append({"params": current_bucket})
                current_bucket = []
                current_size = 0

            current_bucket.append(p)
            current_size += size

        if current_bucket:
            self.buckets.append({"params": current_bucket})

    def _make_hook(self, bucket):
        def hook(grad):
            bucket["ready_count"] += 1

            # if all param grads are ready
            if bucket["ready_count"] == len(bucket["params"]):
                grads = [p.grad for p in bucket["params"] if p.grad is not None]

                # flatten
                flat = torch._utils._flatten_dense_tensors(grads)

                # async all-reduce
                handle = dist.all_reduce(flat, op=dist.ReduceOp.SUM, async_op=True)
                self.handles.append((handle, flat, grads, bucket))

            return grad
        return hook

    def forward(self, *inputs, **kwargs):
        return self.module(*inputs, **kwargs)

    def finish_gradient_synchronization(self):
        world_size = dist.get_world_size()

        for handle, flat, grads, bucket in self.handles:
            handle.wait()

            # unflatten
            synced = torch._utils._unflatten_dense_tensors(flat, grads)

            for p, g in zip(bucket["params"], synced):
                if p.grad is not None:
                    p.grad.copy_(g / world_size)

            bucket["ready_count"] = 0  # reset

        self.handles.clear()

This strategy strikes a balance: fewer, larger communications per bucket, and maximal overlap with computation.

It has been adopted in frameworks like PyTorch DDP, FlashAttention-2, and Hugging Face’s accelerated training for large models.

Note: Choosing the right bucket size is a trade-off: too small increases overhead, too large delays communication.

Benchmarking: Best Practices

1. Warm-up iterations: Always run a few iterations before timing. GPU kernels and communication may be lazy, and warm-ups ensure accurate benchmarking.
2. Synchronous vs. Asynchronous all-reduce: async_op=True allows backward computation to continue while communication happens; handle.wait() ensures completion when needed.
3. Profiler: Use Nsight Systems or PyTorch profiler to visualize overlap. A well-bucketed, asynchronous DDP run will show compute and communication occurring in parallel.
4. Bucket size tuning: Too small → many communication calls, too large → idle compute waiting for bucket to fill. Common practice: 25–100 MB per bucket, depending on GPU memory and model size.

Conclusion

Efficient DDP training isn’t just about splitting data. By overlapping communication with computation and batching gradients into buckets, it minimizes idle time and maximizes throughput.

Using PyTorch’s backward hooks, async all-reduce, and gradient buckets, we can maximize GPU utilization and minimize idle time.

Next time your distributed training feels slow, it’s worth asking: is your bottleneck really compute — or communication?