伯克利開源Ray Tune：大規模AI實驗利器

在本文中，我們將跟大家分享開發兩個重要軟件庫Ray Tune和Ray Cluster Launcher 的經驗， 這兩個庫現在支持很多流行的開源AI庫，很多BAIR研究人員用它們來執行大規模AI實驗。

隨着AI研究變得越來越注重計算，很多AI研究人員的時間和資源都越來越緊缺。現在，很多研究人員依賴AWS或谷歌計算平臺等雲供應商，以獲取訓練大型模型所需的大量計算資源。

理解研究基礎設施

爲了更全面地理解這些內容，我們首先來看看業界的標準機器學習工作流（圖1）。

圖1 展示了業界典型的機器學習模型開發工作流

典型的“研究”工作流實際是步驟2和3之間的緊密循環，大致如圖2所示。

圖2展示了用於研究的典型機器學習模型開發工作流。該研究工作流通常是業界工作流的一個分支。

該研究工作流在很大程度上是個迭代的過程，通常受到實驗執行步驟（如圖2中的B所示）的制約。通常，一個“實驗”包括多個訓練作業或“試驗”，其中每個試驗都是訓練單個模型的任務。每個試驗可能使用一組不同的配置參數（超參數）或不同的種子來訓練一個模型。

在伯克利，我們看到轉向雲的AI研究人員花費大量時間來編寫自己的實驗執行工具，這些工具整理雲供應商的API以啓動實例、設置依賴項和啓動實驗。

不幸的是，除了開發這些工具所要投入的大量時間以外，這些臨時的解決方法還常常在功能上受到限制：

• 簡化的架構：每個試驗通常在單獨的節點上啓動，沒有任何集中控制邏輯。這讓研究人員難以實施優化技術，如基於羣體的訓練（Population-based Training）或貝葉斯優化（Bayesian Optimization），這些技術需要在不同的運行之間進行協調。
• 缺乏故障處理：如果實例失敗，那麼訓練作業的結果會永遠丟失。研究人員常常在電子表格上跟蹤實時實驗來手動管理故障轉移，但是，這麼做既費時又容易出錯。
• 沒有競價性實例折扣（Spot Instance Discount）：缺乏容錯能力也意味着放棄雲供應商提供的競價性實例折扣（最高可達90%）。

總之，在雲資源上測試和管理分佈式實驗既費力又容易出錯。因此，採用易於使用的框架以在執行和研究之間架起橋樑可以大大加快研究過程。來自BAIR的幾個實驗室成員協作構建了兩個互補的工具，用於在雲中進行AI實驗：

Ray Tune：用於訓練和超函數調優的容錯框架。具體來說，Ray Tune（或簡稱“Tune”）：

1. 協調並行作業以支持並行超參數優化。
2. 自動檢查點，並在機器有故障時恢復訓練作業。
3. 提供了很多最先進的超參數搜索算法，如基於羣體的訓練（ Population-based Training）和HyperBand。

Ray Cluster Launcher：一個實用程序，用於管理AWS、GCP和Kubernetes之間的資源調配和集羣配置。

Ray Tune

爲了克服這些臨時實驗執行工具的缺點，我們構建了Tune，利用了Ray Actor API並添加故障處理來解決問題。

基於Actor的訓練

很多用於超參數優化的技術需要一個框架來監視所有併發訓練作業的指標並控制訓練執行。爲了解決這個問題，Tune使用一種master-worker架構來集中決策，並使用Ray Actor API的分佈式worker來通信。

什麼是Ray Actor API？Ray 提供一個API從一個Python類來創建一個“actor”。這使得類和對象可以在並行和分佈式設置中使用 。

Tune使用一個可訓練的（Trainable）類接口來定義一個actor類，這個類專門用於訓練模型。該接口公開了如_train、_stop、_save和_restore等方法，這些方法允許Tune監視中間的訓練指標並殺死效果不佳的試驗。

class NewTrainable(tune.Trainable):
    def _setup(self, config):
        ...
    def _train(self):
        """Run 1 step of training (e.g., one epoch).
        Returns:
            A dict of training metrics.
        """
        ...
    def _save(self, checkpoint_dir):
        ...
    def _restore(self, checkpoint_path):
        ...

更重要的是，通過利用Actor API，我們可以在Tune中實施並行超參數優化模式，如HyperBand和並行貝葉斯優化，而這些是研究人員以前使用的實驗執行工具無法做到的。

容錯

雲供應商常常以很大的折扣提供“可搶佔的實例”（如，競價型實例）。大幅度的折扣使得研究人員可以顯著地降低其雲計算成本。但是，缺點是雲供應商可以在任何時候終止或停止運轉我們的機器，導致我們丟失訓練進度。

爲了可以使用競價型實例，我們構建了Tune以在集羣不同的機器中自動地建立檢查點並恢復訓練作業，這樣，實驗將對搶佔和集羣大小的調整具有彈性。

# Tune will resume training jobs from the last checkpoint
# Tune將從最新的檢查點恢復訓練作業
# even if machines are removed.
#即使機器已被刪除。
analysis = tune.run(
    NewTrainable,
    checkpoint_freq=5,  # Checkpoint every 5 epochs
    config={"lr": tune.grid_search([0.001, 0.01, 0.1])},
)

它是如何工作的？

如果在某個節點丟失的時候，有個訓練作業仍在其上執行且該訓練作業（試驗）的檢查點存在，那麼，Tune將一直等到有可用資源可用，以再次執行該試驗。

如果該試驗被放置在一個不同的節點上，那麼，Tune將自動把之前的檢查點文件推送到那個節點並恢復狀態，即使失敗，該試驗也可以從最新的檢查點恢復。

Ray Cluster Launcher

上面，我們描述了爲自動化集羣設置過程而整理雲供應商API的痛苦之處。但是，即使有了用於分解集羣的工具，研究人員還是必須經過繁瑣的工作流才能進行實驗：

爲了簡化一下這個過程，我們構建了Ray Cluster Launcher，這個工具會調配資源並自動調整資源大小，並且在AWS EC2、GCP和Kubernetes上啓動一個Ray集羣。然後，我們把以上用於進行一個實驗的步驟抽象成一個簡短的配置文件和一個命令：

# CLUSTER.yaml
cluster_name: tune-default
provider: {type: aws, region: us-west-2}
auth: {ssh_user: ubuntu}
min_workers: 0
max_workers: 2
# Deep Learning AMI (Ubuntu) Version 21.0
head_node: {
    InstanceType: c4.2xlarge,
    ImageId: ami-0b294f219d14e6a82}
worker_nodes: {
    InstanceType: c4.2xlarge,
    ImageId: ami-0b294f219d14e6a82}
setup_commands: # Set up each node.
    - pip install ray numpy pandas
file_mounts: {
  '/home/ubuntu/files':'my_files/',
}

下面的命令用來啓動一個集羣、上傳並運行一個用於分佈式超參數調整的腳本，然後關閉該集羣。

$ ray submit CLUSTER.yaml --start --stop tune_experiment.py \
--args="--address=auto"

研究人員現在使用Ray Tune和Ray Cluster Launcher同時在數十臺GPU機器上啓動數百個並行作業。Ray Tune的分佈式實驗文檔頁面展示瞭如何執行該操作。

總結

在過去的一年中，我們一直和BAIR不同的團隊合作，以更好地讓研究人員利用雲。我們必須讓Ray Tune和Ray Cluster Launcher變得足夠通用，以支持大量研究代碼庫，與此同時，要把上手門檻降到人人可以在幾分鐘內嘗試一下的程度。

# An example Ray Tune script for PyTorch.
import torch.optim as optim
from ray import tune
from ray.tune.examples.mnist_pytorch import (
    get_data_loaders, ConvNet, train, test)
class TrainMNIST(tune.Trainable):
    def _setup(self, config):
        self.train_loader, self.test_loader = get_data_loaders()
        self.model = ConvNet()
        self.optimizer = optim.SGD(
            self.model.parameters(), lr=config.get("lr", 0.01))
    def _train(self):
        train(self.model, self.optimizer, self.train_loader)
        acc = test(self.model, self.test_loader)
        return {"mean_accuracy": acc}
    def _save(self, checkpoint_dir):
        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
        torch.save(self.model.state_dict(), checkpoint_path)
        return checkpoint_path
    def _restore(self, checkpoint_path):
        self.model.load_state_dict(torch.load(checkpoint_path))
analysis = tune.run(
    TrainMNIST,
    stop={"training_iteration": 50},
    config={"lr": tune.grid_search([0.001, 0.01, 0.1])})
print("Best hyperparameters: ", analysis.get_best_config(
    metric="mean_accuracy"))
# Get a dataframe for analyzing trial results.
df = analysis.dataframe()

Tune已經成長爲一個流行的超參數調整開源項目。很多其他流行的研究項目也在使用它，從基於羣體的數據增強（ Population-based Data Augmentation）項目到用於AllenNLP的超參數調整（ Hyperparameter Tuning）以及AnalyticsZoo的AutoML。

BAIR的很多開源項目現在都依賴Ray Tune和Ray Cluster Launcher的組合來編排和執行分佈式實驗，其中包括rail-berkeley的softlearning、HumanCompatibleAI的對抗策略（adversarial-policies），以及流項目（flow-project）的流。

來，動手嘗試一下Ray Tune和Ray Cluster Launcher吧！

相關鏈接

• BAIR以前關於Ray項目的博文：
  • Ray: A Distributed System for AI
  • Scaling Multi-Agent Reinforcement Learning
  • Functional RL with Keras and Tensorflow Eager
• Ray網站
• Ray GitHub頁面
• Ray的文檔：
  • Tune
  • Ray Cluster Launcher
  • Ray-project landing page
  • Installation

原文鏈接：

Large Scale Training at BAIR with Ray Tune