Skip to content
项目
群组
代码片段
帮助
当前项目
正在载入...
登录 / 注册
切换导航面板
J
jinchat-server
概览
概览
详情
活动
周期分析
版本库
存储库
文件
提交
分支
标签
贡献者
分支图
比较
统计图
问题
0
议题
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
CI / CD
CI / CD
流水线
作业
日程表
图表
维基
Wiki
代码片段
代码片段
成员
成员
折叠边栏
关闭边栏
活动
图像
聊天
创建新问题
作业
提交
问题看板
Open sidebar
aigc-pioneer
jinchat-server
Commits
86c4147b
Unverified
提交
86c4147b
authored
4月 18, 2023
作者:
imClumsyPanda
提交者:
GitHub
4月 18, 2023
浏览文件
操作
浏览文件
下载
差异文件
Merge pull request #128 from Viscount/dev
1. 参考ChatGLM-6B代码实现模型多卡部署
上级
4df8bb92
edbe155c
隐藏空白字符变更
内嵌
并排
正在显示
1 个修改的文件
包含
55 行增加
和
8 行删除
+55
-8
chatglm_llm.py
models/chatglm_llm.py
+55
-8
没有找到文件。
models/chatglm_llm.py
浏览文件 @
86c4147b
...
...
@@ -5,6 +5,8 @@ from transformers import AutoTokenizer, AutoModel
import
torch
from
configs.model_config
import
LLM_DEVICE
from
typing
import
Dict
,
Tuple
,
Union
,
Optional
DEVICE
=
LLM_DEVICE
DEVICE_ID
=
"0"
if
torch
.
cuda
.
is_available
()
else
None
CUDA_DEVICE
=
f
"{DEVICE}:{DEVICE_ID}"
if
DEVICE_ID
else
DEVICE
...
...
@@ -17,6 +19,36 @@ def torch_gc():
torch
.
cuda
.
ipc_collect
()
def
auto_configure_device_map
(
num_gpus
:
int
)
->
Dict
[
str
,
int
]:
# transformer.word_embeddings 占用1层
# transformer.final_layernorm 和 lm_head 占用1层
# transformer.layers 占用 28 层
# 总共30层分配到num_gpus张卡上
num_trans_layers
=
28
per_gpu_layers
=
30
/
num_gpus
# bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError
# windows下 model.device 会被设置成 transformer.word_embeddings.device
# linux下 model.device 会被设置成 lm_head.device
# 在调用chat或者stream_chat时,input_ids会被放到model.device上
# 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError
# 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上
device_map
=
{
'transformer.word_embeddings'
:
0
,
'transformer.final_layernorm'
:
0
,
'lm_head'
:
0
}
used
=
2
gpu_target
=
0
for
i
in
range
(
num_trans_layers
):
if
used
>=
per_gpu_layers
:
gpu_target
+=
1
used
=
0
assert
gpu_target
<
num_gpus
device_map
[
f
'transformer.layers.{i}'
]
=
gpu_target
used
+=
1
return
device_map
class
ChatGLM
(
LLM
):
max_token
:
int
=
10000
temperature
:
float
=
0.01
...
...
@@ -51,19 +83,34 @@ class ChatGLM(LLM):
def
load_model
(
self
,
model_name_or_path
:
str
=
"THUDM/chatglm-6b"
,
llm_device
=
LLM_DEVICE
):
llm_device
=
LLM_DEVICE
,
device_map
:
Optional
[
Dict
[
str
,
int
]]
=
None
,
**
kwargs
):
self
.
tokenizer
=
AutoTokenizer
.
from_pretrained
(
model_name_or_path
,
trust_remote_code
=
True
)
if
torch
.
cuda
.
is_available
()
and
llm_device
.
lower
()
.
startswith
(
"cuda"
):
self
.
model
=
(
AutoModel
.
from_pretrained
(
model_name_or_path
,
trust_remote_code
=
True
)
.
half
()
.
cuda
()
)
# 根据当前设备GPU数量决定是否进行多卡部署
num_gpus
=
torch
.
cuda
.
device_count
()
if
num_gpus
<
2
and
device_map
is
None
:
self
.
model
=
(
AutoModel
.
from_pretrained
(
model_name_or_path
,
trust_remote_code
=
True
,
**
kwargs
)
.
half
()
.
cuda
()
)
else
:
from
accelerate
import
dispatch_model
model
=
AutoModel
.
from_pretrained
(
model_name_or_path
,
trust_remote_code
=
True
,
**
kwargs
)
.
half
()
# 可传入device_map自定义每张卡的部署情况
if
device_map
is
None
:
device_map
=
auto_configure_device_map
(
num_gpus
)
self
.
model
=
dispatch_model
(
model
,
device_map
=
device_map
)
else
:
self
.
model
=
(
AutoModel
.
from_pretrained
(
...
...
编写
预览
Markdown
格式
0%
重试
或
添加新文件
添加附件
取消
您添加了
0
人
到此讨论。请谨慎行事。
请先完成此评论的编辑!
取消
请
注册
或者
登录
后发表评论