DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False)¶
Implements distributed data parallelism that is based on
torch.distributedpackage at the module level.
This container parallelizes the application of the given module by splitting the input across the specified devices by chunking in the batch dimension. The module is replicated on each machine and each device, and each such replica handles a portion of the input. During the backwards pass, gradients from each node are averaged.
The batch size should be larger than the number of GPUs used locally.
See also: Basics and Use nn.parallel.DistributedDataParallel instead of multiprocessing or nn.DataParallel. The same constraints on input as in
Creation of this class requires that
torch.distributedto be already initialized, by calling
DistributedDataParallelis proven to be significantly faster than
torch.nn.DataParallelfor single-node multi-GPU data parallel training.
Here is how to use it: on each host with N GPUs, you should spawn up N processes, while ensuring that each process individually works on a single GPU from 0 to N-1. Therefore, it is your job to ensure that your training script operates on a single given GPU by calling:
where i is from 0 to N-1. In each process, you should refer the following to construct this module:
>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...') >>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)
In order to spawn up multiple processes per node, you can use either
Please refer to PyTorch Distributed Overview for a brief introduction to all features related to distributed training.
ncclbackend is currently the fastest and highly recommended backend to be used with Multi-Process Single-GPU distributed training and this applies to both single-node and multi-node distributed training
This module also supports mixed-precision distributed training. This means that your model can have different types of parameters such as mixed types of fp16 and fp32, the gradient reduction on these mixed types of parameters will just work fine. Also note that
ncclbackend is currently the fastest and highly recommended backend for fp16/fp32 mixed-precision training.
If you use
torch.saveon one process to checkpoint the module, and
torch.loadon some other processes to recover it, make sure that
map_locationis configured properly for every process. Without
torch.loadwould recover the module to devices where the module was saved from.
This module works only with the
Constructor, forward method, and differentiation of the output (or a function of the output of this module) is a distributed synchronization point. Take that into account in case different processes might be executing different code.
This module assumes all parameters are registered in the model by the time it is created. No parameters should be added nor removed later. Same applies to buffers.
This module assumes all parameters are registered in the model of each distributed processes are in the same order. The module itself will conduct gradient all-reduction following the reverse order of the registered parameters of the model. In other words, it is users’ responsibility to ensure that each distributed process has the exact same model and thus the exact same parameter registration order.
This module allows parameters with non-rowmajor-contiguous strides. For example, your model may contain some parameters whose
torch.contiguous_formatand others whose format is
torch.channels_last. However, corresponding parameters in different processes must have the same strides.
This module doesn’t work with
torch.autograd.grad()(i.e. it will only work if gradients are to be accumulated in
.gradattributes of parameters).
If you plan on using this module with a
ncclbackend or a
gloobackend (that uses Infiniband), together with a DataLoader that uses multiple workers, please change the multiprocessing start method to
forkserver(Python 3 only) or
spawn. Unfortunately Gloo (that uses Infiniband) and NCCL2 are not fork safe, and you will likely experience deadlocks if you don’t change this setting.
Forward and backward hooks defined on
moduleand its submodules won’t be invoked anymore, unless the hooks are initialized in the
You should never try to change your model’s parameters after wrapping up your model with DistributedDataParallel. In other words, when wrapping up your model with DistributedDataParallel, the constructor of DistributedDataParallel will register the additional gradient reduction functions on all the parameters of the model itself at the time of construction. If you change the model’s parameters after the DistributedDataParallel construction, this is not supported and unexpected behaviors can happen, since some parameters’ gradient reduction functions might not get called.
Parameters are never broadcast between processes. The module performs an all-reduce step on gradients and assumes that they will be modified by the optimizer in all processes in the same way. Buffers (e.g. BatchNorm stats) are broadcast from the module in process of rank 0, to all other replicas in the system in every iteration.
If you are using DistributedDataParallel in conjunction with the Distributed RPC Framework, you should always use
torch.distributed.autograd.backward()to compute gradients and
torch.distributed.optim.DistributedOptimizerfor optimizing parameters.
>>> import torch.distributed.autograd as dist_autograd >>> from torch.nn.parallel import DistributedDataParallel as DDP >>> from torch import optim >>> from torch.distributed.optim import DistributedOptimizer >>> from torch.distributed.rpc import RRef >>> >>> t1 = torch.rand((3, 3), requires_grad=True) >>> t2 = torch.rand((3, 3), requires_grad=True) >>> rref = rpc.remote("worker1", torch.add, args=(t1, t2)) >>> ddp_model = DDP(my_model) >>> >>> # Setup optimizer >>> optimizer_params = [rref] >>> for param in ddp_model.parameters(): >>> optimizer_params.append(RRef(param)) >>> >>> dist_optim = DistributedOptimizer( >>> optim.SGD, >>> optimizer_params, >>> lr=0.05, >>> ) >>> >>> with dist_autograd.context() as context_id: >>> pred = ddp_model(rref.to_here()) >>> loss = loss_func(pred, loss) >>> dist_autograd.backward(context_id, loss) >>> dist_optim.step()
Using DistributedDataParallel in conjuction with the Distributed RPC Framework is experimental and subject to change.
module (Module) – module to be parallelized
device_ids (list of python:int or torch.device) – CUDA devices. This should only be provided when the input module resides on a single CUDA device. For single-device modules, the
i``th :attr:`module` replica is placed on ``device_ids[i]. For multi-device modules and CPU modules, device_ids must be None or an empty list, and input data for the forward pass must be placed on the correct device. (default: all devices for single-device modules)
output_device (int or torch.device) – device location of output for single-device CUDA modules. For multi-device modules and CPU modules, it must be None, and the module itself dictates the output location. (default: device_ids for single-device modules)
broadcast_buffers (bool) – flag that enables syncing (broadcasting) buffers of the module at beginning of the forward function. (default:
process_group – the process group to be used for distributed data all-reduction. If
None, the default process group, which is created by
`torch.distributed.init_process_group`, will be used. (default:
bucket_cap_mb – DistributedDataParallel will bucket parameters into multiple buckets so that gradient reduction of each bucket can potentially overlap with backward computation.
bucket_cap_mbcontrols the bucket size in MegaBytes (MB) (default: 25)
find_unused_parameters (bool) – Traverse the autograd graph of all tensors contained in the return value of the wrapped module’s
forwardfunction. Parameters that don’t receive gradients as part of this graph are preemptively marked as being ready to be reduced. Note that all
forwardoutputs that are derived from module parameters must participate in calculating loss and later the gradient computation. If they don’t, this wrapper will hang waiting for autograd to produce gradients for those parameters. Any outputs derived from module parameters that are otherwise unused can be detached from the autograd graph using
check_reduction – when setting to
True, it enables DistributedDataParallel to automatically check if the previous iteration’s backward reductions were successfully issued at the beginning of every iteration’s forward function. You normally don’t need this option enabled unless you are observing weird behaviors such as different ranks are getting different gradients, which should not happen if DistributedDataParallel is correctly used. (default:
~DistributedDataParallel.module (Module) – the module to be parallelized
>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...') >>> net = torch.nn.DistributedDataParallel(model, pg)
A context manager to disable gradient synchronizations across DDP processes. Within this context, gradients will be accumulated on module variables, which will later be synchronized in the first forward-backward pass exiting the context.
>>> ddp = torch.nn.DistributedDataParallel(model, pg) >>> with ddp.no_sync(): ... for input in inputs: ... ddp(input).backward() # no synchronization, accumulate grads ... ddp(another_input).backward() # synchronize grads