🤗 Accelerate 如何透過 PyTorch 執行超大型模型
載入並執行大型模型
Meta AI 和 BigScience 最近開源了超大型語言模型，這些模型無法適應大多數消費級硬體的記憶體（RAM 或 GPU）。在 Hugging Face，我們使命的一部分是使這些大型模型也易於訪問，因此我們開發了工具，即使您沒有超級計算機，也可以執行這些模型。本部落格文章中選擇的所有示例都在免費的 Colab 例項上執行（RAM 和磁碟空間有限），如果您有更多的磁碟空間，請不要猶豫選擇更大的檢查點。
以下是如何執行 OPT-6.7B
import torch
from transformers import pipeline
# This works on a base Colab instance.
# Pick a larger checkpoint if you have time to wait and enough disk space!
checkpoint = "facebook/opt-6.7b"
generator = pipeline("text-generation", model=checkpoint, device_map="auto", torch_dtype=torch.float16)
# Perform inference
generator("More and more large language models are opensourced so Hugging Face has")
我們稍後會解釋每個引數的作用，但首先我們只考慮 PyTorch 中傳統的模型載入流程：它通常包括
- 建立模型
- 在記憶體中載入其權重（通常稱為
state_dict的物件） - 將這些權重載入到建立的模型中
- 將模型移動到裝置上進行推理
儘管這在過去幾年中執行良好，但超大型模型使這種方法面臨挑戰。這裡選擇的模型有 67 億個引數。在預設精度下，這意味著僅第 1 步（建立模型）就需要大約 26.8GB 的 RAM（float32 中的 1 個引數佔用 4 位元組記憶體）。這甚至無法適應您在 Colab 上獲得的 RAM。
然後，第 2 步將在記憶體中載入模型的第二個副本（因此預設精度下會再佔用 26.8GB RAM）。如果您嘗試像這樣載入最大的模型，例如 BLOOM 或 OPT-176B（兩者都有 1760 億個引數），您將需要 1.4 TB 的 CPU RAM。這有點過分了！而所有這些只是為了在第 4 步將模型移動到一個（或多個）GPU 上。
顯然我們需要更智慧的東西。在這篇部落格文章中，我們將解釋 Accelerate 如何利用 PyTorch 功能載入和執行超大型模型的推理，即使它們不適合 RAM 或一個 GPU。簡而言之，它將上述過程更改為這樣
- 建立一個空的（例如，沒有權重）模型
- 決定每一層將去向何處（當有多個裝置可用時）
- 在記憶體中載入其部分權重
- 將這些權重載入到空模型中
- 將權重移動到裝置上進行推理
- 從第 3 步開始重複，載入下一個權重，直到所有權重載入完畢
建立空模型
PyTorch 1.9 引入了一種新的裝置，稱為 meta 裝置。這使我們能夠在不附加任何資料的情況下建立張量：meta 裝置上的張量只需要一個形狀。只要您在 meta 裝置上，就可以建立任意大的張量，而無需擔心 CPU（或 GPU）RAM。
例如，以下程式碼將在 Colab 上崩潰
import torch
large_tensor = torch.randn(100000, 100000)
因為這個大型張量需要 4 * 10**10 位元組（預設精度是 FP32，所以張量的每個元素佔用 4 位元組），即 40GB RAM。然而，在 meta 裝置上相同的操作可以正常工作
import torch
large_tensor = torch.randn(100000, 100000, device="meta")
如果您嘗試顯示此張量，PyTorch 將列印以下內容
tensor(..., device='meta', size=(100000, 100000))
正如我們之前所說，這個張量沒有關聯的資料，只有一個形狀。
您可以直接在 meta 裝置上例項化模型
large_model = torch.nn.Linear(100000, 100000, device="meta")
但是對於現有模型，這種語法將要求您重寫所有建模程式碼，以便每個子模組都接受並傳遞一個 device 關鍵字引數。由於這對於 Transformers 庫中的 150 個模型來說不切實際，我們開發了一個上下文管理器，可以為您例項化一個空模型。
以下是例項化 BLOOM 空版本的方法
from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM
config = AutoConfig.from_pretrained("bigscience/bloom")
with init_empty_weights():
model = AutoModelForCausalLM.from_config(config)
這適用於任何模型，但您會得到一個無法直接使用的 shell：某些操作已針對元裝置實現，但並非所有操作都已實現。例如，您可以使用上面定義的 large_model 進行輸入，但不能使用 BLOOM 模型。即使使用它，輸出也將是元裝置的張量，因此您將獲得結果的形狀，但僅此而已。
作為這項工作的進一步發展，PyTorch 團隊正在開發一個新的 FakeTensor 類，它有點像元裝置上的張量，但帶有裝置資訊（除了形狀和資料型別之外）
由於我們知道每個權重的形狀，因此我們可以知道一旦我們完全載入預訓練張量，它們將消耗多少記憶體。因此，我們可以決定如何將模型拆分到 CPU 和 GPU 上。
計算裝置對映
在開始載入預訓練權重之前，我們需要知道將它們放在何處。這樣，我們每次將權重放置在正確位置時，都可以釋放 CPU RAM。這可以透過元裝置上的空模型完成，因為我們只需要知道每個張量的形狀及其資料型別即可計算它將佔用多少記憶體空間。
Accelerate 提供了一個函式，可以從空模型自動確定裝置對映。它將嘗試最大化所有可用 GPU 的使用，然後是 CPU RAM，最後標記不適合磁碟解除安裝的權重。讓我們使用 OPT-13b 看看。
from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM
config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
model = AutoModelForCausalLM.from_config(config)
device_map = infer_auto_device_map(model)
這將返回一個將模組或權重對映到裝置的字典。例如，在具有一個 Titan RTX 的機器上，我們得到以下結果
{'model.decoder.embed_tokens': 0,
'model.decoder.embed_positions': 0,
'model.decoder.final_layer_norm': 0,
'model.decoder.layers.0': 0,
'model.decoder.layers.1': 0,
...
'model.decoder.layers.9': 0,
'model.decoder.layers.10.self_attn': 0,
'model.decoder.layers.10.activation_fn': 0,
'model.decoder.layers.10.self_attn_layer_norm': 0,
'model.decoder.layers.10.fc1': 'cpu',
'model.decoder.layers.10.fc2': 'cpu',
'model.decoder.layers.10.final_layer_norm': 'cpu',
'model.decoder.layers.11': 'cpu',
...
'model.decoder.layers.17': 'cpu',
'model.decoder.layers.18.self_attn': 'cpu',
'model.decoder.layers.18.activation_fn': 'cpu',
'model.decoder.layers.18.self_attn_layer_norm': 'cpu',
'model.decoder.layers.18.fc1': 'disk',
'model.decoder.layers.18.fc2': 'disk',
'model.decoder.layers.18.final_layer_norm': 'disk',
'model.decoder.layers.19': 'disk',
...
'model.decoder.layers.39': 'disk',
'lm_head': 'disk'}
Accelerate 評估出嵌入層和解碼器直到第 9 個塊都可以放在 GPU（裝置 0）上，然後第 10 個塊的一部分需要放在 CPU 上，以及直到第 17 層的所有後續權重。然後第 18 層在 CPU 和磁碟之間拆分，隨後的所有層都必須解除安裝到磁碟
實際使用此裝置對映以後將無法工作，因為構成此模型的層具有殘差連線（塊的輸入被新增到塊的輸出中），因此給定層的所有內容都應該在同一裝置上。我們可以透過傳遞不應使用 no_split_module_classes 關鍵字引數拆分的模組名稱列表來指示 Accelerate
device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"])
這將返回
'model.decoder.embed_tokens': 0,
'model.decoder.embed_positions': 0,
'model.decoder.final_layer_norm': 0,
'model.decoder.layers.0': 0,
'model.decoder.layers.1': 0,
...
'model.decoder.layers.9': 0,
'model.decoder.layers.10': 'cpu',
'model.decoder.layers.11': 'cpu',
...
'model.decoder.layers.17': 'cpu',
'model.decoder.layers.18': 'disk',
...
'model.decoder.layers.39': 'disk',
'lm_head': 'disk'}
現在，每個層始終在同一裝置上。
在 Transformers 中，當在 from_pretrained() 方法或 pipeline 中使用 device_map 時，這些要保留在同一裝置上的塊類會自動提供，因此您無需擔心它們。請注意，您有以下 device_map 選項（僅當您有多個 GPU 時才相關）
"auto"或"balanced"：Accelerate 將拆分權重，以便每個 GPU 都被同等使用；"balanced_low_0"：Accelerate 會將權重進行拆分，使得每個 GPU 被同等使用，但第一個 GPU 除外，它將嘗試使其權重儘可能少（當您希望在某個 GPU 上處理模型輸出時，例如使用generate函式時，這會很有用）；"sequential"：Accelerate 將按順序填充 GPU（因此最後的 GPU 可能完全不會被使用）。
您也可以傳入自己的 device_map，只要它遵循我們之前看到的格式（字典層/模組名稱到裝置）。
最後，請注意，您收到的 device_map 結果取決於所選的資料型別（因為不同型別的浮點數佔用不同的空間）。提供 dtype="float16" 將給出不同的結果
device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")
在這種精度下，我們可以將模型安裝到 GPU 上的第 21 層
{'model.decoder.embed_tokens': 0,
'model.decoder.embed_positions': 0,
'model.decoder.final_layer_norm': 0,
'model.decoder.layers.0': 0,
'model.decoder.layers.1': 0,
...
'model.decoder.layers.21': 0,
'model.decoder.layers.22': 'cpu',
...
'model.decoder.layers.37': 'cpu',
'model.decoder.layers.38': 'disk',
'model.decoder.layers.39': 'disk',
'lm_head': 'disk'}
現在我們知道每個權重應該去哪裡，我們可以逐步將預訓練權重載入到模型中。
分片 state dicts
傳統上，PyTorch 模型儲存在一個完整的檔案中，該檔案包含從引數名到權重的對映。這個對映通常被稱為 state_dict。以下是 PyTorch 文件中關於儲存和載入的摘錄：
# Save the model weights
torch.save(my_model.state_dict(), 'model_weights.pth')
# Reload them
new_model = ModelClass()
new_model.load_state_dict(torch.load('model_weights.pth'))
這對於引數少於 10 億的模型來說效果很好，但對於更大的模型來說，這在 RAM 中非常耗費資源。BLOOM 模型有 1760 億個引數；即使將權重儲存為 bfloat16 以節省空間，整體仍表示為 352GB。雖然訓練該模型的超級計算機可能擁有這麼多記憶體，但要求推理時也擁有這麼多記憶體是不現實的。
這就是為什麼 Hugging Face Hub 上的大型模型不是以包含所有權重的一個大檔案儲存和共享，而是以多個檔案儲存和共享。例如，如果您訪問 BLOOM 模型頁面，您會看到有 72 個名為 pytorch_model_xxxxx-of-00072.bin 的檔案，每個檔案都包含模型權重的一部分。使用這種格式，我們可以將 state dict 的一部分載入到記憶體中，將權重放入模型中，將它們移動到正確的裝置上，然後在進入下一個部分之前丟棄此 state dict 部分。我們不再需要足夠的 RAM 來容納整個模型，而只需要足夠的 RAM 來獲取最大的檢查點部分，我們稱之為分片，在 BLOOM 的情況下是 7.19GB。
我們將像 BLOOM 這樣儲存在多個檔案中的檢查點稱為分片檢查點，我們已將其格式標準化如下
- 一個檔案（名為
pytorch_model.bin.index.json）包含一些元資料和一個引數名稱到檔名的對映，指示在哪裡可以找到每個權重 - 所有其他檔案都是標準的 PyTorch 狀態字典，它們只包含模型的一部分而不是整個模型。您可以在此處檢視索引檔案的內容。
要將此類分片檢查點載入到模型中，我們只需迴圈遍歷各個分片。Accelerate 提供了一個名為 load_checkpoint_in_model 的函式，如果您克隆了 Hub 中的一個倉庫，它將為您完成此操作，或者您可以直接使用 Transformers 的 from_pretrained 方法，該方法將為您處理下載和快取
import torch
from transformers import AutoModelForCausalLM
# Will error
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float16)
如果自動計算的裝置對映要求將某些權重解除安裝到磁碟上，因為您的 GPU 和 CPU RAM 不足，則您將收到一個錯誤，指示您需要傳遞一個資料夾，其中將解除安裝到磁碟的權重儲存在該資料夾中
ValueError: The current `device_map` had weights offloaded to the disk. Please provide an
`offload_folder` for them.
新增此引數應能解決錯誤
import torch
from transformers import AutoModelForCausalLM
# Will go out of RAM on Colab
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
checkpoint, device_map="auto", offload_folder="offload", torch_dtype=torch.float16
)
請注意，如果您嘗試載入一個需要除 CPU 解除安裝之外還需要磁碟解除安裝的超大型模型，那麼當檢查點的最後一個分片載入時，您可能會耗盡 RAM，因為模型停留在 CPU 上的部分會佔用空間。如果是這種情況，請使用 offload_state_dict=True 選項，在所有權重載入時臨時解除安裝停留在 CPU 上的模型部分，並在所有權重處理完畢後將其重新載入到 RAM 中
import torch
from transformers import AutoModelForCausalLM
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
checkpoint, device_map="auto", offload_folder="offload", offload_state_dict = True, torch_dtype=torch.float16
)
這將在 Colab 中執行，但會非常接近使用所有可用 RAM，當您嘗試生成預測時，它將耗盡 RAM。為了獲得可用的模型，我們需要將另一層解除安裝到磁碟。我們可以透過獲取上一節中計算出的 device_map，稍作調整，然後將其傳遞給 from_pretrained 呼叫來實現
import torch
from transformers import AutoModelForCausalLM
checkpoint = "facebook/opt-13b"
device_map["model.decoder.layers.37"] = "disk"
model = AutoModelForCausalLM.from_pretrained(
checkpoint, device_map=device_map, offload_folder="offload", offload_state_dict = True, torch_dtype=torch.float16
)
在多個裝置上執行拆分模型
我們尚未涉及的最後一部分是 Accelerate 如何使您的模型能夠在其權重分佈在多個 GPU、CPU RAM 和磁碟資料夾中時執行。這透過鉤子非常簡單地完成。
鉤子 是 PyTorch API，它在每次正向傳播呼叫之前執行函式
我們無法直接使用此功能，因為它們只支援正向傳播中具有常規引數且沒有關鍵字引數的模型，但我們採用了相同的思想。模型載入後，dispatch_model 函式將向每個模組和子模組新增鉤子，這些鉤子在每次正向傳播之前和之後執行。它們將
- 確保模組的所有輸入都在與權重相同的裝置上；
- 如果權重已解除安裝到 CPU，則在正向傳播之前將其移動到 GPU 0，並在之後立即移回 CPU；
- 如果權重已解除安裝到磁碟，則在正向傳播之前將其載入到 RAM，然後載入到 GPU 0，並在之後立即釋放此記憶體。
整個過程總結在以下影片中
這樣，即使您的 GPU RAM 和 CPU RAM 不足，您的模型也可以載入和執行。您唯一需要的是磁碟空間（以及大量的耐心！）。雖然這個解決方案在您擁有多個 GPU 時相當幼稚（不涉及巧妙的管道並行性，只是按順序使用 GPU），但它仍然為 BLOOM 帶來了相當不錯的結果。它允許您在較小的設定上執行模型（儘管速度較慢）。
要了解有關 Accelerate 大型模型推理的更多資訊，請參閱文件。
