Update chatglm_llm.py (#242)

This commit is contained in:
imClumsyPanda 2023-05-05 15:14:17 +08:00 committed by GitHub
parent 51d0b7f855
commit e629526589
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions

View File

@ -144,12 +144,12 @@ class ChatGLM(LLM):
config=model_config, **kwargs)
if LLM_LORA_PATH and use_lora:
from peft import PeftModel
model_auto = PeftModel.from_pretrained(model, LLM_LORA_PATH)
model = PeftModel.from_pretrained(model, LLM_LORA_PATH)
# 可传入device_map自定义每张卡的部署情况
if device_map is None:
device_map = auto_configure_device_map(num_gpus)
self.model = dispatch_model(model_auto.half(), device_map=device_map)
self.model = dispatch_model(model.half(), device_map=device_map)
else:
self.model = self.model.float().to(llm_device)