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: