Here’s a snippet showing what happen’s with “gossip”-based parameter aggregation:
for param in model.parameters():
    # A container to hold async requests
    # ---- MARKER A ----
    requests = []

    # We've to exchange params with peers
    to_send = param.data
    if args.gpu:
        to_send = to_send.cpu()

    receive_buffers = [
        (peer, torch.zeros(param.data.shape))
        for peer in peers
    ]

    for peer, buffer in receive_buffers:
        requests.append(dist.isend(
            tensor=to_send,
            dst=peer
        ))
        requests.append(dist.irecv(
            tensor=buffer,
            src=peer
        ))

    # Wait for all the requests to complete
    for r in requests:
        r.wait()
    # ---- MARKER B ----

    # Then compute the average
    for _, buffer in receive_buffers:
        if args.gpu:
            buffer = buffer.cuda(device_id)
        param.data += buffer
    param.data /= len(receive_buffers) + 1

    # ---- MARKER C ----

This is called after every x updates to the parameters. Based on logs, the error seems to occur between ---- MARKER A ---- and ---- MARKER B ---- consistently. (Edit: I just observed this happen between  ---- MARKER B ---- and ---- MARKER C ----)
If you’re curious about the model architecture (this is for the MNIST dataset):
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.l1 = nn.Linear(28 * 28, 1024)
        self.l2 = nn.Linear(1024, 1024)
        self.l3 = nn.Linear(1024, 1024)
        self.out = nn.Linear(1024, 10)
        for name, param in self.named_parameters():
            if 'weight' in name:
                nn.init.kaiming_normal(param)

    def forward(self, x):
        # input
        x = x.view(-1, 28 * 28)
        # reshaped
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.l1(x)
        x = F.relu(x)
        # l1
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.l2(x)
        x = F.relu(x)
        # l2
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.l3(x)
        x = F.relu(x)
        # l3
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.out(x)
        x = F.log_softmax(x, dim=1)
        # output
        return x

Please let me know if you’d like more information.