Performs all_gather operation on the provided tensors. *** Warning ***: torch.distributed.all_gather has no gradient.
(tensor)
| 44 | |
| 45 | @torch.no_grad() |
| 46 | def concat_all_gather(tensor): |
| 47 | """ |
| 48 | Performs all_gather operation on the provided tensors. |
| 49 | *** Warning ***: torch.distributed.all_gather has no gradient. |
| 50 | """ |
| 51 | tensor = tensor.contiguous() |
| 52 | tensors_gather = [ |
| 53 | torch.ones_like(tensor) |
| 54 | for _ in range(torch.distributed.get_world_size()) |
| 55 | ] |
| 56 | torch.distributed.all_gather(tensors_gather, tensor, async_op=False) |
| 57 | |
| 58 | output = torch.cat(tensors_gather, dim=0) |
| 59 | return output |
| 60 | |
| 61 | |
| 62 |