Add GPT-NeoX (StableLM)

[seanmor5]

Should support StableLM. It's rebased off of Llama for rotary embeddings.

@jonatanklosko I wonder if you might be able to take a look at the sliced qkv calculation? When debugging the values I was getting for query are correct, but for key they are way off.

[jonatanklosko]

@seanmor5 yeah, the dense units for GPTNeoX QKV implementation are ordered as (num_attention_heads, 3, head_size), while for example for BLIP it's (3, num_attention_heads, head_size), so we need to split the kernel differently. I pushed a fix.

The remaining difference is the attention block, GPTNeoX has use_parallel_residual option, which uses a different formulation in the attention block, such that attention and MLP are computed in parallel, as in GPT-J. I made a quick tweak locally and after that all tests pass. I will refactor it into an option later :)

[seanmor5]

Beautiful! I believe open assistant is based on this model too? So theoretically this is StableLM and OpenAssistant?

[jonatanklosko]

@seanmor5 the Pythia models are GPTNeoX yeah, the new ones used in huggingface chat are LLaMA I think and they are stored in a very fancy xored way.

[jonatanklosko]

The variations started to accumulate, so I added a separate function to orchestrate the order and flow of layers. That's the best thing I came up with so far, maybe we will arrive at a better abstraction at some point, but it's internal so it's fine as long as it works :D