Hugging Face Accelerate：使用 PyTorch 執行大型模型

Hugging Face Accelerate 允許使用者載入並執行超出消費者硬體可用 RAM 或 GPU 記憶體的超大型語言模型（LLM）。透過利用特定的 PyTorch 功能與修改過的載入管線，Accelerate 會將模型權重分散至可用的硬體資源，包括多個 GPU、CPU RAM 與磁碟儲存空間。

克服模型載入的記憶體限制

傳統的 PyTorch 模型載入遵循線性流程：建立模型、將權重載入記憶體中的 state_dict、將這些權重載入模型，然後將模型移動至裝置。對於超大型模型而言，此流程在 RAM 方面的成本極高。例如，一個擁有 67 億參數、使用 float32 精度的模型，僅在初始模型建立時就需要約 26.8GB 的 RAM，載入 state_dict 又需要額外的 26.8GB，總計超過 53GB，甚至在模型進入 GPU 前就已耗盡。

為了解決此問題，Accelerate 實作了更具記憶體效率的管線：

建立一個空模型（不含權重）。
決定 device map，以指定每層所在的裝置。
以小部分（分片）載入權重。
將這些權重載入空模型。
將權重移動至指定的裝置以進行推論。
對所有剩餘的權重重複上述流程。

透過 Meta 裝置進行高效模型初始化

Accelerate 利用 PyTorch 的「meta」裝置（於 PyTorch 1.9 引入）來實例化模型，而不分配實際資料的記憶體。位於 meta 裝置上的張量僅儲存形狀與資料類型資訊，使得可以建立任意大的模型而不消耗 CPU 或 GPU 的 RAM。

由於為 Transformers 套件中的每個模型重寫以支援 device 關鍵字實務上不可行，Hugging Face 開發了 init_empty_weights() 上下文管理器。它允許任何模型在 meta 裝置上以「外殼」形式實例化，提供必要的結構資訊以計算記憶體需求，而不需載入實際權重。

自動裝置映射與資源分配

Accelerate 使用 infer_auto_device_map 函式自動將模型權重分配至可用硬體。系統會依以下順序優先使用資源：GPU、CPU RAM，最後是磁碟卸載。

裝置映射設定

根據使用情境，使用者可以選擇不同的映射策略：

"auto" 或 "balanced"：將權重平均分配至所有可用的 GPU。
"balanced_low_0"：將權重平均分配至 GPU，但盡量減少第一個 GPU（GPU 0）的負載，適用於第一個 GPU 需要用於模型輸出（例如文字生成）時。
"sequential"：依序填滿 GPU，可能導致後續 GPU 未被使用。

為避免系統將單一層拆分至多個裝置（會破壞殘差連接），Accelerate 允許使用者指定 no_split_module_classes（例如 ["OPTDecoderLayer"]）。

使用分片檢查點降低 RAM 開銷

對大多數硬體而言，載入單一巨大的 state_dict 檔案是不切實際的；例如，BLOOM 模型（1760 億參數）僅以 bfloat16 載入權重就需要 352GB 的 RAM。為了緩解此問題，Hugging Face 使用 分片檢查點。

在分片檢查點中：

pytorch_model.bin.index.json 檔案將每個參數名稱映射至特定的分片檔案。
權重分散於多個標準的 PyTorch state dict 檔案（例如 BLOOM 有 72 個檔案）。

此方法確保系統只需足夠的 RAM 以容納最大的單一分片（例如 BLOOM 為 7.19GB），而不必一次載入整個模型。若 GPU 與 CPU RAM 不足，Accelerate 可將權重卸載至磁碟上指定的 offload_folder。offload_state_dict=True 選項則可在處理其他分片時，暫時將位於 CPU 的模型部分卸載，以進一步降低 RAM 使用量。

透過動態 Hook 執行

為了在多個裝置上執行分割模型，Accelerate 採用了受 PyTorch hook 啟發的機制。dispatch_model 函式會在每個模組與子模組上掛載於前向傳遞前後執行的 hook。這些 hook 執行以下操作：

確保所有模組的輸入與權重位於相同的裝置上。
在前向傳遞前立即將權重從 CPU 移至 GPU 0，並在傳遞結束後立即返回 CPU。
在前向傳遞前將權重從磁碟載入至 RAM 再載入 GPU 0，傳遞結束後立即釋放記憶體。

雖然此方法是以順序方式使用 GPU，而非採用複雜的 pipeline 並行，但它使得在顯著較小的硬體環境下也能執行巨型模型。