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
This commit is contained in:
parent
de27ddd81a
commit
e0db44dcb6
1 changed files with 9 additions and 1 deletions
|
|
@ -505,7 +505,15 @@ class _BaseAutoModelClass:
|
||||||
|
|
||||||
if bigdl_lcmu_enabled:
|
if bigdl_lcmu_enabled:
|
||||||
with ContextManagers(init_contexts):
|
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)
|
model = model_class(config, *model_args, **kwargs)
|
||||||
else:
|
else:
|
||||||
model = model_class(config, *model_args, **kwargs)
|
model = model_class(config, *model_args, **kwargs)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue