Distributed

Distributed PyTorch starts from one idea, worth stating before any API: every rank runs the exact same script. There's no separate driver process orchestrating workers from outside; process 0 runs the training script, and so does process 7, and what tells each one who it is and how many others exist is init_process_group, backed by c10d, PyTorch's collective communication layer. Call it once per process and every rank gets a rank number and a world size, the two facts everything else in this chapter is built on.

That shape, the same script running once per participant rather than one controller dispatching to many workers, is called multi-controller, and it isn't unique to PyTorch. The xla path's McJAX chapter describes the identical architecture from the other framework's side: PyTorch's default distributed world is multi-controller too, arrived at independently because the same problem, many devices, one program, tends to produce the same answer.

Every rank runs your script; the collectives are the meeting points.

DistributedDataParallel is the simpler of the two strategies to reason about: every rank holds a full copy of the model, runs its own forward and backward on its own shard of data, and then all-reduces the gradients so every replica ends the step with the identical update. DDP doesn't wait for the whole backward pass to finish before it starts communicating either. Gradients are bucketed, and each bucket's all-reduce fires the moment that bucket's gradients are ready, overlapping communication with the backward computation still running behind it.

FSDP takes the opposite trade. Instead of replicating the full model on every rank, it shards parameters, gradients, and optimizer state across the group, and gathers a parameter back to its full size only when a layer is about to use it. DDP spends memory, a full copy everywhere, to keep the programming model simple. FSDP spends complexity and extra communication to keep memory down. Neither is strictly better: the choice comes down to which resource you have less of.

Once a job needs more than one axis of parallelism, data-parallel replicas along one dimension, tensor-parallel shards along another, PyTorch names the layout with DeviceMesh, and the tensor parallel APIs built on top of it place a module's weights across that mesh by axis. The idea underneath, a named grid of devices with a spec saying which axis shards which dimension, is the same mesh the jax path's chapter 10 teaches. PyTorch arrives at it from the mutation side of the fence; the destination is the same.

None of this needs multiple machines to observe. The snippet below ran on this machine as a single process standing in for a world of one, over gloo, PyTorch's CPU-friendly collective backend. all_reduce on a world of one can't change any values, there's nobody else to reduce with, but the collective genuinely executes: the process group stands up, the operation dispatches, and the tensor comes back unchanged because unchanged is the correct answer for a group of size one, not because nothing happened.

import torch
import torch.distributed as dist

dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29511", rank=0, world_size=1)
x = torch.ones(4)
dist.all_reduce(x)     # a world of one: unchanged, but the collective ran
print(x)               # tensor([1., 1., 1., 1.])
dist.destroy_process_group()