Jetmoe 8b Chat

🚀 JetMoE: 0.1 millionドルでLLaMA2の性能を達成

JetMoE-8Bは、0.1百万ドル未満のコストで訓練され、Meta AIのLLaMA2-7Bを上回る性能を発揮します。また、完全にオープンソースで学術界にやさしいモデルで、推論時のアクティブパラメータが少なく、計算コストを大幅に削減します。

🚀 クイックスタート

以下は、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は、0.1百万ドル未満のコストで訓練され、多額の訓練リソースを持つMeta AIのLLaMA2-7Bを上回る性能を発揮します。
2. 完全オープンソース：公開データセットのみを使用して訓練され、コードもオープンソースです。独自のリソースは必要ありません。
3. 低計算コスト：推論時には2.2Bのアクティブパラメータのみを使用するため、計算コストを大幅に削減します。

📦 インストール

インストールに関する具体的な手順は提供されていません。

💻 使用例

基本的な使用法

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)

📚 ドキュメント

モデルの詳細

JetMoE-8Bは24個のブロックから構成されています。各ブロックには、Mixture of Attention heads (MoA) とMixture of MLP Experts (MoE) の2つのMoEレイヤーがあります。各MoAとMoEレイヤーには8つのエキスパートがあり、各入力トークンに対して2つのエキスパートがアクティブになります。総計80億個のパラメータを持ち、アクティブパラメータは2.2Bです。JetMoE-8Bは、公開データセットからの1.25Tのトークンで訓練され、学習率は5.0 x 10^-4、グローバルバッチサイズは4Mトークンです。

訓練の詳細

私たちの訓練方法は、MiniCPMの2段階訓練法に従っています。フェーズ1では、一定の学習率で線形ウォームアップを行い、RefinedWeb、Pile、Githubデータなどの大規模オープンソース事前訓練データセットからの1兆トークンで訓練します。フェーズ2では、指数関数的な学習率減衰を使用し、フェーズ1のデータセットと追加の高品質オープンソースデータセットからの2500億トークンで訓練します。

技術レポート

詳細については、JetMoE Technical Reportを参照してください。

JetMoEモデルインデックス

🔧 技術詳細

JetMoE-8Bの技術的な詳細は、上記のモデルの詳細と訓練の詳細で説明されています。

📄 ライセンス

このプロジェクトは、Apache-2.0ライセンスの下で公開されています。

その他の情報

ベンチマーク

私たちは、Open LLM leaderboardと同じ評価方法を使用しています。MBPPコードベンチマークについては、LLaMA2とDeepseek-MoE論文と同じ評価方法を使用しています。結果は以下の通りです。

著者

このプロジェクトは、Yikang Shen、Zhen Guo、Tianle Cai、Zengyi Qinによって貢献されています。技術的な問い合わせは、Yikang Shenに、メディアや協力に関する問い合わせは、Zengyi Qinにお問い合わせください。

協力

もしあなたに素晴らしいアイデアがあるが、より多くのリソース（GPU、データ、資金など）が必要な場合、Zengyi Qinを通じてMyShell.aiにご連絡ください。MyShell.aiは協力を歓迎し、質の高いオープンソースプロジェクトを積極的に支援しています。