JetMoE:以10万美元达到LLaMA2的性能
关键信息
-
JetMoE-8B的训练成本不到10万美元1,但性能超过了拥有数十亿美元训练资源的Meta AI的LLaMA2-7B。这表明大型语言模型的训练可能比人们之前认为的要便宜得多。
-
JetMoE-8B完全开源且对学术界友好,原因如下:
- 它仅使用公开数据集进行训练,且代码开源。不需要任何专有资源。
- 它可以在非常有限的计算预算下进行微调(如消费级GPU),大多数实验室都能负担得起。
-
JetMoE-8B在推理过程中只有22亿个活跃参数,这大大降低了计算成本。与推理计算量相似的模型(如Gemma-2B)相比,JetMoE-8B始终表现更佳。
1 我们使用了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')
模型详情
有关模型详情和训练详情,请参阅技术报告 https://arxiv.org/pdf/2404.07413.pdf。