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

关键信息

JetMoE-8B的训练成本不到10万美元<sup>1</sup>，但性能超过了拥有数十亿美元训练资源的Meta AI的LLaMA2-7B。这表明大型语言模型的训练可能比人们之前认为的要便宜得多。
JetMoE-8B完全开源且对学术界友好，原因如下：
- 它仅使用公开数据集进行训练，且代码开源。不需要任何专有资源。
- 它可以在非常有限的计算预算下进行微调（如消费级GPU），大多数实验室都能负担得起。
JetMoE-8B在推理过程中只有22亿个活跃参数，这大大降低了计算成本。与推理计算量相似的模型（如Gemma-2B）相比，JetMoE-8B始终表现更佳。

<sup>1</sup> 我们使用了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。MyShell.ai欢迎合作，并积极支持高质量的开源项目。

基准测试

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

模型	激活参数	训练词元	开放LLM排行榜平均	ARC	Hellaswag	MMLU	TruthfulQA	WinoGrande	GSM8k	MBPP	HumanEval
Shot				25	10	5	0	5	5	3	0
指标				acc_norm	acc_norm	acc	mc2	acc	acc	Pass@1	Pass@1
LLaMA2-7B	7B	2T	51.0	53.1	78.6	46.9	38.8	74	14.5	20.8	12.8
LLaMA-13B	13B	1T	51.4	56.2	80.9	47.7	39.5	76.2	7.6	22.0	15.8
DeepseekMoE-16B	2.8B	2T	51.1	53.2	79.8	46.3	36.1	73.7	17.3	34.0	25.0
Gemma-2B	2B	2T	46.4	48.4	71.8	41.8	33.1	66.3	16.9	28.0	24.4
JetMoE-8B	2.2B	1.25T	53.0	48.7	80.5	49.2	41.7	70.2	27.8	34.2	14.6

模型	MT-Bench 得分
GPT-4	9.014
GPT-3.5-turbo	7.995
Claude-v1	7.923
JetMoE-8B-chat	6.681
Llama-2-13b-chat	6.650
Vicuna-13b-v1.3	6.413
Wizardlm-13b	6.353
Llama-2-7b-chat	6.269

令我们惊讶的是，尽管训练成本和计算量更低，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')