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

news2024/12/23 13:03:15

Llama 2模型中最大也是最好的模型有700亿个参数。一个fp16参数的大小为2字节。加载Llama 270b需要140 GB内存(700亿* 2字节)。

只要我们的内存够大，我们就可以在CPU上运行上运行Llama 2 70B。但是CPU的推理速度非常的慢，虽然能够运行，速度我们无法忍受。

能否在高端消费级GPU，如NVIDIA RTX 3090或4090，上运行呢，如果我们将Llama 2 70b量化到4位精度，仍然需要35 GB的内存(700亿* 0.5字节)，如果有2个GPU，那么肯定是可以的。或者通过GPTQ量化，可以在不影响模型性能的情况下将精度进一步降低到3位。一个3位参数在内存中占0.375字节。Llama 2 70b量化为3比特后仍重26.25 GB，一个4090还是装不下。那么把精度降低到2位呢。他肯定可以使用24gb的VRAM加载，但根据之前对2位量化的研究，模型的性能会显著下降。

为了避免在模型的性能上损失太多，可以将模型的重要层或部分量化到更高的精度，而将不太重要的部分量化到更低的精度。这被称作混合精度量化。

在本文中，我将展示如何使用ExLlamaV2以混合精度量化模型。我们将看到如何将Llama 2 70b量化到低于3位的平均精度。

Llama 2的混合精度量化

为了量化混合精度的模型，我们需要安装ExLlamaV2。

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

我们的目标是在消费级gpu上运行模型。

对于Llama 2 70b，我们的目标是使用24gb的VRAM，NVIDIA RTX3090/4090 gpu

对于Llama 2 13B，我们的目标是12gb的VRAM。这样RTX3060/3080/4060/4080都可以使用，并且它可以运行在免费的谷歌Colab和T4 GPU上。

如何使用ExLlamaV2

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

在量化过程中，ExLlamaV2会输出测试的结果:

 -- 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脚本一样简单:

 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

脚本的主要参数如下:

input model (-i):以“safetensors”格式包含模型的本地目录。
用于校准的数据集(-c):我们需要一个用于校准量化的数据集。它必须以“parquet”格式存储在本地。
output directory (-cf):存放量化模型的本地目录。
量化的目标精度(-b):模型将以混合精度进行量化，混合精度将平均为目标精度。我们选择以3位精度为目标。

需要注意的是：

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

上面这个这个量化过程在谷歌Colab PRO上耗时2小时5分钟。在整个过程中，它消耗的VRAM不超过5 GB，但CPU RAM的峰值消耗为20 GB。

因为T4相当慢，所以如果使用V100或4090速度会更快。这里不确定在量化过程中使用了多少GPU。可能是CPU速度比GPU对量化时间的影响更大。

量化参数计算

如果要量化Llama 2 70b，我们应该预估一个以多大的精度为目标，才能使量子化的Llama 270b适合24 GB的VRAM?所以在给定硬件的情况下，可以用以下方法来确定模型的精度。

假设我们有24gb的VRAM。因为有一些推理的内存开销。所以我们以22 GB的量化模型大小为目标。

首先，我们需要将22gb转换成bits:

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

这样模型的最大极限为1.76e+11bits(b)。Llama 270b有7e+10个参数§要量化。我们的目标精度是bpw。

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

因此每个参数的平均精度为2.51bits 。

把它四舍五入到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个小时。

使用ExLlamaV2在GPU上运行Llama2 70b

ExLlamaV2还提供了运行混合精度量化模型的脚本。

chat.py脚本将把模型作为聊天机器人运行，并且可以提供交互。还可以使用test_inference.py简单地测试模型。

我们就是用这个脚本检查模型速度和内存消耗:

 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。在推理实验中，它恰好占用了24gb，也就是正好可以使用24G的GPU

为什么需要留出一些内存

内存中的模型实际占用22.15 GB，但推理本身也消耗额外的内存。例如，提示进行编码并将其存储在内存中。如果设置较大的最大序列长度或进行批处理解码，推理也会消耗更多内存。

所以如果使用24gb的GPU，可能会在推理过程中得到CUDA内存不足的错误，这是因为你的GPU会有一些来自操作系统的显存占用，比如（比如Ubuntu桌面消耗大约1.5 GB的VRAM，而Xfce则不到300M)。所以为了保证运行稳定，可以设置更低的bpw。例如2.4甚至2.3，这样给GPU更多的VRAM，可以保证运行稳定。