Hugging Face Accelerate Library Release

Hugging Face 推出了 🤗 Accelerate，這是一個旨在讓 PyTorch 使用者能夠在各種硬體配置上執行其原始訓練腳本的函式庫，包括單一 GPU、多 GPU 集群和 TPU，而無需為了分散式訓練或混合精度而重寫樣板代碼。

Simplified Distributed Training and Mixed Precision

Accelerate 讓開發者在移除管理設備放置和分散式設置的複雜性之餘，仍能保持對訓練迴圈的完全控制。透過在標準 PyTorch 腳本中添加少量代碼，使用者可以從單一設備設置轉換到分散式環境，而無需手動實作 DistributedDataParallel 或 DistributedSampler。

Key benefits include:

  • Minimal Code Changes: Transitioning a script to Accelerate typically requires only a few modifications to the training loop.
  • Unified API: The same functions work across different distributed setups, eliminating the need for device-specific if-statements.
  • Hardware Agnostic: Scripts remain compatible with CPUs, single GPUs, and distributed configurations.

Technical Implementation and Core API

Accelerate 透過抽象化核心 PyTorch 對象的初始化和準備工作來運作。主要介面是 Accelerator 類別。

Initialization

accelerator = Accelerator() 分析環境以確定分散式訓練運行的類型，並執行必要的初始化。使用者可以透過在初始化期間傳遞 cpu=True 或 fp16=True 來顯式強制進行 CPU 訓練或混合精度訓練。

The prepare Method

accelerator.prepare() 方法是該函式庫的核心組件。它封裝了三種主要的對象類型，使其與分散式環境相容性：

  • Models: Wraps the model in the appropriate container (eg, DistributedDataParallel) and handles device placement. Models can be retrieved via accelerator.unwrap_model(model) for saving or accessing specific methods.
  • Optimizers: Wraps the optimizer to handle mixed precision operations and manage the device placement of the state dict.
  • DataLoaders: Wraps the dataloader to ensure each process only retrieves relevant indices from the sampler. This removes the requirement for users to manually implement a DistributedSampler and works with any sampler provided to the dataloader.

Backward Pass

accelerator.backward(loss) 取代了標準的 loss.backward()，以納入混合精度和其他特殊整合所需的必要步驟。

Distributed Evaluation

Accelerate 支援單一進程和分散式評估。對於僅應在主進程上運行的任務，使用者可以利用 if accelerator.is_main_process():。

對於分散式評估，該函式庫提供了 accelerator.gather()，它會收集所有進程中的預測值和標籤的張量。由於準備好的評估 dataloader 可能會為了確保跨進程的一致性批次大小而返回額外的元素，使用者必須將收集到的結果截斷至原始數據集長度以確保準確性。

Deployment and Launching

Accelerate 包含一個 CLI 工具，以簡化在不同硬體配置上執行腳本的執行方式：

  1. Configuration: accelerate config 會啟動一個問卷，以建立一個具有預設訓練設置的配置文件。
  2. Execution: accelerate launch path_to_script.py 使用保存的預設值來運行腳本。

Future Roadmap

Hugging Face 計畫將 Accelerate 進行擴展，以支援 fairscale, deepspeed, 和 AWS SageMaker 特有的數據並行和模型並行。

Sources