@albanD is the right way to think about it then is that the ddp_model(x) is what “blocks” so that the next lines are executed with the right coordinated gradients? (the strange thing is that the demo basic code I pasted is usually the fn that a process would execute on it’s own which is weird because then how do you make sure that each process samples the same x if they have their own dataloader?..I am assuming that the code demo_basic(randk, world_size) is running is not the actual code and it get’s compiled to some other magic that makes it work properly).
Thanks for the help!
Source of example: Getting Started with Distributed Data Parallel — PyTorch Tutorials 1.7.1 documentation 6
class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)