From f53bc7d9a0fa6998113bb478376f34c2d9bb75a9 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Tue, 9 Jan 2024 14:49:13 +0800 Subject: [PATCH] fix #2127 Former-commit-id: ebee4f6a2a20a39ad7fbb87edfa63e244da7b5e6 --- src/llmtuner/model/utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/llmtuner/model/utils.py b/src/llmtuner/model/utils.py index 374808ee..824a1e79 100644 --- a/src/llmtuner/model/utils.py +++ b/src/llmtuner/model/utils.py @@ -27,14 +27,12 @@ def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel": if ( torch.cuda.device_count() > 1 and isinstance(model, PreTrainedModel) - and getattr(model.config, "model_type", None) != "chatglm" + and model._no_split_modules is not None + and model.config.model_type != "chatglm" ): from accelerate import dispatch_model from accelerate.utils import infer_auto_device_map, get_balanced_memory - if getattr(model, "_no_split_modules", None) is None: - raise ValueError("The model class needs to implement the `_no_split_modules` attribute.") - kwargs = {"dtype": model.dtype, "no_split_module_classes": model._get_no_split_modules("auto")} max_memory = get_balanced_memory(model, **kwargs) # Make sure tied weights are tied before creating the device map.