Openmoe Base

🚀 OpenMoE-Base

OpenMoEは、オープンソースのMixture-of-Experts (MoE) 大規模言語モデルのファミリーをリリースするプロジェクトです。このベースモデルは、1280億トークンで訓練され、主にデバッグ目的で使用されました。モデルアーキテクチャの有効性を検証した後、さらなる訓練は行っていません。そのため、性能はあまり良くない可能性があり、チェックポイントは実用的なアプリケーションには適していません。より良い性能は、8Bまたは34Bバージョンで確認できます。

[Github] | [Colab Demo] | [Huggingface] | [Discord] | [Twitter] | [Blog]

✨ 主な機能

• オープンソースのMixture-of-Experts (MoE) 大規模言語モデルのファミリーをリリース。
• 訓練データ、戦略、モデルアーキテクチャ、重みなどをコミュニティと共有。

📦 インストール

PyTorchを使用した推論

私たちのPyTorch実装は、Colossal AI でサポートされています。セットアップを簡単にするために、フォークしたバージョンを直接インストールできます。

# Pythonバージョン: 3.10.12
# ColossalAIのインストール
git clone --branch my_openmoe https://github.com/Orion-Zheng/ColossalAI.git
pip install ./ColossalAI
python -m pip install -r ./ColossalAI/examples/language/openmoe/requirements.txt

💻 使用例

基本的な使用法

from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM

model_path = "ckpts/openmoe-8b-chat"
config = AutoConfig.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True, 
    device_map='auto'
    )
query = 'Question: How do I kill a process? Answer:'
prompt = f'''<<SYS>>
You are a helpful, respectful and honest assistant.
<</SYS>>

<s>[INST] {query} [/INST]'''

inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
sample = model.generate(**inputs, max_new_tokens=32)
print(tokenizer.decode(sample[0]))

📚 ドキュメント

モデルの重み

現在、合計3つのモデルがリリースされています：OpenMoE-base、OpenMoE-8B/8B-Chat、およびOpenMoE-34B（2000億トークンで）。

以下の表は、1.1兆トークンで訓練を完了した8B/8B-Chatモデルを示しています。

また、研究目的ですべての中間チェックポイント（base、8B、34B）も提供しています。

その他の情報

• ベースモデルは、1280億トークンで訓練され、主にデバッグ目的で使用されました。モデルアーキテクチャの有効性を検証した後、さらなる訓練は行っていません。そのため、性能はあまり良くない可能性があり、チェックポイントは実用的なアプリケーションには適していません。より良い性能は、8Bまたは34Bバージョンで確認できます。
• 4つのMoEレイヤーと32のエキスパートを持つOpenMoE-8Bは、1.1兆トークンで訓練されています。SFTバージョンも、OpenMoE-8B-1.1Tをwildchat データセットのGPT-4サブセットで微調整した後にリリースされています。200B、400B、600B、800B、1Tトークンの中間チェックポイントは、MoEアーキテクチャの訓練ダイナミクスを研究するために使用できます。
• 私たちはまだOpenMoE-34Bを訓練中です。これは、8つのMoEレイヤーと32のエキスパートを持つMoEモデルです。2000億トークンで訓練された中間チェックポイントをhuggingfaceでリリースしました。最新のチェックポイントに興味がある場合は、Fuzhaoにメールを送ってください（f.xue@u.nus.edu）。

推論に関する注意事項

• GPUがない場合でも、Colabでモデルを体験できます（注：これには10ドルのColab Proプランが必要です）。こちら から直接ColabでOpenMoE-8B-Chatを実験できます。
• OpenMoE-8Bを実行するには、float32で約49GB、bfloat16で約23GBのメモリが必要です。Colabの CPU High-RAM（float32）ランタイムまたは A100-40GB（bfloat16）ランタイムで実行できますが、どちらもColab Proが必要です。float16精度は推奨されません。場合によっては性能が低下することがあります。
• OpenMoE-34Bを実行するには、bfloat16で約89GB、float32で約180GBのメモリが必要です。複数のデバイスで推論を実行するか、モデルの重みをRAMにオフロードするには、こちら のスクリプトを参照してください。
• より詳細な環境設定スクリプトは、こちら で見つけることができます。また、Dockerを使用する場合は、こちら のDockerfileを参照してください。注：huggingfaceのチェックポイント を使用してjaxチェックポイントを変換しない場合は、t5xとJaxの依存関係は必要ありません。

jaxチェックポイントの変換

また、jaxチェックポイントの変換を示すColab チュートリアル も提供しています。

📄 ライセンス

私たちのコードは、Apache 2.0ライセンスの下にあります。モデルはThe RedpajamaとThe Stackデータセットで訓練されているため、モデルの使用に関しては、これら2つのデータセットのライセンスを確認してください。

👥 作者

このプロジェクトは、現在以下の作者によって貢献されています。

Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, Yang You

🙏 謝辞

このプロジェクトの計算リソースは、Google TPU Research Cloud(TRC) から慷慨に提供されました。TRCには、私たちの研究の成功に不可欠な支援をしてくれたことに心から感謝しています。また、ColossalAIチーム のPyTorch実装に対する多大な支援、特に Xuanlei Zhao と Wenhao Chen のおかげで、GPUでのOpenMoEの訓練と推論が可能になりました。本当に感謝しています。

📖 引用

このリポジトリのモデルとコードを使用する場合は、以下のように引用してください。

@misc{openmoe2023,
  author = {Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou and Yang You},
  title = {OpenMoE: Open Mixture-of-Experts Language Models},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/XueFuzhao/OpenMoE}},
}