[源碼解析] 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.

  1. 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.savetorch.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,然後會依據配置重(比如內部建立異常節點黑名單)新建立環,繼續訓練。

      img

    • TE定義了一個monitor方法,定時調用來監控本地進程異常,轉換爲內部狀態數值,進行處理,如果有一個worker出現了問題,則該node上的agent會重啓本node的所有worker進行新一輪rendezvous,因爲是新一輪 rendezvous,所以其他節點也會重啓其worker,然後大家一起繼續訓練。

      elastic_monitor

4.2 TE 問題

下面是關於一些TE內部的問題,我們後續分析會逐步解答這些問題。

  • RANK 和 WORLD_SIZE 這些字段不再需要手動設置,如何做到?
  • 如何在不同的節點間確定 RANK?RANK 0 的實例會作爲 master 的角色存在?
  • worker 失敗之後,如何實現重啓worker操作?
  • TE 發現了新worker 之後,如何處理?
  • 每個代理上有一個 rendezvous,這些 rendezvous 有master,slave 概念嗎?有一個master專門記錄當前集羣狀態嘛?
  • 如何支持動態地增加或減少參與訓練的 worker 數量?

0x05 PyTorch分佈式系列

PyTorch分佈式其他文章如下:

深度學習利器之自動微分(1)

深度學習利器之自動微分(2)

[源碼解析]深度學習利器之自動微分(3) --- 示例解讀

[源碼解析]PyTorch如何實現前向傳播(1) --- 基礎類(上)

[源碼解析]PyTorch如何實現前向傳播(2) --- 基礎類(下)

[源碼解析] PyTorch如何實現前向傳播(3) --- 具體實現

[源碼解析] Pytorch 如何實現後向傳播 (1)---- 調用引擎

[源碼解析] Pytorch 如何實現後向傳播 (2)---- 引擎靜態結構

[源碼解析] Pytorch 如何實現後向傳播 (3)---- 引擎動態邏輯

[源碼解析] PyTorch 如何實現後向傳播 (4)---- 具體算法

[源碼解析] PyTorch 分佈式(1)------歷史和概述

[源碼解析] PyTorch 分佈式(2) ----- DataParallel(上)

[源碼解析] PyTorch 分佈式(3) ----- DataParallel(下)

[源碼解析] PyTorch 分佈式(4)------分佈式應用基礎概念

[源碼解析] PyTorch分佈式(5) ------ DistributedDataParallel 總述&如何使用

[源碼解析] PyTorch分佈式(6) ---DistributedDataParallel -- 初始化&store

[源碼解析] PyTorch 分佈式(7) ----- DistributedDataParallel 之進程組

[源碼解析] PyTorch 分佈式(8) -------- DistributedDataParallel之論文篇

[源碼解析] PyTorch 分佈式(9) ----- DistributedDataParallel 之初始化

[源碼解析] PyTorch 分佈式(10)------DistributedDataParallel 之 Reducer靜態架構

[源碼解析] PyTorch 分佈式(11) ----- DistributedDataParallel 之 構建Reducer和Join操作

[源碼解析] PyTorch 分佈式(12) ----- DistributedDataParallel 之 前向傳播

[源碼解析] PyTorch 分佈式(13) ----- DistributedDataParallel 之 反向傳播

[源碼解析] PyTorch 分佈式 Autograd (1) ---- 設計

[源碼解析] PyTorch 分佈式 Autograd (2) ---- RPC基礎

[源碼解析] PyTorch 分佈式 Autograd (3) ---- 上下文相關

[源碼解析] PyTorch 分佈式 Autograd (4) ---- 如何切入引擎

[源碼解析] PyTorch 分佈式 Autograd (5) ---- 引擎(上)

[源碼解析] PyTorch 分佈式 Autograd (6) ---- 引擎(下)

[源碼解析] PyTorch分佈式優化器(1)----基石篇

[源碼解析] PyTorch分佈式優化器(2)----數據並行優化器

[源碼解析] PyTorch分佈式優化器(3)---- 模型並行

[源碼解析] PyTorch 分佈式(14) --使用 Distributed Autograd 和 Distributed Optimizer

[源碼解析] PyTorch 分佈式(15) --- 使用分佈式 RPC 框架實現參數服務器

[源碼解析] PyTorch 分佈式(16) --- 使用異步執行實現批處理 RPC

[源碼解析] PyTorch 分佈式(17) --- 結合DDP和分佈式 RPC 框架

[源碼解析] PyTorch 分佈式(18) --- 使用 RPC 的分佈式管道並行

0xFF 參考

vivo AI計算平臺彈性分佈式訓練的探索和實踐

pytorch/elastic 分析

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章