🚀 JetMoE：仅用10万美元达成LLaMA2性能

JetMoE-8B是一款极具性价比的大语言模型。它以不到10万美元的训练成本，实现了超越拥有数十亿训练资源的Meta AI的LLaMA2-7B的性能。该模型完全开源，对学术界友好，且推理时仅需22亿活跃参数，大幅降低了计算成本。

🚀 快速开始

若要加载模型，你需要安装 此包：

pip install -e .

然后，你可以使用以下代码加载模型：

from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, AutoModelForSequenceClassification
from jetmoe import JetMoEForCausalLM, JetMoEConfig, JetMoEForSequenceClassification

AutoConfig.register("jetmoe", JetMoEConfig)
AutoModelForCausalLM.register(JetMoEConfig, JetMoEForCausalLM)
AutoModelForSequenceClassification.register(JetMoEConfig, JetMoEForSequenceClassification)

tokenizer = AutoTokenizer.from_pretrained('jetmoe/jetmoe-8b')
model = AutoModelForCausalLM.from_pretrained('jetmoe/jetmoe-8b')

✨ 主要特性

1. 低成本高性能：JetMoE-8B训练成本不到10万美元¹，却超越了Meta AI的LLaMA2-7B，证明大语言模型训练可以比人们之前认为的便宜得多。
2. 开源友好：仅使用公共数据集进行训练，代码开源，无需专有资源，且可以在大多数实验室能够承担的有限计算预算（如消费级GPU）下进行微调。
3. 低计算成本：推理时仅具有22亿活跃参数，大幅降低了计算成本。与具有类似推理计算量的模型（如Gemma-2B）相比，JetMoE-8B始终表现更优。

¹ 我们使用96×H100 GPU集群进行了2周的训练，成本约为8万美元。

📚 详细文档

基准测试

我们采用与Open LLM排行榜相同的评估方法。对于MBPP代码基准测试，我们采用与LLaMA2和Deepseek-MoE论文相同的评估方法。结果如下：

令人惊讶的是，尽管训练成本和计算量较低，JetMoE-8B的表现甚至优于LLaMA2-7B、LLaMA-13B和DeepseekMoE-16B。与具有类似训练和推理计算量的模型（如Gemma-2B）相比，JetMoE-8B表现更优。

模型详情

JetMoE-8B有24个块。每个块有两个MoE层：注意力头混合层（MoA）和MLP专家混合层（MoE）。每个MoA和MoE层有8个专家，每个输入令牌激活2个专家。它总共有80亿个参数，其中22亿为活跃参数。JetMoE-8B在来自公开可用数据集的1.25T令牌上进行训练，学习率为5.0 x 10^-4，全局批量大小为4M令牌。

训练详情

我们的训练方法遵循MiniCPM的两阶段训练方法。第一阶段使用恒定学习率和线性预热，在来自大规模开源预训练数据集（包括RefinedWeb、Pile、Github数据等）的1万亿令牌上进行训练。第二阶段使用指数学习率衰减，在来自第一阶段数据集和额外高质量开源数据集的2500亿令牌上进行训练。

🔧 技术细节

更多详细信息，请参考 JetMoE技术报告。

📄 许可证

本项目采用Apache-2.0许可证。

🔗 相关链接

• 项目网站：https://research.myshell.ai/jetmoe
• HuggingFace：https://huggingface.co/jetmoe/jetmoe-8b
• Lepton AI在线演示：https://www.lepton.ai/playground/chat?model=jetmoe-8b-chat
• 技术报告：https://arxiv.org/pdf/2404.07413.pdf

👥 作者

本项目由 Yikang Shen、Zhen Guo、Tianle Cai 和 Zengyi Qin 贡献。如需技术咨询，请联系 Yikang Shen。如需媒体和合作咨询，请联系 Zengyi Qin。

🤝 合作

如果您有好的想法但需要更多资源（GPU、数据、资金等），欢迎通过 Zengyi Qin 联系 MyShell.ai。MyShell.ai 乐于合作，并积极支持高质量的开源项目。

📋 JetMoE模型索引

🙏 致谢

我们感谢 Shengding Hu 在第二阶段数据混合方面提供的宝贵建议。我们也感谢 Exabits 在设置GPU集群方面的帮助，以及 Lepton AI 在设置聊天演示方面的支持。