[RFC] Exposing additional XLA collective communication primitives.

[hjm-aws]

🚀 Feature

Add the ability to translate the following Collective Communication ops to native
XLA instructions:

• all_gather
• reduce_scatter
• collective_permute
• send
• recv

At the moment all_gather is implemented with all_reduce in PyTorch/XLA. The other 4
operators are not yet implemented. We are proposing to change the current all_gather
implementation to generate native XLA all_gather instruction, and to add native
XLA support for reduce_scatter, collective_permute, send, and recv operations.

We are intending to submit PRs for this proposal.

Motivation

Adding the ability to distributed training libraries (such as DeepSpeed) and
Apps (such as MegatronLM) to utilize underlying XLA devices/accelerators to
train very large models by calling torch.distributed API.

Currently MegatronLM and DeepSpeed use torch.distributed API to execute Collective
Communication primitives, but torch.distributed does not support XLA. To enable
MegatronLM and DeepSpeed to run on XLA, we are proposing to utilize PT/XLA, which is
a bridge from PT to XLA, to provide the ability to translate the collective
communication primitives listed above to native XLA instructions.

Pitch

We would like users to be able to call the following APIs to insert the corresponding
XLA instructions:

import torch_xla.core.xla_model as xm

xm.all_gather(tensor)
xm.reduce_scatter(xm.REDUCE_SUM, tensor, groups=[[0,1]])
xm.collective_permute(tensor, pairs=[[0,2],[1,3]])
xm.send(tensor, channel)
tensor = xm.recv(shape, channel)

Alternatives

User could implement all these ops with all_reduce, but the efficiency would be too
low for them to be practical.

Additional context

[jysohn23]

FYI, we had added support for ReduceScatter collective (#3075) but had to roll back due to new TF version breaking our internal builds. We should be rolling forward soon.

[jysohn23]

Also, we do expose the CollectivePermute collective: https://github.com/pytorch/xla/blob/master/torch_xla/csrc/ops/collective_permute.cpp

[hjm-aws]

Thank you Daniel for your comment.

Regarding ReduceScatter, yes we noticed that it was only recently added to the XLA side in the last 2 - 3 months. It's great to hear that you already have an existing commit that you can roll forward. Let's ignore ReduceScatter in this RFC.

Regarding CollectivePermute, yes that's a mistake on my part to include it in the RFC. Let's ignore it.

Regarding the rest of the RFC (AllGather, Send, and Recv), are you OK with the proposal?

One more thing I would like to clarify with you is that, I noticed that the CPU IR emitter on the XLA side doesn't have proper support for Send and Recv (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/ir_emitter.cc). Is it our responsibility to implement the CPU IR emitter for these primitives, as a part of the RFC? BTW, just FYI, I believe ReduceScatter is also not implemented in the CPU IR emitter.

[JackCaoG]

@sanjoy for CPU IR emitter question.

[sanjoy]

@d0k for XLA CPU.

[d0k]

Collective ops on CPU are primarily being used to test XLA without accelerators present. Adding support for more collective primitives to XLA:CPU would be nice, but it is optional when your main target is an accelerator and not a requirement for this RFC.

[hjm-aws]

Thank you Benjamin for your clarification!

[mweiss17]

Hello,
Any chance we'll be seeing an autograd version of all_to_all? I've implemented a Switch Transformer in torch-xla, but in order to have model parallelism I kinda need the all-to-all. I expect the all-gather to be quite inefficient for this application.

[JackCaoG]

@mweiss17 Are you talking about something similar to

xla/torch_xla/core/functions.py

Line 55 in 88474a3

class AllGather(torch.autograd.Function):

?

[mweiss17]

Yep, pretty much. I wrote a version like that this weekend, but haven't tested it well. If you think it's worthwhile I could test it and make a PR.

[JackCaoG]

@mweiss17 Sounds good! We will be happy to take your contribution.

[hjm-aws]

@mweiss17 Hi, can you clarify which op you want to contribute for, all-gather or all-to-all? Thanks!

[mweiss17]

Hi @hjm-aws, I'm adding an autograd all-to-all class to torch_xla/core/functions.py. Those functions seem to not have tests (if they do, let me know!), so I'm going to add a new test file as well.