| 455 | # the call method is: |
| 456 | # features_gather = SyncFunction.apply(features) |
| 457 | class SyncFunction(torch.autograd.Function): |
| 458 | |
| 459 | @staticmethod |
| 460 | def forward(ctx, tensor): |
| 461 | ctx.batch_size = tensor.shape[0] |
| 462 | |
| 463 | gathered_tensor = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())] |
| 464 | |
| 465 | torch.distributed.all_gather(gathered_tensor, tensor) |
| 466 | gathered_tensor = torch.cat(gathered_tensor, 0) |
| 467 | |
| 468 | return gathered_tensor |
| 469 | |
| 470 | @staticmethod |
| 471 | def backward(ctx, grad_output): |
| 472 | |
| 473 | grad_input = grad_output.clone() |
| 474 | torch.distributed.all_reduce(grad_input, op=torch.distributed.ReduceOp.SUM, async_op=False) |
| 475 | |
| 476 | idx_from = torch.distributed.get_rank() * ctx.batch_size |
| 477 | idx_to = (torch.distributed.get_rank() + 1) * ctx.batch_size |
| 478 | |
| 479 | return grad_input[idx_from:idx_to] |
| 480 | |
| 481 | |
| 482 | # come from |
nothing calls this directly
no outgoing calls
no test coverage detected