【ChatGLM2-6B】我在3090显卡上部署了清华大学开源中英双语对话模型
站长
· 阅读数 3
前言
盼望已久的国内开源项目终于来了，这次可算是对于现在仍然没有ChatGtp人员的福音。在这里我将着重记录如何部署清华大学开源的ChatGLM2-6B项目。
性能简介
ChatGLM2-6B 是开源中英双语对话模型 ChatGLM-6B 的第二代版本，在保留了初代模型对话流畅、部署门槛较低等众多优秀特性的基础之上，ChatGLM2-6B 引入了如下新特性：
- 更强大的性能：在同尺寸开源模型中具有较强的竞争力
- 更长的上下文：允许更多轮次的对话
- 更高效的推理：INT4 量化下，6G 显存支持的对话长度由 1K 提升到了 8K。
- 更开放的协议：ChatGLM2-6B 权重对学术研究完全开放，在获得官方的书面许可后，亦允许商业使用。如果您发现我们的开源模型对您的业务有用，我们欢迎您对下一代模型 ChatGLM3 研发的捐赠。
硬件选择
在README中的性能一栏中我们发现，作者团队是使用的硬件为A100-SXM4-80G 这对于我们个人来说是很难实现的。如采用CPU推理的话，这将不现实，我们无法忍受如此漫长的推理等待时间。那么GPU则是我们必须面对需要解决的问题了。根据清华团队在readme中所述：
使用官方实现，batch size = 1，max length = 2048，bf16 精度，测试硬件为 A100-SXM4-80G，软件环境为 PyTorch 2.0.1
我们需要满足Pytorch 2.0 在pytorch官网上查阅如下可知我们需要的CUDA为11.7 或 11.8满足算力以及CUDA的版本要求可知常用显卡中3090符合要求。
| 量化等级 | 编码 2048 长度的最小显存 | 生成 8192 长度的最小显存 |
|---|---|---|
| FP16 / BF16 | 13.1 GB | 12.8 GB |
| INT8 | 8.2 GB | 8.1 GB |
| INT4 | 5.5 GB | 5.1 GB |
部署安装
首先需要下载本仓库：
git clone https://github.com/THUDM/ChatGLM2-6B
cd ChatGLM2-6B
第二步： 使用 pip 安装依赖：pip install -r requirements.txt，其中 transformers 库版本推荐为 4.30.2，torch 推荐使用 2.0 以上的版本，以获得最佳的推理性能。
第三步： 考虑到本地加载模型大多数人找不到相关文件夹带来的种种BUG，因此我在这里将不介绍本地加载。创建本地test_chat.py文件并运行等待系统下载模型文件（模型较大，请耐心等待）
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
model = model.eval()
response, history = model.chat(tokenizer, "你好", history=[])
print(response)
使用尝鲜
等待了半天，激动人心的时刻即将来临！直接运行程序进行提问：
python cli_demo.py
