Hugging Face Transformers 透過 DeepSpeed 與 FairScale 整合 ZeRO
Hugging Face Transformers v4.2.0 引入了對 DeepSpeed 與 FairScale 的實驗性支援，透過整合 Zero Redundancy Optimizer (ZeRO) 來優化 GPU 記憶體使用量，讓使用者能夠訓練更大的模型並增加批次大小 (batch size)。此整合允許使用者直接透過 Trainer API，使用 --sharded_ddp (FairScale) 與 --deepspeed (DeepSpeed) 命令列參數來使用這些優化功能。
多 GPU 訓練的 ZeRO 記憶體優化
與標準的 Distributed Data Parallel (DDP) 基準相比，透過 DeepSpeed 或 FairScale 整合 ZeRO 可以顯著減少訓練與評估時間，同時增加最大可能的批次大小 (BS)。在一個使用 t5-large 模型於兩張 24GB Titan RTX GPU 上的基準測試中，觀察到以下效能提升：
| Method | Max Batch Size | Train Time | Eval Time |
|---|---|---|---|
| Baseline (DDP) | 16 | 30.9458 | 56.3310 |
| fp16 | 20 | 21.4943 | 53.4675 |
| sharded_ddp (FairScale) | 30 | 25.9085 | 47.5589 |
| sharded_ddp + fp16 | 30 | 17.3838 | 45.6593 |
| DeepSpeed (no CPU offload) | 40 | 10.4007 | 34.9289 |
| DeepSpeed (with CPU offload) | 50 | 20.9706 | 32.1409 |
DeepSpeed 在批次大小與訓練速度方面展現了最高的增益，而 FairScale 則被認為更容易部署，因為它只需要單一命令列參數，不需要設定檔。
大型模型之單 GPU 訓練
DeepSpeed 可透過 CPU offloading 技術，訓練原本會超出單張 GPU 記憶體容量的模型。在一個使用 t5-3b 模型於單張 24GB RTX-3090 顯卡上的測試中，標準的單 GPU 設定即使在批次大小為 1 的情況下也會發生 Out of Memory (OOM) 錯誤。透過使用 DeepSpeed，該模型得以在批次大小為 20 的情況下成功訓練，且系統僅在批次大小達到 30 時才發生 OOM。
ZeRO 的技術機制
ZeRO (Zero Redundancy Optimizer) 透過在資料平行訓練中加入分散式資料儲存來優化記憶體。ZeRO 不會在所有 GPU 上複製整個模型狀態，而是將參數、梯度與優化器狀態 (optimizer states) 分散在可用的 GPU 之間。
分散式分片 (Distributed Partitioning)
每張 GPU 僅儲存參數、梯度與優化器狀態的一個分片 (shard)。在執行時，每張 GPU 會即時從其他參與的 GPU 獲取特定層所需的資料，確保資料儲存完全沒有重疊。
ZeRO-Offload
ZeRO-Offload 將特定的處理與記憶體需求從 GPU 移至主機 CPU，這對於將龐大的模型（如 t5-3b）放入有限的 GPU 硬體中至關重要。
記憶體碎片管理
DeepSpeed 解決了 GPU 記憶體碎片化問題——即儘管總可用記憶體充足，但因為沒有足夠大的連續區塊而導致 OOM 錯誤——透過獨立管理 GPU 記憶體來分離長期與短期配置。
部署與整合
ZeRO 優化不需要修改模型架構；僅需修改訓練程式碼。Hugging Face Trainer 的使用者可以透過以下標記來實作這些功能：
--sharded_ddp: 啟動 FairScale 整合。--deepspeed: 啟動 DeepSpeed 整合（需要一個 JSON 設定檔）。
提到的未來增強功能包括 DeepSpeed Sparse Attention、1-bit Adam，以及預期將在 FairScale 與 DeepSpeed 中支援模型參數分片 (model parameter sharding)。