近日，阿裡雲機器學習PAI關于深度學習模型高效的分布式訓練架構的論文《 Whale: Efficient Giant Model Training over Heterogeneous GPUs 》被計算機系統領域國際頂級學術會議USENIX ATC'22接收。
Whale是阿裡雲機器學習PAI平台自研的分布式訓練架構，開源後的名稱是EPL(Easy Parallel Library)，Whale通過對不同并行化政策進行統一抽象、封裝，在一套分布式訓練架構中支援多種并行政策，并進行顯存、計算、通信等全方位的優化，來提供易用、高效的分布式訓練架構。Whale提供簡潔易用的接口，使用者隻需添加幾行annotation即可組合各種混合并行政策。同時Whale提供了基于硬體感覺的自動化分布式并行政策，感覺不同硬體的算力、存儲等資源，根據這些資源來合理的切分模型，均衡不同硬體上的計算量，最大化計算效率。借助Whale的硬體感覺負載均衡算法，Bert-Large、Resnet50和GNMT模型 在異構GPU訓練上提速1.2至1.4倍。同時，使用Whale架構， 萬億M6模型使用480 張 V100在3天内完成預訓練。相比此前業界訓練同等規模的模型，節省算力資源超80%，且訓練效率提升近11倍。進一步使用512 GPU在10天内即訓練出具有可用水準的10萬億模型。
背景和挑戰
随着近些年深度學習的火爆，模型的參數規模也增長迅速，OpenAI資料顯示：
- 2012年以前，模型計算耗時每2年增長一倍，和摩爾定律保持一緻；
- 2012年後，模型計算耗時每3.4個月翻一倍，遠超硬體發展速度；
特别最近兩年模型參數規模飛速增長，谷歌、英偉達、阿裡、智源研究院都釋出了萬億參數模型，有大廠也釋出了百億、千億參數模型。随着模型參數規模增大，模型效果也在逐漸提高，但同時也為訓練架構帶來更大的挑戰。目前已經有一些分布式訓練架構，例如：Horovod、Tensorflow Estimator、PyTorch DDP等支援資料并行，Gpipe、PipeDream、PipeMare等支援流水并行，Mesh Tensorflow、FlexFlow、OneFlow、MindSpore等支援算子拆分，但當我們要訓練一個超大規模的模型時會面臨一些挑戰：
- 目前的分布式訓練架構隻支援少量的并行政策，缺乏一個統一的抽象來支援所有的并行政策及其混合政策。
- 實作複雜的并行政策需要大量的模型代碼改造和對底層系統的了解，大大增加了使用者的使用難度。
- 由于叢集中異構GPU計算能力和顯存的差異，靜态的并行化政策無法充分利用異構資源實作高效訓練。
破局
為了應對目前分布式訓練的挑戰，我們研發了高效、通用、硬體感覺的分布式訓練架構Whale。Whale抽象并定義了兩個分布式原語(replicate和split) 。使用者可以通過在模型上添加幾行原語标記，即來表達和轉換各種并行化政策及其組合，極大降低了分布式架構的使用門檻。Whale runtime将使用者的标記資訊融合到計算圖中，自動完成模型的并行化優化。同時Whale提供了基于硬體感覺的自動化分布式并行政策，優化在異構GPU叢集上分布式訓練性能。Whale的設計很好地平衡了模型使用者的幹預和系統優化機會，讓每一個算法工程師都能輕松高效訓練分布式大模型任務。
技術架構
Whale架構如下圖所示，主要分為以下幾個子產品：
- 接口層：使用者的模型程式設計接口基于TensorFlow，同時Whale提供了易用的并行化政策表達接口，讓使用者可以組合使用各種混合并行政策；
- 中間表達層：将使用者模型和并行政策轉成化内部表達，通過TaskGraph、VirtualDevices和政策抽象來表達各種并行政策；
- 并行化引擎層：基于中間表達，Whale會對計算圖做政策探索，結合硬體資源進行顯存/計算/通信優化，并自動生成分布式計算圖。
- Runtime執行引擎：将分布式執行圖轉成TFGraph，再調用TF 的Runtime來執行；
并行化政策表達
Whale通過strategy annotation的方式來劃分模型為多個TaskGraph，并在此基礎上進行并行化。 Whale有兩類strategy：replicate 和 split。通過這兩種并行化接口，可以表達出各種不同的并行化政策，例如：
- 資料并行: 下面這個例子是一個資料并行的例子，每個模型副本用一張卡來計算。如果使用者申請了8張卡，就是一個并行度為8的資料并行任務。
- 流水并行：在下面的例子裡，模型被切分成2個 TaskGraph, "stage0"和"stage1"，使用者可以通過配置pipeline.num_micro_batch參數來設定pipeline的micro batch數量。 在這個例子裡，"stage_0"和"stage_1"組成一個模型副本，共需要2張GPU卡。如果使用者申請了8張卡，Whale會自動在pipeline外嵌套一層并行度為4的資料并行（4個pipeline副本并行執行）。
- 算子拆分并行：在以下例子中，Whale會對split scope下的模型定義做拆分，并放置在不同的GPU卡上做并行計算。
- 同時Whale支援對上述并行政策進行組合和嵌套，來組成各種混合并行政策，更多示例可以參考開源代碼的文檔和示例。
Parallel Planner
Paraller Planner是Whale runtime的重要一環，它的職責是生成一個高效的分布式執行plan。Parallel Planner的流程包含 (a) Paraller Planner的輸入包含使用者模型、使用者标記(可選)、硬體計算資源和其他使用者配置配置。 (b) 将實體計算資源映射成VirtualDevice, 使用者無需擔心如何将算子放置在分布式實體裝置上。(c) 模型被分割成TaskGraph子圖。因為Whale允許對不同的TaskGraph應用不同的分布式政策，是以在TaskGraph之間可能存在輸入/輸出shape不比對。在這種情況下，Paraller Planner将自動在兩個TaskGraphs之間插入相應的橋接層。
硬體感覺的負載均衡算法
當模型訓練資源包含異構硬體（比如混合了V100和T4），硬體感覺的負載均衡算法可以提高在異構資源下的訓練效率。Whale設計了兩種平衡政策：Intra-TaskGraph和Inter-TaskGraph的平衡。
（1）對于Intra-TaskGraph的計算平衡，Whale會profile出模型的算力FLOP，并按照機器的計算能力按比例配置設定對應的計算負載，以此均衡每個step不同卡型上的模型計算時間。對于資料并行，我們可以通過調整不同副本上的batch大小來實作計算負載的均衡（保持全局batch不變）。對于算子拆分，我們可以通過不均勻的次元拆分來實作不同卡上子子產品的計算負載均衡。
（2）當多個TaskGraph在多個異構GPU上執行流水并行，我們需要Inter-TaskGraph的計算平衡來提升整體的計算效率。由于流水并行的計算特點，前面的TaskGraph會比後面的TaskGraph多緩存一些前向計算的結果，是以對于Transformer類模型均衡layer切分的情況下，前面的TaskGraph會有更大的顯存壓力。是以在TaskGraph放置的時候，Whale會優先将前面的TaskGraph放置在顯存容量更大的GPU卡上。與此同時，Whale會按照GPU卡的算力去切分模型，使得切分後的TaskGraph計算load和GPU算力成正比。
應用示例
借助Whale架構，我們4行代碼實作M6模型資料并行+專家并行的混合并行訓練。
如下圖所示，MoElayer采用專家并行，其他layer采用資料并行：
并首次在480 V100 上，3天内完成萬億M6模型的預訓練。相比此前業界訓練同等規模的模型，此次僅使用480張V100 32G GPU就成功訓練出萬億模型M6，節省算力資源超80%，且訓練效率提升近11倍。進一步使用512 GPU在10天内即訓練出具有可用水準的10萬億模型。
結語
Whale 通過統一的抽象支援各種并行化政策，并提出簡潔的并行化原語來提高架構的易用性。同時，Whale提供了異構硬體感覺的自動分布式計算圖優化，實作高效的分布式模型訓練。我們希望Whale能夠成為一個大規模深度學習訓練的基石，進一步推進模型算法創新和系統優化，使大模型訓練技術能夠廣泛應用在實際生産環境中。Whale已經開源(https://github.com/alibaba/EasyParallelLibrary)，歡迎大家來試用和共建。
- 論文名稱：Whale: Efficient Giant Model Training over Heterogeneous GPUs
- 論文作者：賈賢豔，江樂，王昂，肖文聰，石子骥，張傑，李昕元，陳浪石，李永，鄭祯，劉小勇，林偉
- 論文連結：https://arxiv.org/pdf/2011.09208.pdf
- 開源連結：https://github.com/alibaba/easyparallellibrary
參考文獻：
[1] Junyang Lin, An Yang, Jinze Bai, Chang Zhou, Le Jiang, Xianyan Jia, Ang Wang, Jie Zhang, Yong Li, Wei Lin, et al. M6-10t: A sharing-delinking paradigm for effi- cient multi-trillion parameter pretraining. arXiv preprint arXiv:2110.03888, 2021.