llama: add Grok support

[fakerybakery]

Hi,
Please add support for Grok.
Thanks!

Relevant links:

• https://github.com/xai-org/grok
• https://x.ai/blog/grok-os
• https://twitter.com/grok/status/1769441648910479423
• [NEW] Official Upload (thx to @dranger003) for linking: https://huggingface.co/xai-org/grok-1

[bachittle]

Falcon-180b was able to run. Pretty sure a 314b parameter will also be able to run, if at least quantized. Still, I would have no hardware to run either. 😢

[dranger003]

https://huggingface.co/xai-org/grok-1

[nonetrix]

A lot of people are saying the weights released are fp8 or int8 and not the full fp16 for some reason, could this perhaps make it challenging to add?

[netrunnereve]

Considering how this model is too big for most of us to run at home, I wonder how well it performs if we use less experts and run it as a 30Bx4 or something. Or maybe it's possible to merge and distill the expert models into a single 30B model and do some finetuning. If it ends up performing like Yi 30B then we have a win here as this is Apache licensed.

Keep in mind Mixtral 8x7B was built off modified versions of the original Mistral 7B to create the 8 experts. There's a chance the Twitter team did something similar as well.

[stduhpf]

A lot of people are saying the weights released are fp8 or int8 and not the full fp16 for some reason, could this perhaps make it challenging to add?

On the readme on HuggingFace they say it's int8.

[FNsi]

Would like to try this in 1.58 bit 😁

[nonetrix]

Wouldn't that require a retrain and the code is not out either

[Konard]

@fakerybakery https://huggingface.co/alpindale/grok-1 link is broken, please update the issue description.

[nonetrix]

Seems like it was deleted completely since they uploaded it officially I assume

[david-jk]

Has someone had a look at the architecture yet to get an idea of how hard it is to implement?
How much different is the architecture compared to e.g. Mixtral 8x7B?

[fakerybakery]

@fakerybakery https://huggingface.co/alpindale/grok-1 link is broken, please update the issue description.

Fixed, thx! The author removed it since the official weights have been uploaded

[Noeda]

It's not particularly complicated. I started porting the Jax version to a simplified PyTorch version in hopes I can get it to run at least a little faster on a big Mac Studio and also make it more readable. Maybe about halfway done. All the components look like it's made of standard stuff. It's vanilla attention stuff, some MoE layers, same RoPE we all know and love etc.

I haven't seen MoE layers before in any code so can't judge if that's very different from Mixtral. But I can say it's not particularly more complicated or less complicated than most other models I've seen.

Below is a very very rough sketch I made a few hours ago of classes that would roughly correspond to what you'd want to write class Module(nn.Module) for. (don't trust it's correct). I don't have a nice visualization graph. Entry to the code starts from LanguageModel and it's at the bottom of the Grok file .py file. It uses Jax but also some kind of framework that I'm not familiar with (Haiku). I think it's functional style where parameters are not defined like in Torch where you do nn.Parameter() but instead the code is run once, computation graph is made, and then parameter vectors are registered on the course of doing that. Or something along those lines.

The weights come in 8-bit precision, and QuantizedWeight8bit is used for that. So almost all of the 300GB is in the QuantizedWeight8bit one.

class QuantizedWeight8bit:
    weight: parameter [int8]
    scales: parameter [bfloat16]
class KVMemory:
    k: parameter [jax.Array]
    v: parameter [jax.Array]
    step: parameter [jax.Array]
class Memory:
    layers: list [KVMemory]
class Router:
    num_experts: int
    num_selected_experts: int
class MoELayer: ...
class MHABlock: ...
class RMSNorm: ...
class RotaryEmbedding + rotate_half (EleutherAI RoPE looks like)
class MultiHeadAttention: ...
class DenseBlock: ...
class DecoderLayer: ...
class InOutEmbed: looks like weight tying here, just linear layer for first and last layer for embeddings.
class Transformer:
class LanguageModel: this is the top level class that includes everything else

Vocabulary size is 128*1024 = 131072. Tokenizer looks like is a BPE (uses sentencepiece library).

(I'm not working on .gguf support at the moment. But eventually will jump into it if it seems like no one is picking it up).

[schmorp]

@Noeda if you manage to make a working gguf I'll happily make imatrix quants, if no other changes in llama.cpp are needed.

[nonetrix]

Hoping we get Grok 0 weights too, they initially closed my issue as for them as "not planned" but reopened it eventually so maybe?? Idk, I'm hoping Grok 0 is just like 33B model and most people could easily load that and many could probably fine tune it. I imagine if Grok 1 is just a MoE of Grok 0 it would be really trivial to implement it as well