Repository: pytorch
Issue: `torch.distributed.nn.functional.all_gather`: Tensors must be contiguous
Labels: oncall: distributed,module: bootcamp,triaged,better-engineering,pt_distributed_rampup,module: c10d

### 🐛 Describe the bug

In https://github.com/PyTorchLightning/pytorch-lightning/issues/10445, we propose removing our own AllGather implementation in favor of `torch.distributed.nn.functional.all_gather`.

However, when working on the PR: https://github.com/PyTorchLightning/pytorch-lightning/pull/11531 we encountered a difference between implementations. This following snippet reproduces the problem on a GPU machine:

```python
import os

import torch
import torch.distributed.nn


# DIRECTLY COPIED FROM
# https://github.com/pytorch/pytorch/blob/6c8e516a8034d7ab615025ac57b5bc2e7680efb1/torch/distributed/nn/functional.py#L275-L290
class _AllGather(torch.autograd.Function):
    @staticmethod
    def forward(ctx, group, tensor):
        ctx.group = group
        out_tensor_list = [
            torch.empty_like(tensor) for i in range(torch.distributed.get_world_size(group=group))
        ]
        torch.distributed.all_gather(out_tensor_list, tensor, group=group)
        return tuple(out_tensor_list)

    @staticmethod
    def backward(ctx, *grad_outputs):
        gxs = torch.distributed.nn.functional._AlltoAll.apply(ctx.group, *grad_outputs)
        gx = torch.sum(torch.stack(gxs), dim=0)
        return (None, gx)


if __name__ == "__main__":
    orig_environ = os.environ.copy()
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12321"
    os.environ["RANK"] = "0"
    os.environ["WORLD_SIZE"] = "1"
    torch.distributed.init_process_group("nccl")
    try:
        tensor = torch.rand(2, 2, requires_grad=True, device="cuda")
        gathered_tensor = _AllGather.apply(None, tensor)
        gathered_tensor = torch.stack(gathered_tensor)
        gathered_tensor.sum().backward()
    finally:
        torch.distributed.destroy_process_group()
        os.environ.clear()
        os.environ.update(orig_environ)
```

With the error:

```python
  File "kkbug.py", line 19, in backward
    gxs = torch.distributed.nn.functional._AlltoAll.apply(ctx.group, *grad_outputs)
  File "/home/carlos/venv/lib/python3.8/site-packages/torch/distributed/nn/functional.py", line 244, in forward
    dist.all_to_all(out_tensor_list, list(tensors), group=group)
  File "/home/carlos/venv/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 2661, in all_to_all
    work = default_pg.alltoall(output_tensor_list, input_tensor_list, opts)
RuntimeError: Tensors must be contiguous
```

From my investigation, it seems that NCCL all-to-all requires tensors to be contiguous. This is not a problem with the Gloo pg backend.

The tensors returned in `_AllGather.forward`'s `out_tensor_list` are contiguous. However, they are not by the time they are passed as inputs to `_AllGather.backward`.
So the `_AlltoAll.apply` call in backward will fail with the message above.

Is this intended?
If `forward` outputs are contiguous, shouldn't `backward` inputs also be?
If not, shouldn't `_AllGather` make them contiguous before `alltoall`?

In our custom implementation, we circumvent this issue by calling `cat`
https://github.com/PyTorchLightning/pytorch-lightning/blob/d515bcac969c2a485ada673e302bfac51f142331/pytorch_lightning/utilities/distributed.py#L218
which has the side effect of making the tensor contiguous.

Thank you for your help!

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @pietern @SciPioneer @awaelchli 

### Versions

PyTorch version: 1.10.2+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 26 2021, 20:14:08)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-53-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090
GPU 3: NVIDIA GeForce RTX 3090
GPU 4: NVIDIA GeForce RTX 3090
GPU 5: NVIDIA GeForce RTX 3090
GPU 6: NVIDIA GeForce RTX 3090
GPU 7: NVIDIA GeForce RTX 3090

Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] pytorch-lightning==1.6.0.dev0
[pip3] torch==1.10.2+cu113
[pip3] torchmetrics==0.5.0
[pip3] torchvision==0.11.3+cu113
[conda] Could not collect