[Primitive][shard] Use autograd function for all sync ops

[comaniac]

Description

We found that the behavior of PyTorch backward hooks is a bit weird, so for safety, this PR changes all use cases of backward hooks to be forward pre-hook + autograd function. Specifically:

Before

register_backward_hook(dist.all_reduce)

Now

class _ReduceBackwardGradient(torch.autograd.Function):
    def forward(...):
        # no-op

    def backward(...):
        # all-reduce gradient

register_forward_pre_hook(allreduce_backward_gradient)

Accordingly, _ReduceBackwardGradient is added. In addition, this PR also adds unit tests for all supported sync ops.

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @szhengac @chhzh123

[comaniac]

Thanks @szhengac