ExLlama

一个独立的 Python/C++/CUDA 实现的 Llama 模型，用于处理 4 位 GPTQ 权重，旨在现代 GPU 上实现快速且内存高效。

免责声明：该项目正在开发中，但仍处于进行时！

硬件要求

我是在 RTX 4090 和 RTX 3090-Ti 上进行开发的。NVIDIA 的 30 系列及更高版本的 GPU 应该都能很好地支持，但 Pascal 或更早的显卡由于 FP16 支持较差，性能可能不会太好。 对于旧款 GPU，目前 AutoGPTQ 或 GPTQ-for-LLaMa 是更好的选择。理论上 ROCm 也受支持（通过 HIP），但我目前没有 AMD 设备来进行测试和优化。

依赖项

• Python 3.9 或更高版本
• torch 已在 2.0.1 和 2.1.0（夜间版）上测试过，使用 cu118
• safetensors 0.3.2
• sentencepiece
• ninja

此外，仅限 Web UI：

• flask
• waitress

Linux/WSL 前置条件

pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu118

Windows 前置条件

要在 Windows 上运行（不使用 WSL）：

1. 安装 MSVC 2022。你可以选择安装完整的 Visual Studio 2022 IDE，或者只安装 Build Tools for Visual Studio 2022 包（确保在安装程序中勾选了“Desktop development with C++”选项），两者都可以。
2. 安装适当版本的 PyTorch ，选择其中一个 CUDA 版本。我目前正在使用夜间构建版本进行开发，但稳定版本（2.0.1）也应该可以工作。
3. 安装 CUDA 工具包，11.7 和 11.8 都似乎可行，只需确保与 PyTorch 的计算平台版本匹配即可。
4. 为了获得最佳性能，启用硬件加速的 GPU 调度功能。

使用方法

克隆仓库，安装依赖项，并运行基准测试：

git clone https://github.com/turboderp/exllama
cd exllama

pip install -r requirements.txt

python test_benchmark_inference.py -d <模型文件路径> -p -ppl

CUDA 扩展会在运行时加载，因此无需单独安装。它将在首次运行时编译并缓存到 ~/.cache/torch_extensions/ 目录下，这可能需要一些时间。如果一开始没有任何反应，请稍等一分钟让其编译完成。

聊天机器人示例：

python example_chatbot.py -d <模型文件路径> -un "Jeff" -p prompt_chatbort.txt

Python 模块

jllllll 目前维护了一个可安装的 Python 模块 这里，可能更适合将 ExLlama 集成到其他项目中。

Web UI

我还为它制作了一个简单的 Web UI。别看里面的 JavaScript 代码，它基本上都是由 ChatGPT 写的，可能会让你做噩梦。不过它确实能用，而且挺有趣的，尤其是多机器人模式：

要运行它：

pip install -r requirements-web.txt

python webui/app.py -d <模型文件路径>

请注意，默认情况下会话数据存储在 ~/exllama_sessions/ 目录下。如果你需要更改存储位置，可以使用 -sd 参数。

Docker

为了提高安全性并简化部署，也可以将 Web UI 运行在一个隔离的 Docker 容器中。注意：当前的 Docker 镜像仅支持 NVIDIA GPU。

要求

• Docker
• NVIDIA Container Toolkit

建议以 无根模式 运行 Docker。

构建

构建 Docker 镜像最简单的方式是使用 Docker Compose。首先，在 .env 文件中将 MODEL_PATH 和 SESSIONS_PATH 变量设置为宿主机上的实际目录。然后运行：

docker compose build

也可以手动构建镜像：

docker build -t exllama-web .

注意：默认情况下，Docker 容器内的服务是由非 root 用户运行的。因此，在容器入口点脚本 (entrypoint.sh) 中，会将绑定挂载目录（默认 docker-compose.yml 文件中的 /data/model 和 /data/exllama_sessions）的所有权更改为该非 root 用户。若要禁用此功能，可以在使用 docker compose 时将 .env 文件中的 RUN_UID=0 设置为 0，或在手动构建镜像时执行以下命令：

docker build -t exllama-web --build-arg RUN_UID=0 .

运行

使用 Docker Compose：

docker compose up

现在可以通过主机的 http://localhost:5000 访问 Web UI。

配置可以在 docker-compose.yml 中查看，并可通过创建 docker-compose.override.yml 文件进行修改。

手动运行：

docker run --gpus all -p 5000:5000 -v <模型目录路径>:/data/model/ -v <会话目录路径>:/data/exllama_sessions --rm -it exllama-web --host 0.0.0.0:5000

目前的结果

新实现

模型	大小	分组大小	激活函数	序列长度	显存	提示速度	最优	最差	悲观度
Llama	7B	128	否	2,048 tokens	5,194 MB	13,918 tokens/s	173 tokens/s	140 tokens/s	6.45
Llama	13B	128	否	2,048 tokens	9,127 MB	7,507 tokens/s	102 tokens/s	86 tokens/s	5.60
Llama	33B	128	否	2,048 tokens	20,795 MB	2,959 tokens/s	47 tokens/s	40 tokens/s	4.60
Llama	33B	128	是	2,048 tokens	20,795 MB	2,784 tokens/s	45 tokens/s	37 tokens/s	4.55
Llama	33B	32	是	1,550 tokens 1	21,486 MB	2,636 tokens/s	41 tokens/s	37 tokens/s	4.52
Koala	13B	128	是	2,048 tokens	9,127 MB	5,529 tokens/s	93 tokens/s	79 tokens/s	6.73
WizardLM	33B	-	是	2,048 tokens	20,199 MB	2,313 tokens/s	47 tokens/s	40 tokens/s	5.75
OpenLlama	3B	128	是	2,048 tokens	3,128 MB	16,419 tokens/s	226 tokens/s	170 tokens/s	7.81

¹ 若无显存溢出，则无法达到完整序列长度

所有测试均在原装 RTX 4090 / 12900K 上进行，运行桌面环境时，其他一些应用也在使用显存。

“提示” 速度是指在所列序列长度基础上减去 128 个标记后的推理速度。“最差” 是完整上下文最后 128 个标记的平均速度（最坏情况），而 “最优” 则列出空序列中前 128 个标记的速度（最好情况）。

显存使用量由 PyTorch 报告，不包括 PyTorch 自身的开销（CUDA 内核、内部缓冲区等）。无论如何，这部分开销都难以准确预测。最好的办法是优化模型对显存的使用，最好在 24 GB 的显卡上将显存占用控制在 20 GB 左右，以确保为桌面环境和 Torch 的所有内部组件留出空间。

困惑度仅用于验证模型是否正常工作。所使用的数据集是 WikiText 中的一个特定小样本，因此得分与其他 Llama 基准测试不可比较，仅可用于比较不同 Llama 模型之间的性能。

双 GPU 结果

以下基准测试来自 4090 + 3090-Ti，配置为 -gs 17.2,24：

模型	大小	分组大小	激活函数	序列长度	显存	提示速度	最优	最差	悲观度
Llama	65B	128	是	2,048 tokens	39,804 MB	1,109 tokens/s	20 tokens/s	18 tokens/s	4.20
Llama	65B	32	是	2,048 tokens	43,424 MB	1,037 tokens/s	17 tokens/s	16 tokens/s	4.11
Llama-2	70B	128	是	2,048 tokens	40,680 MB	914 tokens/s	17 tokens/s	14 tokens/s	4.15
Llama-2	70B	32	是	2,048 tokens	36,815 MB	874 tokens/s	15 tokens/s	12 tokens/s	4.10

请注意，由于 Llama 和 Llama 2 的预训练数据集不同，它们的困惑度分数可能并不完全具有可比性。

待办事项

待办事项列表已移至 此处。

兼容性

此处列出了目前已确认可用的模型。

最新更新

2023-01-09：添加了 rope_theta 参数，以支持（至少部分）CodeLlama。如果您之前使用 alpha = 97 或类似设置，则现在不再需要为 CodeLlama 模型这样做。不过，关于扩展词汇表的部分仍有待进一步调整。

2023-08-09：增加了对分片模型的支持。config.model_path 现在可以接受文件名或文件名列表。如果传入的是模型目录，model_init() 将自动检测多个 .safetensors 文件。请注意各示例中的变化：model_path = glob.glob(st_pattern)[0] 简化为 model_path = glob.glob(st_pattern)。此外，util/shard.py 中提供了一个小脚本，用于拆分大型 .safetensors 文件。该脚本还会生成一个 index.json 文件，作为分片模型的索引，尽管 ExLlama 并不需要它来读取分片。需要注意的是，safetensors 依赖已升级至版本 0.3.2。

2023-08-12：初步、初始且暂定发布了 ExLlamaV2。虽然它目前尚未具备 ExLlamaV1 的全部功能，但在其擅长的领域表现更佳。不妨去看看吧！

ExLlama 快速上手指南

ExLlama 是一个独立的 Python/C++/CUDA 实现，专为在现代化 GPU 上高效运行 4-bit GPTQ 量化的 Llama 模型而设计，具有速度快、显存占用低的特点。

1. 环境准备

系统要求

• GPU: 推荐 NVIDIA 30 系列及更新型号（如 RTX 3090, 4090）。Pascal 架构或更旧的显卡因 FP16 支持较差，性能可能不理想（旧卡建议使用 AutoGPTQ）。
• 操作系统: Linux (推荐), WSL2, 或 Windows。
• Python: 3.9 或更高版本。

前置依赖

确保已安装以下基础库：

• torch (测试版本：2.0.1 或 2.1.0 nightly, CUDA 11.8)
• safetensors (>= 0.3.2)
• sentencepiece
• ninja

可选依赖（仅用于 Web UI）：

• flask, waitress

2. 安装步骤

第一步：安装 PyTorch (CUDA 版本)

根据操作系统选择以下方式之一：

Linux / WSL:

pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu118

(注：国内用户若下载缓慢，可尝试替换为清华源或阿里源对应的 PyTorch wheel 地址)

Windows:

1. 安装 MSVC 2022 (Visual Studio Build Tools)，务必勾选 "Desktop development with C++"。
2. 安装与 PyTorch 版本匹配的 CUDA Toolkit (推荐 11.7 或 11.8)。
3. 前往 PyTorch 官网 获取并安装对应 CUDA 版本的 PyTorch。
4. (可选) 为获得最佳性能，请在 Windows 设置中启用 "硬件加速 GPU 计划"。

第二步：克隆项目并安装依赖

git clone https://github.com/turboderp/exllama
cd exllama

pip install -r requirements.txt

(如需使用 Web UI，额外运行：pip install -r requirements-web.txt)

注意：CUDA 扩展会在首次运行时自动编译并缓存至 ~/.cache/torch_extensions/。首次执行可能需要等待几分钟进行编译，请耐心等待。

3. 基本使用

请确保你已准备好 4-bit GPTQ 格式的 Llama 模型文件（.safetensors 格式）。

运行基准测试 (Benchmark)

测试模型的推理速度和困惑度：

python test_benchmark_inference.py -d <path_to_model_files> -p -ppl

• <path_to_model_files>: 替换为你的模型文件所在目录路径。

运行命令行聊天机器人

体验简单的交互式对话：

python example_chatbot.py -d <path_to_model_files> -un "Jeff" -p prompt_chatbort.txt

启动 Web UI (可选)

启动基于浏览器的图形界面：

python webui/app.py -d <path_to_model_files>

启动后，访问 http://localhost:5000 即可使用。会话记录默认保存在 ~/exllama_sessions/。