The pseudo codes are like below:
def average_gradients(model):
    size = dist.get_world_size()
    for param in model.parameters():
        dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
        param.grad.data /= size

def train():
    for ep in range(1000):
        optimizer.zero_grad()
        a = torch.tensor(xxx)
        # multipulating a to derive tensor b
        c = torch.autograd.grad(b, a, create_graph=True)[0]
        loss = nn.MSEloss((b, b_target)) + nn.MSEloss((c, c_target))
        loss.backward(retain_graph=True)
        average_gradient(net)
        optimizer.step()
    
if __name__ is '__main__':
    processes = []
    size = 4
    for rank in range(size):
        p = Process(target=train,
                    args=(rank, size, lock))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

I set different value of size to test, the main time consuing part is torch.autograd.gard method. The smaller value of size, the faster for this line computing.