Jetmoe 8b

🚀 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 在設置聊天演示方面的支持。