🎉 简介
LLaMA-MoE是一系列开源的基于LLaMA和SlimPajama的专家混合模型(MoE)。 我们通过以下两个步骤构建LLaMA-MoE:
- 将LLaMA的FFNs分区为稀疏专家,并在每层专家中插入top-K门。
- 使用Sheared LLaMA优化的数据采样权重以及SlimPajama过滤的数据集,对初始化的MoE模型进行持续预训练。
🔥 特性
- 轻量级模型:激活的模型参数数量仅为3.0~3.5B,适合部署和研究使用。
- 多种专家构建方法:
- 神经元独立:随机、聚类、共激活图、梯度(Zhang et al., 2022, Zuo et al., 2022)
- 神经元共享:内部、跨层(残差)
- 多种MoE门控策略:
- TopK 噪声门控(Shazeer et al., 2017)
- 开关门控(Fedus et al., 2022)
- 快速持续预训练:
- 集成FlashAttention-v2(Dao, 2023)
- 快速流式数据集加载
- 丰富的监控项:
- 门负载,门重要性
- 步骤损失,令牌损失,平衡损失
- TGS(令牌/GPU/秒),MFU(模型浮点操作利用率)
- 其他可视化工具
- 动态权重采样:
- 自定义的静态采样权重
- Sheared LLaMA的动态批量加载(Xia et al., 2023)
🚀 快速开始
# python>=3.10
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model_dir = "llama-moe/LLaMA-MoE-v1-3_5B-2_8"
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.bfloat16, trust_remote_code=True)
model.eval()
model.to("cuda:0")
input_text = "苏州以"
inputs = tokenizer(input_text, return_tensors="pt")
inputs = inputs.to("cuda:0")
pred = model.generate(**inputs, max_length=50, temperature=0.0)
print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
# 苏州以其美丽的园林而闻名。最著名的是拙政园。这是一个拥有600多年历史的经典中国园林。该园分为三部分。
⚙️ 安装
- 准备conda环境:
conda create -n smoe python=3.11
(如果您的环境名称不是smoe
,您可能需要在启动脚本中更改环境) - 在
~/.bashrc
中添加正确的环境变量(安装flash-attn
时将gcc
设置为较新版本)。例如:export PATH=/mnt/petrelfs/share/cuda-11.8/bin:$PATH export LD_LIBRARY_PATH=/mnt/petrelfs/share/cuda-11.8/lib64:$LD_LIBRARY_PATH export PATH=/mnt/petrelfs/share/gcc-10.1.0/bin:$PATH export LD_LIBRARY_PATH=/mnt/petrelfs/share/gcc-10.1.0/lib64:$LD_LIBRARY_PATH
- 使变量生效:
source ~/.bashrc
- 安装PyTorch(CUDA-11.8):
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
- 安装依赖:
pip install -r requirements.txt
- 安装
flash-attn
:pip install flash-attn==2.0.1 --no-build-isolation
。您可能需要遵循flash-attn安装说明以避免一些错误。 - 安装最新的Git:
conda install git
- 克隆仓库:
git clone git@github.com:pjlab-sys4nlp/llama-moe.git
(如果您没有设置GitHub的ssh密钥,可能无法通过ssh克隆。请参考文档) - 更改当前目录:
cd llama-moe
- 以可编辑模式安装
smoe
:pip install -e .[dev]
- 设置
pre-commit
钩子:pre-commit install
📊 模型性能
模型 | #激活专家 | #专家 | #激活参数 | 基础模型 | SFT模型 |
---|---|---|---|---|---|
LLaMA-MoE-3.0B | 2 | 16 | 3.0B | 🤗 基础 | 🤗 SFT |
LLaMA-MoE-3.5B (4/16) | 4 | 16 | 3.5B | 🤗 基础 | 🤗 SFT |
LLaMA-MoE-3.5B (2/8) | 2 | 8 | 3.5B | 🤗 基础 | 🤗 SFT |
-
基础模型 | 模型 | 平均分 | SciQ | PIQA | WinoGrande | ARC-e | ARC-c (25) | HellaSwag (10) | LogiQA | BoolQ (32) | LAMBADA | NQ (32) | MMLU (5) | | :------------------------------------------------------------------------------------ | :------: | :------: | :------: | :--------: | :------: | :--------: | :------------: | :------: | :--------: | :------: | :------: | :------: | | OPT-2.7B | 50.3 | 78.9 | 74.8 | 60.8 | 54.4 | 34.0 | 61.4 | 25.8 | 63.3 | 63.6 | 10.7 | 25.8 | | Pythia-2.8B | 51.5 | 83.2 | 73.6 | 59.6 | 58.8 | 36.7 | 60.7 | 28.1 | 65.9 | 64.6 | 8.7 | 26.8 | | INCITE-BASE-3B | 53.7 | 85.6 | 73.9 | 63.5 | 61.7 | 40.3 | 64.7 | 27.5 | 65.8 | 65.4 | 15.2 | 27.2 | | Open-LLaMA-3B-v2 | 55.6 | 88.0 | 77.9 | 63.1 | 63.3 | 40.1 | 71.4 | 28.1 | 69.2 | 67.4 | 16.0 | 26.8 | | Sheared-LLaMA-2.7B | 56.4 | 87.5 | 76.9 | 65.0 | 63.3 | 41.6 | 71.0 | 28.3 | 73.6 | 68.3 | 17.6 | 27.3 | | LLaMA-MoE-3.0B | 55.5 | 84.2 | 77.5 | 63.6 | 60.2 | 40.9 | 70.8 | 30.6 | 71.9 | 66.6 | 17.0 | 26.8 | | LLaMA-MoE-3.5B (4/16) | 57.7 | 87.6 | 77.9 | 65.5 | 65.6 | 44.2 | 73.3 | 29.7 | 75.0 | 69.5 | 20.3 | 26.8 | | LLaMA-MoE-3.5B (2/8) | 57.6 | 88.4 | 77.6 | 66.7 | 65.3 | 43.1 | 73.3 | 29.6 | 73.9 | 69.4 | 19.8 | 27.0 |
-
SFT 模型
模型 | MMLU | ARC-c | HellaSeag | TruthfulQA | MT-Bench |
---|---|---|---|---|---|
Sheared LLaMA-2.7B ShareGPT | 28.41 | 41.04 | 71.21 | 47.65 | 3.79 |
Sheared LLaMA-2.7B Deita6K (Our Impl.) | 25.24 | 43.69 | 71.70 | 49.00 | 4.06 |
LLaMA-MoE-v1-3.0B (2/16) | 23.61 | 43.43 | 72.28 | 44.24 | 4.15 |
LLaMA-MoE-v1-3.5B (4/16) | 26.49 | 48.29 | 75.10 | 45.91 | 4.60 |
LLaMA-MoE-v1-3.5B (2/8) | 25.53 | 45.99 | 74.95 | 44.39 | 4.72 |
🚧 专家构建
- 神经元独立
- 独立随机:
bash ./scripts/expert_construction/split/run_split_random.sh
- 独立聚类:
bash ./scripts/expert_construction/split/run_split_clustering.sh
- 独立随机:
- 神经元共享
- 共享内部:
bash ./scripts/expert_construction/split/run_split_gradient.sh
- 共享外部:
bash ./scripts/expert_construction/split/run_split_gradient_residual.sh
- 共享内部:
更多信息请参考 专家构建文档。
🚅 持续预训练
分词
下载 SlimPajama 到 /path_to_data
并将来自不同领域的数据放入单独的文件夹:
/path_to_data/en_arxiv
/path_to_data/en_book
/path_to_data/en_c4
/path_to_data/en_cc
/path_to_data/en_stack
/path_to_data/en_wikipedia
/path_to_data/github
每个文件应以 *.jsonl
结尾,每行如下所示:
{"id": "id-info", "content": "raw text to be tokenized"}
运行以下命令以对每个文件夹中的数据进行分词:
python -m smoe.utils.tokenize \
-f jsonl \
-t /path_to_tokenizer \
-i /path_to_data/en_arxiv \
-o /path_to_data_tokenized/en_arxiv
持续预训练 (CPT)
- 注意: 请手动创建
logs/
文件夹:mkdir -p logs
- 要运行持续预训练,请检查 CPT 文档。
💎 评估
- 有关自然问题 (NQ) 的评估,请参考 opencompass。
- 对于其他任务,请参考 lm-eval-harness。
💬 监督微调 (SFT)
我们提供 SFT 的简单示例以构建聊天机器人。
请参考 SFT 文档 和 /mnt/petrelfs/zhutong/smoe/scripts/sft
了解更多详情。
📑 引用信息
@article{llama-moe,
title={LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training},
author={Tong Zhu and Xiaoye Qu and Daize Dong and Jiacheng Ruan and Jingqi Tong and Conghui He and Yu Cheng},
journal={arXiv preprint arXiv:2406.16554},
year={2024},
url={https://arxiv.org/abs/2406.16554},
}
LLaMA-MoE 团队 w/ ❤️