pytorch在有限的資源下部署大語言模型（以ChatGLM-6B爲例）

Part1知識準備

在PyTorch中加載預訓練的模型時，通常的工作流程是這樣的：

my_model = ModelClass(...)
state_dict =
torch.load(checkpoint_file)

用簡單的話來說，這些步驟是：

用隨機初始化的權重創建模型。
從磁盤上加載模型權重（在一個通常被稱爲狀態字典的字典中）。
在模型中加載這些權重。

雖然這對常規大小的模型來說非常有效，但當我們處理一個巨大的模型時，這個工作流程有一些明顯的侷限性：在第1步，我們在RAM中加載一個完整版本的模型，並花一些時間隨機初始化權重（這將在第3步被丟棄）。在第2步，我們在RAM中加載另一個完整版本的模型，並使用預訓練的權重。如果你正在加載一個具有60億個參數的模型，這意味着你需要爲每個模型的副本提供24GB的RAM，所以總共需要48GB（其中一半用於在FP16中加載模型）。

1使用accelerate

上下文管理器

引入accelerate處理大模型的第一個工具是上下文管理器init_empty_weights()，它可以幫助你在不使用任何RAM的情況下初始化一個模型，這樣，步驟1就可以可以在任何尺寸的模型上進行。以下是它的工作原理：

from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)

例如：

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

初始化一個空的模型，參數略多於100B。這有賴於PyTorch 1.9中引入的元設備(meta device)。在上下文管理器下的初始化過程中，每次創建一個參數時，它都會移動到該設備上。

分佈式檢查點

你的模型有可能大到即使是一個副本也無法裝入RAM。這並不意味着它不能被加載：如果你有一個或幾個GPU，這將有更多的內存可用於存儲你的模型。在這種情況下，如果你的檢查點被分割成幾個較小的文件，我們稱之爲檢查點碎片，效果會更好。

accelerate將處理分片檢查點，只要你遵循以下格式：你的檢查點應該在一個文件夾中，有幾個文件包含部分狀態字典，應該有一個JSON格式的索引，包含一個字典將參數名稱映射到包含其權重的文件。例如，我們可以有一個包含以下內容的文件夾：

first_state_dict.bin
index.json
second_state_dict.bin

與index.json是以下文件：

{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}

first_state_dict.bin包含 "linear1.weight "和 "linear1.bias "的權重。second_state_dict.bin是 "linear2.weight "和 "linear2.bias "的權重。

加載權重

第二個工具是引入了一個函數load_checkpoint_and_dispatch()，它將允許你在你的空模型中加載一個檢查點。這支持完整的檢查點（一個單個文件包含整個狀態描述）以及分片檢查點。它還會在你可用的設備（GPU、CPURAM）上自動分配這些權重，所以如果你正在加載一個分片檢查點，最大的RAM使用量將是最大分片的大小。

from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

checkpoint = "EleutherAI/gpt-j-6B"
config = AutoConfig.from_pretrained(checkpoint)

with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

請注意，在transformer中用from_config加載模型並不綁定權重，這在加載不包含綁定權重的重複鍵的檢查點時可能導致問題。所以你應該在加載檢查點之前綁定權重。

model.tie_weights()

然後加載我們剛剛下載的檢查點：

model = load_checkpoint_and_dispatch(
    model, "sharded-gpt-j-6B", device_map="auto", no_split_module_classes=["GPTJBlock"]
)

通過傳遞device_map="auto"，根據可用的資源，我們告訴模型的每一層放置在哪裏。

首先，我們使用GPU上的最大可用空間。
如果我們仍然需要空間，我們將剩餘的權重存儲在CPU上。
如果沒有足夠的RAM，我們將剩餘的權重作爲內存映射的張量存儲在硬盤上。

no_split_module_classes=["GPTJBlock"] 表示屬於GPTJBlock的模塊不應該在不同的設備上被分割。你應該在這裏設置所有包括某種residual(殘差連接)的塊。

你可以通過hf_device_map來查看accelearte挑選的設備圖。

model.hf_device_map

{'transformer.wte': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 'transformer.h.1': 0,
 'transformer.h.2': 0,
 'transformer.h.3': 0,
 'transformer.h.4': 0,
 'transformer.h.5': 0,
 'transformer.h.6': 0,
 'transformer.h.7': 0,
 'transformer.h.8': 0,
 'transformer.h.9': 0,
 'transformer.h.10': 0,
 'transformer.h.11': 0,
 'transformer.h.12': 0,
 'transformer.h.13': 0,
 'transformer.h.14': 0,
 'transformer.h.15': 0,
 'transformer.h.16': 0,
 'transformer.h.17': 0,
 'transformer.h.18': 0,
 'transformer.h.19': 0,
 'transformer.h.20': 0,
 'transformer.h.21': 0,
 'transformer.h.22': 0,
 'transformer.h.23': 0,
 'transformer.h.24': 1,
 'transformer.h.25': 1,
 'transformer.h.26': 1,
 'transformer.h.27': 1,
 'transformer.ln_f': 1,
 'lm_head': 1}

如果你喜歡明確地決定每層的位置，你也可以自己設計你的設備圖。在這種情況下，上面的命令變成了：

model = load_checkpoint_and_dispatch(model, "sharded-gpt-j-6B", device_map=my_device_map)

運行模型

現在我們已經做到了這一點，我們的模型位於幾個設備之間，也許還有硬盤。但它仍然可以作爲一個普通的PyTorch模型使用：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer("Hello, my name is", return_tensors="pt")
inputs = inputs.to(0)
output = model.generate(inputs["input_ids"])
tokenizer.decode(output[0].tolist())

在幕後，accelerate爲模型添加了鉤子，因此：

在每一層，輸入被放在正確的設備上（因此，即使你的模型分散在幾個GPU上，它也能工作）。
對於卸載在CPU上的權重，就在向前傳遞之前，它們被放在GPU上，並在之後被清理掉。
對於卸載在硬盤上的權重，它們被加載在RAM中，然後在向前傳遞之前被放在GPU上，並在之後被清理掉。這樣，即使你的模型不適合在某個GPU或CPU RAM上運行，你也可以運行推理！

設計一個設備圖

你可以通過以下選項"auto", "balanced", "balanced_low_0", "sequential"讓acclerate處理設備圖的計算，或自己創建一個。如果你想更多地控制每個層應該去哪裏，你可以在一個元設備上的模型上推導出模型的所有尺寸（從而計算出一個設備圖）。

當你沒有足夠的GPU內存來容納整個模型時，所有的選項都會產生相同的結果（也就是把所有能裝的東西都裝到GPU上，然後把重量卸到CPU上，如果沒有足夠的內存，甚至卸到磁盤上）。

當你有比模型大小更多的GPU內存可用時，這裏是每個選項之間的區別：

"auto"和"balanced"在所有可用的GPU上平均分配模型，使你有可能使用大於1的批次大小。
"balanced_low_0 "將模型均勻地分割到所有的GPU上，除了第一個GPU之外，並且只將不適合其他GPU的部分放在GPU 0上。當你需要使用GPU 0對輸出進行一些處理時，這個選項是非常好的，比如使用transformers的生成函數時。
"順序 "將在GPU 0上安裝它可以安裝的東西，然後在GPU 1上移動，以此類推（所以如果不需要，就不會使用最後的GPU）。

首先注意，你可以通過使用max_memory參數（在fer_auto_device_map()和所有使用該參的函數中可用）限制每個GPU上使用的內存。當設置max_memory時，你應該傳遞一個包含GPU標識符（例如0、1等）和 "cpu "鍵的字典，用於你希望用於CPU卸載的最大RAM。這些值可以是一個整數（以字節爲單位），也可以是一個代表數字及其單位的字符串，例如 "10GiB "或 "10GB"。

這裏有一個例子，我們不希望在兩個GPU上各使用超過10GiB，而在模型權重上不超過30GiB的CPU內存：

from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})

當PyTorch發生首次分配時，它會加載CUDA內核，根據GPU的情況，它需要大約1-2GB的內存。因此，你的可用內存總是少於GPU的實際大小。要查看實際使用了多少內存，請執行torch.ones(1).cuda()並查看內存使用情況。因此，當你用max_memory創建內存映射時，確保相應地調整可用的內存，以避免出先OOM。

此外，如果你對你的輸出做一些額外的操作而不把它們放回CPU（例如在transformer的生成方法裏面），如果你把你的輸入放在一個GPU上，這個GPU將比其他的消耗更多的內存（加速器總是把輸出放回輸入的設備）。因此，如果你想優化最大的批處理量，並且你有很多GPU，給第一個GPU較少的內存。例如在8x80 A100設置上使用BLOOM-176B，接近理想的映射是：

max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}

你可以看到，我們給其餘7個GPU的內存比GPU 0多了50%。

如果你選擇自己完全設計設備映射，它應該是一個字典，鍵是你的模型的模塊名稱，值是一個有效的設備標識符（例如GPU是一個整數）或CPU卸載的 "cpu"，磁盤卸載的 "disc"。鍵需要覆蓋整個模型，然後你可以按照你的意願定義你的設備映射：例如，如果你的模型有兩個塊（比方說block1和block2），它們各自包含三個線性層（比方說線性1、線性2和線性3），一個有效的設備映射可以是：

device_map = {"block1": 0, "block2": 1}

另一個有效的可能是：

device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}

另一方面，這個是無效的，因爲它沒有涵蓋模型的每個參數：

device_map = {"block1"：0, "block2.linear1"：1, "block2.linear2"：1}

爲了達到最高的效率，請確保你的設備映射以連續的方式將參數放在GPU上（例如，不要將第一個權重放在GPU 0上，然後將權重放在GPU 1上，最後一個權重再放回GPU 0），以避免在GPU之間進行多次數據傳輸。

限制和進一步發展

我們知道目前API的侷限性：

雖然理論上這隻可以在一個CPU上工作，並有潛在的磁盤卸載，但你至少需要一個GPU來運行這個API。這將在進一步的開發中得到解決。
infer_auto_device_map() (或load_checkpoint_and_dispatch()中的 device_map="auto")試圖在你執行它的時候最大化它所看到的GPU和CPU RAM。雖然PyTorch在有效地管理GPU RAM方面非常出色（當不需要時就會歸還），但對於Python和CPU RAM來說，這並不完全正確。因此，自動計算的設備圖可能對CPU來說過於緊張。如果你因內存不足而出現崩潰，請將一些模塊移到磁盤設備上。
infer_auto_device_map()（或者load_checkpoint_and_dispatch()中的device_map="auto"）是按順序屬性設備的（以避免來回移動東西），所以如果你的第一層比你的GPU的大小大，最後會把所有東西都放在CPU/磁盤上。
load_checkpoint_and_dispatch()和load_checkpoint_in_model()目前沒有對你的狀態描述與你的模型相比的正確性進行任何檢查（這將在未來的版本中被修復），所以如果你試圖加載一個鍵不匹配或丟失的檢查點，你可能會得到一些奇怪的錯誤。
當你的模型被分割到幾個GPU上時，所使用的模型並行性是天真的，沒有經過優化，這意味着在某個時間只有一個GPU在工作，而另一個則處於閒置狀態。
當權重被卸載在CPU/硬盤上時，沒有預取（還沒有，我們會在未來的版本中努力做到這一點），這意味着權重在需要時被放到GPU上，而不是之前。
如果你運行的硬件沒有磁盤和CPU之間的快速通信（如NVM），硬盤卸載可能會非常慢.

Part2部署ChatGLM-6B

基礎環境：

torch==2.0.0+cu118
transformers==4.28.1
accelerate==0.18.0
Tesla T4 15.3G
內存：11.8G

下載相關文件：

git clone https://github.com/THUDM/ChatGLM-6B
cd ChatGLM-6B

git clone --depth=1 https://huggingface.co/THUDM/chatglm-6b THUDM/chatglm-6b
git clone --depth=1 https://huggingface.co/THUDM/chatglm-6b-int4 THUDM/chatglm-6b-int4

pip install -r requirements.txt
pip install gradio
pip install accelerate

正常情況下，我們使用Chat-GLM需要的顯存大於13G，內存沒有評估過，但上述的肯定是不夠的，16G應該可以。

2第一種方案

直接使用量化以後的模型：

from accelerate import infer_auto_device_map, init_empty_weights, load_checkpoint_and_dispatch
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch
import time

tokenizer = AutoTokenizer.from_pretrained("./THUDM/chatglm-6b-int4", trust_remote_code=True)
model = AutoModel.from_pretrained("./THUDM/chatglm-6b-int4", trust_remote_code=True).half().cuda()

model = model.eval()

def predict(input, history=None):
    print(f'predict started: {time.time()}');
    if history is None:
        history = []
    response, history = model.chat(tokenizer, input, history)
    return response, history

while True:
  text = input(">>用戶：")
  response, history = model.chat(tokenizer, input, history)
  print(">>CHatGLM：", response)

GPU使用4.9G，內存使用5.5G。

3第二種方案

使用acclerate，只有一塊GPU。

%cd /content/ChatGLM-6B

from accelerate import infer_auto_device_map, init_empty_weights, load_checkpoint_and_dispatch
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch
import time


tokenizer = AutoTokenizer.from_pretrained("./THUDM/chatglm-6b", trust_remote_code=True)
config = AutoConfig.from_pretrained("./THUDM/chatglm-6b", trust_remote_code=True)
with init_empty_weights():
  model = AutoModel.from_config(config, trust_remote_code=True)

for name, _ in model.named_parameters():
  print(name)
# device_map = infer_auto_device_map(model, no_split_module_classes=["GLMBlock"])
# print(device_map)
device_map = {'transformer.word_embeddings': 0, 'transformer.layers.0': 0, 'transformer.layers.1': 0, 'transformer.layers.2': 0, 'transformer.layers.3': 0, 'transformer.layers.4': 0, 'transformer.layers.5': 0, 'transformer.layers.6': 0, 'transformer.layers.7': 0, 'transformer.layers.8': 0, 'transformer.layers.9': 0, 'transformer.layers.10': 0, 'transformer.layers.11': 0, 'transformer.layers.12': 0, 'transformer.layers.13': 0, 'transformer.layers.14': 0, 'transformer.layers.15': 0, 'transformer.layers.16': 0, 'transformer.layers.17': 0, 'transformer.layers.18': 0, 'transformer.layers.19': 0, 'transformer.layers.20': 0, 'transformer.layers.21': 'cpu', 'transformer.layers.22': 'cpu', 'transformer.layers.23': 'cpu', 'transformer.layers.24': 'cpu', 'transformer.layers.25': 'cpu', 'transformer.layers.26': 'cpu', 'transformer.layers.27': 'cpu', 'transformer.final_layernorm': 'cpu', 'lm_head': 'cpu'}
model = load_checkpoint_and_dispatch(model, "./THUDM/chatglm-6b", device_map=device_map, offload_folder="offload", offload_state_dict=True, no_split_module_classes=["GLMBlock"]).half()

def predict(input, history=None):
    print(f'predict started: {time.time()}');
    if history is None:
        history = []
    response, history = model.chat(tokenizer, input, history)
    return response, history

while True:
  history = None
  text = input(">>用戶：")
  response, history = model.chat(tokenizer, text, history)
  print(">>CHatGLM：", response)

GPU使用9.7G，內存使用5.9G。第一輪輸入你好後GPU使用11.2G。

4第三種方案

使用accelerate，多塊GPU。

環境：windwos下。GPU：4*4090 24G。內存：128G。python>=3.8，torch==2.0+117，transformers==4.28.1，acclerate==0.18.0。

import os
os.environ["cuda_visible_devices"] = "0,1"

from accelerate import infer_auto_device_map, init_empty_weights, load_checkpoint_and_dispatch
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
# import gradio as gr
# import torch
import time


tokenizer = AutoTokenizer.from_pretrained(".\\chatglm-6b\\", trust_remote_code=True)
config = AutoConfig.from_pretrained(".\\chatglm-6b\\", trust_remote_code=True)
with init_empty_weights():
  model = AutoModel.from_config(config, trust_remote_code=True)

for name, _ in model.named_parameters():
  print(name)
# device_map = infer_auto_device_map(model, no_split_module_classes=["GLMBlock"])
# print(device_map)
# device_map = {'transformer.word_embeddings': 0, 'transformer.layers.0': 0, 'transformer.layers.1': 0, 'transformer.layers.2': 0, 'transformer.layers.3': 0, 'transformer.layers.4': 0, 'transformer.layers.5': 0, 'transformer.layers.6': 0, 'transformer.layers.7': 0, 'transformer.layers.8': 0, 'transformer.layers.9': 0, 'transformer.layers.10': 0, 'transformer.layers.11': 0, 'transformer.layers.12': 0, 'transformer.layers.13': 0, 'transformer.layers.14': 0, 'transformer.layers.15': 0, 'transformer.layers.16': 0, 'transformer.layers.17': 0, 'transformer.layers.18': 0, 'transformer.layers.19': 0, 'transformer.layers.20': 0, 'transformer.layers.21': 'cpu', 'transformer.layers.22': 'cpu', 'transformer.layers.23': 'cpu', 'transformer.layers.24': 'cpu', 'transformer.layers.25': 'cpu', 'transformer.layers.26': 'cpu', 'transformer.layers.27': 'cpu', 'transformer.final_layernorm': 'cpu', 'lm_head': 'cpu'}
model = load_checkpoint_and_dispatch(model, ".\\chatglm-6b\\", device_map="balanced", offload_folder="offload", offload_state_dict=True, no_split_module_classes=["GLMBlock"]).half()

def predict(input, history=None):
    print(f'predict started: {time.time()}')
    if history is None:
        history = []
    response, history = model.chat(tokenizer, input, history)
    return response, history

while True:
  history = None
  text = input(">>用戶：")
  response, history = model.chat(tokenizer, text, history)
  print(">>CHatGLM：", response)

注意，這裏我們設置設備映射爲balanced，並只使用前兩塊GPU。顯卡佔用情況：

會發現平均分配了顯存，當然GPU 0分配得更多些、至此，關於如何進行大模型推理就全部完成了。

Part3參考

https://huggingface.co/docs/accelerate/usage_guides/big_modeling
https://github.com/THUDM/ChatGLM-6B/issues/69
https://github.com/THUDM/ChatGLM-6B/issues/200

pytorch在有限的資源下部署大語言模型（以ChatGLM-6B爲例）

Part1知識準備

1使用accelerate

上下文管理器

分佈式檢查點

加載權重

運行模型

設計一個設備圖

限制和進一步發展

Part2部署ChatGLM-6B

2第一種方案

3第二種方案

4第三種方案

Part3參考

怎麼讓英文大語言模型支持中文？（三）進行指令微調

怎麼讓英文大預言模型支持中文？（一）繼續預訓練

怎麼讓英文大預言模型支持中文？（一）構建自己的tokenization

anaconda和python之間的對應關係

【python】linux下安裝python的一般方法

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結