Implement the MPU from Megatron

[StellaAthena]

The Megatron code contains a "MPU" library. MPU stands for "model parallelism unit." The purpose of an MPU is to allow custom tensor slicing across GPUs. DeepSpeed allows you to hook up a MPU, but doesn't provide one. The goal is to convert the MPU from Megatron to GPT-NeoX. This is a modified clone of Megatron: https://github.com/EleutherAI/MegatronPipeline

You may find the (minimalistic) descriptions DeepSpeed provides helpful:
https://www.deepspeed.ai/features/#model-parallelism
https://www.deepspeed.ai/tutorials/megatron/

The full DeepSpeed docs can be found here: https://deepspeed.readthedocs.io/en/latest/index.html

[StellaAthena]

Superseded by codebase refactoring.