Distributed

Active contributors: kwen2501, weifengpy, fduwjj, H-Huang

Purpose

torch.distributed is PyTorch's distributed training stack. It spans low-level collective communication (NCCL/Gloo/UCC/MPI), high-level data and tensor parallelism (DDP, FSDP/FSDP2, DTensor, TP/SP), pipeline parallelism, RPC, multi-node launching (TorchElastic / torchrun), and checkpointing.

Python lives at torch/distributed/; the C++ runtime (process groups, work objects, the c10d backend interface) lives at torch/csrc/distributed/c10d/. The set of subsystems is wide enough that this page is more of a tour than an exhaustive walkthrough.

Directory layout

Path	Contents
torch/distributed/__init__.py	Public surface
torch/distributed/distributed_c10d.py	High-level Python collectives (all_reduce, etc.) — the largest file at ~260K lines
torch/distributed/device_mesh.py	DeviceMesh (~80K lines)
torch/distributed/_functional_collectives.py	Functional / traceable collectives
torch/distributed/fsdp/	FSDP v1
torch/distributed/_composable/fsdp/	FSDP2 (fully_shard)
torch/distributed/_composable/replicate*.py	Replicate / replicate_with_fsdp
torch/distributed/tensor/	DTensor + TP/SP
torch/distributed/pipelining/	Pipeline parallelism
torch/distributed/rpc/	Distributed RPC
torch/distributed/elastic/	Elastic launcher and rendezvous
torch/distributed/checkpoint/	Distributed checkpointing
torch/distributed/algorithms/	DDP comm hooks, ZeroRedundancyOptimizer, etc.
torch/distributed/run.py	torchrun entry point
torch/distributed/launch.py	Older torch.distributed.launch
torch/distributed/_symmetric_memory/	Symmetric memory primitives for cross-rank shared buffers
torch/distributed/flight_recorder/	NCCL flight recorder for debugging hangs
torch/csrc/distributed/c10d/	C++ side: process groups, NCCL/Gloo/UCC/MPI backends
torch/csrc/distributed/rpc/	C++ RPC implementation

Key abstractions

Type	File	Purpose
ProcessGroup	torch/csrc/distributed/c10d/ProcessGroup.hpp	Group of ranks that can communicate
Backend (C++)	torch/csrc/distributed/c10d/Backend.hpp	NCCL/Gloo/UCC/MPI implementation
Work	torch/csrc/distributed/c10d/Work.hpp	Async result of a collective
DeviceMesh	torch/distributed/device_mesh.py	n-dimensional rank grid
DTensor	torch/distributed/tensor/_api.py	Tensor with sharding/placement on a mesh
Placement (Shard/Replicate/Partial)	torch/distributed/tensor/placement_types.py	How a DTensor is distributed
FullyShardedDataParallel (FSDP1)	torch/distributed/fsdp/fully_sharded_data_parallel.py	FSDP module wrapper
fully_shard (FSDP2)	torch/distributed/_composable/fsdp/fully_shard.py	Per-parameter FSDP
DistributedDataParallel (DDP)	torch/nn/parallel/distributed.py	All-reduce data parallelism
PipelineSchedule	torch/distributed/pipelining/schedules.py	Pipeline schedule (1F1B, GPipe, Looped 1F1B, …)

How it works

c10d: process groups

A process group is a set of ranks that can perform collective ops together. dist.init_process_group("nccl", ...) initializes the default PG; subgroups are created via dist.new_group(...) or via DeviceMesh.

graph LR
    PyAPI["dist.all_reduce(t)"] --> DistC10D["distributed_c10d.py"]
    DistC10D -->|via OperatorEntry|Disp[Dispatcher]
    Disp --> CPP[ProcessGroup C++]
    CPP -->|backend = nccl| NCCL[NCCL]
    CPP -->|backend = gloo| Gloo[Gloo]
    CPP -->|backend = ucc| UCC[UCC]
    NCCL --> Net[Network]
    Gloo --> Net
    UCC --> Net

Backends differ in what they support:

• NCCL — GPU collectives, the production choice for CUDA training.
• Gloo — CPU collectives and a fallback for CUDA, used for one-off init and CPU-only setups.
• UCC — Unified Collective Communications.
• MPI — Used in HPC contexts.

Work objects implement .wait() for synchronous semantics and integrate with CUDA streams when needed.

Functional collectives

A second, traceable API at torch/distributed/_functional_collectives.py exposes ATen-op-level versions of every collective (aten::all_reduce, aten::all_gather_into_tensor, etc.) so they can flow through torch.compile and torch.export. These are now the recommended way for compiled-mode distributed code.

DDP

DistributedDataParallel (torch/nn/parallel/distributed.py) replicates a model on every rank. After backward, gradients are bucketed and all-reduced asynchronously to overlap with compute. DDP is mature, simple, and still the default for moderate-scale data parallelism.

FSDP

FSDP shards parameters across ranks instead of replicating them. Each rank holds 1/N of every parameter; before a forward pass the parameter is all-gathered on the fly, used, and freed. Backward similarly gathers and reduce-scatters.

• FSDP1 (torch/distributed/fsdp/) — flat-parameter FSDP. Each FullyShardedDataParallel instance owns a flat buffer of grouped parameters.
• FSDP2 (torch/distributed/_composable/fsdp/fully_shard.py) — per-parameter FSDP. No flat buffers; integrates with DTensor; composes cleanly with TP, PP, and torch.compile. The recommended path going forward.

DTensor

torch/distributed/tensor/ introduces DTensor: a tensor that knows its DeviceMesh and placements (e.g., Shard(0) along mesh dim 0, Replicate() along dim 1). Every ATen op on a DTensor consults the sharding propagation rules in _ops/ to pick a redistribution strategy and emit the right collectives.

DTensor is the foundation for tensor-parallel and sequence-parallel training and is the canonical way to express n-D parallelism in PyTorch 2.x.

Pipelining

torch/distributed/pipelining/ splits a model into sequential stages, places each stage on a different rank, and runs micro-batches with one of several schedules (GPipe, 1F1B, Interleaved 1F1B, Looped, Zero-Bubble). Schedules are encoded as small DSL programs.

RPC

torch/distributed/rpc/ provides rpc_sync, rpc_async, remote, and the RRef (remote reference). Used by classic distributed actor patterns (parameter server, model parallel via RPC). Less central in modern training stacks than DTensor/FSDP/PP.

TorchElastic and torchrun

torch/distributed/elastic/ and torch/distributed/run.py provide the multi-node launcher. torchrun --nnodes=8 --nproc_per_node=8 train.py is the modern entry point. Rendezvous (etcd, c10d store, static) handles rank assignment.

Checkpointing

torch/distributed/checkpoint/ provides distributed save/load: each rank writes its shards in parallel to a shared filesystem; on load the planner re-distributes shards according to the new world size. Successor to the old torch.save(state_dict) for distributed jobs.

Symmetric memory and flight recorder

• torch/distributed/_symmetric_memory/ exposes IPC-shared buffers across ranks, used by experimental low-latency communication kernels.
• torch/distributed/flight_recorder/ reads NCCL's flight recorder ring buffer to diagnose hangs and unmatched collectives.

Integration points

• Autograd. DDP, FSDP, and DTensor all interact with autograd via accumulation hooks and custom Functions.
• torch.compile. Functional collectives + DTensor are designed to flow through Dynamo and Inductor unchanged.
• torch.export / serving. Distributed checkpoint and torch.distributed.tensor interoperate with the export pipeline.
• Mobile/edge. Out of scope: distributed isn't part of mobile builds.

Entry points for modification

• New collective op → declare in aten/src/ATen/native/native_functions.yaml, register in torch/csrc/distributed/c10d/Ops.cpp, expose in _functional_collectives.py, surface a Python wrapper in distributed_c10d.py.
• New process group backend → subclass c10d::Backend and register via the C++ backend registry.
• New DTensor sharding rule → add to torch/distributed/tensor/_ops/.
• New FSDP feature → most active dev is in torch/distributed/_composable/fsdp/.
• For debugging hangs, TORCH_NCCL_DESYNC_DEBUG=1 and TORCH_NCCL_TRACE_BUFFER_SIZE enable the flight recorder.

Key source files

File	Purpose
torch/distributed/distributed_c10d.py	High-level Python collectives
torch/distributed/device_mesh.py	DeviceMesh
torch/distributed/_functional_collectives.py	Traceable collectives
torch/distributed/fsdp/fully_sharded_data_parallel.py	FSDP1
torch/distributed/_composable/fsdp/fully_shard.py	FSDP2
torch/distributed/tensor/_api.py	DTensor
torch/distributed/tensor/placement_types.py	Placement
torch/distributed/pipelining/schedules.py	Pipeline schedules
torch/distributed/run.py	torchrun
torch/distributed/checkpoint/	Distributed checkpoint
torch/csrc/distributed/c10d/ProcessGroup.hpp	C++ ProcessGroup
torch/csrc/distributed/c10d/Backend.hpp	C++ backend interface
torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp	NCCL backend
torch/csrc/distributed/c10d/ProcessGroupGloo.cpp	Gloo backend

For the user-facing "how do I scale up training" feature page see Features / Distributed training.