LLaMA-MoE 使用教程

LLaMA-MoE 使用教程

llama-moe ⛷️ LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training (EMNLP 2024) 项目地址: https://gitcode.com/gh_mirrors/ll/llama-moe

1. 项目介绍

LLaMA-MoE 是基于 LLaMA 和 SlimPajama 的开源 Mixture-of-Expert（MoE）模型系列。该模型通过以下两个步骤构建：

1. 将 LLaMA 的全连接层（FFNs）划分为稀疏专家，并在每一层专家中插入 top-K 门控。
2. 使用 Sheared LLaMA 和 SlimPajama 过滤的数据集，对初始化的 MoE 模型进行持续预训练。

LLaMA-MoE 模型具有以下特点：

• 轻量级模型：激活的模型参数数量仅为 3.0~3.5B，便于部署和研究使用。
• 多种专家构建方法：包括独立神经元、聚类、协同激活图、梯度等。
• 多种 MoE 门控策略：包括 TopK 噪声门控、开关门控等。
• 快速持续预训练：集成 FlashAttention-v2，支持快速数据流加载。

2. 项目快速启动

以下是一个快速启动 LLaMA-MoE 模型的示例代码：

# 确保已安装 torch 和 transformers
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)

# 将模型设置为评估模式，并移动到 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))

3. 应用案例和最佳实践

专家构建

根据项目需求，可以选择不同的专家构建方法。以下是一些构建方法的示例命令：

• 独立随机构建：

  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
  

持续预训练

为了进行持续预训练，需要准备数据集并进行分词，然后运行预训练脚本。

• 数据分词：

  python -m smoe.utils.tokenize -f jsonl -t /path_to_tokenizer -i /path_to_data/en_arxiv -o /path_to_data_tokenized/en_arxiv
  

• 持续预训练（CP）： 预训练的具体命令会根据数据集和模型配置有所不同，请参考官方文档进行配置和运行。

4. 典型生态项目

LLaMA-MoE 作为一种 MoE 模型，可以应用于多种自然语言处理任务，例如文本生成、问答系统等。以下是一些可能与之集成的典型生态项目：

• 文本生成工具：可以集成到内容管理系统或聊天机器人中，用于生成自然语言文本。
• 问答系统：可以用于构建智能问答系统，提供准确的答案。
• 语言理解服务：可以作为服务端组件，为应用提供语言理解能力。

通过以上教程，您可以开始使用 LLaMA-MoE，并根据项目需求进行定制化开发。

llama-moe ⛷️ LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training (EMNLP 2024) 项目地址: https://gitcode.com/gh_mirrors/ll/llama-moe