Final question. Not clear why you have another all reduce at the end

That’s to indicate one process has exited the loop, otherwise the sum would never be smaller than the world_size. It needs to be guarded by a flag, because only the processes that exit early need that. Sth like:
for batch in get_batch():
    ....

if signal.item() >= world_size:
    all_reduce(torch.tensor([0]))