Distributed training

What it is

Scaling a PyTorch training job from one GPU to many is what torch.distributed is for. The framework offers several composable strategies; this page is the user-level orientation. For the implementation see Systems / Distributed.

The strategies

Data parallelism

Strategy	When to use
nn.parallel.DistributedDataParallel (DDP)	Model fits on one GPU. Simple, fast, mature.
_composable.replicate (DDP2)	Modern composable replacement for DDP.
nn.DataParallel	Don't. Single-process multi-GPU; deprecated.

DDP replicates the model on every rank. After backward, gradients are bucketed and all-reduced asynchronously to overlap with compute. The Python entry point is torch/nn/parallel/distributed.py; comm hooks for things like compression and async accumulation live in torch/distributed/algorithms/ddp_comm_hooks/.

Sharded data parallelism

Strategy	When to use
distributed.fsdp.FullyShardedDataParallel (FSDP1)	Model doesn't fit in one GPU. Mature, flat-buffer based.
_composable.fsdp.fully_shard (FSDP2)	Newer per-parameter FSDP. The recommended path going forward.

FSDP shards parameters across data-parallel ranks; each rank only holds 1/N of every parameter, gathering full parameters on the fly. Activation memory dominates instead of parameter memory.

Tensor / sequence parallelism

DTensor (torch/distributed/tensor/) is the modern way to express tensor parallelism. Build a DeviceMesh, annotate parameters with Shard/Replicate/Partial placements, and let the framework insert the right collectives:

from torch.distributed.tensor import DeviceMesh, Shard, Replicate
from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel

mesh = DeviceMesh("cuda", torch.arange(world_size))
parallelize_module(layer, mesh, {
    "qkv": ColwiseParallel(),
    "proj": RowwiseParallel(),
})

Sequence parallelism (SP) is a small variation that shards the sequence dim through a transformer block; the helpers in torch.distributed.tensor.parallel cover the canonical patterns.

Pipeline parallelism

torch/distributed/pipelining/ splits a model into sequential stages, places each on a different group of ranks, and runs micro-batches with one of several schedules: ScheduleGPipe, Schedule1F1B, ScheduleInterleaved1F1B, ScheduleLooped1F1B, ScheduleZBVZeroBubble. The schedule is a small DSL of (forward, backward, microbatch_idx, stage_idx) tuples.

N-D parallelism

The strategies compose. A typical large-LLM setup might be:

• 8-way tensor parallel,
• 4-way pipeline parallel,
• the rest data-parallel via FSDP2,

all expressed on a 3-D DeviceMesh.

Launching jobs

torchrun --nnodes=4 --nproc_per_node=8 --rdzv_backend=c10d --rdzv_endpoint=$MASTER train.py

torchrun is torch/distributed/run.py; it handles process spawning, environment, rendezvous, and restarts. TorchElastic provides the elastic / re-spawn-on-failure semantics for long-running jobs.

Checkpointing

torch.distributed.checkpoint writes per-shard files in parallel to a shared filesystem and re-distributes shards on load:

import torch.distributed.checkpoint as dcp
dcp.save({"model": model, "optim": optim}, checkpoint_id="step-100")
dcp.load({"model": model, "optim": optim}, checkpoint_id="step-100")

It works correctly when the world size or sharding changes between save and load.

Compile + distributed

torch.compile and the modern distributed stack are designed to compose:

• Functional collectives (torch.distributed._functional_collectives) are traceable ATen ops.
• DTensor's collective insertion happens before tracing, so the compiled graph contains explicit collectives.
• FSDP2 + torch.compile is the recommended path for compiled large-model training.

Debugging

Tool	Purpose
TORCH_DISTRIBUTED_DEBUG=DETAIL	Verbose collective logging
TORCH_NCCL_DESYNC_DEBUG=1	Detect collective desync (mismatched calls)
TORCH_NCCL_TRACE_BUFFER_SIZE=1000000	Enable NCCL flight recorder for hang debugging
torch.distributed.flight_recorder	Decoder for the flight recorder dumps
torch.distributed.checkpoint.utils	Checkpoint sanity checks

Where to read next

• Systems / Distributed — implementation.
• The CONTRIBUTING.md in torch/distributed/ (torch/distributed/CONTRIBUTING.md) — contributor guide.
• The PyTorch tutorials at https://pytorch.org/tutorials/distributed/index.html.