將大模型載入到記憶體

在 PyTorch 中載入預訓練模型時，通常的工作流程如下：

import torch

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

通俗地說，這些步驟是：

1. 使用隨機初始化的權重建立模型
2. 從磁碟載入模型權重（通常在一個稱為 state dict 的字典中）
3. 將這些權重載入到模型中

雖然這個工作流程對於常規大小的模型非常有效，但在處理巨大模型時，它有一些明顯的侷限性：在第 1 步，我們在 RAM 中載入一個完整版本的模型，並花一些時間隨機初始化權重（這些權重將在第 3 步被丟棄）。在第 2 步，我們載入另一個完整版本的模型到 RAM 中，其中包含預訓練的權重。如果你正在載入一個擁有 60 億引數的模型，這意味著你需要為每個模型副本準備 24GB 的 RAM，總共需要 48GB（其中一半用於以 FP16 格式載入模型）。

工作流程快速概覽

工作流程：程式碼實踐

例項化一個空模型

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)])

初始化一個引數超過 1000 億的空模型。在幕後，這依賴於 PyTorch 1.9 中引入的元裝置（meta device）。在上下文管理器下進行初始化時，每當建立一個引數，它就會被立即移動到該裝置上。

分片檢查點

可能你的模型太大，以至於單個副本都無法裝入 RAM。這並不意味著它不能被載入：如果你有一個或多個 GPU，就有更多的可用記憶體來儲存你的模型。在這種情況下，最好將你的檢查點分割成幾個較小的檔案，我們稱之為檢查點分片。

只要你遵循以下格式，Accelerate 就會處理分片檢查點：你的檢查點應該在一個資料夾中，包含幾個存有部分狀態字典的檔案，並且應該有一個 JSON 格式的索引，其中包含一個將引數名稱對映到其權重所在檔案的字典。你可以使用 save_model() 輕鬆地對模型進行分片。例如，我們可能有一個包含以下內容的資料夾：

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"` 的權重。

載入權重

Accelerate 提供的第二個工具是一個函式 load_checkpoint_and_dispatch()，它允許你將檢查點載入到空模型中。這支援完整的檢查點（單個檔案包含整個狀態字典）以及分片檢查點。它還會自動將這些權重分派到你可用的裝置上（GPU、CPU RAM），所以如果你載入的是分片檢查點，最大 RAM 使用量將是最大分片的大小。

如果你想在 Transformers 模型上使用大模型推理，請檢視此 文件。

下面是我們如何使用它來載入 GPT2-1.5B 模型。

讓我們下載這個模型的分片版本。

pip install huggingface_hub

from huggingface_hub import snapshot_download
checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)

為了初始化模型，我們將使用 minGPT 庫。

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)

然後，使用以下命令載入我們剛剛下載的檢查點：

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)

透過傳遞 `device_map="auto"`，我們告訴 Accelerate 根據可用資源自動確定模型各層的位置：

• 首先，我們使用 GPU 上的最大可用空間
• 如果仍需要空間，我們將剩餘的權重儲存在 CPU 上
• 如果 RAM 不足，我們將剩餘的權重作為記憶體對映張量儲存在硬碟上

no_split_module_classes

此引數將指示某些名為 `"Block"` 的模組不應跨不同裝置進行分割。你應該在此處設定所有包含某種殘差連線的塊。

The device_map

你可以透過訪問模型的 `hf_device_map` 屬性來檢視 Accelerate 選擇的 `device_map`：

model.hf_device_map

{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}

完全可以為你自己的層建立裝置對映，指定要使用的 GPU 裝置（一個數字）、`"cpu"` 或 `"disk"`，並將其傳入：

device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map
)

執行模型

現在我們已經完成了這些，我們的模型分佈在多個裝置上，甚至可能在硬碟上。但它仍然可以像一個常規的 PyTorch 模型一樣使用：

from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())

在幕後，Accelerate 為模型添加了鉤子，以便：

• 在每一層，輸入都被放置在正確的裝置上（所以即使你的模型分佈在多個 GPU 上，它也能正常工作）
• 對於解除安裝到 CPU 的權重，它們會在前向傳播之前被放置到 GPU 上，並在之後立即清理
• 對於解除安裝到硬碟的權重，它們會先載入到 RAM 中，然後在前向傳播之前被放置到 GPU 上，並在之後立即清理

這樣，即使你的模型無法完全裝入單個 GPU 或 CPU RAM，它仍然可以進行推理！

設計裝置對映

你可以讓 Accelerate 處理裝置對映的計算，方法是將 `device_map` 設定為支援的選項之一（`"auto"`、`"balanced"`、`"balanced_low_0"`、`"sequential"`），或者如果你想對每個層的放置有更多控制，可以自己建立一個。

當你的 GPU 記憶體不足以容納整個模型時，所有選項都會產生相同的結果（即將所有能放下的內容都放在 GPU 上，然後將權重解除安裝到 CPU，甚至在 RAM 不足時解除安裝到磁碟）。

當你可用的 GPU 記憶體超過模型大小時，以下是每個選項的區別：

• `"auto"` 和 `"balanced"` 會將模型均勻地分配到所有可用的 GPU 上，使你能夠使用大於 1 的批處理大小。
• `"balanced_low_0"` 會將模型均勻地分配到除第一個 GPU 外的所有 GPU 上，只將無法容納在其他 GPU 上的部分放在 GPU 0 上。當你需要使用 GPU 0 對輸出進行某些處理時（例如使用 Transformers 模型的 `generate` 函式），這個選項非常好。
• `"sequential"` 會先將能容納的部分放在 GPU 0 上，然後移到 GPU 1，以此類推（所以如果不需要，就不會使用最後的 GPU）。

首先請注意，你可以透過使用 `max_memory` 引數（在 infer_auto_device_map() 和所有使用它的函式中都可用）來限制每個 GPU 上使用的記憶體。設定 `max_memory` 時，你應該傳入一個包含 GPU 識別符號（例如 `0`、`1` 等）和 `"cpu"` 鍵的字典，用於表示你希望用於 CPU 解除安裝的最大 RAM。值可以是整數（以位元組為單位）或表示帶有單位的數字的字串，例如 `"10GiB"` 或 `"10GB"`。

下面是一個例子，我們不希望在兩個 GPU 上分別使用超過 10GiB，並且模型權重使用的 CPU RAM 不超過 30GiB：

from accelerate import infer_auto_device_map

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

此外，如果你對輸出進行額外操作而沒有將它們放回 CPU（例如在 Transformers 的 `generate` 方法內），並且你將輸入放在了 GPU 上，那麼該 GPU 將比其他 GPU 消耗更多記憶體（Accelerate 總是將輸出放回輸入的裝置）。因此，如果你想最佳化最大批處理大小並且有多個 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%。

如果你選擇完全自己設計 `device_map`，它應該是一個字典，鍵是模型的模組名稱，值是有效的裝置識別符號（例如，GPU 的整數）或用於 CPU 解除安裝的 `"cpu"`、用於磁碟解除安裝的 `"disk"`。鍵需要覆蓋整個模型，然後你可以隨心所欲地定義你的裝置對映：例如，如果你的模型有兩個塊（假設是 `block1` 和 `block2`），每個塊包含三個線性層（假設是 `linear1`、`linear2` 和 `linear3`），一個有效的裝置對映可以是：

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}

僅 CPU 解除安裝

如果你想將模型解除安裝到 CPU，可以使用 cpu_offload()。這樣，模型的所有引數都將被解除安裝，並且只保留模型狀態字典的一個副本。在前向傳播期間，引數將從該狀態字典中提取，並根據需要放置到執行裝置上，然後再次解除安裝。

cpu_offload(model, execution_device)

你也可以使用 cpu_offload_with_hook()。這個函式會將模型解除安裝到 CPU，並在執行時將其放回執行裝置。與 cpu_offload() 的區別在於，模型在前向傳播後會保留在執行裝置上，只有在呼叫返回的 `hook` 的 `offload` 方法時才會再次解除安裝。此外，cpu_offload_with_hook() 效能更高但節省的記憶體較少。它對於在迴圈中執行模型的流水線很有用。

model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

僅磁碟解除安裝

要執行磁碟解除安裝，你可以使用 disk_offload()。這樣，模型的所有引數都將被解除安裝為給定資料夾中的記憶體對映陣列。在前向傳播期間，將從該資料夾訪問引數，並根據需要將其放置在執行裝置上，然後再次解除安裝。

disk_offload(model, offload_dir, execution_device)

限制與未來發展

我們意識到當前 API 的侷限性：

• infer_auto_device_map()（或在 load_checkpoint_and_dispatch() 中的 `device_map="auto"`）會嘗試最大化執行時可用的 GPU 和 CPU RAM。雖然 PyTorch 在高效管理 GPU RAM 方面非常出色（並在不需要時釋放它），但對於 Python 和 CPU RAM 來說並非完全如此。因此，自動計算的裝置對映可能會對 CPU 造成過大壓力。如果因 RAM 不足而崩潰，請將一些模組移動到磁碟裝置。
• infer_auto_device_map()（或在 load_checkpoint_and_dispatch() 中的 `device_map="auto"`）按順序分配裝置（以避免來回移動），所以如果你的第一層比你擁有的 GPU 大小還大，最終所有內容都會放在 CPU/磁碟上。
• load_checkpoint_and_dispatch() 和 load_checkpoint_in_model() 目前不會對你的狀態字典與模型的正確性進行任何檢查（這將在未來版本中修復），所以如果嘗試載入的檢查點鍵不匹配或缺失，你可能會遇到一些奇怪的錯誤。
• 當你的模型分佈在多個 GPU 上時，所使用的模型並行是簡單且未經最佳化的，這意味著在任何給定時間只有一個 GPU 在工作，而其他 GPU 處於空閒狀態。
• 當權重被解除安裝到 CPU/硬碟時，沒有預取功能（我們將在未來版本中對此進行改進），這意味著權重是在需要時才被放到 GPU 上，而不是提前。
• 如果你的硬體在磁碟和 CPU 之間的通訊速度不快（例如沒有 NVMe），硬碟解除安裝可能會非常慢。