DistributedDataParallel (DDP) wraps a PyTorch model so synchronized data-parallel training can run across processes on one machine or across machines. Each process joins a training group, while DDP coordinates distributed communication, gradient aggregation, and local parameter updates. section 4.2.3
A typical DDP setup has these responsibilities:
- Initialize the process group: provide each process's rank, the total world size, and the master's address and port; processes wait until the group is complete. section 4.2.3
- Place the model on its device and wrap it with DDP, such as DDP(model, device_ids=[rank]). section 4.2.3
- Run the forward pass and call backward; when backward returns, parameter gradients contain synchronized tensors across workers. section 4.2.3
- Tune gradient bucket size when needed: smaller or larger buckets trade off device saturation against network saturation during synchronization. section 4.2.2
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = 'xxxx'
os.environ['MASTER_PORT'] = 'xxx'
# initialize the process group, "gloo" is one of the communication
# backends Pytorch supports, it also supports MPI and NCCL.
# rank is the process’s rank, it's a globally unique id
# for this process. rank=0 means master process.
# world_size is the total number of processes in this training group.
dist.init_process_group("gloo", rank=rank, world_size=world_size)
def cleanup():
dist.destroy_process_group()
Use the setup listing to inspect environment configuration, process-group initialization, supported communication backends, and cleanup. section 4.2.3
import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP # create model and move it to GPU model = DpModel().to(device) # wrap the model with DDP ddp_model = DDP(model, device_ids=[rank]) outputs = ddp_model(data) #1 # compute the loss and sync gradient with other workers. # when 'backward' function returns, the param.grad already # contains synchronized gradient tensor loss_fn(outputs, labels).backward()
Use the model listing to see device placement, DDP wrapping, forward execution, and synchronized gradients after backward. section 4.2.3
