Gradient hooks with DataParallel and DDP

agolynski · April 12, 2021, 5:02pm

Hi,

Seems like this question had come up before. Seems like this limitation is intrinsic to DataParallel, e.g. see Yanli_Zhao answer here:

Also see warnings here

github.com

pytorch/pytorch/blob/main/torch/nn/parallel/data_parallel.py#L67


      
          
          .. warning::
              It is recommended to use :class:`~torch.nn.parallel.DistributedDataParallel`,
              instead of this class, to do multi-GPU training, even if there is only a single
              node. See: :ref:`cuda-nn-ddp-instead` and :ref:`ddp`.
          
          Arbitrary positional and keyword inputs are allowed to be passed into
          DataParallel but some types are specially handled. tensors will be
          **scattered** on dim specified (default 0). tuple, list and dict types will
          be shallow copied. The other types will be shared among different threads
          and can be corrupted if written to in the model's forward pass.
          
          The parallelized :attr:`module` must have its parameters and buffers on
          ``device_ids[0]`` before running this :class:`~torch.nn.DataParallel`
          module.
          
          .. warning::
              In each forward, :attr:`module` is **replicated** on each device, so any
              updates to the running module in ``forward`` will be lost. For example,
              if :attr:`module` has a counter attribute that is incremented in each
              ``forward``, it will always stay at the initial value because the update