[feat] Support TPUs for intra-layer model parallel training

## 🚀 Feature
The current implementation is based on Megatron and only supports GPUs. Now that we’re migrating fairseq to this implementation, we should add TPU support here as well.

## Motivation
Several fairseq users and internal projects would benefit from TPU support. For example, see https://github.com/pytorch/fairseq/issues/2503.

## Pitch
Replace all the CUDA-specific calls with device-agnostic versions that are compatible with PyTorch/XLA.

## Alternatives
Not support TPUs.

## Additional context
I have a preliminary version of the needed changes, but they are based on an old version of Megatron, so would need to be rebased over the (newer) Megatron fork in fairscale.