Openmoe Base

🚀 OpenMoE-Base

OpenMoE 項目旨在激發開源混合專家（MoE）社區的活力！我們正在發佈一系列開源的混合專家（MoE）大語言模型。該基礎模型使用 1280 億個標記進行訓練，主要用於調試目的。在驗證了模型架構的有效性後，我們未進行進一步訓練。因此，其性能可能不太理想，該檢查點不適用於實際應用。我們的 8B 或 34B 版本能展現出更好的性能。

🚀 快速開始

使用 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

然後，你可以在配備 A100 80GB 的機器上使用以下代碼進行推理：

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]))

如果你手頭沒有 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 依賴項。

此外，我們還提供了一個 Colab 教程，演示了 Jax 檢查點的轉換。

✨ 主要特性

• 本項目致力於開源混合專家（MoE）大語言模型，推動該領域的研究發展。
• 提供多種規模的模型，包括 OpenMoE-base、OpenMoE-8B/8B-Chat 和 OpenMoE-34B 等，滿足不同需求。
• 分享訓練數據、策略、模型架構、權重等所有資源，促進社區共同參與。

📦 安裝指南

安裝 ColossalAI

# 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]))

📚 詳細文檔

模型權重

目前，總共發佈了三個模型：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 萬億個標記上進行了訓練。在我們對 OpenMoE-8B-1.1T 在 wildchat 數據集的 GPT-4 子集上進行微調後，也發佈了 SFT 版本。2000 億、4000 億、6000 億、8000 億和 1 萬億個標記的中間檢查點可用於研究 MoE 架構的訓練動態。

我們仍在訓練 OpenMoE-34B，這是一個具有 8 個 MoE 層和 32 個專家的 MoE 模型。我們在 Huggingface 上發佈了在 2000 億個標記上訓練的中間檢查點。如果你對最新的檢查點感興趣，請隨時給 Fuzhao 發郵件（f.xue@u.nus.edu）。

📄 許可證

我們的代碼遵循 Apache 2.0 許可證。

由於模型是在 The Redpajama 和 The Stack 數據集上訓練的，請檢查這兩個數據集的許可證以確定你的模型使用方式。

👥 作者

本項目目前由以下作者貢獻： 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，使 OpenMoE 在 GPU 上的訓練和推理成為現實。

📚 引用

如果您使用了本倉庫中的模型和代碼，請引用本倉庫。

@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}},
}