From 6720189f3fc59db985d404a40d4d014781c82678 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Sat, 2 Dec 2023 00:37:53 +0800 Subject: [PATCH] fix #1642 Former-commit-id: b69763ff9237a450194c0de6377745ba3086a4ce --- src/llmtuner/model/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llmtuner/model/utils.py b/src/llmtuner/model/utils.py index a9138e7e..42bef35b 100644 --- a/src/llmtuner/model/utils.py +++ b/src/llmtuner/model/utils.py @@ -25,7 +25,7 @@ def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel": if getattr(model, "quantization_method", None): # already set on current device return model - if torch.cuda.device_count() > 1: + if torch.cuda.device_count() > 1 and getattr(model.config, "model_type", None) != "chatglm": from accelerate import dispatch_model from accelerate.utils import infer_auto_device_map, get_balanced_memory