RuntimeError: function AllGatherGradBackward returned an incorrect number of gradients (expected 2, got 1)

[ArvinZhuang]

🐛 Bug

When using self.all_gather in training_step to gather tensor with gradient function to compute and return loss throwing
RuntimeError: function AllGatherGradBackward returned an incorrect number of gradients (expected 2, got 1)

I think the bug is that the forward(ctx, tensor, group=group.WORLD) in distributed.AllGatherGrad function has two arguments but the backward(ctx, *grad_output) only returns one output.

class AllGatherGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor, group=group.WORLD):
        ctx.group = group

        gathered_tensor = [
            torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())
        ]

        torch.distributed.all_gather(gathered_tensor, tensor, group=group)
        gathered_tensor = torch.stack(gathered_tensor, dim=0)

        return gathered_tensor

    @staticmethod
    def backward(ctx, *grad_output):
        grad_output = torch.cat(grad_output)

        torch.distributed.all_reduce(
            grad_output,
            op=torch.distributed.ReduceOp.SUM,
            async_op=False,
            group=ctx.group
        )

        return grad_output[torch.distributed.get_rank()]

The error can be solved by changing return grad_output[torch.distributed.get_rank()] to return grad_output[torch.distributed.get_rank()], None

Environment

• PyTorch Version: 1.7.1
• OS: Linux
• How you installed PyTorch: pip
• Python version: 3.7
• lightning version: 1.1.8