(ctx, grad_output)
| 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