多GPU支持, 模型文件夹没有index.json会自动保存模型到multi_gpu_model_cache_dir以支持多GPU

api.py

@@ -1,6 +1,10 @@
+import datetime
+import json
+
+import uvicorn
 from fastapi import FastAPI, Request
-from transformers import AutoTokenizer, AutoModel
-import uvicorn, json, datetime
+
+from utils import load_mode_and_tokenizer
 
 app = FastAPI()
 
@@ -30,6 +34,4 @@ async def create_item(request: Request):
 if __name__ == '__main__':
     uvicorn.run('api:app', host='0.0.0.0', port=8000, workers=1)
 
-tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
-model.eval()
+model, tokenizer = load_mode_and_tokenizer("THUDM/chatglm-6b", num_gpus=1)

cli_demo.py

@@ -1,10 +1,9 @@
 import os
 import platform
-from transformers import AutoTokenizer, AutoModel
-from chatglm_parallel import load_model_on_gpus
 
-tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = load_model_on_gpus("THUDM/chatglm-6b", num_gpus=2)
+from utils import load_mode_and_tokenizer
+
+model, tokenizer = load_mode_and_tokenizer("THUDM/chatglm-6b", num_gpus=1)
 
 os_name = platform.system()
 clear_command = 'cls' if os_name == 'Windows' else 'clear'

utils.py

@@ -0,0 +1,66 @@
+import os
+from typing import Dict, Tuple, Union
+
+from accelerate import load_checkpoint_and_dispatch
+from transformers import AutoModel, AutoTokenizer
+
+
+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
+
+    device_map = {'transformer.word_embeddings': 0,
+                  'transformer.final_layernorm': num_gpus - 1, 'lm_head': num_gpus - 1}
+
+    used = 1
+    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.layers.{i}'] = gpu_target
+        used += 1
+
+    return device_map
+
+
+def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike],
+                       multi_gpu_model_cache_dir: Union[str, os.PathLike] = "./temp_model_dir",
+                       num_gpus: int = 2, **kwargs):
+    model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs)
+    model = model.eval()
+
+    device_map = auto_configure_device_map(num_gpus)
+    try:
+        model = load_checkpoint_and_dispatch(
+            model, checkpoint_path, device_map=device_map, offload_folder="offload", offload_state_dict=True).half()
+    except ValueError:
+        # index.json not found
+        print(f"index.json not found, auto fixing and saving model to {multi_gpu_model_cache_dir} ...")
+
+        assert multi_gpu_model_cache_dir is not None, "using auto fix, cache_dir must not be None"
+        model.save_pretrained(multi_gpu_model_cache_dir, max_shard_size='2GB')
+        model = load_checkpoint_and_dispatch(
+            model, multi_gpu_model_cache_dir, device_map=device_map,
+            offload_folder="offload", offload_state_dict=True).half()
+
+        print(f"loading model successfully, you should use checkpoint_path={multi_gpu_model_cache_dir} next time")
+
+    return model
+
+
+def load_mode_and_tokenizer(checkpoint_path: Union[str, os.PathLike],
+                            multi_gpu_model_cache_dir: Union[str, os.PathLike] = "./temp_model_dir",
+                            num_gpus: int = 1, **kwargs) -> Tuple[AutoModel, AutoTokenizer]:
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs)
+    if num_gpus < 2:
+        model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda()
+        model = model.eval()
+    else:
+        model = load_model_on_gpus(checkpoint_path, multi_gpu_model_cache_dir, num_gpus, **kwargs)
+    return model, tokenizer

web_demo.py

@@ -1,9 +1,7 @@
-from transformers import AutoTokenizer
 import gradio as gr
-from chatglm_parallel import load_model_on_gpus
+from utils import load_mode_and_tokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = load_model_on_gpus("THUDM/chatglm-6b", num_gpus=2)
+model, tokenizer = load_mode_and_tokenizer("THUDM/chatglm-6b", num_gpus=1)
 
 MAX_TURNS = 20
 MAX_BOXES = MAX_TURNS * 2

web_demo2.py

@@ -1,8 +1,7 @@
 from transformers import AutoModel, AutoTokenizer
 import streamlit as st
 from streamlit_chat import message
-from chatglm_parallel import load_model_on_gpus
-
+from utils import load_mode_and_tokenizer
 
 st.set_page_config(
     page_title="ChatGLM-6b 演示",
@@ -12,8 +11,7 @@
 
 @st.cache_resource
 def get_model():
-    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-    model = load_model_on_gpus("THUDM/chatglm-6b", num_gpus=2)
+    model, tokenizer = load_mode_and_tokenizer("THUDM/chatglm-6b", num_gpus=1)
     return tokenizer, model