Enable SequenceParallel in te.LayerNormMLP layer

[vchiley]

Given mosaicml#432 we can relatively easily enable SequenceParallel training.

This branch is just for testing for now; if it works well, we can consider merging it.

Note: I'm pretty sure model ckpt using tp, then load without tp will be a headache and will require tooling be built.

[vchiley]

On a single node of 8 GPUs

MPT 125M training with / without Sequence Parallel is nearly identical with TP size in [2, 4, 8]

BUT, since the 125M model is small, its always slower to use TP vs the standard model even if mbs=1

If we bump up the model size to 7B, at mbs=1, a TP world size of 2 and 4 are slightly faster than the baseline

At mbs=2, the advantage does away.

If we bump up the model size to 13B, at mbs=1, again a TP world size of 2 and 4 are slightly faster than the baseline

At mbs=2, a TP world size of 2 is almost as fast, but again the advantage goes away

Note this is done on one node and does not factor in interconnect speed.
wandb here

[vchiley]

Issue: FSDP operates on the module level. The module has [layer_norm_weight, layer_norm_bias, fc1_weight, fc1_bias, fc2_weight, fc2_bias] parameters, where FSDP needs to tree [layer_norm_weight, layer_norm_bias, fc2_bias] the standard sharded way, and [fc1_weight, fc1_bias, fc2_weight] need to be treated standard TP sharded way...