Source code for torch.nn.parallel.distributed_cpu

import torch
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
import torch.distributed as dist
from torch.nn.modules import Module
from collections import defaultdict
from torch.autograd import Variable
import torch.utils.hooks

[docs]class DistributedDataParallelCPU(Module): r"""Implements distributed data parallelism for CPU at the module level. This module supports the ``mpi`` and ``gloo`` backends. 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 such replica handles a portion of the input. During the backwards pass, gradients from each node are averaged. This module could be used in conjunction with the DistributedSampler, (see :class:``) which will load a subset of the original dataset for each node with the same batch size. So strong scaling should be configured like this: n = 1, batch size = 12 n = 2, batch size = 64 n = 4, batch size = 32 n = 8, batch size = 16 Creation of this class requires the distributed package to be already initialized in the process group mode (see :func:`torch.distributed.init_process_group`). .. warning:: 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 node might be executing different code. .. warning:: This module assumes all parameters are registered in the model by the time it is created. No parameters should be added nor removed later. .. warning:: This module assumes all gradients are dense. .. warning:: This module doesn't work with :func:`torch.autograd.grad` (i.e. it will only work if gradients are to be accumulated in ``.grad`` attributes of parameters). .. warning:: Forward and backward hooks defined on :attr:`module` and its submodules won't be invoked anymore, unless the hooks are initialized in the :meth:`forward` method. .. note:: Parameters are broadcast between nodes in the __init__() function. The module performs an all-reduce step on gradients and assumes that they will be modified by the optimizer in all nodes in the same way. Args: module: module to be parallelized Example:: >>> torch.distributed.init_process_group(world_size=4, init_method='...') >>> net = torch.nn.DistributedDataParallelCPU(model) """ def __init__(self, module): super(DistributedDataParallelCPU, self).__init__() self.module = module self.sync_parameters() def allreduce_params(): if self.needs_reduction: self.needs_reduction = False buckets = defaultdict(list) for param in self.module.parameters(): if param.requires_grad and param.grad is not None: tp = type( buckets[tp].append(param) for bucket in buckets.values(): grads = [ for param in bucket] coalesced = _flatten_dense_tensors(grads) dist.all_reduce(coalesced) coalesced /= dist.get_world_size() for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)): buf.copy_(synced) for param in list(self.module.parameters()): @torch.utils.hooks.unserializable_hook def allreduce_hook(*unused): Variable._execution_engine.queue_callback(allreduce_params) if param.requires_grad: param.register_hook(allreduce_hook) def sync_parameters(self): for param in self.module.parameters(): dist.broadcast(, 0) def forward(self, *inputs, **kwargs): self.needs_reduction = True return self.module(*inputs, **kwargs)


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources