ChatGlm-6B多张显卡的web_demo.py

from transformers import AutoModel, AutoTokenizer
import gradio as gr
import mdtex2html
from utils import load_model_on_gpus ＃要加的包
import os ＃要加的包

os.environ["CUDA_VISIBLE_DEVICES"]='0,1' #把显卡加到全局变量中
tokenizer = AutoTokenizer.from_pretrained("chatglm-6b", trust_remote_code=True)
model = load_model_on_gpus("chatglm-6b",num_gpus=2)＃使用两块显卡
# model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
model = model.eval()

在运行上面的代码需要保证你的显卡的驱动已经安装，需要验证以下代码反回为true

import torch
print(torch.cuda.is_available())
True

Print(torch.cuda.device_count())
2

以下是我运行ChatGlm的pip list

(glm) root@gpt:/home/gpt# pip list
Package            Version
------------------ ------------
accelerate         0.19.0
aiofiles           23.1.0
aiohttp            3.8.4
aiosignal          1.3.1
altair             4.2.2
anyio              3.6.2
async-timeout      4.0.2
attrs              23.1.0
certifi            2022.12.7
charset-normalizer 2.1.1
click              8.1.3
cmake              3.25.0
contourpy          1.0.7
cpm-kernels        1.0.11
cycler             0.11.0
entrypoints        0.4
fastapi            0.95.1
ffmpy              0.3.0
filelock           3.9.0
fonttools          4.39.3
frozenlist         1.3.3
fsspec             2023.5.0
gradio             3.28.3
gradio_client      0.2.1
h11                0.14.0
httpcore           0.17.0
httpx              0.24.0
huggingface-hub    0.14.1
idna               3.4
Jinja2             3.1.2
jsonschema         4.17.3
kiwisolver         1.4.4
latex2mathml       3.75.5
linkify-it-py      2.0.2
lit                15.0.7
Markdown           3.4.3
markdown-it-py     2.2.0
MarkupSafe         2.1.2
matplotlib         3.7.1
mdit-py-plugins    0.3.3
mdtex2html         1.2.0
mdurl              0.1.2
mpmath             1.2.1
multidict          6.0.4
networkx           3.0
numpy              1.24.1
orjson             3.8.12
packaging          23.1
pandas             2.0.1
Pillow             9.3.0
pip                23.0.1
protobuf           4.23.0
psutil             5.9.5
pydantic           1.10.7
pydub              0.25.1
Pygments           2.15.1
pyparsing          3.0.9
pyrsistent         0.19.3
python-dateutil    2.8.2
python-multipart   0.0.6
pytz               2023.3
PyYAML             6.0
regex              2023.5.5
requests           2.28.1
semantic-version   2.10.0
sentencepiece      0.1.99
setuptools         66.0.0
six                1.16.0
sniffio            1.3.0
starlette          0.26.1
sympy              1.11.1
tokenizers         0.13.3
toolz              0.12.0
torch              2.0.1+cu118
torchaudio         2.0.2+cu118
torchvision        0.15.2+cu118
tqdm               4.65.0
transformers       4.27.1
triton             2.0.0
typing_extensions  4.4.0
tzdata             2023.3
uc-micro-py        1.0.2
urllib3            1.26.13
uvicorn            0.22.0
websockets         11.0.3
wheel              0.38.4
yarl               1.9.2