Jetmoe 8b Chat

🚀 JetMoE：僅用10萬美元達到LLaMA2性能

JetMoE-8B模型以不到10萬美元的訓練成本，超越了擁有數十億美元訓練資源的Meta AI的LLaMA2-7B。它完全開源，對學術界友好，推理時僅22億活躍參數，大幅降低計算成本。

🚀 快速開始

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

import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
# Initialize the model and tokenizer
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)
# Check if a GPU is available and move the model to GPU if it is
if torch.cuda.is_available():
    model = model.cuda()
    print("Using GPU:", torch.cuda.get_device_name(torch.cuda.current_device()))
else:
    print("GPU is not available, using CPU instead.")
# Encode input context
messages = [
    {
        "role": "system",
        "content": "You are a friendly chatbot",
    },
    {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
print(tokenized_chat)
# If using a GPU, move the input IDs to the GPU
if torch.cuda.is_available():
    input_ids = tokenized_chat.cuda()
# Generate text
output = model.generate(input_ids, max_length=500, num_return_sequences=1, no_repeat_ngram_size=2)
# If the output is on the GPU, move it back to CPU for decoding
if torch.cuda.is_available():
    output = output.cpu()
# Decode the generated text
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_text)

✨ 主要特性

1. 低成本高性能：JetMoE-8B訓練成本不到10萬美元，卻超越了Meta AI的LLaMA2-7B。這表明大語言模型的訓練成本可以比人們之前認為的低得多。
2. 完全開源且友好：僅使用公共數據集進行訓練，代碼開源，無需專有資源。並且可以在大多數實驗室都能負擔得起的有限計算預算（如消費級GPU）下進行微調。
3. 低計算成本：推理時僅22億活躍參數，與Gemma-2B等推理計算量相似的模型相比，性能更優。

📚 詳細文檔

基準測試

我們使用與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⁻⁴，全局批量大小為4M令牌。

訓練詳情

我們的訓練方法遵循MiniCPM的兩階段訓練方法。第一階段使用恆定學習率和線性預熱，在來自大規模開源預訓練數據集（包括RefinedWeb、Pile、Github數據等）的1萬億令牌上進行訓練。第二階段使用指數學習率衰減，在第一階段數據集和額外高質量開源數據集的2500億令牌上進行訓練。

JetMoE模型索引

🔧 技術細節

如需更多詳細信息，請參考JetMoE技術報告。

📄 許可證

本項目採用Apache-2.0許可證。

👥 作者信息

本項目由Yikang Shen、Zhen Guo、Tianle Cai和Zengyi Qin貢獻。技術諮詢請聯繫Yikang Shen。媒體和合作諮詢請聯繫Zengyi Qin。

🤝 合作機會

如果您有好的想法但需要更多資源（GPU、數據、資金等），歡迎通過Zengyi Qin聯繫MyShell.ai。MyShell.ai樂於合作，並積極支持高質量的開源項目。

🙏 致謝

我們感謝Shengding Hu在第二階段數據混合方面提供的寶貴建議。我們也感謝Exabits在搭建GPU集群方面的幫助，以及Lepton AI在搭建聊天演示方面的支持。

🔗 相關鏈接

• 官網：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