weiweiw
/
Langchain-Chatchat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add device_map config to support chatglm2-6b (#734) 

chatglm-6b和chatglm2-6b的参数命名不一致，本次提交旨在解决chatglm2-6b device_map 创建的问题。在chatglm_auto_configure_device_map 函数中新增了chatglm2-6b device_map 创建的相关代码。
This commit is contained in:
Jingsong-Yan 2023-06-30 21:08:38 +08:00 committed by GitHub
parent 51ed739d6c
commit 421ce3da3a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
models/loader/loader.py
View File

@ -257,10 +257,21 @@ class LoaderCheckPoint:
# 在调用chat或者stream_chat时,input_ids会被放到model.device上 # 在调用chat或者stream_chat时,input_ids会被放到model.device上
# 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError # 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError
# 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上 # 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上
encode = ""
if 'chatglm2' in self.model_name:
device_map = {
f"{layer_prefix}.embedding.word_embeddings": 0,
f"{layer_prefix}.rotary_pos_emb": 0,
f"{layer_prefix}.output_layer": 0,
f"{layer_prefix}.encoder.final_layernorm": 0,
f"base_model.model.output_layer": 0
}
encode = ".encoder"
else:
device_map = {f'{layer_prefix}.word_embeddings': 0, device_map = {f'{layer_prefix}.word_embeddings': 0,
f'{layer_prefix}.final_layernorm': 0, 'lm_head': 0, f'{layer_prefix}.final_layernorm': 0, 'lm_head': 0,
f'base_model.model.lm_head': 0, } f'base_model.model.lm_head': 0, }
used = 2 used = 2
gpu_target = 0 gpu_target = 0
for i in range(num_trans_layers): for i in range(num_trans_layers):
@ -268,7 +279,7 @@ class LoaderCheckPoint:
gpu_target += 1 gpu_target += 1
used = 0 used = 0
assert gpu_target < num_gpus assert gpu_target < num_gpus
device_map[f'{layer_prefix}.layers.{i}'] = gpu_target device_map[f'{layer_prefix}{encode}.layers.{i}'] = gpu_target
used += 1 used += 1
return device_map return device_map