使用 Transformers 量化 Meta AI LLaMA2 中文版大模型

本篇文章聊聊如何使用 HuggingFace 的 Transformers 来量化 Meta AI 出品的 LLaMA2 大模型，让模型能够只使用 5GB 左右显存就能够运行。

写在前面

在前两篇文章《使用 Docker 快速上手官方版 LLaMA2 开源大模型》和《使用 Docker 快速上手中文版 LLaMA2 开源大模型》中，我们聊过了如何快速上手和使用新鲜出炉的 Meta AI LLaMA2 大模型。

经过实际测试，不论是原版模型（英文），还是中文版模型（双语），我们都需要 13～14 GB 显存，才能够将其运行起来。

为了能够让更多同学能够玩起来 LLaMA2 模型，我尝试使用HuggingFace 的 Transformers 对模型进行了量化，量化后的模型只需要 5GB 左右显存即可运行。

LLaMA2 Chat 项目

完整的代码和模型，我已经上传到了 GitHub 和 HuggingFace，感兴趣的同学可以自取。

准备工作

本文中所有的方法，你都可以参考并在非 Docker 容器中使用。

为了简单省事，可以参考前两篇文章，可以快速的搞定原版或者中文版的 LLaMA2 模型运行环境和 Docker 镜像。如果你本地环境完备，那么忽略 Docker 相关的命令，直接在 Bash 中执行各种具体的程序命令即可。

接下来，我们以使用 LLaMA2 中文模型镜像为例，进行模型的量化操作。

在前文中，我们使用下面的命令快速启动一个 LLaMA2 中文模型应用：

因为要对量化的模型进行保存，我们首先对上面的命令进行简单调整，添加一个参数：

这里我们多添加一个参数 -v `pwd`/soulteary:/app/soulteary ，来将当前执行命令的目录下的 soulteary 目录映射到容器内的 /app/soulteary，用来保存未来的量化后的模型文件。

执行命令后，我们将进入环境完备的 Docker 容器内部的交互式命令行环境中。

使用 Transformers 对 LLaMA2 进行量化

这里，我们只使用 HuggingFace 出品的 Transformers 就能够完成一切所需的工作，不需要引入其他的开源项目。

Transformers 量化模型的核心配置

Transformers 的量化功能实现是调用了bitsandbytes。想正确调用这个函数库进行量化，则需要在 AutoModelForCausalLM.from_pretrained 方法中完成 quantization_config 的参数配置。

在 Transformers 的 utils/quantization_config.py#L37 源代码中，我们能够直观的看到函数的运行方式和参数定义，最简单的 4BIT 量化的配置如下：

这里的 bnb_4bit_quant_type 之所以设置为 nf4，是因为在 HuggingFace 的 QLoRA 大模型量化实践中，使用 nf4 （NormalFloat）这种新的数据类型，能够在不牺牲性能的前提下，尽可能节省内存消耗。

而 bnb_4bit_compute_dtype 之所以设置为 torch.bfloat16 则是因为 HuggingFace 的另外一篇说明，我们可以使用这种新的数据格式，来减少传统 FP32 的“空间浪费” 和避免 FP32 转换 FP16 存在的潜在的溢出问题。

编写模型量化程序

综上所述，不难写出一段简单的不到三十行的程序，来完成对于 LLaMA2 模型的量化（相关程序，我上传到了 soulteary/docker-llama2-chat/llama2-7b-cn-4bit/quantization_4bit.py）：

对模型执行量化操作

我们将上面的内容保存为 quantization_4bit.py，放置于和 LLaMA2 模型目录 meta-llama 或 LinkSoul 同级的目录中，然后使用 python quantization_4bit.py 执行程序，即可开始模型的量化工作：

稍等片刻，就能够在当前程序目录中找到自动创建，并保存来新的模型对目录 soulteary/Chinese-Llama-2-7b-4bit/ 啦：

补齐模型运行文件

模型量化计算是结束了，但是此时的模型还不能使用，因为缺少了 tokenizer 相关的程序文件。类似 LLaMA2 官方版本和中文版本全部兼容，这里的量化版的模型和量化前的模型，也是全部兼容的。

解决这个问题非常简单，我们只需要将量化前的模型中的文件复制到新模型目录即可：

调整模型程序

前文中提到，这里量化的程序和原版程序没有使用上的区别，所以多数程序都可以保持原样。不过因为是新的模型文件，还是要进行几处简单的调整的。

更新模型运行程序

前文中提到，这里量化的程序和原版程序没有使用上的区别，所以多数程序都可以保持原样。为了能够让模型正确的通过 4BIT 方式加载和运行，我们需要调整两处内容：

我们需要调整前两篇文章中相关项目使用的 model.py 中的 model_id 变量，以及在 AutoModelForCausalLM.from_pretrained 调用中加上 load_in_4bit=True：

这部分完整的代码在 soulteary/docker-llama2-chat/llama2-7b-cn-4bit/model.py 可以找到。

运行模型应用

模型应用程序，我上传到了soulteary/docker-llama2-chat/llama2-7b-cn-4bit，因为和前两篇文章中没有什么区别，就不展开了。

如果你选择不在容器内运行，直接使用 python app.py ，模型程序就会快速的运行起来。

构建新的容器镜像

构建 4BIT 的镜像，和之前的文章中一样，执行脚本，等待镜像构建完成即可：

如果你之前跟着前两篇文章走过一遍，那么这个应该操作能够在 1～2 秒内完成。

使用容器启动模型应用

使用容器启动应用和之前的文章也并没有什么区别，执行命令，调用下面的脚本即可：

等待日志中出现 Running on local URL: http://0.0.0.0:7860，我们就能够正常进行使用和测试啦。

运行起来的量化后的中文 LLaMA2 项目

显存资源使用

显存资源一直是大家都比较关注的部分，模型启动大概需要 5G 出头的显存资源。

使用一段时间后，依旧还在 6GB 以内，是不是感觉还凑合？

最后

并不是所有的同学都是人手一张或几张 4090 或者 A100，所以即使量化会带来一些效果的下降，但总归比因为显存不足无法跑起来模型，不能一起玩要好呀。

况且，即使效果下降，依旧是适合做非常多场景下的使用的。后面的文章里，我们再做展开。

工程的艺术就在于 “trade-off”，前一阵线下偶然听到一位新朋友提起，有一种将心底藏了很久的东西唤醒的感觉。

--EOF

本文使用「署名 4.0 国际 (CC BY 4.0)」许可协议，欢迎转载、或重新修改使用，但需要注明来源。 署名 4.0 国际 (CC BY 4.0)

本文作者: 苏洋

创建时间: 2023年07月22日

统计字数: 8491字

阅读时间: 17分钟阅读

本文链接: https://soulteary.com/2023/07/22/quantizing-meta-ai-llama2-chinese-version-large-models-using-transformers.html