Add Mamba Architecture

[haileyschoelkopf]

closes #1148

This PR adds Mamba to NeoX, along with flags for turning on/off the selective scan + conv1d + full mamba_inner_fn kernels.

For now, does not support parallelism, but want to investigate adding Tensor Parallel to this.

[haileyschoelkopf]

This seems to train well without parallelism, but am having bugs in a conversion script I wrote (gibberish output). I'll be checking for differences in output between a single instantiated layer for this versus the mamba_ssm module to try to diagnose further

[haileyschoelkopf]

Worked after state-spaces/mamba#211 !

Got performance from a 160m (trained with Pythia config, untied embed + unembed) on-par with the Mamba-130m results in paper!

Will cleanup code slightly and add sample configs, then mark ready for review. This also pairs with a DeeperSpeed PR I'll make that should allow for holding specified parameters in fp32 despite Deepspeed trying to cast everything to 16 bit.

I want to check out adding tensor parallelism for Mamba too, but will do that later.

[Quentin-Anthony]

Worked after state-spaces/mamba#211 !

Got performance from a 160m (trained with Pythia config, untied embed + unembed) on-par with the Mamba-130m results in paper!

Will cleanup code slightly and add sample configs, then mark ready for review. This also pairs with a DeeperSpeed PR I'll make that should allow for holding specified parameters in fp32 despite Deepspeed trying to cast everything to 16 bit.

I want to check out adding tensor parallelism for Mamba too, but will do that later.

Awesome! Great work. I have some TP ideas that we can discuss on discord.

[haileyschoelkopf]

Ready for initial review!

Pairs with EleutherAI/DeeperSpeed#61 , which I'd appreciate feedback on if the approach there is acceptable.

[Quentin-Anthony]

Note for the future: In addition to the attention_config, we should add a block_config where the user can choose on a per-block basis where to place individual blocks for MLP(and its variants like MoE)/Mamba/Attention(and its variants). The attention_config would then just allow the user to choose between attention variants for any attention blocks in the broader block_config. A similar mlp_config could be useful.

E.g.

"block_config": ["mamba", "attention", "mamba", "attention", "mlp"],
"attention_config": [[["flash"], 2]],
"mlp_config": [[["moe"], 1],

One drawback from this strategy is that we're adding a lot of annoying book-keeping for the user.

It's already a bit confusing that we're putting the attention-free mamba block under attention_config, but I think it's fine for this PR.

[Quentin-Anthony]

In the future, we should add support for the triton RMSNorm kernel introduced by mamba. Noting here and adding a TODO for later.

https://github.com/state-spaces/mamba/blob/v1.2.0/mamba_ssm/ops/triton/layernorm.py