Grok-1 MoE support
Author: AlpinDale
Pull request: no
Labels: name: New model

### Model description

X-AI recently released [grok-1](https://huggingface.co/xai-org/grok-1), a massive MoE model, with a total parameter count of 314B across 8 experts, 2 active at a time. Would be interesting if we could have support for the weights in transformers, would make integrations everywhere else much easier to handle.

### Open source status

- [X] The model implementation is available
- [X] The model weights are available

### Provide useful links for the implementation

They have a reference implementation in JAX, but similar to meta's original release of llama, the inference code is highly unoptimized, and has a hard-coded number of GPUs as a requirement.

https://github.com/xai-org/grok-1

Comments:
Can't wait to have it in Transformers 😬

However, this model is too gigantic for even an advanced server. Yet, from what I saw in its official code, it's highly sub-optimal (it uses the "vanillest" implementation of Attention, not even the Flash one!). I believe it will benefit greatly for pre-existing and future optimizations thanks for the Transformers library modularity. Examples include:

1. XFormers Flash Attention
2. Torch 2 Scaled dot product attention
3. [Ring Attention (from LWM - Large World Model), to scale the context to millions of tokens](https://github.com/lucidrains/ring-attention-pytorch)
4. [Linear "Rebased" Flash Attention](https://github.com/corl-team/rebased) - eliminates the quadratic attention cost and requires less rework than switching to Mamba, more compatible to existing architecture 
5. Hybrid attention head sequence parallelism (e.g. FastSeq) including ZeRO (see [Accelerated Transformer](https://github.com/hpcaitech/Open-Sora/blob/main/docs/acceleration.md#accelerated-transformer))
6. Layer offloading
7. Bitsandbytes/AutoGPTQ Quantization

And the best thing is that all these points (except for 1/2 and 3/4) can be combined!

With these optimizations we can hope to shrink the gap between the current requirements and the consumer/low spectrum server requirement

---

Mirror of https://github.com/xai-org/grok-1/discussions/242

I wrote a quick implementation yesterday at [Huggingface](https://huggingface.co/keyfan/grok-1-hf/tree/main) which can be loaded using `AutoModelForCausalLM` with `trust_remote_code=True`. It can also be loaded with Bitsandbytes 4&8 bit. The weight conversion code is [here](https://gist.github.com/chu-tianxiang/ec310e15d56949fd0f351cb5f65ee7a1) if you have downloaded the origin weight.
I compared the logits with the JAX implementation and I think they're close enough. But there're probably many subtle differences to be fixed since I don't understand JAX at all.
Besides, I don't think XFormers or SDPA can be easily adapted, because Grok-1 adds an additional nonlinear transformation to the attention logits.

@chu-tianxiang I tried to load your model on 8xA10 with 4bit, I was getting this error.

```
ValueError:
Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit
the quantized mode l. If you want to dispatch the model on the CPU or the disk while keeping
these modules in 32-bit, you need to set load_in_8bit_fp32_cpu_offload=True and pass a custom
device_map to from_pretrained. Check
https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu
for more details.
```

regardless, I did manage to load on cpu and uploaded a 4bit version here, in case someone wants to test it out.
https://huggingface.co/eastwind/grok-1-hf-4bit

I don't have access to A100s to test. Is there a way to load it on like 32xA10 by using Ray for example?