Which version of pytorch support the function torch.distributed.ring_exchange()?

1579534387 · February 22, 2023, 2:05am

ptrblck · February 22, 2023, 5:51am

I don’t think this method is exposed but is used in dist.batch_isend_irecv.
CC @kwen2501 to correct me if I’m wrong.

kwen2501 · February 23, 2023, 9:34pm

torch.distributed does not have an API called ring_exchange().

But you can express a ring-exchange pattern using torch.distributed.batch_isend_irecv

Here is an example:

github.com

pytorch/pytorch/blob/master/torch/distributed/distributed_c10d.py#L1436-L1445


      
          >>> send_tensor = torch.arange(2) + 2 * rank
          >>> recv_tensor = torch.randn(2)
          >>> send_op = dist.P2POp(dist.isend, send_tensor, (rank + 1)%world_size)
          >>> recv_op = dist.P2POp(dist.irecv, recv_tensor, (rank - 1 + world_size)%world_size)
          >>> reqs = batch_isend_irecv([send_op, recv_op])
          >>> for req in reqs:
          >>>     req.wait()
          >>> recv_tensor
          tensor([2, 3])     # Rank 0
          tensor([0, 1])     # Rank 1