[原始碼解析] PyTorch 分散式之彈性訓練(1) --- 總體思路

[原始碼解析] PyTorch 分散式之彈性訓練(1) --- 總體思路

0x00 摘要

在前面的文章之中，我們已經學習了PyTorch 分散式的基本模組，介紹了官方的幾個例子，我們接下來會介紹PyTorch的彈性訓練，本文是第一篇，介紹其歷史和設計理念，也會與Horovod做一下對比。

注：後續會對Horovod文章進行系統整理，屆時對本文進行更新，會加入更多對比分析。

0x01 痛點

因為機器學習的模型越來越龐大，單個GPU視訊記憶體早已無法容納模型引數，所以一般都是使用大量節點或者叢集進行訓練，隨著訓練規模擴大，硬體薄弱或設計原因會導致單點故障概率隨之增加，這就帶來了一些問題或者痛點，比如：

痛點 1：缺少容災功能。
- 問題點：單個節點故障往往會導致整個訓練job結束。雖然框架提供了checkpoint功能，但是頻繁呼叫會導致效能問題，所以依然會丟失一段時間的訓練成果，並且還得繼續進行任務排隊。
- 理想狀態：單個節點失敗不會影響整體訓練，在節點故障時候，自動剔除該節點，同時訓練繼續平滑進行。
痛點 2：缺少彈性算力感知和動態訓練擴縮容機制。
- 問題點：使用者只能在提交任務時候確定所需要的固定靜態資源，無法對叢集資源進行實時動態感知，導致叢集資源利用率低。
- 理想狀態：應該在有少量空閒機器時候就開始訓練，當有更多資源時候，彈性任務同上層排程系統可以和i進行配合，從而能有效檢測到這些潛在資源，在訓練過程中可以自動增加worker數量。當本任務有空閒算力時候，會自動釋放資源。而且在worker數量變化時，不會中斷訓練任務，做到平滑過渡。
痛點 3：叢集資源配置/排程機制不靈活
- 問題點：目前不支援動態配置worker，不支援高優先順序搶佔例項。因此當資源不足時，無法按需為其他高優先順序業務騰出資源，只能等待任務自己主動終止或者出錯終止。
- 理想狀態：訓練任務可以被搶佔，可以主動騰出資源，可以在不同用途/配置的機器間進行漂移。

0x02 難點

我們接下來看看實現彈性訓練需要面對哪些挑戰和難點，這裡只從工程角度來看，不考慮資料切分/學習率/batch size 調整等問題。

難點1 ：需要一個節點/程式之間彼此發現的機制。

節點/訓練程式自動進入或者退出時候，其他節點/訓練程式如何感知。

難點2：如何處理成員變更

當發現有成員變更之後，如何處理。

難點3：如何捕獲單個程式訓練失敗。

如何在單個節點上管理所有訓練程式，從而當某個程式發生錯誤時候，可以捕獲其失敗，或者重試或者重啟該程式。

難點4：如何與現有訓練程式碼整合。

如何以最小工作量與現有程式碼進行整合，不引入複雜的抽象，或者對使用者最大限度遮蔽底層實現。

0x03 TorchElastic

我們接下來看看 PyTorch 彈性機制總體狀況。PyTorch 彈性機制是合併自 TorchElastic https://github.com/pytorch/elastic，所以我們就從 TorchElastic 看起。

TorchElastic（TE）是從 PyTorch 1.9 正式引入的，我們從兩個地方看彈性訓練的i歷史。

3.1 歷史

3.1.1 PyTorch 1.7

Release note 裡面提到了 TE 已經被加入了 PyTorch Docker 之中：

[Stable] TorchElastic now bundled into PyTorch docker image

Torchelastic提供了"torch.distributed.launch" CLI的一個嚴格超集，並新增了容錯和彈性功能。如果使用者對容錯不感興趣，他們可以通過設定"max_restarts=0"來獲得準確的功能/行為，並增加自動分配"RANK"和"MASTER_ADDR"埠的便利性（而不是在"torch.distributed.launch"中手動指定）。

3.1.2 PyTorch 1.9

從TE repository 可以看到，TE現在已經被合併到主版本 PyTorch 1.9 之中。

IMPORTANT: This repository is deprecated.

TorchElastic has been upstreamed to PyTorch 1.9 under torch.distributed.elastic. Please refer to the PyTorch documentation here.

3.2 設計理念

PET經歷了兩個版本，從 https://github.com/pytorch/elastic/blob/master/design/torchelastic/0.2.0/design_doc.md 可以看到其設計理念。

3.2.1 基本功能

PyTorch Elastic Trainer (PET) 提供了一個可以用容錯和彈性方式跨叢集來訓練模型的框架。PET 通過兩種方式提供這些功能：

當 PyTorch worker 程式丟擲某類可重試錯誤時，它會被 PET 捕獲並重試訓練過程。
只要worker的數量維持在開始工作時指定的範圍內，新worker就可以隨時離開或加入到現有訓練job的程式池。當成員發生變化時，所有worker會重新集合（re-rendezvous）以建立一個新的程式組，並從以前的良好狀態之中恢復訓練。

為了與 PET 整合，PyTorch 使用者需要對其訓練邏輯進行以下更改：

使用者需要使 PET 來控制他們的訓練迴圈。
- 本質上，使用者提供了一個“內部訓練”迴圈，該迴圈被 PET 包裹在一個可重試的迴圈中。
- PET迴圈是可重試的迴圈，其負責建立或重新建立過程組，以及將使用者的訓練恢復到良好狀態。
在新worker加入程式池時，使用者需要指定狀態是什麼以及如何把狀態施加到一個新worker之上。

3.2.2 新設計概述

PET v0.2 從 v0.1 之中獲取了不少經驗，下面講講 v0.2的設計理念。

動態範圍

在 PET v.0.2 中，我們不再嘗試恢復訓練函式中的錯誤。相反，PET 嘗試維護工作程式的數量，使它們保持在作業所需的 [ min , max ] 範圍內。應用編寫者負責從現有可用還原點檔案載入和重新啟動。與 v0.1 不同，PET v0.2 不強制指定如何管理checkpoints。應用編寫者可以任意使用torch.save 和 torch.load 或更高層次的框架如PyTorch Lightening 進行處理。

本地代理

PET v0.2 使用了一個名為elastic-agent的新程式，每個節點有一個獨立的elastic-agent。每個代理程式只負責管理該節點的一組本地工作程式，並與本作業其他節點上的彈性代理一起協調來確定程式組成員身份的變化。具體如下圖所示：

圖片來源：https://github.com/pytorch/elastic/raw/master/design/torchelastic/0.2.0/torchelastic_diagram.jpg

成員變更

成員變更的處理方式如下：當一個工作程式失敗時，管理它的彈性代理會殺死該節點上的所有worker，然後與其他代理建立一個集合操作（rendezvous），並使用新的集合資訊來重啟worker。但是，當代理以非零錯誤程式碼退出時，應該由上層排程模組（例如 Kubernetes）來重新啟動代理（同理，此代理將重新啟動它負責的所有worker）。相同的恢復機制也適用於節點級故障。編排工具（諸如 Kubernetes ）會排程作業以便job可以使用最小數目的代理副本執行，然後每個代理將依次編排使用者的訓練指令碼。

圖片來源：https://github.com/pytorch/elastic/raw/master/design/torchelastic/0.2.0/torchelastic_agent_diagram.jpg

相容

要採用 PET v0.2，應用程式只需讓其入口點或main函式與PyTorch distributed launcher相容。我們期望通過分散式啟動器啟動的分散式訓練作業可以通過彈性代理無縫啟動，無需更改或最小化程式碼更改。唯一的區別是在後一種情況下，應用程式將能夠在出現某些故障的情況下依然取得進展。

3.2.3 bare-bones

新的PET設計是想成為一個“bare-bones”：它從簡單性和健壯性兩方面權衡了應用程式可恢復的粒度。

將來，TE 希望為檢查點機制提供更多更方便的API，開發人員可以選擇使用這些API來實現更高效的重啟語義。

因為 PET 是 “bare-bones”，所以對使用者如何處理也給了一些指導性意見，比如 checkpoint 的處理。

一旦發生故障或成員變更，所有幸存的worker將立即被殺掉。所以使用者需要手動地處理 checkpoint，定期儲存你的工作進度，來保證重啟後訓練能夠繼續下去。檢查點的頻率應取決於使用者job對於失敗的容忍度。建議使用者指令碼採用如下結構進行處理：

  def main():
    load_checkpoint(checkpoint_path)
    initialize()
    train()

  def train():
    for batch in iter(dataset):
      train_step(batch)

      if should_checkpoint:
        save_checkpoint(checkpoint_path)

3.3 小結

不難發現，TE的設計理念主要就是回答了之前提到的4個難點。

難點1 ：需要一個節點/程式之間彼此發現的機制。

TE的答案是：當成員發生變化時，所有worker會重新集合（re-rendezvous）以建立一個新的程式組。rendezvous就是這個發現機制。

難點2：如何處理成員變更

TE的答案是：當一個工作程式失敗時，管理它的彈性代理會殺死該節點上的所有worker，然後與其他代理建立一個集合操作（rendezvous），並使用新的集合資訊來重啟worker。但是，當代理以非零錯誤程式碼退出時，應該由上層排程模組（例如 Kubernetes）來重新啟動代理（同理，此代理將重新啟動它負責的所有worker）。

難點3：如何捕獲單個程式訓練失敗，如何在單個節點上管理所有訓練程式。

TE的答案是：每個代理程式只負責管理該節點的一組本地工作程式，並與本作業其他節點上的彈性代理一起協調來確定程式組成員身份的變化。

難點4：如何與現有訓練程式碼整合。

TE的答案是：應用程式只需讓其入口點或main函式與PyTorch distributed launcher相容。

0x04 問題

4.1 VS Horovod

因為我們已經有了 Horovod 彈性訓練的基礎，所以我們就用 Horovod 為基準，提出一系列問題，然後去 PyTorch 探尋比對。

如何管理本地訓練程式？
- Horovod 通過後臺 Driver 程式來管理本地訓練程式。
- TE 通過後臺 Agent 程式來管理本地訓練程式。
如何儲存狀態？
- Horovod 提供了內建實現，在每次訓練間隙，使用 state.commit() 完成checkpoint。
- TE 需要用自己實現儲存/載入 checkpoint。
如何發現新節點？
- Horovod 讓使用者自己實現節點發現的邏輯，這需要使用者提供一個discovery_hosts.sh，其中指定了正在參與訓練的節點。Horovod 會定期執行這個指令碼來發現當前節點。
- TE 利用分散式一致性中介軟體 ETCD 或者自帶的 C10D後端（基於TcpStore）來解決節點之間互相發現的問題。
如何捕獲異常？
- Horovod 捕獲集合通訊異常/節點異常/擴縮容，轉換為Horovod自己的Exception，然後會依據配置重（比如內部建立異常節點黑名單）新建立環，繼續訓練。
- TE定義了一個monitor方法，定時呼叫來監控本地程式異常，轉換為內部狀態數值，進行處理，如果有一個worker出現了問題，則該node上的agent會重啟本node的所有worker進行新一輪rendezvous，因為是新一輪 rendezvous，所以其他節點也會重啟其worker，然後大家一起繼續訓練。