Buckets and the reducer

Wrapping a module in DDP does something before any training happens: it broadcasts rank 0's parameters and buffers to every other rank. That is not documentation, it is observable. Seed each rank differently so the models genuinely differ, print the weights, wrap, and print again.

Rank 1 walks in holding [0.364, -0.312] and walks out holding rank 0's [-0.005, 0.379]. Nothing in the training script asked for that. It happens in the constructor because every later step of DDP assumes the replicas started identical, and a gradient all-reduce keeps replicas in sync only if they were in sync to begin with.

The same assumption is why the DDP docstring warns that parameters must be registered in the same order on every rank. Ranks are matched up by position in the parameter list, not by name, so a model whose module construction order varies by rank will happily reduce one layer's gradient into another's.

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29551"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    torch.manual_seed(rank)                    # deliberately different models
    net = nn.Linear(2, 1, bias=False)
    print(f"rank {rank}: before ddp {net.weight.detach().flatten().tolist()}")
    DDP(net)                                   # the wrap is the broadcast
    print(f"rank {rank}: after ddp  {net.weight.detach().flatten().tolist()}")
    dist.destroy_process_group()

# rank 1: before ddp [0.36434608697891235, -0.3121015429496765]
# rank 1: after ddp  [-0.0052939653396606445, 0.37932294607162476]
# rank 0: before ddp [-0.0052939653396606445, 0.37932294607162476]
# rank 0: after ddp  [-0.0052939653396606445, 0.37932294607162476]

The clearest way to see what DDP does to a gradient is to compute the same gradient twice, once outside the wrapper and once inside, on ranks holding different data. A Linear(3, 1) with no bias, fed a row of ones on rank 0 and a row of twos on rank 1, has an analytic gradient: the input row itself.

Alone, rank 0 gets [1.0, 1.0, 1.0] and rank 1 gets [2.0, 2.0, 2.0], which is exactly what a per-rank backward should produce. Inside DDP both ranks get [1.5, 1.5, 1.5]. The all-reduce summed the two gradients and the reducer divided by the world size, so what lands in .grad is the mean over the global batch rather than the sum.

That division is why a DDP step and a single-process step on the concatenated batch agree, and it is also the reason a learning rate tuned on one process usually survives the move. What changes is the effective batch size, which is now the per-rank batch times the world size.

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29536"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    torch.manual_seed(0)
    net = nn.Linear(3, 1, bias=False)
    x = torch.full((1, 3), float(rank + 1))

    solo = nn.Linear(3, 1, bias=False)
    solo.load_state_dict(net.state_dict())
    solo(x).sum().backward()
    print(f"rank {rank}: alone  {solo.weight.grad.tolist()}")

    ddp = DDP(net)
    ddp(x).sum().backward()
    print(f"rank {rank}: in ddp {net.weight.grad.tolist()}")

# rank 1: alone  [[2.0, 2.0, 2.0]]
# rank 1: in ddp [[1.5, 1.5, 1.5]]
# rank 0: alone  [[1.0, 1.0, 1.0]]
# rank 0: in ddp [[1.5, 1.5, 1.5]]

DDP does not invent its grouping. It calls dist._compute_bucket_assignment_by_size, hands it the parameter list and a list of size limits, and gets back a list of index lists. That function is importable, needs no process group, and runs in a single process, which makes the bucketing rule something you can check instead of infer.

Run it on three Linear(512, 512) layers with a 1MB cap and the answer is [[0], [1, 2], [3, 4], [5]]. Parameter 0 is a weight of exactly 1048576 bytes, which fills the first bucket on its own. Parameters 1 and 2 are a 2048-byte bias followed by the next weight, so that bucket ends up at 1050624 bytes, a little over the cap: the fill is greedy, and a parameter that pushes a bucket past the limit still goes in before the next bucket opens.

Two size limits go in, not one. The first is dist._DEFAULT_FIRST_BUCKET_BYTES, 1048576 on this build, and it exists so that the parameters defined first, whose gradients arrive last in the backward, get a small bucket rather than spilling into a large one that then has to wait. Every later bucket uses bucket_cap_mb, which defaults to 25MB.

DDP then reverses the list before handing it to the reducer. The comment on that line says why in one clause: the reversal approximates the order gradients are produced in, on the assumption that layers are used in the forward pass in the order they were defined.

import sys
import torch
import torch.distributed as dist
import torch.nn as nn

m = nn.Sequential(nn.Linear(512, 512), nn.ReLU(),
                  nn.Linear(512, 512), nn.ReLU(),
                  nn.Linear(512, 512))
ps = list(m.parameters())
sparse = [False] * len(ps)

idx, lim = dist._compute_bucket_assignment_by_size(
    ps, [dist._DEFAULT_FIRST_BUCKET_BYTES, 1024 * 1024], sparse)
print("two-limit :", idx)
idx2, _ = dist._compute_bucket_assignment_by_size(ps, [sys.maxsize], sparse)
print("one-limit :", idx2)
print("bytes     :", [p.numel() * p.element_size() for p in ps])

# two-limit : [[0], [1, 2], [3, 4], [5]]
# one-limit : [[0, 1, 2, 3, 4, 5]]
# bytes     : [1048576, 2048, 1048576, 2048, 1048576, 2048]

Chapter 8 describes bucketed all-reduce firing as each bucket fills, which is what a steady-state step does. The first iteration is the exception, and the reason is written into _ddp_init_helper at v2.2.2: when static_graph is true or find_unused_parameters is false, the bucket size limit list is [sys.maxsize], which produces exactly one bucket.

The comment above that branch explains the hazard it is avoiding. Before the first backward has run, DDP only knows the order parameters were registered in, which for a model with control flow can be nothing like the order gradients actually become ready. Bucketing on a guess can fire a bucket's all-reduce early on one rank and late on another, and two ranks issuing collectives in different orders is the deadlock from lesson one.

After the first backward, DDP has evidence instead of a guess. It records the order gradients arrived in and rebuilds the buckets against that order. Run twelve steps on the same three-layer model and the logging data says so directly: has_rebuilt_buckets = 1, gradients became ready in the order 5, 4, 3, 2, 1, 0, and the rebuilt buckets are 5 4, 3 2, 1 0.

Read those index pairs against the parameter list and the rebuild is legible. Index 5 is the last layer's bias and index 4 is its weight, so each rebuilt bucket is one layer's bias and weight together, in the order the backward produced them, at 1050624 bytes each. The initial assignment had split those same six parameters as [[0], [1, 2], [3, 4], [5]], which pairs each weight with the next layer's bias, a grouping that only made sense before anything had been measured.

Iteration one buys correctness with one big all-reduce. Iteration two onward buys overlap with the order it just measured.

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29535"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(512, 512), nn.ReLU(),
                          nn.Linear(512, 512), nn.ReLU(),
                          nn.Linear(512, 512))
    ddp = DDP(model, bucket_cap_mb=1)
    x = torch.full((8, 512), float(rank + 1))
    for _ in range(12):
        ddp.zero_grad()
        ddp(x).sum().backward()
    if rank == 0:
        d = ddp._get_ddp_logging_data()
        for k in sorted(d):
            if "bucket" in k or "rebuilt" in k or k == "iteration":
                print(k, "=", d[k])

# bucket_cap_bytes = 1048576
# bucket_sizes = 3151872
# gradient_as_bucket_view = 0
# has_rebuilt_buckets = 1
# iteration = 10
# prev_iteration_grad_ready_order_indices = 5, 4, 3, 2, 1, 0
# rebuilt_bucket_sizes = 1050624, 1050624, 1050624
# rebuilt_per_bucket_param_indices = 5 4, 3 2, 1 0

Five constructor arguments change the ledger above, and each one trades a different resource. bucket_cap_mb sets how coarse the overlap is: smaller buckets start communicating sooner and send more, larger buckets send more efficiently and start later. find_unused_parameters makes DDP traverse the autograd graph every iteration to discover which parameters got no gradient, which costs a graph walk per step and is what you need when a model skips branches.

gradient_as_bucket_view points each parameter's .grad at a slice of the bucket rather than at its own tensor, which removes one full copy of the gradients from memory. static_graph promises the model's graph never changes, which lets DDP keep the first iteration's analysis forever. broadcast_buffers decides whether buffers, batch-norm running statistics being the usual case, get re-broadcast from rank 0 at every forward.

The communication itself is replaceable. DDP comm hooks let you install a function that runs instead of the plain all-reduce for each bucket, which is how bf16 gradient compression and PowerSGD are implemented. That is a mechanism this machine cannot demonstrate honestly at any useful scale, so it stays a pointer to the hook reference rather than a measurement.

One knob lives in the training loop rather than the constructor, and chapter 4's accumulation lesson already owns it: no_sync suppresses the reduction for every micro-batch but the last. Read it there; it belongs to the accumulation story, not this one.