DDP Communication Hooks

The contract

Register with tensorplay.nn.parallel.DistributedDataParallel.register_comm_hook():

ddp_model.register_comm_hook(state, hook)

• state is any object you choose — it is passed back to every hook call untouched. Use it for per-process context (a communicator, a compression state, a step counter).

• hook is called as hook(state, bucket) once per gradient bucket, when the bucket is full, on the backward path. It must return the reduced buffer: either the tensor itself, or a Future-like object whose wait() (or value()) yields it. The reduced values are copied back into the parameters’ .grad.

The bucket argument is a GradBucket:

Two rules from the registration check itself: a hook can be registered only once (a second call raises RuntimeError), and it must be registered before the first backward pass. Registering any hook also switches DDP off its C++ fast reducer and onto the Python hook path — the flexibility costs the no-GIL hot path, so hooks are for when you actually need custom communication.

A minimal hook

A hook that logs each bucket and does the plain average:

import tensorplay as tp
import tensorplay.distributed as dist
from tensorplay.futures import Future


def logging_avg_hook(state, bucket):
    buf = bucket.buffer()
    print(f"[rank {state}] bucket {bucket.index()} reducing "
          f"{buf.numel()} elements")
    dist.all_reduce(buf, op=dist.ReduceOp.AVG)
    fut = Future()
    fut.set_result(buf)
    return fut

Registration and use, in the usual two-process gloo setup:

def example(rank, world_size):
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    model = tp.nn.Linear(10, 10)
    ddp_model = tp.nn.parallel.DistributedDataParallel(model)
    ddp_model.register_comm_hook(rank, logging_avg_hook)

    out = ddp_model(tp.randn(4, 10))
    out.square().mean().backward()
    # every parameter's .grad is now the average across the two processes

The Future return is the asynchronous form — a hook that kicks off a non-blocking collective can return a future that resolves when the collective completes, and backwards of later buckets overlap with it. Returning the tensor directly is the synchronous form and is also accepted.

Composing with the shipped hooks

The compression wrappers in tensorplay.distributed.algorithms.ddp_comm_hooks are built on exactly this contract — fp16_compress_wrapper(inner_hook) runs inner_hook on a half-precision copy of the bucket, so a custom hook composed with them stays inside the same mechanism. Reading their source is the fastest way to see the state/bucket discipline in real use.

Where to go next

• the distributed algorithms page — every shipped hook and its knobs.

• the DDP note — how buckets are built and ordered in the first place.