We provide convenience methods for some of the communication collectives, one main reason for that is they work in both distributed as well as non-distributed code. However in the none-distributed path we still were doing some work such as moving the data from cpu to gpu or vice versa. We can always avoid doing this work, reduce and gather ops should be noops if we only have one process.

Another simplification is to remove the custom autograd function for gather. PyTorch added support for gradients in all_gather some time ago: pytorch/pytorch#40762, thus we can just use the normal all_gather functions