Rate this Page
★★★★★

Getting Started with Distributed Data Parallel#

Created On: Apr 23, 2019 | Last Updated: Sep 23, 2025 | Last Verified: Nov 05, 2024

Author:Shen Li

Edited by:Joe Zhu,Chirag Pandya

Note

edit View and edit this tutorial ingithub.

Prerequisites:

DistributedDataParallel(DDP) is a powerful module in PyTorch that allows you to parallelize your model acrossmultiple machines, making it perfect for large-scale deep learning applications.To use DDP, you’ll need to spawn multiple processes and create a single instance of DDP per process.

But how does it work? DDP uses collective communications from thetorch.distributedpackage to synchronize gradients and buffers across all processes. This means that each process will haveits own copy of the model, but they’ll all work together to train the model as if it were on a single machine.

To make this happen, DDP registers an autograd hook for each parameter in the model.When the backward pass is run, this hook fires and triggers gradient synchronization across all processes.This ensures that each process has the same gradients, which are then used to update the model.

For more information on how DDP works and how to use it effectively, be sure to check out theDDP design note.With DDP, you can train your models faster and more efficiently than ever before!

The recommended way to use DDP is to spawn one process for each model replica. The model replica can spanmultiple devices. DDP processes can be placed on the same machine or across machines. Note that GPU devicescannot be shared across DDP processes (i.e. one GPU for one DDP process).

In this tutorial, we’ll start with a basic DDP use case and then demonstrate more advanced use cases,including checkpointing models and combining DDP with model parallel.

Note

The code in this tutorial runs on an 8-GPU server, but it can be easilygeneralized to other environments.

Comparison betweenDataParallel andDistributedDataParallel#

Before we dive in, let’s clarify why you would consider usingDistributedDataParalleloverDataParallel, despite its added complexity:

  • First,DataParallel is single-process, multi-threaded, but it only works on asingle machine. In contrast,DistributedDataParallel is multi-process and supportsboth single- and multi- machine training.Due to GIL contention across threads, per-iteration replicated model, and additional overhead introduced byscattering inputs and gathering outputs,DataParallel is usuallyslower thanDistributedDataParallel even on a single machine.

  • Recall from theprior tutorialthat if your model is too large to fit on a single GPU, you must usemodel parallelto split it across multiple GPUs.DistributedDataParallel works withmodel parallel, whileDataParallel does not at this time. When DDP is combinedwith model parallel, each DDP process would use model parallel, and all processescollectively would use data parallel.

Basic Use Case#

To create a DDP module, you must first set up process groups properly. More details canbe found inWriting Distributed Applications with PyTorch.

importosimportsysimporttempfileimporttorchimporttorch.distributedasdistimporttorch.nnasnnimporttorch.optimasoptimimporttorch.multiprocessingasmpfromtorch.nn.parallelimportDistributedDataParallelasDDP# On Windows platform, the torch.distributed package only# supports Gloo backend, FileStore and TcpStore.# For FileStore, set init_method parameter in init_process_group# to a local file. Example as follow:# init_method="file:///f:/libtmp/some_file"# dist.init_process_group(#    "gloo",#    rank=rank,#    init_method=init_method,#    world_size=world_size)# For TcpStore, same way as on Linux.defsetup(rank,world_size):os.environ['MASTER_ADDR']='localhost'os.environ['MASTER_PORT']='12355'# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__# such as CUDA, MPS, MTIA, or XPU.acc=torch.accelerator.current_accelerator()backend=torch.distributed.get_default_backend_for_device(acc)# initialize the process groupdist.init_process_group(backend,rank=rank,world_size=world_size)defcleanup():dist.destroy_process_group()

Now, let’s create a toy module, wrap it with DDP, and feed it some dummyinput data. Please note, as DDP broadcasts model states from rank 0 process toall other processes in the DDP constructor, you do not need to worry aboutdifferent DDP processes starting from different initial model parameter values.

classToyModel(nn.Module):def__init__(self):super(ToyModel,self).__init__()self.net1=nn.Linear(10,10)self.relu=nn.ReLU()self.net2=nn.Linear(10,5)defforward(self,x):returnself.net2(self.relu(self.net1(x)))defdemo_basic(rank,world_size):print(f"Running basic DDP example on rank{rank}.")setup(rank,world_size)# create model and move it to GPU with id rankmodel=ToyModel().to(rank)ddp_model=DDP(model,device_ids=[rank])loss_fn=nn.MSELoss()optimizer=optim.SGD(ddp_model.parameters(),lr=0.001)optimizer.zero_grad()outputs=ddp_model(torch.randn(20,10))labels=torch.randn(20,5).to(rank)loss_fn(outputs,labels).backward()optimizer.step()cleanup()print(f"Finished running basic DDP example on rank{rank}.")defrun_demo(demo_fn,world_size):mp.spawn(demo_fn,args=(world_size,),nprocs=world_size,join=True)

As you can see, DDP wraps lower-level distributed communication details andprovides a clean API as if it were a local model. Gradient synchronizationcommunications take place during the backward pass and overlap with thebackward computation. When thebackward() returns,param.grad alreadycontains the synchronized gradient tensor. For basic use cases, DDP onlyrequires a few more lines of code to set up the process group. When applying DDP to moreadvanced use cases, some caveats require caution.

Skewed Processing Speeds#

In DDP, the constructor, the forward pass, and the backward pass aredistributed synchronization points. Different processes are expected to launchthe same number of synchronizations and reach these synchronization points inthe same order and enter each synchronization point at roughly the same time.Otherwise, fast processes might arrive early and timeout while waiting forstragglers. Hence, users are responsible for balancing workload distributionsacross processes. Sometimes, skewed processing speeds are inevitable due to,e.g., network delays, resource contentions, or unpredictable workload spikes. Toavoid timeouts in these situations, make sure that you pass a sufficientlylargetimeout value when callinginit_process_group.

Save and Load Checkpoints#

It’s common to usetorch.save andtorch.load to checkpoint modulesduring training and recover from checkpoints. SeeSAVING AND LOADING MODELSfor more details. When using DDP, one optimization is to save the model inonly one process and then load it on all processes, reducing write overhead.This works because all processes start from the same parameters andgradients are synchronized in backward passes, and hence optimizers should keepsetting parameters to the same values.If you use this optimization (i.e. save on one process but restore on all), make sure no process startsloading before the saving is finished. Additionally, whenloading the module, you need to provide an appropriatemap_locationargument to prevent processes from stepping into others’ devices. Ifmap_locationis missing,torch.load will first load the module to CPU and then copy eachparameter to where it was saved, which would result in all processes on thesame machine using the same set of devices. For more advanced failure recoveryand elasticity support, please refer toTorchElastic.

defdemo_checkpoint(rank,world_size):print(f"Running DDP checkpoint example on rank{rank}.")setup(rank,world_size)model=ToyModel().to(rank)ddp_model=DDP(model,device_ids=[rank])CHECKPOINT_PATH=tempfile.gettempdir()+"/model.checkpoint"ifrank==0:# All processes should see same parameters as they all start from same# random parameters and gradients are synchronized in backward passes.# Therefore, saving it in one process is sufficient.torch.save(ddp_model.state_dict(),CHECKPOINT_PATH)# Use a barrier() to make sure that process 1 loads the model after process# 0 saves it.dist.barrier()# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__# such as CUDA, MPS, MTIA, or XPU.acc=torch.accelerator.current_accelerator()# configure map_location properlymap_location={f'{acc}:0':f'{acc}:{rank}'}ddp_model.load_state_dict(torch.load(CHECKPOINT_PATH,map_location=map_location,weights_only=True))loss_fn=nn.MSELoss()optimizer=optim.SGD(ddp_model.parameters(),lr=0.001)optimizer.zero_grad()outputs=ddp_model(torch.randn(20,10))labels=torch.randn(20,5).to(rank)loss_fn(outputs,labels).backward()optimizer.step()# Not necessary to use a dist.barrier() to guard the file deletion below# as the AllReduce ops in the backward pass of DDP already served as# a synchronization.ifrank==0:os.remove(CHECKPOINT_PATH)cleanup()print(f"Finished running DDP checkpoint example on rank{rank}.")

Combining DDP with Model Parallelism#

DDP also works with multi-GPU models. DDP wrapping multi-GPU models is especiallyhelpful when training large models with a huge amount of data.

classToyMpModel(nn.Module):def__init__(self,dev0,dev1):super(ToyMpModel,self).__init__()self.dev0=dev0self.dev1=dev1self.net1=torch.nn.Linear(10,10).to(dev0)self.relu=torch.nn.ReLU()self.net2=torch.nn.Linear(10,5).to(dev1)defforward(self,x):x=x.to(self.dev0)x=self.relu(self.net1(x))x=x.to(self.dev1)returnself.net2(x)

When passing a multi-GPU model to DDP,device_ids andoutput_devicemust NOT be set. Input and output data will be placed in proper devices byeither the application or the modelforward() method.

defdemo_model_parallel(rank,world_size):print(f"Running DDP with model parallel example on rank{rank}.")setup(rank,world_size)# setup mp_model and devices for this processdev0=rank*2dev1=rank*2+1mp_model=ToyMpModel(dev0,dev1)ddp_mp_model=DDP(mp_model)loss_fn=nn.MSELoss()optimizer=optim.SGD(ddp_mp_model.parameters(),lr=0.001)optimizer.zero_grad()# outputs will be on dev1outputs=ddp_mp_model(torch.randn(20,10))labels=torch.randn(20,5).to(dev1)loss_fn(outputs,labels).backward()optimizer.step()cleanup()print(f"Finished running DDP with model parallel example on rank{rank}.")if__name__=="__main__":n_gpus=torch.accelerator.device_count()assertn_gpus>=2,f"Requires at least 2 GPUs to run, but got{n_gpus}"world_size=n_gpusrun_demo(demo_basic,world_size)run_demo(demo_checkpoint,world_size)world_size=n_gpus//2run_demo(demo_model_parallel,world_size)

Initialize DDP with torch.distributed.run/torchrun#

We can leverage PyTorch Elastic to simplify the DDP code and initialize the job more easily.Let’s still use the Toymodel example and create a file namedelastic_ddp.py.

importosimporttorchimporttorch.distributedasdistimporttorch.nnasnnimporttorch.optimasoptimfromtorch.nn.parallelimportDistributedDataParallelasDDPclassToyModel(nn.Module):def__init__(self):super(ToyModel,self).__init__()self.net1=nn.Linear(10,10)self.relu=nn.ReLU()self.net2=nn.Linear(10,5)defforward(self,x):returnself.net2(self.relu(self.net1(x)))defdemo_basic():torch.accelerator.set_device_index(int(os.environ["LOCAL_RANK"]))acc=torch.accelerator.current_accelerator()backend=torch.distributed.get_default_backend_for_device(acc)dist.init_process_group(backend)rank=dist.get_rank()print(f"Start running basic DDP example on rank{rank}.")# create model and move it to GPU with id rankdevice_id=rank%torch.accelerator.device_count()model=ToyModel().to(device_id)ddp_model=DDP(model,device_ids=[device_id])loss_fn=nn.MSELoss()optimizer=optim.SGD(ddp_model.parameters(),lr=0.001)optimizer.zero_grad()outputs=ddp_model(torch.randn(20,10))labels=torch.randn(20,5).to(device_id)loss_fn(outputs,labels).backward()optimizer.step()dist.destroy_process_group()print(f"Finished running basic DDP example on rank{rank}.")if__name__=="__main__":demo_basic()

One can then run atorch elastic/torchrun commandon all nodes to initialize the DDP job created above:

torchrun--nnodes=2--nproc_per_node=8--rdzv_id=100--rdzv_backend=c10d--rdzv_endpoint=$MASTER_ADDR:29400elastic_ddp.py

In the example above, we are running the DDP script on two hosts and we run with 8 processes on each host. That is, weare running this job on 16 GPUs. Note that$MASTER_ADDR must be the same across all nodes.

Heretorchrun will launch 8 processes and invokeelastic_ddp.pyon each process on the node it is launched on, but user also needs to apply clustermanagement tools like slurm to actually run this command on 2 nodes.

For example, on a SLURM enabled cluster, we can write a script to run the command aboveand setMASTER_ADDR as:

exportMASTER_ADDR=$(scontrolshowhostname${SLURM_NODELIST}|head-n1)

Then we can just run this script using the SLURM command:srun--nodes=2./torchrun_script.sh.

This is just an example; you can choose your own cluster scheduling tools to initiate thetorchrun job.

For more information about Elastic run, please see thequick start document.