Merge pull request THUDM#23 from qinjx/main

Multiple GPUs support. 多显卡支持

Author: duzx16

Diff for: README.md

@@ -165,6 +165,13 @@ cd ChatGLM2-6B
 git clone https://huggingface.co/THUDM/chatglm2-6b
 ```
 
+如果你从 Hugging Face Hub 上下载 checkpoint 的速度较慢，可以只下载模型实现
+```Shell
+GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/THUDM/chatglm2-6b
+```
+然后从[这里](https://cloud.tsinghua.edu.cn/d/674208019e314311ab5c/)手动下载模型参数文件，并将下载的文件替换到本地的 `chatglm2-6b` 目录下。
+
+
 将模型下载到本地之后，将以上代码中的 `THUDM/chatglm2-6b` 替换为你本地的 `chatglm2-6b` 文件夹的路径，即可从本地加载模型。
 
 模型的实现仍然处在变动中。如果希望固定使用的模型实现以保证兼容性，可以在 `from_pretrained` 的调用中增加 `revision="v1.0"` 参数。`v1.0` 是当前最新的版本号，完整的版本列表参见 [Change Log](https://huggingface.co/THUDM/chatglm2-6b#change-log)。

Diff for: api.py

@@ -52,5 +52,9 @@ async def create_item(request: Request):
 if __name__ == '__main__':
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
     model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+    # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
+    # model_path = "THUDM/chatglm2-6b"
+    # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    # model = load_model_on_gpus(model_path, num_gpus=2)
     model.eval()
     uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)

Diff for: cli_demo.py

@@ -3,9 +3,14 @@
 import signal
 from transformers import AutoTokenizer, AutoModel
 import readline
+from utils import load_model_on_gpus
 
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
 model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+# 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
+# model_path = "THUDM/chatglm2-6b"
+# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+# model = load_model_on_gpus(model_path, num_gpus=2)
 model = model.eval()
 
 os_name = platform.system()
@@ -17,7 +22,7 @@ def build_prompt(history):
     prompt = "欢迎使用 ChatGLM2-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序"
     for query, response in history:
         prompt += f"\n\n用户：{query}"
-        prompt += f"\n\nChatGLM-6B：{response}"
+        prompt += f"\n\nChatGLM2-6B：{response}"
     return prompt
 
 
@@ -39,21 +44,18 @@ def main():
             os.system(clear_command)
             print("欢迎使用 ChatGLM2-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序")
             continue
-        count = 0
+        print("\nChatGLM：", end="")
+        current_length = 0
         for response, history, past_key_values in model.stream_chat(tokenizer, query, history=history,
                                                                     past_key_values=past_key_values,
                                                                     return_past_key_values=True):
             if stop_stream:
                 stop_stream = False
                 break
             else:
-                count += 1
-                if count % 8 == 0:
-                    os.system(clear_command)
-                    print(build_prompt(history), flush=True)
-                    signal.signal(signal.SIGINT, signal_handler)
-        os.system(clear_command)
-        print(build_prompt(history), flush=True)
+                print(response[current_length:], end="", flush=True)
+                current_length = len(response)
+        print("")
 
 
 if __name__ == "__main__":

Diff for: openai_api.py

@@ -158,6 +158,10 @@ async def predict(query: str, history: List[List[str]], model_id: str):
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
     model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+    # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
+    # model_path = "THUDM/chatglm2-6b"
+    # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    # model = load_model_on_gpus(model_path, num_gpus=2)
     model.eval()
 
     uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)

Diff for: utils.py

@@ -0,0 +1,59 @@
+import os
+from typing import Dict, Tuple, Union, Optional
+
+from torch.nn import Module
+from transformers import AutoModel
+
+
+def auto_configure_device_map(num_gpus: int) -> Dict[str, int]:
+    # transformer.word_embeddings 占用1层
+    # transformer.final_layernorm 和 lm_head 占用1层
+    # transformer.layers 占用 28 层
+    # 总共30层分配到num_gpus张卡上
+    num_trans_layers = 28
+    per_gpu_layers = 30 / num_gpus
+
+    # 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都放到第一张卡上
+    # 本文件来源于https://github.com/THUDM/ChatGLM-6B/blob/main/utils.py
+    # 仅此处做少许修改以支持ChatGLM2
+    device_map = {
+        'transformer.embedding.word_embeddings': 0,
+        'transformer.encoder.final_layernorm': 0,
+        'transformer.output_layer': 0,
+        'transformer.rotary_pos_emb': 0,
+        'lm_head': 0
+    }
+
+    used = 2
+    gpu_target = 0
+    for i in range(num_trans_layers):
+        if used >= per_gpu_layers:
+            gpu_target += 1
+            used = 0
+        assert gpu_target < num_gpus
+        device_map[f'transformer.encoder.layers.{i}'] = gpu_target
+        used += 1
+
+    return device_map
+
+
+def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2,
+                       device_map: Optional[Dict[str, int]] = None, **kwargs) -> Module:
+    if num_gpus < 2 and device_map is None:
+        model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda()
+    else:
+        from accelerate import dispatch_model
+
+        model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half()
+
+        if device_map is None:
+            device_map = auto_configure_device_map(num_gpus)
+
+        model = dispatch_model(model, device_map=device_map)
+
+    return model

Diff for: web_demo.py

@@ -1,9 +1,14 @@
 from transformers import AutoModel, AutoTokenizer
 import gradio as gr
 import mdtex2html
+from utils import load_model_on_gpus
 
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
 model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+# 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
+# model_path = "THUDM/chatglm2-6b"
+# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+# model = load_model_on_gpus(model_path, num_gpus=2)
 model = model.eval()
 
 """Override Chatbot.postprocess"""

Diff for: web_demo2.py

@@ -14,6 +14,10 @@
 def get_model():
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
     model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+    # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
+    # model_path = "THUDM/chatglm2-6b"
+    # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    # model = load_model_on_gpus(model_path, num_gpus=2)
     model = model.eval()
     return tokenizer, model