LLaMA-MoE学习资料汇总 - 基于LLaMA的轻量级混合专家模型

LLaMA-MoE:基于LLaMA的轻量级混合专家模型

LLaMA-MoE是一个基于LLaMA和SlimPajama的开源混合专家(Mixture-of-Experts, MoE)模型系列。该项目旨在构建一个轻量级且人人可用的MoE模型。本文将为大家介绍LLaMA-MoE的主要特点、安装使用方法以及相关学习资源。

项目简介

LLaMA-MoE通过以下两个步骤构建:

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

主要特点

1. 轻量级模型:激活的模型参数仅为3.0-3.5B,便于部署和研究使用。
2. 多种专家构建方法:包括神经元独立和神经元共享两大类。
3. 多种MoE门控策略:TopK Noisy Gate和Switch Gating。
4. 快速持续预训练:集成FlashAttention-v2,支持快速流式数据集加载。
5. 丰富的监控项:包括门控负载、损失、TGS、MFU等。
6. 动态权重采样:支持自定义静态采样权重和Sheared LLaMA的动态批处理加载。

快速开始

以下是使用LLaMA-MoE模型的简单示例:

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 = "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))

安装指南

1. 准备conda环境:conda create -n smoe python=3.11
2. 安装PyTorch: pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
3. 安装依赖: pip install -r requirements.txt
4. 安装flash-attn: pip install flash-attn==2.0.1 --no-build-isolation
5. 克隆仓库: git clone git@github.com:pjlab-sys4nlp/llama-moe.git
6. 安装smoe: pip install -e .[dev]

详细的安装步骤请参考项目README。

模型性能

LLaMA-MoE提供了多个版本的模型,包括3.0B和3.5B参数规模。以下是部分模型在各项任务上的表现:

完整的性能数据请查看项目README。

相关资源

• 📃 技术报告
• 🤗 模型权重
• 🚧 专家构建文档
• 🚅 持续预训练文档
• 💬 监督微调(SFT)文档

总结

LLaMA-MoE项目为研究人员和开发者提供了一个轻量级且易于使用的混合专家模型。通过本文介绍的资源,读者可以快速了解项目的核心概念,并开始使用LLaMA-MoE进行实验和开发。欢迎大家探索这个有趣的项目,为NLP领域的发展贡献力量。