From 5f4f8dcedc63e08cd7b66f9845736a4f21498319 Mon Sep 17 00:00:00 2001 From: hzg0601 Date: Fri, 21 Jul 2023 09:49:18 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3loader.py=E4=B8=AD=E4=BB=A5?= =?UTF-8?q?=E5=90=8D=E7=A7=B0=E5=8C=B9=E9=85=8D=E6=A8=A1=E5=9E=8B=E5=BC=8F?= =?UTF-8?q?=E6=9C=AA=E6=8E=92=E9=99=A4chatglm2=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- models/loader/loader.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/models/loader/loader.py b/models/loader/loader.py index ae4b408..986f5c4 100644 --- a/models/loader/loader.py +++ b/models/loader/loader.py @@ -148,7 +148,7 @@ class LoaderCheckPoint: trust_remote_code=True).half() # 可传入device_map自定义每张卡的部署情况 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) elif 'moss' in self.model_name.lower(): 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, max_memory=max_memory, 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) else: