rwkv.cpp
这是一个将BlinkDL/RWKV-LM移植到ggerganov/ggml的项目。
除了常规的FP32外,它还支持FP16、量化为INT4、INT5和INT8的推理。这个项目主要针对CPU,但也支持cuBLAS。
本项目提供了一个C语言库rwkv.h和一个方便使用的Python封装。
RWKV是一种大型语言模型架构,其中最大的模型拥有140亿参数。与具有O(n^2)
注意力复杂度的Transformer不同,RWKV只需要前一步的状态就能计算logits。这使得RWKV在处理长上下文时非常适合CPU。
RWKV v5是RWKV架构的一次重大升级,使其在质量上可以与Transformer竞争。本项目支持RWKV v5模型。
RWKV v6是RWKV架构的进一步改进,具有更高的质量。本项目支持RWKV v6模型。
通过merge_lora_into_ggml.py脚本支持加载Blealtan格式的LoRA检查点。
质量和性能
如果你将rwkv.cpp
用于任何严肃的用途,请在代表性数据集上测试所有可用格式的困惑度和延迟,并决定哪种权衡最适合你。
总的来说,RWKV v5
模型与RWKV v4
模型一样快,在延迟和内存消耗上有细微差异,但质量远高于v4
。因此,推荐使用RWKV v5
。
下表仅供参考。测量在具有AVX2的4核/8线程x86 CPU上进行,使用4个线程。模型为RWKV v4 Pile 169M
和RWKV v4 Pile 1.5B
。
格式 | 困惑度 (169M) | 延迟,毫秒 (1.5B) | 文件大小,GB (1.5B) |
---|---|---|---|
Q4_0 | 17.507 | 76 | 1.53 |
Q4_1 | 17.187 | 72 | 1.68 |
Q5_0 | 16.194 | 78 | 1.60 |
Q5_1 | 15.851 | 81 | 1.68 |
Q8_0 | 15.652 | 89 | 2.13 |
FP16 | 15.623 | 117 | 2.82 |
FP32 | 15.623 | 198 | 5.64 |
使用cuBLAS
测量在Intel i7 13700K和NVIDIA 3060 Ti 8 GB上进行。模型为RWKV-4-Pile-169M
,12层卸载到GPU。
显示的是每个令牌的延迟(毫秒)。
格式 | 1线程 | 2线程 | 4线程 | 8线程 | 24线程 |
---|---|---|---|---|---|
Q4_0 | 7.9 | 6.2 | 6.9 | 8.6 | 20 |
Q4_1 | 7.8 | 6.7 | 6.9 | 8.6 | 21 |
Q5_1 | 8.1 | 6.7 | 6.9 | 9.0 | 22 |
格式 | 1线程 | 2线程 | 4线程 | 8线程 | 24线程 |
---|---|---|---|---|---|
Q4_0 | 59 | 51 | 50 | 54 | 94 |
Q4_1 | 59 | 51 | 49 | 54 | 94 |
Q5_1 | 77 | 69 | 67 | 72 | 101 |
注意:由于cuBLAS仅支持ggml_mul_mat()
,我们仍需使用少量CPU资源来执行剩余操作。
使用hipBLAS
测量在CPU AMD Ryzen 9 5900X和GPU AMD Radeon RX 7900 XTX上进行。模型为RWKV-novel-4-World-7B-20230810-ctx128k
,32层卸载到GPU。
显示的是每个令牌的延迟(毫秒)。
格式 | 1线程 | 2线程 | 4线程 | 8线程 | 24线程 |
---|---|---|---|---|---|
f16 | 94 | 91 | 94 | 106 | 944 |
Q4_0 | 83 | 77 | 75 | 110 | 1692 |
Q4_1 | 85 | 80 | 85 | 93 | 1691 |
Q5_1 | 83 | 78 | 83 | 90 | 1115 |
注意:与cuBLAS一样,hipBLAS仅支持ggml_mul_mat()
,我们仍需使用少量CPU资源来执行剩余操作。
使用方法
1. 克隆仓库
要求:git。
git clone --recursive https://github.com/saharNooby/rwkv.cpp.git
cd rwkv.cpp
2. 获取rwkv.cpp库
选项2.1. 下载预编译库
Windows / Linux / MacOS
查看Releases,下载适合你的操作系统和CPU的ZIP文件,将rwkv
库文件解压到仓库目录中。
Windows用户:要检查你的CPU是否支持AVX2或AVX-512,请使用CPU-Z。
选项2.2. 自行构建库
为获得最佳性能,推荐使用此选项,因为库将专门为你的CPU和操作系统构建。
Windows
要求:CMake或来自anaconda的CMake,Visual Studio 2019的构建工具。
cmake .
cmake --build . --config Release
如果一切正常,bin\Release\rwkv.dll
文件应该会出现。
Windows + cuBLAS
请参阅docs/cuBLAS_on_Windows.md获取详细指南。
Windows + hipBLAS
请参阅docs/hipBLAS_on_Windows.md获取详细指南。
Linux / MacOS
要求:CMake(Linux:sudo apt install cmake
,MacOS:brew install cmake
,anaconoda:cmake包)。
cmake .
cmake --build . --config Release
Anaconda和M1用户:请在运行cmake .
后验证CMAKE_SYSTEM_PROCESSOR: arm64
——如果检测到x86_64
,请编辑CMakeLists.txt
文件,在# Compile flags
下添加set(CMAKE_SYSTEM_PROCESSOR "arm64")
。
如果一切正常,librwkv.so
(Linux)或librwkv.dylib
(MacOS)文件应该会出现在仓库的基本文件夹中。
Linux / MacOS + cuBLAS
cmake . -DRWKV_CUBLAS=ON
cmake --build . --config Release
如果一切正常,librwkv.so
(Linux)或librwkv.dylib
(MacOS)文件应该会出现在仓库的基本文件夹中。
3. 获取RWKV模型
要求:Python 3.x 和 PyTorch。
第一步,从 Hugging Face 下载一个模型,比如 这个。
第二步,使用以下命令将其转换为 rwkv.cpp
格式:
# Windows
python python\convert_pytorch_to_ggml.py C:\RWKV-4-Pile-169M-20220807-8023.pth C:\rwkv.cpp-169M.bin FP16
# Linux / MacOS
python python/convert_pytorch_to_ggml.py ~/Downloads/RWKV-4-Pile-169M-20220807-8023.pth ~/Downloads/rwkv.cpp-169M.bin FP16
可选步骤,将模型量化为上表中的一种量化格式:
# Windows
python python\quantize.py C:\rwkv.cpp-169M.bin C:\rwkv.cpp-169M-Q5_1.bin Q5_1
# Linux / MacOS
python python/quantize.py ~/Downloads/rwkv.cpp-169M.bin ~/Downloads/rwkv.cpp-169M-Q5_1.bin Q5_1
4. 运行模型
使用命令行
要求:Python 3.x 和 numpy。如果使用 Pile
或 Raven
模型,还需要 tokenizers。
生成一些文本,运行:
# Windows
python python\generate_completions.py C:\rwkv.cpp-169M-Q5_1.bin
# Linux / MacOS
python python/generate_completions.py ~/Downloads/rwkv.cpp-169M-Q5_1.bin
与机器人聊天,运行:
# Windows
python python\chat_with_bot.py C:\rwkv.cpp-169M-Q5_1.bin
# Linux / MacOS
python python/chat_with_bot.py ~/Downloads/rwkv.cpp-169M-Q5_1.bin
编辑 generate_completions.py 或 chat_with_bot.py 来更改提示和采样设置。
在自己的代码中使用
简短的脚本 inference_example.py 演示了在 Python 中使用 rwkv.cpp
的方法。
要在 C/C++ 中使用 rwkv.cpp
,请包含头文件 rwkv.h。
要在任何其他语言中使用 rwkv.cpp
,请参见下面的 绑定 部分。如果缺少你的语言,你可以尝试使用该语言提供的工具绑定到 C API。
绑定
这些项目封装了 rwkv.cpp
,以便在其他语言/框架中更易使用。
- Golang: seasonjs/rwkv
- Node.js: Atome-FE/llama-node
兼容性
ggml
发展迅速,偶尔可能会破坏与旧文件格式的兼容性。
rwkv.cpp
会尽最大努力解释为什么无法加载模型文件,并向用户提供下一步可用的选择。
仅供参考,以下是支持旧格式的 rwkv.cpp
最新版本列表。这些版本将不提供任何支持。
Q4_2
,量化格式的旧布局Q4_3
,Q4_1_O
另请参阅 docs/FILE_FORMAT.md 了解 rwkv.cpp
模型文件的版本号及其更新日志。
贡献
请遵循 docs/CODE_STYLE.md 中描述的代码风格。