From e0db44dcb625b1dc1b0f4eb854eabe32b7642ce8 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 24 Jan 2024 13:20:46 +0800 Subject: [PATCH] fix unexpected keyword argument 'device' (#9982) * add device for chatglm3 only * add comment for this change * fix style * fix style * fix style again.. * finally fixed style --- python/llm/src/bigdl/llm/transformers/model.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/llm/src/bigdl/llm/transformers/model.py b/python/llm/src/bigdl/llm/transformers/model.py index 265a3300..93afdbe9 100644 --- a/python/llm/src/bigdl/llm/transformers/model.py +++ b/python/llm/src/bigdl/llm/transformers/model.py @@ -505,7 +505,15 @@ class _BaseAutoModelClass: if bigdl_lcmu_enabled: with ContextManagers(init_contexts): - kwargs["device"] = "meta" + if config.architectures is not None and config.architectures[0] in \ + ["ChatGLMModel", "ChatGLMForConditionalGeneration"]: + + """ + ChatGLMModel uses skip_init by default, which will force modules placed on cpu + if the device is not specified. This will further cause replaced linear + allocating memory on cpu. + """ + kwargs["device"] = "meta" model = model_class(config, *model_args, **kwargs) else: model = model_class(config, *model_args, **kwargs)