服务器环境:
显卡驱动:Driver Version: 530.30.02
CUDA版本:CUDA Version: 12.1
显卡:NVIDIA GeForce RTX 3090共4张
注意:最好把显卡驱动升级到530,CUDA版本之前使用11.7有问题。
docker-compose.yml内容如下:
version: "3.2"
services:
chatglm2-6b:
container_name: chatglm2-6b
environment:
TZ: Asia/Shanghai
image: woshikid/chatglm2-6b:int4
ports:
- "8001:8000"
command:
[
"/bin/sh",
"-c",
"python openai_api.py",
]
restart: always
shm_size: "8g"
networks:
- general
deploy:
resources:
reservations:
devices:
- device_ids: ['2','3']
capabilities: ["gpu"]
driver: "nvidia"
networks:
general:
external: true
使用的镜像说明地址为:
Docker Hubhttps://hub.docker.com/r/woshikid/chatglm2-6b
用的是int4量化模型,实际对话过程中发现并没有使用显卡!
进入容器,修改openai_api.py文件,修改末尾的代码,将代码:
if __name__ == "__main__":
tokenizer = AutoTokenizer.from_pretrained("/chatglm2-6b-int4", trust_remote_code=True)
model = AutoModel.from_pretrained("/chatglm2-6b-int4", trust_remote_code=True).float()
# 多显卡支持,使用下面两行代替上面一行,将num_gpus改为你实际的显卡数量
#from utils import load_model_on_gpus
#model = load_model_on_gpus("/chatglm2-6b-int4", num_gpus=2)
model.eval()
改为:
if __name__ == "__main__":
tokenizer = AutoTokenizer.from_pretrained("/chatglm2-6b-int4", trust_remote_code=True)
#model = AutoModel.from_pretrained("/chatglm2-6b-int4", trust_remote_code=True).float()
# 多显卡支持,使用下面两行代替上面一行,将num_gpus改为你实际的显卡数量
from utils import load_model_on_gpus
model = load_model_on_gpus("/chatglm2-6b-int4", num_gpus=2)
model.eval()
运行时出现错误,RuntimeError: Library cudart is not initialized。
这是因为CUDA库文件损坏或者没安装导致的,这是最大的坑,执行下面代码安装:
apt update
apt install nvidia-cuda-toolkit
默认安装了cuda11.8版本,所以宿主的CUDA版本一定要>=cuda11.8,否则安装后仍然提示错误:RuntimeError: Unexpected error from cudaGetDeviceCou。
把chatglm2-6b接入到fastgpt,秒级反应,效果不错。
参考: