Thanks for the tips! It took me some while to understand and implement async_op.
I would like to point out a problem when I ran my own code above.
I changed my code to
is_next = torch.Tensor([self.data_loader.peek()]).cuda(self.gpu)
col_handle = dist.all_reduce(is_next, op=dist.ReduceOp.BAND, async_op)
...
col_handle.wait()
if not is_next: break
and tried it with SPSG with 2 processes. The final value of is_next is [2] rather than [True] or [1]. It seems that dist.ReduceOp.BAND adds up input tensors rather than doing a regular AND. Therefore I changed the first line into:
is_next = torch.Tensor([self.data_loader.peek()]).bool().cuda(self.gpu)
The Error Message says all_reduce does not support this Tensor type for now. In order to achieve my goal, I use dist.ReduceOp.MIN instead. Here’s my final code that actually runs smoothly without imbalanced for-loop counts blocking the synchornization process.
for inputs, labels in self.data_loader:
    is_next = torch.Tensor([self.data_loader.peek()]).cuda(self.gpu)
    col_handle = dist.all_reduce(is_next, op=dist.ReduceOp.MIN, async_op=True)
    # forward and backward and step and stuff                    
    col_handle.wait()
    if not is_next: break