许可证：apache-2.0
基础模型：jetmoe/jetmoe-8b
标签：

• alignment-handbook
• generated_from_trainer
  数据集：
• HuggingFaceH4/ultrachat_200k
• HuggingFaceH4/airoboros-3.2
• HuggingFaceH4/Code-Feedback
• HuggingFaceH4/orca-math-word-problems-200k
• HuggingFaceH4/SystemChat
• HuggingFaceH4/capybara
  模型索引：
• 名称：jetmoe-8b-sft
  结果：[]

 

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

核心信息

1. JetMoE-8B训练成本不足10万美元¹，却超越Meta AI耗资数十亿美元训练的LLaMA2-7B。这表明大语言模型训练可以比人们预想的便宜得多。

2. JetMoE-8B是完全开源且学术友好的，因为：

   • 仅使用公开数据集训练，代码开源，无需专有资源
   • 可在大多数实验室负担得起的有限算力（如消费级GPU）上进行微调

3. JetMoE-8B推理时仅激活22亿参数，大幅降低计算成本。与Gemma-2B等推理算力相近的模型相比，JetMoE-8B性能持续更优。

¹ 我们使用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排行榜相同评估方法，代码基准MBPP评估与LLaMA2、Deepseek-MoE论文方法一致：

令人惊讶的是，尽管训练成本和算力更低，JetMoE-8B仍优于LLaMA2-7B、LLaMA-13B和DeepseekMoE-16B。与Gemma-2B等训练/推理算力相近的模型相比，JetMoE-8B表现更佳。

模型使用

快速使用JetMoE-8B-chat的示例：

import torch  
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM  
# 初始化模型和分词器  
model_name = "jetmoe/jetmoe-8b-chat"  
tokenizer = AutoTokenizer.from_pretrained(model_name)  
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, attn_implementation="eager", trust_remote_code=True)  
# 检查GPU可用性  
if torch.cuda.is_available():  
    model = model.cuda()  
    print("使用GPU:", torch.cuda.get_device_name(torch.cuda.current_device()))  
else:  
    print("未检测到GPU，使用CPU运行")  
# 构建对话输入  
messages = [  
    {"role": "system", "content": "你是一个友好的聊天机器人"},  
    {"role": "user", "content": "人类一顿饭能吃多少架直升机？"},  
]  
tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")  
# GPU推理时转移输入  
if torch.cuda.is_available():  
    input_ids = tokenized_chat.cuda()  
# 生成文本  
output = model.generate(input_ids, max_length=500, num_return_sequences=1, no_repeat_ngram_size=2)  
# 转移回CPU解码  
if torch.cuda.is_available():  
    output = output.cpu()  
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)  
print(generated_text)  

模型架构

JetMoE-8B包含24个模块，每个模块含混合注意力头（MoA）和混合专家MLP（MoE）两层。每层MoA/MoE有8个专家，每token激活2个专家。总参数量80亿，激活参数量22亿。模型在1.25T公开数据集token上训练，学习率5.0×10^-4，全局批大小400万token。

训练方案

采用MiniCPM两阶段训练法：

• 阶段1：恒定学习率+线性预热，在RefinedWeb、Pile、Github等1T token开源数据上训练
• 阶段2：指数学习率衰减，在阶段1数据+额外高质量开源数据的250B token上训练

技术报告

详见JetMoE技术报告。

模型索引

致谢

感谢胡圣鼎对阶段2数据配比的建议，感谢Exabits协助搭建GPU集群，感谢Lepton AI支持聊天演示部署。