🎉 简介

LLaMA-MoE 是一系列基于 LLaMA 和 SlimPajama 开源的专家混合（MoE）模型。我们通过以下两个步骤构建了 LLaMA-MoE：

1. 将 LLaMA 的前馈网络（FFN）划分为稀疏专家，并为每一层专家插入 top-K 门控机制。
2. 使用来自 Sheared LLaMA 的优化数据采样权重以及从 SlimPajama 过滤出的数据集，对初始化的 MoE 模型进行持续预训练。

🔥 特性

1. 轻量级模型：激活的模型参数仅为 3.0~3.5B，非常适合部署和研究使用。
2. 多种专家构建方法：
   1. 神经元独立型：随机、聚类、协同激活图、梯度法（Zhang et al., 2022, Zuo et al., 2022）
   2. 神经元共享型：内部共享、外部共享（残差连接）
3. 多种 MoE 门控策略：
   1. TopK 噪声门控（Shazeer et al., 2017）
   2. Switch 门控（Fedus et al., 2022）
4. 快速持续预训练：
   1. 集成 FlashAttention-v2（Dao, 2023）
   2. 快速流式数据加载
5. 丰富的监控指标：
   1. 门控负载、门控重要性
   2. 步骤损失、令牌损失、平衡损失
   3. TGS（每秒处理的令牌数/显卡）、MFU（模型浮点运算利用率）
   4. 其他可视化工具
6. 动态权重采样：
   1. 自定义静态采样权重
   2. 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))

# 苏州以其美丽的园林而闻名。其中最著名的是拙政园。它是一座有着六百多年历史的古典中国园林。这座园林被分为三个

⚙️ 安装

1. 准备 conda 环境：conda create -n smoe python=3.11（如果您的环境名称不是 smoe，可能需要在启动脚本中更改环境）
2. 在 ~/.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
   

3. 使变量生效：source ~/.bashrc
4. 安装 PyTorch（CUDA-11.8）：pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
5. 安装依赖项：pip install -r requirements.txt
6. 安装 flash-attn：pip install flash-attn==2.0.1 --no-build-isolation。您可能需要遵循 flash-attn 的安装说明 以避免一些错误。
7. 安装最新版 Git：conda install git
8. 克隆仓库：git clone git@github.com:pjlab-sys4nlp/llama-moe.git（如果您未将 SSH 密钥设置到 GitHub，则可能无法通过 SSH 克隆。请参阅 文档 了解相关操作。）
9. 更改当前目录：cd llama-moe
10. 以 可编辑模式 安装 smoe：pip install -e .[dev]
11. 设置 pre-commit 钩子：pre-commit install

📊 模型性能

模型	激活专家数	专家总数	激活参数量	基础模型	SFT 模型
LLaMA-MoE-3.0B	2	16	3.0B	🤗 base	🤗 SFT
LLaMA-MoE-3.5B (4/16)	4	16	3.5B	🤗 base	🤗 SFT
LLaMA-MoE-3.5B (2/8)	2	8	3.5B	🤗 base	🤗 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": "待分词的原始文本"}

运行以下命令对各文件夹中的数据进行分词：

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 文档。

💎 评估

• 关于 Natural Questions (NQ) 数据集上的评估，请参考 opencompass。
• 对于其他任务，请参考 lm-eval-harness。

💬 有监督微调 (SFT)

我们提供了用于构建聊天机器人的简单 SFT 示例。 更多详情请参阅 SFT 文档以及 /mnt/petrelfs/zhutong/smoe/scripts/sft 目录。

📑 引用

@article{llama-moe,
  title={LLaMA-MoE: 基于 LLaMA 的混合专家模型及其持续预训练},
  author={Tong Zhu 和 Xiaoye Qu 和 Daize Dong 和 Jiacheng Ruan 和 Jingqi Tong 和 Conghui He 和 Yu Cheng},
  journal={arXiv 预印本 arXiv:2406.16554},
  year={2024},
  url={https://arxiv.org/abs/2406.16554},
}

---

LLaMA-MoE 团队 敬上 ❤️

LLaMA-MoE 快速上手指南

LLaMA-MoE 是一个基于 LLaMA 架构的开源混合专家（MoE）模型系列。其激活参数量仅为 3.0B~3.5B，兼具高性能与低部署成本，适合研究与轻量级应用。

1. 环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐)
• Python 版本: >= 3.10 (推荐 3.11)
• CUDA 版本: 11.8 (用于安装 FlashAttention)
• 编译器: GCC 较新版本 (推荐 10.1.0+)
• 硬件: NVIDIA GPU (支持 CUDA)

前置依赖检查： 请确保已安装 conda 和 git。若需编译 flash-attn，请确认环境变量中已正确配置 CUDA 和 GCC 路径。

2. 安装步骤

请按顺序执行以下命令完成环境搭建：

2.1 创建 Conda 环境

conda create -n smoe python=3.11
conda activate smoe

2.2 配置环境变量 (可选但推荐)

如果您的系统默认 GCC 或 CUDA 版本过旧，请在 ~/.bashrc 中添加如下配置（路径需根据实际安装位置调整）：

export PATH=/path/to/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/path/to/cuda-11.8/lib64:$LD_LIBRARY_PATH
export PATH=/path/to/gcc-10.1.0/bin:$PATH
export LD_LIBRARY_PATH=/path/to/gcc-10.1.0/lib64:$LD_LIBRARY_PATH

使配置生效：

source ~/.bashrc

2.3 安装 PyTorch (CUDA 11.8)

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

国内加速建议：如遇下载缓慢，可使用清华源： pip3 install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple (需确保 wheel 兼容 CUDA 11.8，否则建议使用官方源或手动下载 whl)。

2.4 克隆项目并安装依赖

# 克隆仓库
git clone git@github.com:pjlab-sys4nlp/llama-moe.git
cd llama-moe

# 安装项目依赖
pip install -r requirements.txt

# 安装 flash-attn (必须指定版本且禁用构建隔离以避免错误)
pip install flash-attn==2.0.1 --no-build-isolation

# 以可编辑模式安装 smoe 包
pip install -e .[dev]

# 初始化 pre-commit 钩子
pre-commit install

3. 基本使用

安装完成后，您可以直接使用 Hugging Face transformers 加载模型进行推理。以下是最简单的文本生成示例：

# python>=3.10

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 指定模型目录 (可从 HuggingFace 自动下载或填写本地路径)
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)

# 设置模型为评估模式并移至 GPU
model.eval()
model.to("cuda:0")

# 准备输入文本
input_text = "Suzhou is famous of"
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))

预期输出示例：

Suzhou is famous of its beautiful gardens. The most famous one is the Humble Administrator's Garden. It is a classical Chinese garden with a history of more than 600 years...

---

注：更多高级功能（如专家构建、持续预训练、SFT 微调）请参考项目官方文档。