使用 ExLlamaV2 在消费级 GPU 上运行 Llama 2 70B

Llama 2 系列中最大、最好的模型拥有 700亿 个参数。一个 fp16 参数占 2 个字节。加载 Llama 2 70B 需要 140 GB 内存（700 亿 * 2 字节）。

Llama 2 70B 明显小于 Falcon 180B。

Llama 2 70B 可以完全适合单个消费级 GPU 吗？

这是个很有挑战性的问题。高端消费类 GPU（例如 NVIDIA RTX 3090 或 4090）具有 24 GB 的显存VRAM。如果将 Llama 2 70B 量化到 4-bit 精度，仍然需要 35 GB 显存（700 亿 * 0.5 字节）。该模型可以安装到 2 个消费级 GPU 中。

通过 GPTQ 量化，可以进一步将精度降低到 3-bit，而不会损失太多模型的性能。 3-bit 参数在内存中的大小为 0.375 字节。 Llama 2 70B 量化为 3-bit 后仍需显存 26.25 GB，还是不适合一款消费级 GPU。

继续降低精度，将精度降低到 2-bit，需要现场 24 GB，但根据之前关于 2-bit 量化的研究，模型的性能会显着下降。

为了避免模型性能损失太多，可以将模型的重要层或部分量化为更高的精度，将不太重要的部分量化为更低的精度，将模型以混合精度方式进行量化。

ExLlamaV2（MIT 许可证）实现混合精度量化。

在本文中，将展示如何使用 ExLlamaV2 以混合精度方式量化模型。具体地说，将了解如何将 Llama 2 70B 量化到低于 3-bit 的平均精度。

什么是 Llama 2

Meta AI 的 Llama 2 摆脱了前身 Llama 的阴影，在聊天机器人领域树立新基准方面迈出了重大一步。它的前身 Llama 通过根据提示生成文本和代码来激起波澜，就像它的聊天机器人同行一样。

Llama 2 扮演两个不同的角色 Llama 2 和 Llama 2-Chat。Llama 2-Chat 特别适合进行双向对话。它们根据其复杂程度进一步分为不同的版本，范围从 7b 个参数到高达 70b 个参数的模型。

混合精度 Llama 2 的量化

要以混合精度量化模型并运行它们，需要安装 ExLlamaV2 。

从源码安装：

git clone https://github.com/turboderp/exllamav2
cd exllamav2
pip install -r requirements.txt

目标是在消费级 GPU 上运行模型。

• Llama 2 70B：目标是 24 GB 显存。 NVIDIA RTX3090/4090 GPU 可以工作。如果使用 Google Colab，则无法在免费的 Google Colab 上运行它。只有 Google Colab PRO 的 A100 有足够的显存。
• Llama 2 13B：目标是 12 GB 显存。许多具有至少 12 GB VRAM 的 GPU 都可用。 RTX3060/3080/4060/4080 就是其中的型号。它可以在带有 T4 GPU 的免费 Google Colab 上运行。

如何使用 ExLlamaV2 以混合精度进行量化

ExLlamaV2 使用的量化算法与 GPTQ 类似。但 ExLlamaV2 没有选择一种精度类型，而是在测量量化误差时为每一层尝试不同的精度类型。所有的尝试和相关的错误率都会被保存。然后，给定用户提供的目标精度，ExLlamaV2 算法将通过为每层模块选择领先的量化精度来量化模型，平均以最低的错误率达到目标精度。

在量化期间，ExLlamaV2 输出所有尝试：

Llama 2 13B 第10层 up_proj 模块的量化尝试：

-- Linear: model.layers.10.mlp.up_proj
 -- 0.05:3b/0.95:2b 32g s4         2.18 bpw    rfn_error: 0.21867
 -- 0.25:3b/0.75:2b 32g s4         2.38 bpw    rfn_error: 0.20617
 -- 0.25:4b/0.75:2b 32g s4         2.63 bpw    rfn_error: 0.20230
 -- 0.1:4b/0.4:3b/0.5:2b 32g s4    2.73 bpw    rfn_error: 0.18449
 -- 0.1:4b/0.9:3b 32g s4           3.23 bpw    rfn_error: 0.10229
 -- 0.2:6b/0.8:3b 32g s4           3.73 bpw    rfn_error: 0.09791
 -- 1.0:3b 128g s4                 3.03 bpw    rfn_error: 0.11354
 -- 1.0:3b 32g s4                  3.13 bpw    rfn_error: 0.10491
 -- 0.05:4b/0.95:3b 32g s4         3.18 bpw    rfn_error: 0.10363
 -- 0.4:4b/0.6:3b 32g s4           3.53 bpw    rfn_error: 0.09272
 -- 0.6:4b/0.4:3b 64g s4           3.66 bpw    rfn_error: 0.08835
 -- 1.0:4b 128g s4                 4.03 bpw    rfn_error: 0.05756
 -- 1.0:4b 32g s4                  4.13 bpw    rfn_error: 0.05007
 -- 0.1:5b/0.9:4b 32g s4           4.23 bpw    rfn_error: 0.04889
 -- 0.1:6b/0.9:4b 32g s4           4.33 bpw    rfn_error: 0.04861
 -- 1.0:5b 128g s4                 5.03 bpw    rfn_error: 0.02879
 -- 0.1:6b/0.9:5b 32g s4           5.23 bpw    rfn_error: 0.02494
 -- 0.05:8b/0.05:6b/0.9:5b 32g s4  5.33 bpw    rfn_error: 0.02486
 -- 0.4:6b/0.6:5b 32g s4           5.53 bpw    rfn_error: 0.02297
 -- 0.1:8b/0.3:6b/0.6:5b 32g s4    5.73 bpw    rfn_error: 0.02280
 -- 1.0:6b 128g s4                 6.03 bpw    rfn_error: 0.01503
 -- 1.0:6b 32g s4                  6.13 bpw    rfn_error: 0.01471
 -- 0.1:8b/0.9:6b 128g s4          6.23 bpw    rfn_error: 0.01463
 -- 1.0:8b 32g s4                  8.13 bpw    rfn_error: 0.00934
 -- Time: 19.57 seconds

可以看到，正如预期的那样，随着量化精度（bpw，即每权重位数）的增加，错误率降低。

使用 ExLlamaV2 进行量化就像运行 convert.py 脚本一样简单，其中 convert.py 位于ExLlamaV2 的根目录下：

python convert.py \
    -i ./Llama-2-13b-hf/ \
    -o ./Llama-2-13b-hf/temp/ \
    -c test.parquet \
    -cf ./Llama-2-13b-hf/3.0bpw/ \
    -b 3.0

ExLlamaV2 不支持 Hugging Face 库，它期望模型和校准数据集存储在本地。

该脚本的主要参数如下：

• 输入模型 -i：包含 safetensors 格式模型的本地目录。
• 用于校准的数据集 -c：需要一个用于校准量化的数据集。它必须以 parquet 格式存储在本地。
• 输出目录 -cf：保存量化模型的本地目录。
• 量化的目标精度 -b：模型将以混合精度进行量化，该混合精度将平均为目标精度。在这里，选择以 3-bit 精度为目标。

这次量化花了 2小时5分钟。使用具有 T4 GPU 和高 CPU RAM 的 Google Colab PRO。整个过程中没有消耗超过 5 GB 的 VRAM，但 CPU RAM 的峰值消耗达到了 20 GB。

T4 相当慢。使用 Google Colab V100 或 RTX GPU 可以缩短量化时间。 注意：不清楚量化期间 GPU 的使用量。 CPU 速度对量化时间的影响可能比 GPU 更大。

要量化 Llama 2 70B，可以执行相同的操作。

应该以什么精度为目标，以便量化的 Llama 2 70B 适合 24 GB VRAM？

可以使用以下方法来确定给定硬件的模型的精度。

假设有 24 GB 的 VRAM。还应该始终预期推理会产生一些内存开销。因此，目标是 22 GB 的量化模型大小。

首先，需要将 22 GB 转换为 bits：

22 GB = 2.2e+10 bytes = 1.76e+11 bits (since 1 byte = 8 bits)

有 1.76e+11 bits (b) 可用。 Llama 2 70B 有 7e+10 个要量化的参数 p。目标精度是 bpw。

bpw = b/p
bpw = 176 000 000 000 / 70 000 000 000 = 2.51

因此可以承受每个参数 2.51 bits 的平均精度。将其四舍五入为 2.5 bits。 为了将 Llama 2 70B 量化到 2.5 bits 的平均精度，运行：

python convert.py \
    -i ./Llama-2-70b-hf/ \
    -o ./Llama-2-70b-hf/temp/ \
    -c test.parquet \
    -cf ./Llama-2-70b-hf/2.5bpw/ \
    -b 2.5

这种量化在具有 24 GB GPU 的消费类硬件上也是可行的。最多可能需要 15 小时。如果想使用 Google Colab 进行此操作，请注意，必须将原始模型存储在 Google Colab 硬盘之外，因为使用 A100 GPU 时原始模型太小。

使用 ExLlamaV2 在 GPU 上运行 Llama 2 70B

ExLlamaV2 提供运行混合精度量化模型所需的一切。

有一个 chat.py 脚本将模型作为聊天机器人运行以供交互使用。还可以使用 test_inference.py 简单地测试模型。这就是检查模型速度和内存消耗的方法。

为了测试以 2.5 bpw 量化的 Llama 2 70B，运行：

python test_inference.py -m ./Llama-2-70b-2.5bpw/ -p "Once upon a time,"

• -p 为测试提示符。

这应该需要几分钟（在 A100 GPU 上需要 8 分钟）。 ExLlamaV2 使用torch.compile。根据 PyTorch 文档：

torch.compile 通过将 PyTorch 代码 JIT 编译到优化的内核中，使 PyTorch 代码运行得更快，同时只需要最少的代码更改。

这个编译很耗时但是会被缓存。

如果运行 test_inference.py，同样只需要 30 秒。

该模型本身需要显存 22.15 GB。在推理实验中，它正好占用了 24 GB。它几乎不适合消费级 GPU。

为什么它不只消耗 22.15 GB 显存？

显存中的模型实际上占用了 22.15 GB，但推理本身也消耗了额外的显存。例如，必须对提示进行编码并将其存储在显存中。此外，如果设置更高的最大序列长度或进行批量解码，推理将消耗更多显存。

使用 Google Colab 的 A100 进行这个实验。如果使用 24 GB 的 GPU，则在推理过程中可能会出现 CUDA 内存不足错误，特别是如果还使用 GPU 运行操作系统图形用户界面（例如，Ubuntu 桌面消耗大约 1.5 GB 的显存） 。

为了给一些利润，目标是较低的 bpw。 2.4 甚至 2.3 会留下几 GB 的 VRAM 可用于推理。

ExLlamaV2 模型的速度也非常快。观察到生成速度在 15 到 30 个令牌/秒之间。为了给一个比较点，当使用 GPTQ（一个小 10 倍的模型）对量化为 4-bit 的 Llama 2 7B 进行基准测试时，使用 Hugging Face 转换器进行生成，可以达到大约 28 个令牌/秒的速度。

总结

混合精度的量化很直观，在影响较小的地方积极降低模型的精度。在单个消费级 GPU 上运行 Llama 2 70B 等大型模型是可能的。

请务必评估使用不同目标精度量化的模型。虽然较大的模型更容易量化而不会造成太大的性能损失，但总有一个精度，在该精度下，量化的模型将变得比未量化但参数较少的模型更差，例如，Llama 2 70B 2-bit 可能比 4-bit 差很多，但仍然更大。

翻译：towardsdatascience.com/run-llama-2...