作者：王林、颯洋

導讀

最近阿里雲機器學習PAI平臺和達摩院智慧計算實驗室一起釋出“低碳版”巨模型M6-10T，模型引數已經從萬億躍遷到10萬億，規模遠超業界此前釋出的萬億級模型，成為當前全球最大的AI預訓練模型。同時做到了業內極致的低碳高效，使用512 GPU在10天內即訓練出具有可用水平的10萬億模型。相比之前釋出的大模型GPT-3，M6實現同等引數規模，能耗僅為其1%。M6模型訓練使用的正是阿里雲機器學習PAI平臺自研的分散式訓練框架EPL(Easy Parallel Library，原名whale)。EPL透過對不同並行化策略進行統一抽象、封裝，在一套分散式訓練框架中支援多種並行策略，並進行視訊記憶體、計算、通訊等全方位的最佳化，來提供易用、高效的分散式訓練框架。

EPL背後的技術框架是如何設計的？作為開發者我們可以怎麼使用EPL？未來有哪些規劃？今天一起來深入瞭解。

EPL是什麼

EPL(Easy Parallel Library)是阿里最近開源的，統一多種並行策略的、靈活易用的自研分散式深度學習訓練框架。

專案背景

隨著近些年深度學習的火爆，模型的引數規模也增長迅速，OpenAI資料顯示：

• 2012年以前，模型計算耗時每2年增長一倍，和摩爾定律保持一致；
• 2012年後，模型計算耗時每3.4個月翻一倍，遠超硬體發展速度；

特別最近一年模型引數規模飛速增長，谷歌、英偉達、阿里、智源研究院都發布了萬億引數模型，有大廠也釋出了百億、千億引數模型。隨著模型引數規模增大，模型效果也在逐步提高，但同時也為訓練框架帶來更大的挑戰。當前已經有一些分散式訓練框架，例如：Horovod、Tensorflow Estimator、PyTorch DDP等支援資料並行，Gpipe、PipeDream、PipeMare等支援流水並行，Mesh Tensorflow、FlexFlow、OneFlow、MindSpore等支援運算元拆分，但當我們要訓練一個超大規模的模型時會面臨一些挑戰：

• 如何簡潔易用：
• 
  
  • 接入門檻高：使用者實現模型分散式版本難度大、成本高，需要有領域專家經驗才能實現高效的分散式並行策略；
  • 最優策略難：隨著研究人員設計出越來越靈活的模型，以及越來越多的並行加速方法，如果沒有自動並行策略探索支援，使用者很難找到最適合自身的並行策略；
  • 遷移代價大：不同模型適合不同的混合並行策略，但切換並行策略時可能需要切換不同的框架，遷移成本高；
• 如何提高價效比：
• 
  
  • 業界訓練萬億規模模型需要的資源：英偉達 3072 A100、谷歌 2048 TPU v3，資源成本非常高；
  • 如何降本增效，組合使用各種技術和方法來減少需要的資源，提高訓練的速度；

為了應對當前分散式訓練的挑戰，我們研發了分散式訓練框架EPL，將不同並行化策略進行統一抽象、封裝，在一套分散式訓練框架中支援多種並行策略。EPL同時提供簡潔易用的介面，使用者只需新增幾行annotation即可完成並行策略的配置，模型程式碼不需要改動。在使用者無感的情況下，EPL透過進行各種視訊記憶體、計算、通訊最佳化，打造高效的分散式訓練框架。

主要特性

• 多種並行策略統一：在一套分散式訓練框架中支援多種並行策略（資料/流水/運算元/專家並行）和其各種組合、巢狀使用；
• 介面靈活易用：使用者只需新增幾行程式碼就可以使用EPL豐富的分散式並行策略，模型程式碼無需修改；
• 自動並行策略探索：運算元拆分時自動探索拆分策略，流水並行時自動探索模型切分策略；
• 分散式效能更優：提供了多維度的視訊記憶體最佳化、計算最佳化，同時結合模型結構和網路拓撲進行排程和通訊最佳化，提供高效的分散式訓練；

開源地址

EPL(Easy Parallel Library)的開源地址是：

https://github.com/alibaba/EasyParallelLibrary

我們同時提供了model zoo，歡迎大家試用：

https://github.com/alibaba/FastNN

EPL主要技術特點

EPL透過豐富並行化策略，簡單易用的介面，多維度的視訊記憶體最佳化技術，和最佳化的計算通訊加速技術，讓每一個演算法工程師都能輕鬆訓練分散式大模型任務。

• 豐富的並行化策略：EPL提供了多種並行化策略及其組合策略，包含資料並行，流水並行，運算元拆分並行及並行策略的組合巢狀。豐富的策略選擇使得不同的模型結構都能找到最適合自己的分散式訓練方式。
• 易用性：使用者的模型程式設計介面和訓練介面均基於TensorFlow，使用者只需在已有的單機單卡模型上做簡單的標記即可實現不同的分散式策略。EPL設計了兩種簡單的策略介面(replicate/split)來表達分散式策略及混合並行。分散式策略標記的方式讓使用者無需學習新的模型程式設計介面，僅需幾行程式碼即可實現和轉換分散式策略，極大降低了分散式框架的使用門檻。
• 視訊記憶體最佳化：EPL提供了多維度的視訊記憶體最佳化技術，包含自動重算技術(Gradient Checkpoint)，ZeRO資料並行視訊記憶體最佳化技術，CPU Offload技術等，幫助使用者用更少的資源訓練更大的模型。
• 通訊最佳化技術：EPL深度最佳化了分散式通訊庫，包括硬體拓撲感知，通訊執行緒池，梯度分組融合，混合精度通訊、梯度壓縮等技術。

技術架構

EPL框架如下圖所示，主要分為以下幾個模組：

• 介面層：使用者的模型程式設計介面基於TensorFlow，同時EPL提供了易用的並行化策略表達介面，讓使用者可以組合使用各種混合並行策略；
• 中間表達層：將使用者模型和並行策略轉成化內部表達，透過TaskGraph、VirtualDevices和策略抽象來表達各種並行策略；
• 並行化引擎層：基於中間表達，EPL會對計算圖做策略探索，進行視訊記憶體/計算/通訊最佳化，並自動生成分散式計算圖。
• Runtime執行引擎：將分散式執行圖轉成TFGraph，再呼叫TF 的Runtime來執行；

並行化策略表達

EPL透過strategy annotation的方式來劃分模型為多個TaskGraph，並在此基礎上進行並行化。 EPL有兩類strategy：replicate 和 split。透過這兩種並行化介面，可以表達出各種不同的並行化策略，例如：

1. 資料並行: 下面這個例子是一個資料並行的例子，每個模型副本用一張卡來計算。如果使用者申請了8張卡，就是一個並行度為8的資料並行任務。

2. 流水並行：在下面的例子裡，模型被切分成2個 TaskGraph, "stage0"和"stage1"，使用者可以透過配置pipeline.num_micro_batch引數來設定pipeline的micro batch數量。 在這個例子裡，"stage_0"和"stage_1"組成一個模型副本，共需要2張GPU卡。如果使用者申請了8張卡，EPL會自動在pipeline外巢狀一層並行度為4的資料並行（4個pipeline副本並行執行）。

3. 運算元拆分並行：在以下例子中，EPL會對split scope下的模型定義做拆分，並放置在不同的GPU卡上做平行計算。

4. 同時EPL支援對上述並行策略進行組合和巢狀，來組成各種混合並行策略，更多示例可以參考開原始碼的文件和示例。

視訊記憶體最佳化

當模型增長，GPU的視訊記憶體常常成為訓練大模型的瓶頸。EPL提供了多維度的視訊記憶體最佳化技術，極致最佳化了訓練視訊記憶體消化。

• 重算 Recomputation (Gradient Checkpoint)：正常的DNN前向過程中會生成activation，這部分activation會在後向過程中用於梯度的計算。因此在梯度生成之前，前向的activation會一直存留在視訊記憶體中。activation大小和模型結構以及batch size相關，通常佔比都非常高。Gradient Checkpoint (GC) 透過保留前向傳播過程中的部分activation，在反向傳播中重算被釋放的activation，用時間換空間。GC中比較重要的一部分是如何選擇合適的checkpoint點，在節省視訊記憶體，保證效能的同時，又不會影響收斂性。EPL提供了自動GC功能，讓使用者可以一鍵開啟GC最佳化功能。
• ZeRO：在資料並行的場景下，每個卡上會存放一個模型副本，optimizer state等，這些資訊在每張卡上都是一樣，存在很大的冗餘量。當模型變大，很容易超出單卡的視訊記憶體限制。在分散式場景下，我們可以透過類似DeepSpeed ZeRO的思路，將optimizer state和gradient分片存在不同的卡上，從而減少單卡的persistent memory佔用。
• 視訊記憶體最佳化的AMP(Auto Mixed Precision)：在常規的AMP裡，需要維護一個FP16的weight buffer，對於引數量比較大的模型，也是不小的開銷。EPL提供了一個視訊記憶體最佳化的AMP版本，FP16只有在用的時候才cast，從而節約視訊記憶體。
• Offload: Offload將訓練的儲存空間從視訊記憶體擴充套件到記憶體甚至磁碟，可以用有限的資源來訓練大模型。

同時，EPL支援各種視訊記憶體最佳化技術的組合使用，達到視訊記憶體的極致最佳化。我們在T5模型上開啟GC+ZeRO+視訊記憶體最佳化的AMP技術，在效能保持不變的情況下，視訊記憶體降低2.6倍。

應用場景

EPL適合不同場景的模型，在阿里內部已經支援影像、推薦、語音、影片、自然語言、多模態等業務場景，同時也可以支援不同規模的模型，最大我們完成了10萬億規模的M6模型訓練，下面以M6和Bert模型為例進行介紹。

萬億/10萬億 M6模型預訓練

訓練一個萬億/10萬億引數模型的算力需求非常大，為了降低算力需求，EPL中實現了MoE(Mixture-of-Experts)結構，MoE的主要特點是稀疏啟用，使用Gating(Router)來為輸入選擇Top-k的expert進行計算（k常用取值1、2），從而大大減少算力需求。

EPL支援專家並行(Expert Parallelism, EP)，將experts拆分到多個devices上，降低單個device的視訊記憶體和算力需求。同時資料並行有利於提升訓練的併發度，因此採用資料並行+專家並行組合的混合並行策略來訓練M6模型：MoE layer採用專家並行，其他layer採用資料並行。

EPL中提供簡潔易用的介面來進行模型的混合並行訓練，只需要增加幾行annotation來配置並行策略，模型本身不需要任何修改。例如，M6模型採用資料並行+專家並行的策略，只需要增加如下圖的annotation：

同時為了節約訓練資源，提高訓練效率，我們採用了EPL的視訊記憶體最佳化技術和計算通訊加速技術，包含自動 Gradient Checkpointing節省activation視訊記憶體佔用，CPU Offload技術用於最佳化optimizer states和weight的視訊記憶體佔用，採用DP+EP混合並行策略降低算力需求，和結合混合精度、編譯最佳化等技術提高訓練效率等。

藉助EPL框架，首次在480 V100 上，3天內完成萬億M6模型的預訓練。相比此前業界訓練同等規模的模型，此次僅使用480張V100 32G GPU就成功訓練出萬億模型M6，節省算力資源超80%，且訓練效率提升近11倍。進一步使用512 GPU在10天內即訓練出具有可用水平的10萬億模型。

流水並行加速Bert Large模型訓練

對於Bert Large模型，結構圖如下圖所示：

由於Bert Large模型對視訊記憶體消耗較大，Nvidia V100 16G顯示卡上batch size常常只有2-8左右（具體值和Embedding大小、Sequence Length等有關）。Batch size太小會導致演算法收斂波動大，收斂效果差的問題。同時透過資料並行模式訓練通訊佔比較高，分散式加速效果不理想。

分析Bert Large模型，由24層重複結構的encoder組成，可以使用流水並行進行加速。這裡我們將Bert Large中的Encoder Layer 1~8層、Encoder Layer 9~16層，Encoder Layer 17~24層分別放在不同的卡上進行訓練，並行化後的計算圖如下圖所示：

此時每張卡訓練時的視訊記憶體開銷會減少，從而可以增大batch size以提升收斂加速。另外，對於模型過大單卡視訊記憶體無法放下導致無法訓練的場景，透過Layer間拆分的模型並行方式來進行分散式訓練。透過epl.replicate介面可以實現模型的stage劃分，同時透過流水並行的執行排程來提升並行化效能，如下圖所示：

上述例子是一個流水micro batch mumber為5的情況。透過流水並行最佳化後的時間軸可以看出，在同一個時間上多張卡可以平行計算。當5個micro batch結束後，每張卡會將梯度進行本地的累計之後再進行update。與單純的模型並行相比，透過流水的交替執行，提高了GPU的利用率。EPL還透過採用Backward-Preferred排程最佳化策略來提升流水並行效能，降低GPU空閒時間和視訊記憶體開銷。

為能夠獲得更高的水平擴充套件，EPL還支援在流水並行外巢狀資料並行來提升訓練吞吐。EPL會自動推導巢狀的資料並行的並行度。在32卡GPU規模下，使用EPL的流水並行+資料並行，相較於資料並行得到了66%的提升。

Roadmap

我們決定建設開源生態主要有如下的考慮：

• EPL發源於阿里雲內部的業務需求，在服務內部業務的過程中，我們積累了大量的經驗和理解，很好的支援了大規模、多樣性的業務場景，在我們自身隨著業務需求的迭代逐漸完善的同時，也希望能夠開源給社群，將自身積累的經驗和理解也同樣回饋給社群，希望和深度學習訓練框架的開發者或深度學習從業者之間有更多更好的交流和共建，為這個行業貢獻我們的技術力量。
• 我們希望能夠藉助開源的工作，收到更多真實業務場景下的使用者反饋，以幫助我們持續完善和迭代，併為後續的工作投入方向提供輸入。
• 同時我們希望藉助開源的工作，能吸引一些志同道合的同學、公司或組織來參與共建，持續完善深度學習生態。

後續我們計劃以兩個月為單位釋出Release版本。EPL近期的Roadmap如下：

• 持續的效能最佳化和穩定性改進；
• 通用運算元拆分功能；
• 自動拆分策略探索的基礎版；
• 自動流水並行策略探索；

此外，在中長期，我們在下面幾個探索性的方向上會持續投入精力，也歡迎各種維度的反饋和改進建議以及技術討論，同時我們十分歡迎和期待對開源社群建設感興趣的同行一起參與共建。

• 全自動的模型並行策略探索；
• 高效的策略探索演算法和精準的CostModel評估；
• eager model下的並行策略探索；
• 更多新硬體的支援、適配和協同最佳化；
• 高效的運算元最佳化和整合、極致的視訊記憶體最佳化、軟硬一體的通訊最佳化；

參考文獻

1. Whale: Scaling Deep Learning Model Training to the Trillions https://arxiv.org/abs/2011.09208
2. https://openai.com/blog/ai-and-compute/
3. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. https://arxiv.org/abs/1910.02054
4. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding https://arxiv.org/abs/1810.04805