Merge branch 'master' into dev

README.md

@@ -193,6 +193,6 @@ Web UI 可以实现如下功能：
   - [ ] 实现调用 API 的 Web UI Demo
 
 ## 项目交流群
-![二维码](img/qr_code_15.jpg)
+![二维码](img/qr_code_16.jpg)
 
 🎉 langchain-ChatGLM 项目交流群，如果你也对本项目感兴趣，欢迎加入群聊参与讨论交流。

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)