今天我爲大家介紹飛槳大模型分佈式訓練技術，內容分爲以下幾個部分：
首先，我會介紹大模型訓練面臨的重點難題；然後，爲大家介紹飛槳在大模型訓練領域的特色分佈式訓練技術和優化方案；最後，伴隨着代碼示例和大模型套件，爲大家展示如何使用飛槳框架訓練大模型。
1. 背景與挑戰
近年來，大模型由於良好的模型效果和廣闊的應用前景，逐漸受到業界的廣泛重視。主流的 AI 科技公司也站在了大模型研究的前沿，模型參數量的規模呈現快速增長的趨勢。從 2018 年 1 億參數規模的模型增長至今已達千億參數量規模。
大模型的出現給模型訓練帶來極大的挑戰。即使使用 A800、H800 這樣的 GPU，單張 GPU 的算力和顯存都是遠遠無法滿足大模型訓練需求的。爲了保證大模型可訓練，並提高整體訓練吞吐，需要用到模型並行 + 數據並行等技術。
這張圖展示的是大模型分佈式訓練技術的發展歷程。
對於十億及以下的模型，單卡往往就能放下全量模型參數和梯度，傳統的數據並行即可即可覆蓋其應用場景。當模型規模到了百億量級以後，需要使用分組參數切片的方式將模型參數、梯度和優化器狀態切分到各個卡上，保證單機可放下。當模型規模到了千億以後，則需要同時使用模型並行、數據並行等多種並行技術混合進行高效訓練。在這個階段裏，分佈式並行技術從單一的基礎並行策略演進爲多種並行策略的組合。
當模型規模到了萬億級別以後，稠密模型已經難以高效訓練，從而衍生出稀疏專家模型，也伴隨着 MoE 等並行策略的演進與迭代。
下面，我將爲大家介紹飛槳在大模型訓練領域的特色分佈式訓練技術。
2. 飛槳大模型特色分佈式訓練技術
飛槳在 2021 年提出了業界首個 4D 混合並行訓練方案，將數據並行、張量模型並行、流水線並行、分組參數切片並行有機結合起來，實現百億、千億參數稠密大模型的高效訓練。
在最內層，我們進行張量模型並行以及 Sharding 的參數的切分，並使用流水線並行將張量模型並行和 Sharding 並行串聯起來。在最外層，使用數據並行技術來提高整體訓練吞吐。這種技術應用在了文心 ERNIE 千億模型訓練上，相較於 3D 混合並行方案獲得了 24%~44% 的提速效果。
而在稀疏專家模型領域，飛槳在 2022 年提出了高效可擴展的 MoE 訓練架構，並結合 4D 混合並行技術，解決稀疏專家模型訓練效率低下的難題。
首先，MoE 訓練面臨的一個重要問題是巨大的參數量，參數量達到萬億級別，但是這些參數是高度稀疏的，MoE 的路由模塊每次只會激活部分參數。針對這個問題，飛槳提出了通用多級參數存儲的方案，將短期不使用到的參數先 offload 到 CPU 內存、甚至 offload 到 SSD 中去；當即將需要使用時，再從 CPU 內存或 SSD 中預先加載到顯存中，保證模型訓練和參數預加載 overlap 起來。
其次，MoE 訓練面臨的另一個重要問題是分佈式通信慢，在 MoE 的路由模塊裏存在 All-to-All 通信，嚴重影響模型訓練效率。飛槳提出了硬件拓撲自動感知技術，感知當前集羣的拓撲信息，對 All-to-All 通信進行分層處理，提高分佈式通信效率。這些技術最終應用在了百度在 2022 年發佈的 VIMER-UFO 視覺大模型中，性能較 PyTorch 有 66% 的提升，達到業內領先水平。
這頁 PPT 展示的是蛋白質結構預測領域飛槳的混合並行加速方案。蛋白質結構預測模型中存在以下幾個重點難題。其一，蛋白質結構預測中往往會引入 MSA 的維度，MSA 維度往往會特別大，導致單卡無法放下全量模型。其二，在模型訓練過程中，模型中有多個分支結構，在傳統的方案裏這幾個分支結構是串行執行的，計算效率較低。其三，模型訓練過程中會產生許多很小的張量，數千個小張量的計算和通信都會造成較大的性能瓶頸。
爲了解決這些問題，飛槳創新提出了分支並行方案，讓模型中的多個分支結構運行在不同設備上，並在必要時候進行聚合。與此同時，我們融合了分支並行、數據並行和動態軸並行形成了 3D 混合並行方案，並使用梯度融合的方式解決數千個小張量計算和通信慢的問題，相比於 AlphaFold 2 獲得了 106% 的性能提升。同時，我們使用該方案在千卡國產硬件上訓練蛋白質預測模型，整體訓練時長從 7 天縮短到了 2.6 天，極大提高了訓練效率。
這頁 PPT 展示的是飛槳在超長序列訓練領域的解決方案。目前許多大模型採用的是 Transformer 的結構，Transformer 結構裏的 Self Attention 模塊的顯存佔用是與序列長度的平方成正比的，這給超長序列訓練帶來極大的挑戰。針對這個問題，飛槳集合了 FlashAttention 方案，可將 Attention 模塊的顯存佔用從平方複雜度降低到線性複雜度。
但是，即使將 Attention 模塊的顯存佔用從平方複雜度降低到線性複雜度，整個網絡裏 Activation 的顯存佔用仍然與序列長度成正比，當序列長度達到 16K、32K 時會面臨巨大的顯存壓力。而在多種並行策略中，基本只有張量模型並行才能節省 Activation 的顯存佔用，但當張量模型並行大約 8 後，就會面臨跨機的問題，跨機張量模型並行會造成性能的嚴重下滑。
針對這個問題，飛槳獨創了 Segment Parallel 的策略切分序列維度，且可以與數據並行、張量模型並行、序列並行、流水線並行等方案自由組合，達到近線性擴展。
以上介紹的是在某些特定場景下的分佈式並行策略。
但是，實際大模型訓練過程是非常複雜的，如何根據當前模型結構和集羣環境找出最佳的分佈式並行策略是十分困難的。針對這個問題，飛槳提出了端到端自適應分佈式訓練框架，我們通過建立統一的邏輯計算視圖表徵大模型的模型結構和算力要求，建立統一的異構資源視圖表徵算力環境和網絡環境，結合分佈式代價模型，自動搜索出適合當前模型和資源狀況的最優並行策略，並進行資源的彈性調度。
以上介紹的是飛槳特色大模型分佈式並行策略，下面我爲大家介紹飛槳的大模型端到端性能優化方案。
3. 飛槳大模型端到端性能優化方案
下圖展示的是飛槳的大模型端到端性能優化方案，涵蓋數據讀取、模型實現、高性能算子實現和分佈式策略優化等模型訓練全鏈條。
在原始數據讀取方面，飛槳提供了多進程數據讀取 + 每個進程獨享數據隊列的方式加速數據讀取過程，防止 Python 端 GIL 鎖以及隊列資源搶佔導致數據加載低效等問題。
在模型訓練過程中，往往還需要對原始數據進程預處理。飛槳在 CV 領域提供了 FlyCV 庫，提供高效的圖像預處理能力；在 NLP 領域提供了 FastTokenizer，性能達到業內領先水平。
這裏是我們針對變長輸入模型數據預處理的優化技術。在變長輸入下，各設備間的計算負載是不均等的，需要對輸入數據進行負載均衡處理保證計算的均衡性。飛槳提出了 GPU + CPU 混合處理的方式對輸入數據進行數據交換，實現負載均衡，且讓負載均衡過程與模型訓練過程完全 overlap，提高大模型整體訓練速度。
在模型實現層面，飛槳也做了許多優化工作。例如，針對變長輸入模型，飛槳提出了 FlashAttention + 分組多流的優化思路，在 FlashAttention 的基礎上對變長輸入進行分組併發計算，提高整體訓練吞吐。除此以外，我們還針對 Transformer 類模型進行了多粒度的選擇性重計算策略，用戶可以根據當前顯存佔用情況選取最佳的重計算策略，達到計算性能和顯存佔用之間的平衡。
在高性能算子優化方面，針對 Transformer 類大模型，飛槳實現了全環節的算子融合策略，涵蓋 Attention、Linear、FFN、Optimizer 等模型的每一個環節，發揮極致的硬件性能。
在分佈式策略優化方面，除了常規的反向計算與梯度通信的 overlap 外，飛槳還創新提出了跨 step 的前向計算與參數廣播的 overlap，讓上一個 step 的參數廣播與下一個 step 的前向計算併發執行，在 GPT 類大模型上可獲得 11% 左右的提速。
結合上述在大模型上的性能優化工作，我們使用飛槳框架在國際權威 AI 基準測試 MLPerf Training v2.0 和 MLPerf Training v2.1 上提交了 BERT 模型的單機和多機性能數據。在同等 GPU 配置下，飛槳框架性能排名世界第一。
同樣地，我們在 GPT 等典型大模型上也是達到了業內領先的效果。下圖展示的 GPT 百億模型 Megatron-LM 和飛槳的性能對比結果。飛槳相較 Megatron-LM 有 9%~11% 的提速。
以上介紹的是飛槳在大模型分佈式訓練上的特色技術。下面我將用幾個實踐案例，爲大家介紹如何使用飛槳框架來訓練大模型。
4. 飛槳大模型應用實踐
首先介紹的是飛槳的大模型套件 PaddleFleetX。PaddleFleetX 的定位是端到端大模型套件，它打通了大模型預訓練-有監督精調-壓縮-推理等開發部署全流程，降低大模型的入門門檻。這裏貼的是 PaddleFleetX 的 github 地址，感興趣的聽衆可以嘗試使用。（https://github.com/PaddlePaddle/PaddleFleetX）
下面我首先介紹下如何使用飛槳實現混合並行的模型。首先介紹的是張量模型並行的例子。在 GPT 的 FFN 結構裏有 2 個 Linear 層。變成張量模型並行的代碼，只需要把這兩個 Linear 層對應地替換成飛槳的張量模型並行 API 即可。與此同時，在模型組網結束後，調用 distributed_model、distributed_optimizer 等 API 將單機模型轉換爲分佈式模型即可。
這裏介紹的是一個流水線模型並行的代碼例子。與其他並行方式不同，流水線是做的層間的切分，不同設備分到不同的層，因此組網寫法上會有一定的特殊性。首先，我們需要定義若干的 LayerDesc，每個 LayerDesc 對應網絡中的每一層。然後，同樣地，將由 LayerDesc 組出來的網絡調用 distributed_model、distributed_optimizer 等 API 將單機模型轉換爲分佈式模型即可。最後調用 train_batch 這個 API 進行流水線並行調度。
上述講的 2 個例子偏手動並行，即用戶組網的時候需要感知這是個分佈式程序。這裏給大家展示的是一個單卡組網 + 簡單的分佈式標記即可自動轉換爲分佈式程序的例子。
這裏涉及的是飛槳的自動並行技術。用戶在組網的時候寫的是單卡的程序，通過調用自動並行的 shard_tensor 等 API 標記部分 tensor 應做什麼樣的分佈式切分。最後調用 Engine 接口，內部會自動將用戶的單卡組網轉換爲分佈式組網，並進行訓練。
以上是我今天分享的全部內容。
本文分享自微信公衆號 - 百度開發者中心（baidudev）。
如有侵權，請聯繫 [email protected] 刪除。
本文參與“OSC源創計劃”，歡迎正在閱讀的你也加入，一起分享。