Torch XLA Model all_gather does not work with tensors of different sizes along dimension 0

[ajayvohra2005]

🐛 Bug

Torch XLA Model all_gather works with tensors of same size along dim=0, but if tensor sizes are different along dim=0, it hangs.

To Reproduce

Save this code in test_all_gather.py

import torch
import torch_xla.core.xla_model as xm
import torch_xla.runtime as xr
import torch_xla.distributed.xla_backend as xb
import torch.distributed


def test_all_gather():

    same = [512, 512, 512, 512, 512, 512, 512, 512]

    different = [416, 536, 560, 544, 576, 512, 592, 360]
    torch.distributed.init_process_group(backend="xla", init_method="xla://")       

    rank = torch.distributed.get_rank()
    device = xm.xla_device()
    input = torch.randn((same[rank], 16), dtype=torch.float32, device=device)
    
    all_inputs = xm.all_gather(input, dim=0, groups=[[0,1,2,3,4,5,6,7]], pin_layout=False)
    print(f"!!!!!! rank: {rank}, all_inputs: {all_inputs}")
    
    input = torch.randn((different[rank], 16), dtype=torch.float32, device=device)
    
    all_inputs = xm.all_gather(input, dim=0, groups=[[0,1,2,3,4,5,6,7]], pin_layout=False)
    
    print(f"!!!!!! rank: {rank}, all_inputs: {all_inputs}")
    torch.distributed.destroy_process_group()
    
if __name__ == "__main__":
    test_all_gather()

torchrun --nproc_per_node=8 test_all_gather.py

Expected behavior

It should gather all the tensors from all the devices along dim=0

Environment

Docker image
us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.5.0_3.10_cuda_12.4

Additional context

According to this documentation for all_gather https://pytorch.org/docs/stable/distributed.html uneven tensor sizes are supported.

[zpcore]

Thanks for pointing this out. Currently we don't support dist ops with uneven sizes. Basically XLA needd to know the gathered output tensor shape before hand. One work around we can do is to pad those tensor to even sizes before all_gather.

By the way, we now support dist ops directly like:

xla/test/pjrt/test_collective_ops_tpu.py

Line 206 in 9ae017e

dist.all_gather(output_tensor, input, None)

e.g.:

output_tensor = [
    torch.empty((512, 16), dtype=torch.float32).to(device)
    for _ in range(xr.world_size())
    ]
all_inputs = torch.randn((512,16), dtype=torch.float32, device=device)
dist.all_gather(output_tensor, all_inputs)

, and we suggest launch multiprocess using

xla/torch_xla/torch_xla.py

Line 212 in 9ae017e

def launch(

This should handle both torchrun call or xmp.spawn.

[miladm]

@zpcore what's the plan to bring in this support?

cc @ysiraichi

[ysiraichi]

If this usage is not supported, I believe we should mention it somewhere (e.g. its documentation).

[ajayvohra2005]

I am already using the workaround. I opened the issue to avoid using the workaround. At the very least, may I request it should be clearly documented so as to avoid costly debugging effort to debug the hang.

[zpcore]

Thanks for the report, I will update documents for this issue.