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

解决loader.py中以名称匹配模型式未排除chatglm2的问题

This commit is contained in:
hzg0601 2023-07-21 09:49:18 +08:00
parent a5294748a4
commit 5f4f8dcedc
1 changed files with 1 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
models/loader/loader.py
View File

@ -148,7 +148,7 @@ class LoaderCheckPoint:
trust_remote_code=True).half() trust_remote_code=True).half()
# 可传入device_map自定义每张卡的部署情况 # 可传入device_map自定义每张卡的部署情况
if self.device_map is None: if self.device_map is None:
if 'chatglm' in self.model_name.lower(): if 'chatglm' in self.model_name.lower() and not "chatglm2" in self.model_name.lower():
self.device_map = self.chatglm_auto_configure_device_map(num_gpus) self.device_map = self.chatglm_auto_configure_device_map(num_gpus)
elif 'moss' in self.model_name.lower(): elif 'moss' in self.model_name.lower():
self.device_map = self.moss_auto_configure_device_map(num_gpus, checkpoint) self.device_map = self.moss_auto_configure_device_map(num_gpus, checkpoint)
@ -164,13 +164,6 @@ class LoaderCheckPoint:
dtype=torch.float16 if not self.load_in_8bit else torch.int8, dtype=torch.float16 if not self.load_in_8bit else torch.int8,
max_memory=max_memory, max_memory=max_memory,
no_split_module_classes=model._no_split_modules) no_split_module_classes=model._no_split_modules)
# 对于chaglm和moss意外的模型应使用自动指定，而非调用chatglm的配置方式
# 其他模型定义的层类几乎不可能与chatglm和moss一致，使用chatglm_auto_configure_device_map
# 百分百会报错，使用infer_auto_device_map虽然可能导致负载不均衡，但至少不会报错
# 实测在bloom模型上如此
# self.device_map = infer_auto_device_map(model,
# dtype=torch.int8,
# no_split_module_classes=model._no_split_modules)
model = dispatch_model(model, device_map=self.device_map) model = dispatch_model(model, device_map=self.device_map)
else: else: