license: apache-2.0

 

JetMoE：以10万美元成本达到LLaMA2性能

核心亮点

1. 成本革命：JetMoE-8B训练成本不足10万美元¹，却超越Meta AI耗资数十亿美元训练的LLaMA2-7B，证明大模型训练成本可远低于行业认知。

2. 学术友好：完全开源的设计理念体现在：

   • 纯公共数据集训练：代码开源，无需任何专有资源
   • 低门槛微调：消费级GPU即可完成，普通实验室也能负担

3. 推理高效：仅激活22亿参数，计算成本显著降低。相比Gemma-2B等同类推理计算量的模型，性能持续领先。

¹ 使用96张H100 GPU集群训练2周，成本约8万美元。

官网：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

核心团队

项目由沈亦康、郭振、蔡天乐和覃增怡共同贡献。技术咨询请联系沈亦康，媒体与合作事宜请联系覃增怡。

合作邀约

如有创新构想但需要资源支持（GPU/数据/资金等），欢迎通过覃增怡联系MyShell.ai。我们持续赋能优质开源项目，期待与您共创未来。

性能基准

采用Open LLM leaderboard标准评估体系，代码能力评估与LLaMA2、Deepseek-MoE论文方法一致：

令人惊喜的是，在更低训练成本和计算量下，JetMoE-8B性能全面超越LLaMA2-7B、LLaMA-13B和DeepseekMoE-16B。与Gemma-2B等同类计算量模型相比，JetMoE-8B始终保持优势。

模型调用

安装依赖包：

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')

架构细节

• 24层模块化设计：每层包含注意力头混合模块(MoA)和专家混合模块(MoE)
• 动态激活：每个token激活2/8个专家，80亿总参数中仅激活22亿
• 训练配置：1.25T公开数据集token，学习率5.0×10^-4，全局batch size 400万token

训练策略

采用MiniCPM两阶段训练法：

• 第一阶段：1T token基础训练（RefinedWeb/Pile/Github等开源数据集），恒定学习率+线性预热
• 第二阶段：250B token精炼训练（新增高质量数据集），指数衰减学习率

技术报告

详见JetMoE技术报告

模型索引

致谢

特别感谢：

• 胡盛鼎对第二阶段数据配比的建议
• Exabits的GPU集群支持
• Lepton AI的对话演示技术支持