search

pytorch / xla

Enabling PyTorch on XLA Devices (e.g. Google TPU)
https://pytorch.org/xla
Other
2.46k stars 467 forks source link

Support dist.all_to_all_single #8064

Closed zpcore closed 4 days ago

zpcore commented 5 days ago

Support to use torch.distributed.all_to_all_single to both dynamo and nondynamo case.

Note that there is a function signature mismatch between torch's all_to_all_single and xla op's AllToAll. To leverage the AllToAll op, we doesn't support specifying the input_split_sizes and output_split_sizes at this time. Check test_collective_ops_tpu.py for the usage.

will-cromar commented 4 days ago

Thanks!