Hey @MRGAO1996
Does your model use any buffers (e.g., running mean in BatchNorm)?
As you have wrapped the second forward with with torch.no_grad():, rank 0 will skip the code in if torch.is_grad_enabled() and ... branch (see DDP forward below), but the behavior of the next forward will be different on two ranks, as rank 0 would skip _sync_params() but rank 1 would execute that. But this method should only make a difference when your model has buffers.

  
      github.com
  
  
    pytorch/pytorch/blob/d5ae4a07ef5b2e77cf51737fb0a3aafc2e71231d/torch/nn/parallel/distributed.py#L562-L590 2

def forward(self, *inputs, **kwargs):
    if self.require_forward_param_sync:
        self._sync_params()

    if self.device_ids:
        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            output = self.module(*inputs[0], **kwargs[0])
        else:
            outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
            output = self.gather(outputs, self.output_device)
    else:
        output = self.module(*inputs, **kwargs)

    if torch.is_grad_enabled() and self.require_backward_grad_sync:
        self.require_forward_param_sync = True
        # We'll return the output object verbatim since it is a freeform
        # object. We need to find any tensors in this object, though,
        # because we need to figure out which parameters were used during
        # this forward pass, to ensure we short circuit reduction for any


  This file has been truncated. show original