Hey @epignatelli
When did you take the snapshot of self.k? DDP will broadcast all buffers from rank 0 process to other processes right before calling Mode.forward. See the code below. So given the above code, the buffer should be consistent across all processes before the self.callback is launched.

  
      github.com
  
  
    pytorch/pytorch/blob/79cfd85987c128b672850b64d6680719e4eda552/torch/nn/parallel/distributed.py#L695-L713

# module buffer sync
if self.broadcast_buffers and len(self.modules_buffers[0]) > 0:
    # Synchronize buffers across processes.
    # The process with rank 0 is considered the authoritative copy.
    self._distributed_broadcast_coalesced(
        self.modules_buffers[0],
        self.broadcast_bucket_size)
    # only do intra-node buffer sync for replicated single-device
    # CUDA modules
    if self.device_ids and len(self.device_ids) > 1:
        # intra-node buffer sync
        result = comm.broadcast_coalesced(
            self.modules_buffers[0],
            self.device_ids,
            self.broadcast_bucket_size)
        for tensors, module_buffers in zip(result[1:],
                                           self.modules_buffers[1:]):
            for tensor, buffer in zip(tensors, module_buffers):
                buffer.set_(tensor)