Fix 70b oom (#9738)

* add default value to bigdl llm

* fix model oom
This commit is contained in:
Chen, Zhentao 2023-12-21 10:40:52 +08:00 committed by GitHub
parent 6c3e698bf1
commit b06a3146c8

View file

@ -45,9 +45,12 @@ class BigDLLM(AutoCausalLM):
keys = list(kwargs.keys()) keys = list(kwargs.keys())
for k in keys: for k in keys:
if k not in self.AutoCausalLM_ARGS: if k not in self.AutoCausalLM_ARGS:
self.bigdl_llm_kwargs[k] = kwargs[k] self.bigdl_llm_kwargs[k] = kwargs.pop(k)
kwargs.pop(k)
self.bigdl_llm_kwargs['use_cache'] = kwargs.get('use_cache', True)
self.bigdl_llm_kwargs['optimize_model'] = kwargs.get('optimize_model', True)
AutoModelForCausalLM.from_pretrained = partial(AutoModelForCausalLM.from_pretrained, **self.bigdl_llm_kwargs) AutoModelForCausalLM.from_pretrained = partial(AutoModelForCausalLM.from_pretrained, **self.bigdl_llm_kwargs)
kwargs['trust_remote_code'] = kwargs.get('trust_remote_code', True) kwargs['trust_remote_code'] = kwargs.get('trust_remote_code', True)
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)