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