I see, that’s the uneven DDP input batches issue. See the previous tracking issue here 20. We have just finished design discussion on this, and plan to create a ddp.join() context manager to solve this problem. @rvarm1 will post a design RFC later.
For short-term unblock, you can either make sure that all DDP instances are running through the same number of batches or detect that proactively in the application, sth like:
for batch in get_batch():
    signal = torch.tensor([1])
    work = all_reduce(signal, async_op=True)
    loss = model(batch).sum()
    work.wait()
    if signal.item() < world_size:
        break
    loss.backward()


all_reduce(torch.tensor([0]))