使用 PyTorch DDP、Accelerate 與 Transformers Trainer 的分散式訓練

Hugging Face 詳細說明了在 PyTorch 中實作分散式訓練的三個抽象層級，從原生的分散式資料平行（Distributed Data Parallelism，DDP）到 Accelerate 函式庫，最後到 Transformers Trainer API。這樣的層級讓開發者可以在對訓練迴圈的細緻控制與高階的分散式設定自動化之間做選擇。

原生 PyTorch 分散式資料平行 (DDP)

PyTorch DDP 允許在多個 GPU 上進行訓練，無論是單機（多 GPU）還是跨多台機器（多節點）。它的運作方式是將模型複製到每個 GPU，並在 loss.backward() 呼叫期間對所有副本的梯度進行平均，以確保各設備之間的權重一致性。

實作原生 DDP 需要多個手動設定步驟：

• Process Group Setup：開發者必須定義 setup 與 cleanup 函式，以初始化與銷毀 process group，並指定用於通訊的 MASTER_ADDR 與 MASTER_PORT。
• Model Wrapping：模型必須以 DistributedDataParallel（DDP）模組包裹，且 optimizer 必須基於此包裹後的模型宣告，以確保梯度正確計算。
• Execution：腳本通常使用 torchrun 命令列模組啟動，並指定節點數量與每個節點的進程數。

使用 🤗 Accelerate 簡化分散式

Accelerate 是一個輕量的 pytorch.distributed 包裝器，讓相同的程式碼能在單一 GPU、多個 GPU 或 TPU 上執行，只需極少的變更。它省去手動設定 process group 的需求，並簡化裝置配置。

技術改進

Accelerate 引入了 Accelerator 類別，透過單一呼叫 accelerator.prepare() 即可處理模型、optimizer 與資料載入器的分散。這取代了手動的 .to(rank) 呼叫與 DDP 包裝。

除了簡化之外，Accelerate 透過自訂 sampler 提升記憶體效能。它不會在各設備上建立多個完整的資料載入器副本，而是確保記憶體中僅保有原始資料集的一份完整拷貝，並將資料子集分配至可用的節點。這可防止在訓練極大資料集時出現記憶體爆炸的情況。

Notebook 整合

對於在 Jupyter Notebook 中工作的使用者，Accelerate 提供了 notebook_launcher 工具。透過傳入訓練函式與欲使用的進程數，即可直接在 notebook 中觸發多 GPU 訓練。

透過 🤗 Trainer 的高階抽象

Transformers 的 Trainer API 提供了最高層級的抽象，移除幾乎所有與分散式訓練相關的樣板程式碼。它會自動處理底層的分散式邏輯，使用者無需撰寫明確的訓練迴圈。

要使用 Trainer，開發者需要定義 TrainingArguments 以管理超參數，並可繼承 Trainer 以實作自訂的 compute_loss 函式。之後，Trainer 會自動管理跨分散式系統的訓練與評估流程。與 Accelerate 範例類似，Trainer 也能與 notebook_launcher 整合，以在互動環境中快速實驗。