diff --git a/README.md b/README.md index d9413d3..442b404 100644 --- a/README.md +++ b/README.md @@ -193,6 +193,6 @@ Web UI 可以实现如下功能: - [ ] 实现调用 API 的 Web UI Demo ## 项目交流群 -![二维码](img/qr_code_15.jpg) +![二维码](img/qr_code_16.jpg) 🎉 langchain-ChatGLM 项目交流群,如果你也对本项目感兴趣,欢迎加入群聊参与讨论交流。 diff --git a/img/qr_code_16.jpg b/img/qr_code_16.jpg new file mode 100644 index 0000000..febb78b Binary files /dev/null and b/img/qr_code_16.jpg differ diff --git a/models/chatglm_llm.py b/models/chatglm_llm.py index ec2c04e..49ce97c 100644 --- a/models/chatglm_llm.py +++ b/models/chatglm_llm.py @@ -11,7 +11,7 @@ DEVICE_ID = "0" if torch.cuda.is_available() else None DEVICE = f"{DEVICE_}:{DEVICE_ID}" if DEVICE_ID else DEVICE_ -def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: +def auto_configure_device_map(num_gpus: int, use_lora: bool) -> Dict[str, int]: # transformer.word_embeddings 占用1层 # transformer.final_layernorm 和 lm_head 占用1层 # transformer.layers 占用 28 层 @@ -19,14 +19,21 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: num_trans_layers = 28 per_gpu_layers = 30 / num_gpus + # bugfix: PEFT加载lora模型出现的层命名不同 + if LLM_LORA_PATH and use_lora: + layer_prefix = 'base_model.model.transformer' + else: + layer_prefix = 'transformer' + # bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError # windows下 model.device 会被设置成 transformer.word_embeddings.device # linux下 model.device 会被设置成 lm_head.device # 在调用chat或者stream_chat时,input_ids会被放到model.device上 # 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError # 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上 - device_map = {'transformer.word_embeddings': 0, - 'transformer.final_layernorm': 0, 'lm_head': 0} + device_map = {f'{layer_prefix}.word_embeddings': 0, + f'{layer_prefix}.final_layernorm': 0, 'lm_head': 0, + f'base_model.model.lm_head': 0, } used = 2 gpu_target = 0 @@ -35,7 +42,7 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: gpu_target += 1 used = 0 assert gpu_target < num_gpus - device_map[f'transformer.layers.{i}'] = gpu_target + device_map[f'{layer_prefix}.layers.{i}'] = gpu_target used += 1 return device_map @@ -141,16 +148,16 @@ class ChatGLM(LLM): else: from accelerate import dispatch_model - model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True, - config=model_config, **kwargs) + # model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True, + # config=model_config, **kwargs) if LLM_LORA_PATH and use_lora: from peft import PeftModel - model = PeftModel.from_pretrained(model, LLM_LORA_PATH) + model = PeftModel.from_pretrained(self.model, LLM_LORA_PATH) # 可传入device_map自定义每张卡的部署情况 if device_map is None: - device_map = auto_configure_device_map(num_gpus) + device_map = auto_configure_device_map(num_gpus, use_lora) - self.model = dispatch_model(model.half(), device_map=device_map) + self.model = dispatch_model(self.model.half(), device_map=device_map) else: self.model = self.model.float().to(llm_device)