Xla: Support for Layerwise Model Parallelism

[ultrons]

🚀 Feature

Methods to place network subgraphs/submodules on separate TPU devices, manage forward and backward passes to compute and aggregate the gradients and communication collectives to allow various reduction among these devices. It will enable the support for layer-wise model parallelism. The communication collectives will include operations such as scatter, gather and reduce. Since broadcast can be implemented as scatter operation it is not requested separately.

Motivation

Enable training large capacity models (X-XXB parameters).

Pitch

Both , the state of the art and research frontiers of deep learning for task specific or more generalized learning have been moving towards multi-billion parameter neural networks. Model parallelism has been one the key approaches to enable such large scale model training. With the methods already provided in torch_xla, it appears that we have access to individual gradients values from XLA devices and also .to(device=...) capability to place the subnets and data onto individual devices to some degree. Infeed of the devices and gradient management to support layerwise model parallelism is needed.

[myleott]

This would be especially helpful in implementing something like GPipe. A good open source implementation of this in PyTorch can be found here: https://github.com/kakaobrain/torchgpipe. It'd be great to support something similar with PyTorch-XLA!

[dlibenzi]

Hi!
As we mentioned in other threads, enabling full model parallelism requires deep changes all around the stack (from client side down to TF/XRT - which is on the deprecated path), but we can follow two different tracks to see if we can come up with alternative solutions for the problem.
One, at expense of some performance, we could use all_reduce(SUM) to simulate send/recv ops.
Think like:

def send(t, groups=[]):
  xm.all_reduce('sum', [t], groups=groups)

def recv(shape, device, groups=[]):
  t = torch.zeros(shape, device=device)
  xm.all_reduce('sum', [t], groups=groups)
  return t

In the above case (like every pytorch/xla case, for the same-graph requirements) there needs to be one sender and N-1 receivers for each replica group (where N is the size of the replica group).

We will also expose XLA's AllToAll() op:

https://www.tensorflow.org/xla/operation_semantics#alltoall

I am not sure how easy is to stitch all this together to create a working model able to scale POD size, but it might be doable.

[dlibenzi]

The XLA's AllToAll() exposure should go in today, and we will also be exposing XLA's CollectivePermute():

https://www.tensorflow.org/xla/operation_semantics#collectivepermute

Whether an efficient model parallelism can be built on top of these, given the current architecture remains to be seen.

[dlibenzi]

Thanks to @JackCaoG now even XLA's CollectivePermute is available.

https://github.com/pytorch/xla/blob/39d22ee6ee6af6b022ba90b72cc2cca071316136/torch_xla/core/xla_model.py#L382
https://github.com/pytorch/xla/blob/39d22ee6ee6af6b022ba90b72cc2cca071316136/torch_xla/core/xla_model.py#L412

[jysohn23]

FYI https://github.com/jysohn23/xla/blob/model-parallel-colab/Gather_Scatter_Broadcast_PyTorch_XLA.ipynb

[dlibenzi]

Thanks @jysohn23 !
Some ops will be more efficient using AllToAll or CollectivePermute.

[ultrons]

Awesome. Thanks @dlibenzi , @jysohn23 and @JackCaoG !
@jysohn23 , Cool implementations of collectives in terms of all_reduce. your colab shows examples at sub-group level. You shared with me a similar example of all gather across the cores using the same idea. I don't know where these examples will be integrated but will nice to add your across core all_gather example as well.

[jysohn23]

For applying all above ops globally just call xm.all_reduce (as well as AllToAll and CollectivePermute) with groups=None.

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.