Repository: pytorch
Issue: xpu: torch distributed ops like c10d::allgather are not implemented for XPU backend
Labels: triaged,module: xpu

With:
* (pytorch) https://github.com/pytorch/pytorch/commit/00134d68af2ce50560fa5a74473665ea229e6c9d
* https://github.com/intel/torch-xpu-ops/commit/98f47b621e3d8757ea6c03c4dacbad491dc84014

Running with `gloo` torch distributed backend, the following aten operators are not currently implemented for XPU backend (likely there are more not implemented ops in the same series...):
* This one does not allow manual CPU fallback, `PYTORCH_ENABLE_XPU_FALLBACK=1` will fail:
  - [ ] `c10d::allgather_`
* These allow manual CPU fallback, `PYTORCH_ENABLE_XPU_FALLBACK=1` will pass:
  - [ ] `c10d::_allgather_base_`
  - [ ] `c10d::allgather_into_tensor_coalesced_`

To reproduce the following scripts can be used ran thru `torchrun` (with `--nnodes=1` can be reproduced in the single GPU). **NOTE: some projects might use torch distributed package running even on a single GPU device setting `world_size=1` and step into this issue. One example of such project is [llama-models](https://github.com/meta-llama/llama-models/tree/main) - they use torch distributed running vision models. [Here](https://github.com/meta-llama/llama-models/blob/17107dbe165f48270eebb17014ba880c6eb6a7c9/models/llama3/reference_impl/multimodal/model.py#L66) is a call to the allgather.**

* For `c10d::allgather_`:
```
$ cat run.py
import os
import torch

torch.distributed.init_process_group('gloo')

world_size=int(os.environ["WORLD_SIZE"])
local_rank=int(os.environ['LOCAL_RANK'])
print(f'world_size={world_size}')
print(f'local_rank={local_rank}')

device = torch.device(f'xpu:0')

tensor_list = [torch.zeros(2, dtype=torch.int64, device=device) for _ in range(world_size)]
print(tensor_list)

tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * local_rank
print(tensor)

torch.distributed.all_gather(tensor_list, tensor)
print(tensor_list)

$ PYTORCH_DEBUG_XPU_FALLBACK=1 torchrun --nnodes=1 --nproc-per-node=1 run.py
...
[rank0]: NotImplementedError: The operator 'c10d::allgather_' is not currently implemented for the XPU device.
...
```

* For `c10d::_allgather_base_`:
```
$ car run.py
import os
import torch

torch.distributed.init_process_group('gloo')

world_size=int(os.environ["WORLD_SIZE"])
local_rank=int(os.environ['LOCAL_RANK'])
print(f'world_size={world_size}')
print(f'local_rank={local_rank}')

device = torch.device(f'xpu:0')

tensor_in = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * local_rank
print(tensor_in)

tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device=device)
torch.distributed.all_gather_into_tensor(tensor_out, tensor_in)
print(tensor_out)

$ PYTORCH_DEBUG_XPU_FALLBACK=1 torchrun --nnodes=1 --nproc-per-node=1 run.py
...
[rank0]:[W1204 16:55:39.981637569 RegisterXPU.cpp:46330] Warning: The operator 'c10d::_allgather_base_ on the XPU backend is falling back to run on the CPU. (function xpu_fallback_impl)
...
```

* I don't have simple repro script for `c10d::allgather_into_tensor_coalesced_`. I did see this one running Llama3.2-11B-Vision-Instruct model working with llama-models project, see https://github.com/meta-llama/llama-models/pull/233.

NOTE: above examples for `c10d::allgather_` and `c10d::_allgather_base_` work on CUDA with `gloo` torch distributed backend.

CC: @gujinghui @EikanWang @fengyuan14 @guangyey @jgong5

cc @gujinghui @EikanWang @fengyuan14 @guangyey