LLM Lab Notes 03: Writing Collective Communication Autograd — How AllReduce Joins Backpropagation

Code file: llm_lab/02_nccl/tp_autograd_comm.py

Note 02 said DDP "hides gradient sync inside backward." This note opens that box: how exactly does collective communication get embedded into automatic differentiation? The answer is PyTorch's torch.autograd.Function — the standard implementation pattern for the Row Parallel communication layer in Megatron-LM and for all kinds of custom parallel operators.

1. Experiment Goal

1. Wrap all_reduce in a differentiable operator with torch.autograd.Function;
2. Verify the core tensor-parallelism mechanism: "communicate in forward, communicate in backward too";
3. Confirm with printed gradients that x.grad of the AllReduce.apply(x) input equals the sum of gradients across all GPUs.

2. Code Walkthrough

Full code (tp_autograd_comm.py, 37 lines):

import os
import torch
import torch.distributed as dist

dist.init_process_group("nccl")

rank = dist.get_rank()
local_rank = int(os.environ["LOCAL_RANK"])

torch.cuda.set_device(local_rank)

class AllReduce(torch.autograd.Function):          # ① custom operator
    @staticmethod
    def forward(ctx, x):
        x = x.contiguous()                          # ② ensure contiguous memory
        dist.all_reduce(x)                          # ③ forward: cross-GPU reduction
        return x

    @staticmethod
    def backward(ctx, grad_output):                 # ④ backward must communicate too!
        grad_output = grad_output.contiguous()
        dist.all_reduce(grad_output)
        return grad_output

x = torch.ones(4, device="cuda", requires_grad=True)  # ⑤ input with grad enabled

y = AllReduce.apply(x)                                # ⑥ through the custom operator

loss = y.sum()

loss.backward()                                       # ⑦ backpropagation

if rank == 0:
    print("x.grad =", x.grad)

dist.destroy_process_group()

① torch.autograd.Function

autograd.Function is PyTorch's "custom differentiable operator" interface: implement the static methods forward and backward, and invoke via Class.apply(...). The autograd engine chains forward's output and backward into a computation graph — the backward we write determines how gradients flow back.

② The hidden trap of contiguous()

all_reduce (like most NCCL collectives) requires contiguous memory. If the upstream tensor is a non-contiguous view (from transpose, slicing, etc.), communicating directly errors out or silently corrupts results. contiguous() is the "safety belt" of collective operators — production code has it in nearly every communication operator.

③④ Forward reduction + backward reduction

This is the core insight of the note — worth reading three times:

• Forward: the input x is summed across GPUs, and every GPU receives the same aggregated value;
• Backward: the gradient grad_output must also be summed across GPUs before flowing back to the input.

Why must backward communicate again? Because the forward all_reduce means "multiple GPUs shared the same output." If the output's gradient w.r.t. the total loss is g, then every GPU's input contributed to this output, and each should receive its own share of the gradient. For a SUM reduction, the correct semantics are: each input's gradient = the output gradient. However, in multi-GPU autodiff, every GPU independently computes the same g; without communication the per-GPU gradients would diverge from there on. So backward performs one more AllReduce: keeping gradients consistent across all GPUs.

3. How It Works: This Is the "Glue" of Tensor Parallelism

The AllReduce autograd operator appears decisively in two places:

1. DDP gradient sync (Note 02): as each parameter's gradient becomes ready, it effectively executes a "gradient AllReduce" — the same thing as the backward here.

2. Row Parallel cross-GPU reduction in tensor parallelism (expanded in Note 04): after splitting the weight by rows, each GPU computes a partial sum; forward AllReduce assembles the full output; backward AllReduce keeps each GPU's weight gradient consistent. One operator, communicating in both directions, re-glues the split matrix into a complete computation — this is the f (forward) and b (backward) communication layers in Megatron-LM.

4. Results

Running on 4 GPUs with x = ones(4), after the forward AllReduce y = ones(4) * 4, loss = y.sum(), and loss.backward(), every GPU should print:

x.grad = tensor([4., 4., 4., 4.])

Why 4, not 1: x.grad = output gradient (all ones) × the backward AllReduce result. The backward reduction sums the per-GPU gradients (each all ones) into 4 — direct evidence that "gradients are globally consistent and equal to the aggregated value."

5. Engineering Takeaways

1. Custom parallel operator = communication + autograd: for any custom layer that needs synchronization after partitioning (tensor parallelism, expert parallelism, sequence parallelism's communication points), the implementation pattern is autograd.Function + one collective in forward and one in backward. Understand this note and you understand the skeleton of Megatron/DeepSpeed communication layers.

2. The cost of backward communication: one backward AllReduce means gradient traffic proportional to data size. In tensor parallelism, many small communications kill throughput; engineering mitigates this with operator fusion, communication-compute overlap, and choosing partition directions that avoid communication in backward (e.g., Column Parallel) — which is also the design motivation behind "why Megatron cuts weights in a specific direction" (Note 04).

3. Gradient-correctness validation is the first gate of any parallel change: before shipping any custom parallel layer, run "N-GPU gradient consistency vs single-GPU baseline" (the x.grad print here is the minimal form), then talk performance.

Next: now we can actually split matrices — Note 04 implements tensor parallelism with both Column and Row cuts and reproduces Megatron's FFN: LLM Lab Notes 04: Tensor Parallelism 101.