Mamba + Tensor Parallel Support

[haileyschoelkopf]

This PR adds Mamba + TP support.

Loss curves comparing TP=2 to TP=1 with + without mamba_inner_func_fusion:

Versus when allreduce was missing with inner func fusion turned on:

Also tested that PP seems to work.

[Quentin-Anthony]

LGTM, no comments.

@haileyschoelkopf -- As a final check, I'd like to verify that TP gives the expected memory benefits. Can you link the wandb here so I can take a look?

[haileyschoelkopf]

Yep! https://wandb.ai/eleutherai/mamba-neox-tp-memsavings/workspace?nw=nwuserschoelkopf

this should be public and a clean comparison, lmk if more is needed or for some reason it's not visible (MBS=16 for both TP=1 and TP=2, corrected for w/ grad accum.) Seeing 29-30GB for TP=1, 19GB for TP = 2.

Here's the full wandb of trial runs including initial tests + adding stuff like Mamba's GPT-2 style init, Mamba-160m, and some Pythia-160m baseline curves: https://wandb.ai/eleutherai/mamba-neox?nw=nwuserschoelkopf