PyTorch DDP Deep Dive — Initialization, Buffer Broadcast, the Reducer, and the torch.compile Story

All permalinks in this post are pinned to pytorch/pytorch commit 9b980e3f. Line numbers are verified against that revision.

The trace that started this post

A colleague handed me a profiler trace they were trying to make sense of. The top of the stack looks like a normal nn.Module forward, then something unexpected:

nn.Module: DistributedDataParallel_0
  torch/nn/modules/module.py(1780): _call_impl
  torch/nn/parallel/distributed.py(1660): forward
  torch/nn/parallel/distributed.py(1521): _pre_forward
  torch/nn/parallel/distributed.py(2178): _sync_buffers
  torch/nn/parallel/distributed.py(2197): _sync_module_buffers
  torch/nn/parallel/distributed.py(2207): _default_broadcast_coalesced
  torch/nn/parallel/distributed.py(2133): _distributed_broadcast_coalesced
    <built-in method _broadcast_coalesced of pybind11_builtins...>
    c10d::broadcast_

A broadcast is firing every forward step. That isn’t the all-reduce of gradients you read about in every DDP intro — those fire in backward. So what is being broadcast on every iteration, why, and where in the source does the decision live?

That single trace turns out to surface most of what’s worth knowing about DDP: how it sets itself up at construction time, what work the C++ Reducer does in backward, what runs in forward that isn’t your model, and how torch.compile has to bend around the whole thing. This post walks all of it.

The line numbers in the user’s trace are from a slightly older revision than the one I have checked out — the function names and the relationships between them are unchanged, only line numbers shifted. I’ll use line numbers from the pinned commit (9b980e3f) below.

DDP in one paragraph

DistributedDataParallel is a module wrapper. Each rank owns a full replica of the model. On every iteration each rank computes gradients on its own data shard, then DDP all-reduces the gradients so every replica sees the same averaged grads, and the optimizer steps. Replicas stay bit-identical if (a) they started identical and (b) they receive identical gradients. The whole rest of DDP is the bookkeeping that keeps those two invariants true, plus the engineering to overlap the all-reduces with backward compute.

Part 1 — Initialization

DDP’s __init__ does six things in order. Pulling them out as bullets so they map to source ranges below.

1. Resolve the process group (or device mesh).
2. Collect parameters and buffers that DDP will manage; drop anything in parameters_to_ignore.
3. Configure bucket capacities.
4. One-time broadcast of params and buffers from rank 0 to all other ranks, so replicas start identical.
5. Build the C++ Reducer, which decides bucket assignments and registers a post-accumulate-grad hook on every parameter’s gradient accumulator. These hooks are what drive bucketed allreduce in backward.
6. Various optional setup: mixed precision, in-backward optimizers, static graph, delayed allreduce.

The constructor lives at distributed.py:802. Here’s the load-bearing chunk:

# distributed.py:1015–1040
parameters, expect_sparse_gradient = self._build_params_for_reducer()

if init_sync:
    _verify_param_shape_across_processes(self.process_group, parameters)
    _sync_module_states(                  # one-shot init broadcast
        module=self.module,
        process_group=self.process_group,
        broadcast_bucket_size=self.broadcast_bucket_size,
        src=0,
        params_and_buffers_to_ignore=self.parameters_to_ignore,
        broadcast_buffers=self.broadcast_buffers,
    )
...
self._ddp_init_helper(parameters, expect_sparse_gradient,
                      param_to_name_mapping, static_graph)

_sync_module_states is the init-time broadcast — it flattens params + buffers into a list and calls dist._broadcast_coalesced once (torch/distributed/utils.py:289). After this, replicas are bit-identical. This is not the broadcast in the trace above — that one fires in forward, not in __init__.

What _ddp_init_helper actually does

_ddp_init_helper (distributed.py:1322) computes bucket assignments and then constructs the C++ Reducer:

# distributed.py:1385–1407
self.reducer = dist.Reducer(
    parameters,
    list(reversed(bucket_indices)),
    list(reversed(per_bucket_size_limits)),
    self.process_group,
    expect_sparse_gradient,
    self.bucket_bytes_cap,
    self.find_unused_parameters,
    self.gradient_as_bucket_view,
    param_to_name_mapping,
    self._bucket_config.first_bucket_bytes_cap,
    ...
)

Two non-obvious choices here:

reversed(bucket_indices) — _compute_bucket_assignment_by_size returns buckets in parameter-registration order, but gradients become ready in reverse order during backward (autograd walks the graph from the loss tensor toward the parameters). DDP reverses so bucket index 0 contains the parameters that will be ready first in backward — the ones near the output. Reducer kicks off allreduce on bucket next_bucket_ first and increments — so the bucket index order has to match the autograd readiness order to overlap maximally.

Optional smaller first bucket (_bucket_config.first_bucket_bytes_cap) — the comment in the source explains: “the bucket size limit is specified in the constructor. Additionally, we allow for a single small bucket for parameters that are defined first, such that their gradients don’t spill into a much larger bucket, adding unnecessary latency after gradient computation finishes. Experiments showed 1MB is a reasonable value.” In other words, the first bucket DDP fires in backward is intentionally small so it can launch quickly, get warm comms going, and overlap with the rest of compute.

Reducer registers grad-accumulator post-hooks

This is the piece of init that’s the most interesting. The C++ Reducer constructor walks every parameter and attaches a post-hook on its grad accumulator (reducer.cpp:175–235):

// reducer.cpp:184–209
auto grad_accumulator = torch::autograd::impl::grad_accumulator(variable);

hooks_.emplace_back(
    grad_accumulator->add_post_hook(std::make_unique<
                                    torch::autograd::utils::
                                        LambdaPostHook>(
        [this, variable_index](
            const torch::autograd::variable_list& outputs,
            const torch::autograd::variable_list& /* unused */) {
          this->autograd_hook(variable_index);
          return outputs;
        },
        [this](torch::autograd::CompiledNodeArgs& args) {
          TORCH_CHECK(this->use_python_reducer_,
              "Compiled autograd is not compatible with C++ DDP Reducer, ...");
        })),
    grad_accumulator);

A “grad accumulator” is the autograd Node that accumulates incoming grad-tensors into a leaf parameter’s .grad field. Every leaf parameter has exactly one, lazily created the first time autograd traverses it. The post-hook lets us run code after the accumulator has finished writing .grad for that one parameter. That’s the signal we use to mark a parameter “ready for allreduce”.

Note the second lambda — the compiled-autograd hook — fails with a check that points the user at optimize_ddp="python_reducer". We’ll get to why later.

Part 2 — A full training step

Time to walk one iteration. We’ll use the user’s trace to anchor where in source each line comes from.

Step 1: forward

The forward entry point at distributed.py:1826:

def forward(self, *inputs, **kwargs):
    with torch.autograd.profiler.record_function("DistributedDataParallel.forward"):
        inputs, kwargs = self._pre_forward(*inputs, **kwargs)
        output = (
            self.module.forward(*inputs, **kwargs)
            if self._delay_all_reduce_all_params
            else self._run_ddp_forward(*inputs, **kwargs)
        )
        return self._post_forward(output)

Three things happen: _pre_forward, the user module’s forward, then _post_forward. The pre and post are where DDP does its housekeeping.

_pre_forward — here is the buffer broadcast

distributed.py:1686:

def _pre_forward(self, *inputs, **kwargs):
    if self._use_python_reducer:
        return inputs, kwargs

    if not self._lazy_init_ran and not torch.compiler.is_compiling():
        self._lazy_init()

    ...

    if torch.is_grad_enabled() and self.require_backward_grad_sync:
        self.logger.set_runtime_stats_and_log()
        self.reducer.prepare_for_forward()

    # Notify the join context that this process has not joined, if needed
    work = Join.notify_join_context(self)
    ...

    # Calling _rebuild_buckets before forward computation, ...
    if torch.is_grad_enabled() and self.reducer._rebuild_buckets():
        logger.info("Reducer buckets have been rebuilt in this iteration.")
        self._has_rebuilt_buckets = True

    # sync params according to location (before/after forward) user
    # specified as part of hook, if hook was specified.
    if self._check_sync_bufs_pre_fwd():
        self._sync_buffers()
    ...

The interesting line is if self._check_sync_bufs_pre_fwd(): self._sync_buffers(). This is where the trace in the intro is coming from. _check_sync_bufs_pre_fwd at distributed.py:2317:

def _check_sync_bufs_pre_fwd(self):
    return self.will_sync_module_buffers() and (
        not hasattr(self, "buffer_hook")
        or self.buffer_hook.buffer_comm_hook_location
        == _BufferCommHookLocation.PRE_FORWARD
    )

def will_sync_module_buffers(self):
    return (
        self.require_forward_param_sync
        and self.broadcast_buffers
        and len(self.modules_buffers) > 0
    )

So a buffer broadcast fires on every forward iff:

• broadcast_buffers=True (default) and
• the module actually has buffers (BatchNorm, anything register_buffer‘d) and
• a user-installed buffer_hook hasn’t moved the sync to post-forward.

_sync_buffers itself (distributed.py:2347):

def _sync_buffers(self):
    with torch.no_grad():
        if self._join_config.enable:
            authoritative_rank = self._find_common_rank(self._distributed_rank, True)
        else:
            authoritative_rank = 0
        # Update self.modules_buffers in case any buffers were reassigned.
        self._assign_modules_buffers()
        with torch.autograd._unsafe_preserve_version_counter(
            tuple(self.modules_buffers)
        ):
            self._sync_module_buffers(authoritative_rank)

…then _sync_module_buffers → _default_broadcast_coalesced → _distributed_broadcast_coalesced → dist._broadcast_coalesced — exactly the chain in the trace.

The motivation: parameter equivalence is preserved by gradient allreduce + identical optimizer state. Buffer equivalence is not. BatchNorm’s running_mean and running_var are updated as a side-effect of forward(training=True), with no autograd involvement — and the update sees only this rank’s local mini-batch. So after every forward, BN buffers drift apart across ranks. The cheap fix is to overwrite them on the other ranks with rank 0’s copy at the start of the next forward. That’s what this broadcast does. It’s coalesced (multiple buffer tensors flattened, broadcast in one collective, then unflattened) and the buffer size is 250 * 1024 * 1024 (distributed.py:976).

Why broadcast and not all-reduce? All-reduce would average BN stats. Treating rank 0 as the authoritative source matches what a single-process trained model would see. It also keeps the cost asymmetric: one tensor flying out to N-1 readers via NCCL ring is cheap relative to all-reduce.

Can you turn it off? Yes — pass broadcast_buffers=False to the DDP constructor if your model has no buffers that need synchronizing across ranks, or use SyncBatchNorm which keeps its statistics consistent through a different mechanism.

Subtle bit: _assign_modules_buffers() is called every time before the broadcast, because user modules can reassign buffers in their forward (self.my_buf = torch.zeros(...) instead of self.my_buf.copy_(...)). Without re-scanning, DDP would broadcast into the old tensor and the user code would never see it. See pytorch issue #63916.

_run_ddp_forward and the active-DDP context

After _pre_forward returns, _run_ddp_forward calls the user model — but it does so inside a context manager that sets DistributedDataParallel._active_ddp_module = self. This isn’t for the eager path; it’s a flag torch.compile reads to decide whether to route to DDPOptimizer (we’ll get to it in Part 3). distributed.py:1637–1657:

@contextmanager
@torch._disable_dynamo(recursive=False)
def _inside_ddp_forward(self):
    old = DistributedDataParallel._active_ddp_module
    DistributedDataParallel._active_ddp_module = self
    try:
        yield
    finally:
        DistributedDataParallel._active_ddp_module = old

def _run_ddp_forward(self, *inputs, **kwargs):
    if self._use_python_reducer:
        return self.module(*inputs, **kwargs)
    else:
        with self._inside_ddp_forward():
            return self.module(*inputs, **kwargs)

_post_forward

distributed.py:1757. Three things of note:

• If the user installed a buffer_hook with POST_FORWARD location, the buffer sync happens here instead of in _pre_forward.
• self.reducer.prepare_for_backward(...) is called — it tells the C++ reducer which params should be expected to receive grads. With find_unused_parameters=True, it walks the output’s autograd graph to identify which params were actually used, and the rest are marked ready immediately so backward doesn’t wait on grads that will never arrive.
• If find_unused_parameters or static_graph is on (with some subtleties), the output is routed through _DDPSink.apply(...) — a no-op autograd Function whose only purpose is to give DDP a custom backward node that fires before per-parameter grad accumulators, so it can do bookkeeping. For the common case (find_unused_parameters=False, no static graph) this is skipped.

Step 2: backward

You call loss.backward(). Autograd walks the graph. For every leaf parameter that receives a grad, its grad-accumulator runs, copies the incoming grad into .grad, and then fires the post-hook DDP registered in __init__.

Reducer::autograd_hook(index) is the entry point (reducer.cpp:667). The interesting bit at the end:

// reducer.cpp:730–751
if (static_graph_after_first_iteration()) {
  REDUCER_CHECK(numGradHooksTriggeredMapPerIteration_[index] > 0, logger_, ...);
  if (--numGradHooksTriggeredMapPerIteration_[index] == 0) {
    if (should_rebuild_buckets()) {
      push_rebuilt_params(index);
    }
    mark_variable_ready(index);
  }
} else {
  if (should_rebuild_buckets()) {
    push_rebuilt_params(index);
  }
  mark_variable_ready(index);
}

mark_variable_ready(index) at reducer.cpp:894 does three things in sequence:

1. Copy this parameter’s grad into its slot in the bucket’s flat buffer (mark_variable_ready_dense). The bucket holds one big gradients tensor; each parameter has a slice (bucket_views_in[i]) it gets copy_‘d into. As the copy happens, it also pre-divides by div_factor_ (the world size, modulo join semantics), so the eventual allreduce-sum produces an average.
2. Decrement bucket.pending. If it hits zero, all parameters in this bucket are ready.
3. If the bucket is ready, call mark_bucket_ready (reducer.cpp:1050), which walks forward from next_bucket_ and fires all_reduce_bucket on each contiguous ready bucket.

// reducer.cpp:1050–1074
void Reducer::mark_bucket_ready(size_t bucket_index) {
  TORCH_INTERNAL_ASSERT(bucket_index >= next_bucket_);

  // Buckets are reduced in sequence. Ignore this bucket if
  // it's not its turn to be reduced.
  if (bucket_index > next_bucket_) {
    return;
  }

  // Keep going, until we either:
  // - have kicked off reduction for all buckets, or
  // - found a bucket that's not yet ready for reduction.
  for (; next_bucket_ < buckets_.size() && buckets_[next_bucket_].pending == 0;
       next_bucket_++) {
    ...
    auto& bucket = buckets_[next_bucket_];
    if (!should_skip_all_reduce_bucket(bucket)) {
      all_reduce_bucket(bucket);
      num_buckets_reduced_++;
    }
  }
}

The key insight: DDP doesn’t fire allreduce as soon as any one parameter is ready. It waits for an entire bucket to be ready, then fires one allreduce on the bucket’s flat tensor. That’s the compute/comm overlap engine: while autograd keeps cranking on the rest of the graph, NCCL is asynchronously reducing bucket 0 in the background.

all_reduce_bucket (reducer.cpp:977) doesn’t directly call into NCCL — it calls run_comm_hook, which is run_allreduce_hook by default (plain sum-allreduce) but is overrideable via register_comm_hook (gradient compression, PowerSGD, etc.). The hook returns an intrusive_ptr<Future> stored on the bucket; the future is awaited at the end of backward.

Step 3: finalize_backward

When the last bucket marks ready, the reducer enqueues an autograd engine callback (reducer.cpp:945) that runs after autograd is done:

torch::autograd::Engine::get_default_engine().queue_callback([this] {
  std::lock_guard<std::mutex> lock(this->mutex_);
  ...
  this->finalize_backward();
});

finalize_backward (reducer.cpp:1729) waits on each bucket’s future, then either copies the reduced flat tensor back into the per-parameter .grad views, or — if gradient_as_bucket_view=True — does nothing because .grad already is the bucket view. The optimizer can now step on synchronized gradients.

Step 4: optimizer.step()

DDP isn’t involved. The optimizer sees the synchronized .grad tensors and does its update on every rank in parallel. Since the grads are identical and the parameters started identical, the updated parameters are identical on every rank too. Invariant preserved.

Annotating the trace

Now the original trace reads end-to-end:

Frame	What’s happening
_call_impl(args)	The Python module __call__ machinery calls forward.
forward(args, kwargs)	DDP’s forward opens the profiler scope, then calls _pre_forward.
_pre_forward(args, kwargs)	Lazy init, prepare_for_forward, _rebuild_buckets, then _check_sync_bufs_pre_fwd() → _sync_buffers().
_sync_buffers()	Determine authoritative rank (rank 0 normally, or the lowest non-joined rank under the Join context). Re-scan module.named_buffers() in case the user reassigned any.
_sync_module_buffers(authoritative_rank)	Either a user-registered buffer_hook runs here, or the default broadcast path.
_default_broadcast_coalesced(...)	Fills in defaults (self.modules_buffers, self.broadcast_bucket_size=250MB), forwards.
_distributed_broadcast_coalesced(...)	Thin wrapper around dist._broadcast_coalesced.
dist._broadcast_coalesced (C++)	Coalesces tensors into 250MB-cap buckets, flattens each, process_group->broadcast(flat_tensor, root=0), unflattens, copy_’s back. Up to 2 broadcasts in flight at once.
c10d::broadcast_	The actual NCCL collective.

So the answer to “where do these buffer broadcasts come from” is: DDP’s forward synchronizes module buffers (e.g. BatchNorm running stats) from rank 0 to every other rank on every iteration, because nothing else keeps them aligned. If your module has any registered buffers and you didn’t pass broadcast_buffers=False, you’ll see this in every profiler trace.

Part 3 — DDP + torch.compile

DDP’s bucket-allreduce design assumes the autograd engine can fire per-node hooks while backward is mid-execution. That assumption is precisely the one torch.compile breaks.

Consider what torch.compile(ddp_wrapped_model) does naïvely:

1. Dynamo captures the user’s forward as one FX graph.
2. AOTAutograd extends it into a joint forward+backward graph.
3. Inductor codegens a single fused kernel pipeline for the whole backward.
4. At runtime, autograd sees the joint backward as a single Node — when it runs, all parameters’ grads become ready at the same instant, after the entire fused backward has finished.

The Reducer’s grad-accumulator post-hooks all fire essentially simultaneously after the whole backward kernel returns. Buckets all fill at once. Allreduces all fire at once. There is no overlap between compute and communication. Worse, the allreduces serialize behind a now-complete backward instead of streaming alongside it. For a large model this is catastrophic.

PyTorch ships two ways around this. Both are selected via torch._dynamo.config.optimize_ddp:

optimize_ddp: (
    bool
    | Literal[
        "ddp_optimizer",
        "python_reducer",
        "python_reducer_without_compiled_forward",
        "no_optimization",
    ]
) = True   # default is "ddp_optimizer"

Path A — ddp_optimizer (the default): break the graph at bucket boundaries

The idea is small and clean: if a single fused backward defeats overlap, split it. Compile the model as N smaller graphs, one per DDP bucket. Each subgraph’s backward is its own autograd Node, so the grad-accumulator post-hooks for that subgraph’s parameters fire as soon as that subgraph’s backward finishes — exactly the granularity DDP wants.

DDPOptimizer lives at torch/_dynamo/backends/distributed.py:373. Its docstring is worth reading in full — it explains both the algorithm and the rationale (distributed.py:374–432). The activation point is in convert_frame.py, where Dynamo checks if there’s an active DDP module before compilation:

# convert_frame.py:2480–2502
if torch._dynamo.utils.get_optimize_ddp_mode() == "ddp_optimizer":
    ddp_module = DistributedDataParallel._get_active_ddp_module()
    if ddp_module:
        with compile_lock:
            from torch._dynamo.backends.distributed import DDPOptimizer

            ddp_optimizer = DDPOptimizer(
                bucket_bytes_cap=ddp_module.bucket_bytes_cap,
                backend_compile_fn=self._torchdynamo_orig_backend._torchdynamo_orig_backend,
            )
            ...
            hijacked_callback = (
                self._torchdynamo_orig_backend._clone_with_backend(
                    ddp_optimizer.compile_fn,
                )
            )
            return hijacked_callback(frame, cache_entry, self.hooks, frame_state)

This is the payoff for the _active_ddp_module flag we saw set in _run_ddp_forward: Dynamo only takes the DDPOptimizer path when compilation happens inside a DDP forward.

compile_fn at distributed.py:489 is the meat:

# distributed.py:501–566
# 1: compute the partition map according to DDP bucket logic
buckets = [Bucket()]
processed_modules: set[torch.nn.Module] = set()
for node in reversed(gm.graph.nodes):           # reverse-order walk
    if node.op in ("output", "placeholder"):
        continue

    if (
        buckets[0].size >= self.bucket_bytes_cap
        or len(buckets) == 1
        and buckets[0].size >= self.first_bucket_cap
    ):
        if bucket_has_external_output(buckets[0]):
            buckets.insert(0, Bucket())
        else:
            # extend bucket past its parameter capacity until it
            # contains at least one externally-used output
            ...

    if node.op == "call_function":
        self.add_param_args(buckets[0], node)
    elif node.op == "call_module":
        target_mod = gm.get_submodule(node.target)
        if target_mod not in processed_modules:
            self.add_module_params_to_bucket(target_mod, buckets[0], ...)
    ...
    elif node.op == "get_attr":
        maybe_param = getattr(gm, node.target)
        if isinstance(maybe_param, torch.nn.Parameter) and ...:
            self.add_param(buckets[0], maybe_param, node.target)

    buckets[0].nodes.append(node)

A few things worth pulling out:

The walk is in reverse order, mirroring the order gradients become ready in backward. Each parameter encountered in this reverse walk is added to the current bucket; once that bucket exceeds bucket_bytes_cap, a new bucket is opened. The first bucket can be smaller (first_bucket_cap) — mirroring the C++ Reducer’s “small first bucket for fast warm-up” trick.

bucket_has_external_output (distributed.py:72) — graph-partitioning requires each subgraph to produce something the next subgraph consumes (otherwise its result has no consumer and Inductor optimizes it away). DDPOptimizer won’t close out a bucket until it’s added at least one node that produces a value used by an earlier node in the original graph (which becomes the next subgraph in execution order). If a bucket’s contents are all inplace mutations or self-contained side-effects, the bucket is extended past its parameter capacity until an externally-consumed node lands in it.

The split:

# distributed.py:577–586
partition_map = {}
for idx, b in enumerate(buckets):
    for node in b.nodes:
        partition_map[node] = idx

split_gm = fx.passes.split_module.split_module(
    gm,
    None,
    lambda node: partition_map[node],
)

Hand each partition off to fx.passes.split_module.split_module, which produces a top-level split_gm whose nodes are calls into per-bucket sub-GraphModules.

Per-subgraph AOT+Inductor compile is then driven by SubmodCompiler (distributed.py:181), which is an fx.Interpreter that walks split_gm, calls AOTAutograd’s compile on each submodule, swaps the compiled module back in, and ensures correctly-strided fake tensors flow between submodules so Inductor’s output stride choices in one subgraph are respected by the next. That last piece — the FakeifyFirstAOTInvocationGuard and “fakify_first_call” — is the part that took the most engineering effort; it makes the per-bucket pipeline produce the same output strides as a single-piece compile would. (distributed.py:312–360)

The runtime behavior after this:

forward:
  split_gm calls compiled_submod_0   ─┐
              ↓ output                │ N separate Inductor artifacts;
            compiled_submod_1         │ each is its own AOTAutograd
              ↓                       │ joint fwd/bwd.
            ...                       │
            compiled_submod_{N-1}    ─┘

backward (driven by autograd engine):
  compiled_submod_{N-1}.backward fires
     → its parameters' grads land → grad-accumulator hooks fire
     → Reducer marks that bucket ready → allreduce kicks off
  ... while compiled_submod_{N-2}.backward is already running ...
  → compute and allreduce overlap, as in eager DDP.

The cost: more compilation, less cross-bucket Inductor fusion. For models where the per-bucket loss in fusion is larger than the comm-overlap gain, the tradeoff goes the wrong way — but those are unusual cases, and the default optimize_ddp=True reflects that for most large models, comm overlap dominates.

Path B — python_reducer + compiled autograd

The other approach inverts the problem: instead of breaking the compile to fit DDP’s eager hooks, replace the C++ Reducer with a Python reducer that emits its own all_reduce ops directly into the graph, then let compiled_autograd trace the whole forward+backward as one graph including the allreduces, and let Inductor fuse compute and comms together.

The mode is set by torch._dynamo.config.optimize_ddp = "python_reducer". When DDP detects it (distributed.py:826 and distributed.py:1095–1108):

# distributed.py:826
self._use_python_reducer = (
    torch._dynamo.utils.get_optimize_ddp_mode() == "python_reducer"
)

# distributed.py:1096–1108 — at end of __init__
if self._use_python_reducer:
    torch._inductor.config._fuse_ddp_communication = True
    torch._inductor.config._fuse_ddp_bucket_size = bucket_cap_mb
    torch._dynamo.trace_rules.LEGACY_MOD_INLINELIST.add(
        "torch.nn.parallel.distributed"
    )
    torch._dynamo.trace_rules.get_legacy_mod_inlinelist.cache_clear()
    self._register_accum_grad_hook()

…it does three things:

1. Skip the C++ Reducer. _pre_forward and _post_forward early-return when self._use_python_reducer is true, and the wrapper around self.module(...) in _run_ddp_forward is also skipped. The model is just nn.Module.forward — no Reducer bookkeeping.

2. Register a Python post_accumulate_grad_hook on every parameter (distributed.py:1113):

   def compiled_accum_grad_hook(param, *, param_index: int):
       if not self.require_backward_grad_sync:
           return
       if param.grad is None:
           return
       if self._comm_hooks:
           for hook, state in self._comm_hooks:
               hook(state, (param.grad, param))
       else:
           gradient = param.grad / self.process_group.size()
           gradient = fcol.all_reduce(gradient, "sum", self.process_group)
           param.grad.copy_(gradient)
   

   The hook is plain Python. With compiled_autograd on, it gets traced as part of the backward graph — the fcol.all_reduce becomes a c10d_functional.all_reduce node in the FX graph, not a runtime side-effect.

3. Enable Inductor’s DDP fusion passes: _fuse_ddp_communication = True, with _fuse_ddp_bucket_size = bucket_cap_mb. The pass _fuse_ddp_communication (torch/_inductor/fx_passes/ddp_fusion.py:467) walks the post-grad FX graph and coalesces consecutive small all_reduce calls into a single large one with the same bucket-cap heuristic the C++ Reducer uses.

The net effect: backward is one fused compiled graph, but the graph contains the allreduces, and Inductor schedules them inline with compute. There are no Python-level grad-accumulator post-hook callbacks at runtime — that work was moved to compile time.

This path is gated on compiled_autograd because vanilla autograd would still execute the backward as one Node with all the side effects packaged opaquely; it’s compiled_autograd that walks the autograd graph at compile time and turns it into a separate FX graph that Dynamo can re-trace with full visibility.

Quick comparison

Mode	What gets compiled	Where allreduces live	Comm/compute overlap	When to use
"no_optimization"	One graph per forward, plus eager backward	C++ Reducer hooks, eager	None — backward is one fused Node, hooks all fire at end	Debugging; cheap models where overlap doesn’t matter
"ddp_optimizer" (default)	N small graphs, one per bucket	C++ Reducer hooks, between subgraphs	Yes — subgraph backward emits Node per bucket, hooks fire at bucket boundary	Default, well-tested, no compiled-autograd dependency
"python_reducer"	One fused graph including allreduces	Inlined in FX via Python hook + Inductor DDP-fusion pass	Yes — Inductor schedules comm + compute together	Requires compiled_autograd, can squeeze out more overlap; experimental

"ddp_optimizer" is what you’ll hit by default. "python_reducer" is what to reach for if you’ve already adopted compiled_autograd and want maximum fusion across the entire backward.

Summary

DDP is a small idea (replicas + averaged grads) wrapped in a meaningful amount of plumbing to (a) keep replicas bit-identical given that some state isn’t covered by gradient averaging, and (b) hide allreduce latency under backward compute. Concretely:

• __init__ resolves the process group, broadcasts initial state once, configures buckets, and registers a post-hook on every parameter’s grad accumulator via the C++ Reducer.
• Every forward re-broadcasts buffers from rank 0 — this is what shows up in the user’s trace, and it exists because BatchNorm-style buffers update as a forward side-effect and would otherwise drift across ranks.
• Every backward’s grad hooks fire mark_variable_ready, which fills a bucket’s flat tensor slot; when a bucket is full, all_reduce_bucket kicks off an async allreduce while the rest of backward keeps running. finalize_backward (queued via an autograd engine callback) joins on the futures and writes reduced grads back into .grad.
• torch.compile collides with the bucket-hook design: a single fused backward defeats overlap. DDPOptimizer resolves it by splitting the FX graph along bucket boundaries before AOTAutograd sees it; python_reducer resolves it by replacing the C++ Reducer entirely and letting compiled_autograd + Inductor’s _fuse_ddp_communication pass schedule allreduces inline with compute.

The buffer broadcast that prompted this post sits on the boundary between two ideas: parameters are synchronized by allreduce + identical optimizer steps, but buffers don’t go through autograd, so they have to be synchronized by something else. DDP’s choice — overwrite from rank 0 every forward — is cheap, predictable, and accounts for ~250MB worth of coalesced broadcast per iteration on a typical BN-heavy model. Once you’ve internalized that, every line in the trace at the top of this post makes sense.