Add JetMoE model
Author: yikangshen
Pull request: yes

# What does this PR do?

Add support to JetMoE architecture by Yikang Shen and MyShell AI. 
JetMoE is a new sparsely activated architecture inspired by the [ModuleFormer](https://arxiv.org/abs/2306.04640). Each JetMoE block consists of two MoE layers: a mixture of Attention Heads and a Mixture of MLP Experts. Given the input tokens, JetMoE activates a subset of its experts to process them. This sparse activation schema enables JetMoE to achieve much better training throughput than similar-sized dense models.

## Who can review?

@ArthurZucker and @younesbelkada


Comments:
https://huggingface.co/jetmoe/jetmoe-8b

I will be off next week -- please ping ArthurZucker [adding @ before 😉 ] if you have any questions or if you've addressed the review comments!

Hi @gante @ArthurZucker, I have updated the commit to address the review comments. Please let me know if there are further comments. Thanks!

Reviewing asap

Feel free to ping me again for another review 🤗 

> Feel free to ping me again for another review 🤗

I am fixing last a few bugs, will update the code ping you very soon. Thanks!

Hi @ArthurZucker, I have refactored the MoE and attention code and also made other changes according to your comments. The current check failure seems to be a test environmental issue, and has nothing to do with my changes. Please let me if you have further comments and if I can do something to fix the check failure. Thanks!

cc @younesbelkada for a review here! 

> Thanks so much for this work ! Quite clean already - one point that needs to be addressed IMO is the logic around `single_xxx` & `batched_xxx`, if batched inference works I don't see why it can't work with batch_size == 1? Do you think it's possible to harmonize a bit the current API, by generalising the forward pass of the modules that define these methods, to make the code easier to maintain for us? Thanks !

Thanks for the comments! The `single_xxx` skips many unnecessary ops for this specific scenario, thus making the model 30% faster during batch_size=1 decoding. However, I am also okay with removing these functions if needed.

Hi @younesbelkada, I have updated the code according to your suggestion. Please take a look when you have time, thanks!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30005). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> Thanks so much @yikangshen for addressing my comments, I believe the code is much more cleaner now ! We should be really close merging this in my opinion. I just have one main concern about the harmonisation of expert weights across your JetMoE modules. To harmonise everything let's either use `torch.nn.Linear` for all MoE / MoA expert modules or Module Lists instead of `nn.Parameter` unless you think that there is a valid reason I did not see. The rest should be quite clean in my opinion, just left few nits here and there ! Again thanks a lot for this huge work ! 🚀

Thanks for the comments @younesbelkada . I have fixed most of the issues you mentioned above. As for the expert weights, I hope we can keep the existing format, as it will provide the best compatibility with many other MoE acceleration kernels, including [Megablock](https://github.com/databricks/megablocks), [vllm](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py), and [scattermoe](https://github.com/shawntan/scattermoe). Please let me know what you think; I will make further updates based on your preference.

Thanks for the hard work, reviewing in a bit! 

Feel free to ping me whenever for another review! 🤗 

> Feel free to ping me whenever for another review! 🤗

Thanks @ArthurZucker. I have updated the code according to your suggestions. I hope the extra comments will make the code more clear. 

Thanks! having a look 😉 

Failing test is unrelated, should I merge? 🔥 

> Failing test is unrelated, should I merge? 🔥

Yes! I have tested offline with the following command:
`RUN_SLOW=1 python -m pytest tests/models/jetmoe/test_modeling_jetmoe.py -vv`
and all the tests are passed.

Congrats for this great work! We'll do a release on Thursday! 

> Congrats for this great work! We'll do a release on Thursday!

Thanks a lot for the review and comments! @ArthurZucker @gante @younesbelkada