# Pytorch 如何加速大型模型 [原文](https://huggingface.co/blog/accelerate-large-models)
隨著大型語言模型（Large Language Models, LLMs）的崛起，如 Hugging Face 最近推出的各種 6.7B 或 13B 參數的模型，如何在有限的記憶體（RAM）和儲存空間（Disk Space）上有效地運行這些模型，成為一大挑戰。這篇文章將詳細介紹如何使用 Pytorch 和 Hugging Face Accelerate 套件來優化大型模型的載入與推論。
## 傳統 Pytorch 模型載入流程
在傳統的 Pytorch 模型使用流程中，通常會依序進行以下步驟：
1. **建立模型架構**：定義模型的各個層次和計算圖。
2. **載入參數至記憶體（state\_dict）**：從檔案中讀取模型的參數（權重）。
3. **將參數載入模型**：將這些參數賦值給模型的對應層。
4. **將模型放入裝置（CPU 或 GPU）**：將模型的所有參數移至指定裝置進行推論。
### 傳統方法與磁碟空間的使用
在傳統方法中，模型參數（如 state\_dict）通常儲存在磁碟上（Disk），當我們載入模型時，這些參數會從磁碟載入至記憶體（RAM）。然而，整個模型的參數一旦載入，就會完全駐留在記憶體中，無論是 CPU 還是 GPU。
這種方式在小型模型上運作良好，但對於 6.7B 或 13B 參數的模型來說，這會導致記憶體不足：
* **模型架構建立：** 6.7B 參數 × 4 bytes（Float32）= 26.8 GB 記憶體。
* **載入參數：** 再次需要 26.8 GB 記憶體。
* **搬移至 GPU：** 額外消耗 GPU 記憶體。
這使得在普通硬體（如 Colab）上幾乎無法運行這些大型模型。
## Hugging Face Accelerate：高效的模型載入方式
與傳統方法不同，Hugging Face Accelerate 提供了一種動態且分散的載入方式，這種方法透過以下流程實現：
1. **建立模型（Meta Tensors）**：使用 `meta` device 創建虛擬張量，僅紀錄張量大小而不佔用實際記憶體。
2. **劃分儲存空間（Device Map）**：將不同層次的模型分配至不同裝置（GPU、CPU、Disk）。
3. **載入參數（Offloading）**：依照 Device Map 將參數載入對應裝置。
4. **將參數載入模型並進行推論**：僅在需要時將指定參數加載至指定裝置。
這種方式在處理大型模型時更加靈活，尤其是當記憶體（RAM 和 GPU）有限時，可以將部分參數儲存於磁碟（Disk），並在需要時動態載入。
## 內存優化的關鍵：Meta Tensors
在 Accelerate 之前，Pytorch 一旦建立模型張量（Tensor），無論是否有實際數據，都會立即佔用記憶體。這在大型模型上顯然是不可行的。
```python
import torch
# 這行會消耗超過 40GB 的記憶體，可能導致崩潰
large_tensor = torch.randn(100000, 100000)
# 使用 meta device 不會消耗實際記憶體
large_tensor = torch.randn(100000, 100000, device="meta")
```
### 解決辦法：Context Manager 介入管理初始化
Context Manager 是 Python 中用來管理資源工具的概念。Hugging Face Accelerate 提供 `init_empty_weights` 這個 Context Manager 來讓模型初始化為 meta device（僅紀錄大小而不佔用記憶體）。
```python
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)
```
## Device Map：優化模型儲存位置
Accelerate 用 device map 去規劃資料被載入的位置。
```python
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)
```
這個 device map 在 GPU 會顯示
```
{'model.decoder.embed_tokens': 0,
'model.decoder.layers.0': 0,
'model.decoder.layers.9': 'cpu',
'model.decoder.layers.18': 'disk',
'lm_head': 'disk'}
```
這表示模型的前 9 層在 GPU，接下來的 9 層在 CPU，最後幾層則放在磁碟。
### 殘差網路與分層優化
當模型有殘差連接（Residual Connection）時，這些參數必須位於同一裝置才能計算，因此我們可以指定 `no_split_module_classes` 來確保它們不會被分開：
```python
device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"])
```
結果可能顯示：
```
{'model.decoder.embed_tokens': 0,
'model.decoder.layers.0': 0,
'model.decoder.layers.21': 0,
'model.decoder.layers.22': 'cpu',
'model.decoder.layers.37': 'cpu',
'model.decoder.layers.38': 'disk',
'lm_head': 'disk'}
```
### 多 GPU 環境中的 Device Map
* **auto / balanced**：GPU 使用量均衡。
* **balanced\_low\_0**：第一個 GPU 使用較少，適合後處理。
* **sequential**：填滿一個 GPU 後再使用下一個。
當然也可以手動指定 `device_map` 每層的位置
## Dict 與磁碟空間優化
在 Pytorch 中，模型的參數（如權重和偏置）通常使用 state dict（狀態字典）來儲存，這是一種將模型各層名稱對應至權重張量的字典：
```python
# 儲存模型權重
torch.save(my_model.state_dict(), 'model_weights.pth')
# 重新載入模型
new_model = ModelClass()
new_model.load_state_dict(torch.load('model_weights.pth'))
```
### 為何大型模型無法使用傳統 State Dict
當模型規模增至數十億參數（如 13B 參數的 LLM）時，這種一次性保存和載入方法變得不可行：
* **記憶體限制：** 需要將整個模型一次性載入至記憶體。
* **磁碟空間：** 單一檔案會變得極大，載入和儲存時間過長。
### Hugging Face 的 Sharded State Dict 優化
Hugging Face 解決此問題的方法是將模型權重分割成多個小檔案（Shards）：
* **分割檔案**：`pytorch_model_xxxxx-of-00072.bin`，每個檔案儲存一部分模型參數。
* **索引檔案**：`pytorch_model.bin.index.json`，儲存所有 Shard 的結構和位置。
這種方法允許模型在載入時僅讀取當前需要的 Shard，而非整個模型，極大地降低了記憶體和磁碟壓力。
### 自動分割與快取
在實際操作中，只需使用 Hugging Face 的 `from_pretrained` 方法，這些分割與載入將自動處理：
```python
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
checkpoint, device_map="auto", offload_folder="offload", torch_dtype=torch.float16
)
```
### Offload Folder
在有限資源環境（如 Colab）中，當模型無法完全載入至 RAM 或 GPU 時，常見的錯誤是：
```
ValueError: The current `device_map` had weights offloaded to the disk. Please provide an `offload_folder` for them.
```
這表示你需要指定一個 `offload_folder` 來將部分參數儲存至磁碟。
```python
model = AutoModelForCausalLM.from_pretrained(
checkpoint, device_map="auto", offload_folder="offload", torch_dtype=torch.float16
)
```
### Offload Folder 記憶體不足問題
在最後一個 shard 需要載入時，如果整體模型內存已滿，則無法再分配記憶體。這種情況特別常見於混合 CPU 和磁碟 offload 的情況。解決方法是在 `from_pretrained` 中使用 `offload_state_dict=True`，這將確保每個 shard 在載入後即時釋放，避免內存佔滿：
```python
model = AutoModelForCausalLM.from_pretrained(
checkpoint, device_map="auto", offload_folder="offload", offload_state_dict=True, torch_dtype=torch.float16
)
```
### Offload Folder 推論時記憶體不足問題
在 Colab 等資源有限的環境中，使用以上方式可以載入模型，但如果模型無法生成預測（表示記憶體已滿），可以手動將特定層指定為儲存在磁碟：
```python
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
)
```
## 如何做到：Hook 機制與 Accelerate
前面都在討論架構本身，但還沒解釋到 Hugging Face Accelerate 如何在在大型模型上實現不同裝置（如 CPU、GPU、Disk）間的動態流動參數。
### 傳統 Hook 機制
一般來說會先想到 Hook API，這是在 Pytorch 中一種監控或修改計算圖中張量的方式（Hook 是一種「鉤子」或「監控器」，它允許你在模型運算時插入自訂的邏輯），常用於：
* **Forward Hook**：監控輸入和輸出。
* **Backward Hook**：監控梯度計算。
* **Parameter Hook**：監控和修改模型參數。
```python
import torch
# 簡單的 Hook 範例
def print_gradient(module, grad_input, grad_output):
print("梯度：", grad_output)
model = torch.nn.Linear(10, 5)
model.register_backward_hook(print_gradient)
```
但傳統 Hook 無法控制模型參數在不同裝置之間的流動。
### Hook 與 Hugging Face Accelerate 的關係
在 Hugging Face Accelerate 中，雖然它沒有明確使用 Hook 這個詞，但它採用了類似 Hook 的機制來實現「動態模型分配」：
* 動態分配：在計算過程中，模型層會自動分配至 GPU、CPU 或磁碟（`Dispatch Model`）。
* 即時釋放：計算結束後自動釋放記憶體，避免記憶體不足。
* 這種機制本質上是一種 進階的 Hook 機制，但應用於模型的分配與釋放。
Sign in with Wallet
Connect another wallet