近兩年對
AI
技術架構影響最大的就是大模型。在大模型產生、迭代和演進的過程中，它對底層的基礎設施提出了新的挑戰。
今天的分享主要分爲四個部分，第一部分是從業務視角來介紹大模型帶來的關鍵變化。第二部分分享是在大模型時代，大模型訓練對基礎設施提出的挑戰以及百度智能雲的應對方法。第三部分是結合大模型和平臺建設的需求，講解百度智能雲所做的軟硬結合的聯合優化。第四部分則是百度智能雲對大模型未來的發展、對基礎設施新要求的一些思考。
1. GPT-3 開啓大模型時代
大模型時代是由 GPT-3 開啓的，該模型有如下幾個特點：
第一個特點是該模型的參數有極大的提升，單個模型達到了 1750 億的參數，這也帶來了準確性的大幅提升。從左邊圖中我們可以看到，隨着模型參數越來越多，模型的準確性也在不斷提升。
右邊圖則展示了它更令人震驚的特點：基於預訓練好的 1750 億參數的模型，只需要通過少量樣本的訓練，就可以接近 BERT 使用大樣本訓練後的效果。這在一定程度上反映了模型的規模變大，就可以帶來模型性能和模型通用性上的提升。
除此之外，GPT-3 還在數學計算、閱讀理解、多輪問答等任務上均表現出了一定的通用型，僅通過少量樣本就可以讓模型達到較高的、甚至近似於人類的準確度。
正因此，大模型也給 AI 整體研發模式帶來了新的變化。在未來，我們可能會先預訓練一個大模型，再針對具體的任務，通過少量的樣本做 Fine-tune，就可以得到不錯的訓練效果。而不需要像現在訓練模型時，每個任務都需要從零開始，完整地做迭代和訓練。
百度很早就開始訓練大模型，擁有 2600 億參數的文心大模型就是在 2021 年發佈的。現在隨着Stable Diffusion、AIGC 文生圖，以及最近大火的聊天機器人 ChatGPT 等受到整個社會的關注，讓大家真正意識到大模型時代到了。
各廠商也在佈局大模型相關的產品，Google 此前剛發佈了 Bard，百度也即將在 3 月份推出“文心一言”。
大模型有一個 Scaling Law，如左邊圖所示，隨着模型的參數規模和訓練數據越來越多，效果就會越來越好。
但這裏有個前提，參數一定要足夠大，並且數據集足以支撐整個參數的訓練。而這帶來後果就是計算量呈指數上升。對於一個普通的小模型來說，單機單卡就可以搞定。但對於大模型來說，它對訓練量的要求就需要有大規模的資源去支撐它的訓練。
以 GPT-3 爲例，它是一個有 1750 億參數的模型，需要 3000 億條詞語的訓練才能支撐它達到一個不錯的效果。它的計算量在論文裏面評估是 314 ZFLOPs 。對比 NVIDIA GPU A100，一張卡還是隻是 312 TFLOPS 的 AI 算力，這中間絕對數值相差了 9 個數量級。
所以爲了更好地支持大模型的的計算、訓練和演進，如何去設計、開發基礎設施就成了非常重要的問題。
2. 大模型基礎設施全景圖
這是百度智能雲面向大模型的基礎設施全景圖。這是一個覆蓋了從框架到集羣，軟硬結合的全棧基礎設施。
在大模型中基礎設施不再僅僅涵蓋底層的硬件、網絡等傳統的基礎設施。我們還需要把所有的相關資源都納入到基礎設施的範疇之中。
具體來說，它大概分爲幾個層級：
-
最上面的是模型層，包括內外部發布的模型和一些配套組件。比如百度的飛槳 PaddlePaddle 和 Fleet ，Fleet 是在飛槳上做的分佈式策略。同時在開源社區，比如 PyTorch 有 DeepSpeed/Megatron 等一些基於 PyTorch 框架做的大模型訓練的框架和加速能力。
-
在框架之下我們還會做加速庫的相關能力，包括 AI 算子加速、通信加速等。
-
-
最底下是硬件資源，比如單機單卡、異構芯片、網絡相關的能力。
以上就是整個基礎設施的全景。今天我會先重點從 AI 框架入手，接着再延伸到加速層和硬件層，分享百度智能雲的一些具體工作。
對於傳統的在單卡上進行的小模型訓練，我們使用訓練數據不停做前向、反向梯度更新就可以完成整個訓練。
而大模型的訓練，主要存在兩點挑戰：算力牆和存儲牆。
算力牆指的是如果要完成 GPT-3 這種需要 314 ZFLOPs 算力的模型訓練，而單卡只有 312 TFLOPS 算力時，我們如何通過分佈式的方法去打破單算力時間太長的問題。畢竟一張卡訓練一個模型耗時 32 年，明顯是不可行的。
其次是存儲牆。對大模型來說這是更大的挑戰。當單卡放不下模型的時候，模型一定要有一些切分的方法。
舉個例子，千億級別大模型的存儲（包括參數、訓練中間值等）需要 2 TB 的存儲空間，而單卡的顯存目前最大隻有 80 GB。因此，大模型的訓練需要一些切分策略以解決單卡放不下的問題。
一個很簡單也是大家最熟悉的想法就是數據並行，將訓練樣本切到不同的卡上。整體來說，我們現在觀察到的大模型訓練過程中，主要採用的是同步的數據更新策略。
重點來說說存儲牆問題的解決。其中關鍵就是模型切分的策略和方法。
我們用下圖的例子來做一個說明。對於一個模型來說，它是由很多層組成的，訓練時先做前向，再做反向。比如圖中 0、1、2 三層在一張卡上放不下，我們按層進行切分後，會把這個模型中的某些層放到第一張卡上。比如下圖中，綠色區域代表 GPU 0，紅色區域代表 GPU 1。我們可以將前幾層放到 GPU 0 上，其他幾層放到 GPU 1 上。數據在流動的時候，會先在 GPU 0 上做兩次前向，然後在 GPU 1 上做一次前向和一次反向，再返回到 GPU 0 上做兩次反向。由於這個過程特別像流水線，所以我們把它叫做流水線並行。
但流水線並行有個主要的問題，即流水線氣泡。因爲數據間會存在依賴關係，梯度依賴於前面層的計算，所以在數據流動的過程中，就會產生氣泡，也就造成了空等。針對這樣的問題，我們通過調整不同 mini-batch 的執行策略來減少氣泡的空等時間。
上面是算法工程師或者框架工程師的視角來看該問題。
而從基礎設施工程師的視角來看，我們更關注的是它會給基礎設施帶來何種不同的變化。
這裏重點關注它的通信語義。它在前向和反向之間，它要傳遞自己的激活值和梯度值，這就會帶來額外的
Send/Receive
操作。而
Send/Receive
是點到點的，我們在後文中會提到應對的辦法。
以上就是第一種打破存儲牆的並行模型切分的策略：流水線並行。
第二種切分方法是張量並行，
它解決的是單層參數過大的問題。
模型雖然有很多層，但其中某一層的計算量很大。此時我們期望將這一層的計算量在機器之間或者在卡間進行共同計算。從算法工程師的視角來看，解決的方法就是將不同的輸入切分成幾塊，然後分別用不同的塊分別做部分計算，最後再做聚合。
但從基礎設施工程師的視角來看，我們還是關注這個過程中引入了哪些額外的操作。剛剛提到的場景中，額外操作就是圖中的 f 和 g。f 和 g 分別代表的什麼意思？在做前向的時候， f 是一個不變的操作，通過 f 把 x 透傳過去，後面做部分的計算即可。但在最後做結果聚合的時候，還需要把整個數值繼續往下傳遞。對於這種情況，就要引入 g 的操作。 g 就是 AllReduce 的操作，它在語義上就是將兩個不同的數值做一次聚合，以保證輸出 z 在兩張卡上都能拿到相同數據。
所以從基礎設施工程師的視角來看，就會看到它引入額外的 AllReduce 的通信操作。該操作的通信量還是比較大的，因爲裏面的參數規模比較大。我們在後文中也會提到應對的方法。
第三種切分方法是分組參數切片。
該方法在一定程度上減少了數據並行中的顯存冗餘。傳統的數據運行中，每張卡都會有自己的模型參數和優化器狀態。由於它們在各自的訓練過程中還需要做同步和更新，這些狀態在不同卡上都有完整備份，
對於上述過程，其實就是同樣的數據和參數，在不同的卡上做了冗餘的存儲。因爲大模型對存儲空間要求特別高，所以這樣的冗餘存儲是不可接受的。爲了解決這個問題，我們將模型參數做了切分，每張卡上只保留一部分參數。
在真正需要計算的時候，我們用時間換空間：將其中的參數先做同步，等計算完成之後，再丟棄這些冗餘數據。通過這種方式可以進一步壓縮對顯存的需求，可以更好地在現有的機器上做訓練。
同樣，站在基礎設施工程師的視角來看，我們就需要引入 Broadcast 這樣的通信操作，通信的內容就是這些優化器的狀態與模型的參數。
除了上面提到的顯存優化的方法和策略之外，還有一種方法是減少模型的計算量。
當數據量足夠大的時候，模型的參數越多，模型的精度越好。但是隨着參數的增加，計算量也隨之增加，就需要更多的資源，與此同時計算的時間也會更長。
那如何保證在參數規模不變的情況下減少計算量呢？其中一個解決方案就是條件計算：根據一定的條件（即右圖中的 Gating 層，或者稱爲路由層），去選擇激活其中部分參數。
舉個例子，在右側圖中，我們將參數切分爲三份，樣本根據模型的條件，只激活部分的參數來在專家網絡 2 中進行計算。而專家 1 和專家 3 中的部分參數則不做計算。這樣就可以在保證參數規模的同時，減少其中的計算量。
業界基於上述方法，我們提出了混合專家的模式，
就是將模型抽象爲多個專家，每張卡處理其中不同的樣本。具體來說就是在模型層插入了一些路由的選擇，然後根據這種選擇只激活其中的一部分參數。同時，不同的卡上會保留不同專家的參數。這樣，在樣本分發的過程中就會分到不同的卡上去做計算。
但是站在基礎設施工程師的視角來看，我們發現這個過程中引入了 All2All 的操作。如下方右圖所示，在多個 Device 上，存儲了 0、1、2、3 這樣的樣本。Device 裏面的數值則表示適合被哪張卡來計算，或者說它適合被哪個專家來計算。比如 0 表示適合被 0 號專家 ，也就是 0 號卡來計算，以此類推。每張卡均會去判斷存儲的數據分別適合被誰來計算，比如 1 號卡，它判斷出部分參數適合被 0 號卡計算，另一部分參數適合被 1 號卡計算。然後下一個動作就是把樣本分發到不同的卡上去。
經過如上操作之後， 0 號卡上全爲 0 的樣本，1 號卡上全爲 1 的樣本。以上過程在通信裏就叫 All2All。這個操作從基礎設施工程師的視角來看是一個比較重的操作，我們需要在此基礎上做一些相關的優化。我們在後文中也會有進一步介紹。
在這個模式裏，我們觀察到的一個現象是，如果採用混合專家的模式，它在同樣參數的模型下，訓練的精度不如剛纔提到的各種並行策略、混合疊加策略好，大家需要根據自己的實際情況來酌情選擇。
剛剛介紹了幾種並行策略，接下來分享下百度智能雲的一個內部實戰。
我們用飛槳訓練了一個有 2600 億參數的大模型，堆疊了一些優化過的 Transformer 類的層。我們可以按照橫切豎切的方式給它做一個切分，比如豎切就是採用流水線並行策略將模型按 Transformer 層進行切分。橫切則是採用模型並行策略/張量並行策略對 Transformer 內部的 MetaMul 這種大的矩陣乘的計算做切分。同時我們再輔以數據並行的縱向優化，以及配合在數據並行上做的分組模型參數切分的顯存優化。通過如上四種方式，我們推出了飛槳的 4D 混合並行訓練的框架。
在千億參數模型的訓練配置上，我們採用機內八卡做張量並行，同時配合數據並行進行一些分組參數切分操作。同時還使用多組機器組成流水線並行，以此來承載 2600 億的模型參數。最後，再利用數據並行的方式進行分佈式計算，從而完成模型的月級別訓練。
以上就是我們整個的模型並行參數模型並行策略的一個實戰。
接下來，讓我們再回到基礎設施的視角，評估模型訓練中不同的切分策略對通信和算力的需求。
如表格所示，我們按照千億參數的規模，列出了不同切分方式所需的通信量和計算時間。從整個訓練過程來說，最好的效果是計算過程和通信過程可以完全覆蓋，或者互相重疊。
從這張表格中，我們可以推算出千億參數模型對集羣、硬件、網絡，包括整體通信模式上的需求。根據估算 1750 億參數的模型在 1024 張 A100 卡上使用 3000 億詞語進行訓練，需要 34 天就可以完成完整的端到端訓練。
在單機層面，由於需要在機內做大量的 AllReduce 和 Broadcast 操作，我們希望機內可以支持高性能，高帶寬的連接方式。所以我們在選型中採用了當時最先進的 A100 80G 的套餐，採用 8 張 A100 組成單機。
此外，在外部的網絡連接方式中，最重要的就是拓撲的連接方式。我們希望網卡和 GPU 卡之間能儘可能在同一個 PCIe Switch 下，通過對稱的方式可以更好地減少整個訓練過程中卡之間交互的吞吐瓶頸。同時也要儘量避免它們經過 CPU 的 PCIe Root Port。
首先我們先評估下需求，如果我們對業務的預期是在一個月之內完成模型的端到端訓練，就需要單作業訓練中達到千卡級別，大模型訓練集羣達到萬卡級別。因此，在網絡設計的過程中，我們應該兼顧兩個點：第一，爲了滿足流水線中的 Send/Receive 這種點到點的操作，需要減少 P2P 的延遲。第二，由於 AI 訓練中網絡側流量集中於同號卡 AllReduce 操作，我們還希望它有很高的通信吞吐。
針對這種通信需求。我們設計了右邊所示的三層 CLOS 架構的拓撲。這種拓撲跟傳統方式相比，最重要的是做了八導軌的優化，讓任一同號卡在不同機器中的通信中的跳步數儘可能少。
在 CLOS 架構中，最底下的一層是 Unit。每個 Unit 中有 20 臺機器，我們將每臺機器中的同號 GPU 卡連接到與對應編號的同一組 TOR 上。這樣單個 Unit 內的所有同號卡，只需要一跳就可以完成通信，可以很好地提升同號卡之間的通信。
但是，一個
Unit
中只有
20
臺機器共
160
張卡，這個規模是達不到大模型訓練要求的。所以我們又設計了第二層
Leaf
層。
Leaf
層將不同的 U
nit
中的同號卡連到同一組
Leaf
的交換設備上，解決的依然是同號卡互聯的問題。通過這一層，我們可以將
20
個
Unit
再做一次互聯。到這裏，我們已經可以連出
400
臺機器共
3200
卡的規模。對於這樣一個
3200
卡規模的集羣來說，任意兩張同號卡之間最多就跳到
3
跳即可實現通信。
如果我們想支持異號卡的通信怎麼辦？我們在上面又加了 Spine 層，進而解決了異號卡之間通信的問題。
通過這三層的架構，我們便實現了一個支持 3200 卡間專門針對 AllReduce 操作優化的整體架構。如果是在 IB 的網絡設備上，該架構可以支持 16000 卡的規模，這也是目前 IB 盒式組網的最大規模。
我們將 CLOS 架構跟其他的一些網絡架構做了比較，比如 Dragonfly，Torus 等。跟它們相比，這套架構的網絡帶寬更加充足、節點間的跳步數更加穩定，對於做可預期的訓練性能的估計很有幫助。
3. 軟硬件結合的聯合優化
大模型訓練並不是把硬件買好了，放在那裏就可以完成訓練。我們還需要對硬件和軟件進行聯合優化。
首先我們先來說說計算優化。大模型的訓練整體來看還是一個計算密集型的過程。在計算優化上，目前很多思路和想法都是基於靜態圖的多後端加速。用戶構建的圖，無論是飛槳、PyTorch 還是 TensorFlow，都會先通過圖捕獲把動態圖轉換成靜態圖，然後再讓靜態圖進入到後端進行加速。
下圖是我們整個基於靜態圖的多後端架構，它分爲如下幾部分。
第二個是多後端接入的方法，通過不同的後端提供基於計時的選優能力。
第三個是圖優化，我們針對靜態圖做了一些計算上的優化和圖轉換，從而進一步提升計算效率。
最後我們會通過一些自定義的算子，整體加速大模型的訓練過程。
大模型的訓練架構中，第一部分是圖接入。在 AI 框架裏描述圖的時候，通常分爲靜態圖和動態圖。
靜態圖是用戶在執行之前先把圖構造出來，再結合自己的實際輸入做執行。結合這樣的特點，在計算過程中可以提前做些編譯上的優化或者調度的優化，可以更好地提升訓練性能。
但與之對應的是動態圖的構造過程。用戶隨便寫了一些代碼，在寫的過程中就動態執行了。比如 PyTorch，在用戶寫完一條語句之後，它就會做相關的執行和求值。對用戶來說，這種方式的好處就是很容易開發和調試。但是對於執行器或者對於加速過程來說，由於每次看到都是一部分很小的操作，反而不是很好優化。
爲了解決這個問題，一般的想法就是把動態圖和靜態圖做融合，使用動態圖做開發，再通過靜態圖做執行。我們現在看到的主要有兩種實現的路徑。
第一種就是基於 Python AST 去做靜態的轉換。比如我們拿到用戶寫的 Python 源碼，將其轉換成 Python AST 樹後，再根據 AST 樹去做 CodeGen。在這個過程中使用靜態組圖的方法和 API 就可以將 Python 的動態源碼轉換成靜態圖。
但這個過程中，最大的問題是 Python 語言的靈活性，導致靜態的分析沒辦法很好地理解語義，進而出現動態圖轉靜態圖轉換失敗的情況。比如在靜態分析過程中，它沒辦法推斷動態類型，又比如靜態分析沒辦法推斷 range 的範圍，導致實際轉換過程中經常失敗。所以靜態轉換隻能適用一些簡單的模型場景。
第二條路線就是通過 Tracing 或者 Symbolic Tracing 的方法，去做簡單的執行和模擬。Tracer 記錄過程中遇到的一些計算節點，將這些計算節點記錄下來之後，再通過回放或者重組的方式去後驗構造一整個靜態圖。這種方式的好處是可以通過模擬一些輸入的方法，或者通過一些構造特殊結構的方法來做整體動態圖的捕獲和計算，進而可以較爲成功地捕獲到一條路徑。
但是這個過程其實也存在一些問題。對於依賴輸入的分支或者循環結構來說，由於 Tracer 是通過構造模擬輸入的方式來構造靜態圖，因此 Tracer 只會走到其中一些分支，從而導致安全性上存在問題。
這幾種方法比較下來之後，我們發現在 Python 現有的語言靈活性下，如果想讓動態圖完整地轉換成靜態圖，在現階段基本上是一個不太可能完成的任務。因此我們的重心就轉移到了在雲上如何給用戶提供更安全易用的圖轉化能力。現階段有如下幾種方案。
第一個方案是自研基於 AST 代碼替換的方法。該方法由百度智能雲來提供相應的模型轉換和優化的能力，對用戶而言是無感的。比如，用戶輸入一段源碼，但是其中部分代碼（圖中 XXXX、YYYY 所示）在做靜態圖捕獲、圖優化、算子優化的過程中，我們發現這部分代碼無法將動態圖轉換爲靜態圖、或者代碼有性能優化的空間。那麼我們就會寫一段替換代碼，如中間圖所示。左邊是我們認爲是可以被替換的一段 Python 代碼，右邊是我們替換後的另一段 Python 代碼。然後，我們就會通過 AST 匹配的方法，把用戶的輸入和我們的原始的目標模式做 AST 的轉換，在這上面執行我們子樹的樹匹配算法。
通過這種方式，就可以把我們原始輸入的 XXXX、YYYY 變成 WWWW、HHHH，變成一個可以更好執行的方案，一定程度上提升了動態圖轉靜態圖的成功率，同時提升了算子的性能，並可以做到用戶基本無感的效果。
第二種是社區上的一些方案，尤其是 PyTorch 2.0 提出的 TorchDynamo 方案，這也是我們目前看到的比較適合計算優化的方案。它可以實現部分的圖捕獲，不支持的結構可以 fallback 回 Python。 通過這樣的方式，它可以一定程度上把其中的部分子圖吐給後端，然後後端再針對這些子圖做進一步的計算加速。
當我們將圖整個捕獲完之後，接下來就要真正開始做計算加速了，即後端加速。
我們認爲，在 GPU 計算的時序圖中比較關鍵的其實是訪存時間和計算時間。我們從如下幾個角度去加速這個時間。
第一個是算子融合。算子融合主要的收益就是去掉 kernel launch 的時間、提升計算密度、減少額外的訪存。我們將算子單位訪存上的計算次數，定義爲計算密度。
根據計算密度的不同，我們把算子分爲計算密集型和訪存密集型兩類。舉個例子，像 GEMM就是典型的計算密集型算子，像 Elementwise 就是典型的訪存密集型的算子。我們發現，“計算密集型算子 + 訪存密集型的算子”和“訪存密集型算子 + 訪存密集型算子”之間，是可以做比較好的融合的。
我們的目標是將所有在 GPU 上執行的算子轉化爲偏計算密集性的算子，這樣可以充分利用我們的算力。
左邊是我們的一個例子，比如在 Transformer 的結構中，最重要的 Multihead Attention，是可以做一個很好的融合的。同時右邊還有我們發現一些其他模式，限於篇幅就不一一列舉了。
算子實現的本質問題是如何將計算邏輯和芯片架構相結合，從而更好地去實現整個計算過程。我們目前看到三類方案：
第一類是手寫算子。相關廠商會提供比如 cuBLAS 、cuDNN 等算子庫。它提供的的算子性能是最好的，但是它支持的操作有限，同時對定製開發的支持也比較差。
第二類是半自動化模板，比如 CUTLASS。這種方法做了開源的抽象，讓開發人員可以在上面做二次開發。這也是我們目前實現計算密集型與訪存密集型算子融合中實際使用的方法。
第三個是基於搜索的優化。我們關注了社區上一些像 Halide， TVM 的編譯方法。目前發現該方法在一些算子上是有效果的，但在另外一些算子上則還需要進一步打磨。
在實踐中這三種方法各有優勢，因此我們會通過計時選擇的方式來爲大家提供最好的實現。
第一個是交換機哈希衝突問題的解決。下圖是我們做的一個實驗，我們起了一個 32 卡的任務，每次執行 30 次 AllReduce 操作。下圖是我們測出來通信帶寬，可以看到它有很大的概率是會出現降速。這是在大模型訓練中比較嚴重的問題。
降速的原因是因爲存在哈希衝突。雖然在網絡設計上，交換機沒有收斂比，即我們網絡設計上的帶寬資源都是充足的，但是由於使用了 RoCE 這種基於以太網四元組選路的方法，仍然可能發生網絡側的流量衝突。
比如下圖中的例子，綠色的機器之間要通信，紅色的機器之間也要通信，那麼在選路的過程中，就會因爲哈希衝突導致大家的通信爭搶到了同一帶寬上，導致雖然整體的網絡帶寬是充足的，但是依然會形成局部網絡熱點，導致通信性能降速。
我們的解決方法其實也很簡單。在整個通信過程中有源 IP、 源端口、目的 IP、 目的端口這個四元組。其中源 IP、目的 IP、目的端口固定不變，而源端口是可以調整的。利用這個特點，我們通過不斷地調整源端口來選擇不同的路徑，再通過整體的貪心算法，儘量減少哈希衝突的發生。
通信優化中，除了剛纔說到在 AllReduce 上的一些優化，在 All2All 上也有一定優化空間，尤其是我們這種面向八導軌的專門定製過的網絡。
該網絡在整個 All2All 的操作上會給上層的 Spine 交換機較大的壓力。優化方法是使用了 NCCL 裏面的 Rail-Local All2All，或者 PXN 的優化。原理就是將異號卡之間的通信通過機內高性能的 NVLink 轉換爲同號卡之間的通信。
通過這種方式，我們將機間的原本要上到 Spine 層的所有的網絡通信轉換爲機內通信，這樣只需要做 TOR 層或者 Leaf 層通信就可以實現異號卡通信，在性能上也會有較大的提升。
此外，除了在 RoCE 上做的這些優化，還有一個直接就能拿到效果的就是使能 Infiniband。比如剛纔提到的交換機哈希衝突，它自己的自適應路由就可以搞定。而對於 AllReduce 來說，它還有一些高級的特性如 Sharp，可以很好地把 AllReduce 計算操作的一部分卸載到我們的網絡設備中，以釋放計算單元，同時提升計算性能。通過這樣的方法，我們可以很好地將 AllReduce 的訓練效果再次提升。
剛纔講完了計算和通信的優化，我們接下來再從端到端去看這個問題。
從整個大模型訓練來說，它其實分爲兩大部分，第一部分是模型代碼，第二部分是高性能的網絡。在這兩個不同的層級上，有一個問題亟待解決：經過多種切分策略切分後的模型放置到哪張卡上是最合適的？
我們舉個例子，在做張量並行的時候，我們需要將一個張量的計算切分爲兩塊。由於分塊計算結果之間需要做大量的 AllReduce 操作，因此需要很高的帶寬。
如果我們將一個張量並行切完的兩塊分別放到不同機器的兩張卡上，就會引入網絡通信，造成性能問題。反之，如果我們把這兩塊放到同一機器內，就可以高效完成計算任務，提升訓練效率。因此，放置問題的核心訴求，就是要找到切分的模型和異構硬件間最合適或者說性能最好的映射關係。
在我們早期的模型訓練中，是基於專家經驗知識手工完成映射。比如下圖就是我們跟業務團隊合作時，當我們認爲機內的帶寬好，就建議放到機內。如果我們認爲機間可能有改進時，就建議放到機間。
我們核心的解法就是構建計算、通信的 Cost Model，然後基於 Cost Model 做搜索優化。通過這種方式產出最優的映射。
在整個過程中會先將框架側的模型網絡做抽象和切分，映射成一個計算框架圖。同時會將單機和集羣上的計算能力、通信能力做相關建模，建立集羣的拓撲圖。
當我們有了右圖左邊的模型上的計算和通信需求，以及右圖右邊硬件上的計算和通信能力，我們就可以通過圖算法或者其他一些搜索方法進行模型的拆分和映射，最終拿到右圖下邊的一個最優解。
在實際的大模型訓練過程中，通過這種方式可以讓最後的性能提升 2.1 倍。
4. 大模型發展推動基礎設施演進
最後和大家探討下在未來，大模型會對基礎設施提出哪些新的要求。
目前我們看到的變化有三個。第一個是模型的參數，模型參數還會持續地增長，從 GPT-3 的 1750 億到 PaLM 的 5400 億。關於未來參數增長的最終值，我們可以參考下大概有 60 萬億參數規模的人腦。
第二個是多模態訓練。在未來我們會處理更多的模態數據。不同的模態數據對存儲、計算量、顯存都會帶來更多的挑戰。
第三個是異構資源。在未來我們會有越來越多的異構資源。在訓練過程中，各種類型的算力，如何更好地使用它們也是一個亟待解決的挑戰。
同時從業務角度看，在一個完整的訓練過程中可能還會有不同的類型的作業，可能同時會存在傳統的 GPT-3 訓練、強化學習訓練、和數據標註任務。如何把這些異構的任務更好地放置在我們異構的集羣上將會是一個更大的問題。
我們現在看到的幾個方法，其中一種方法是基於統一視圖的端到端的優化：將整個模型和異構資源做視圖上的統一，基於統一視圖擴展 Cost Model ，可以支持單任務多作業在異構資源集羣下的放置。結合彈性調度的能力，可以更好地感知集羣資源的變化。
上文講到的所有能力，都已經全部融合到了百度百舸的· AI 異構計算平臺中。
本文分享自微信公衆號 - 百度開發者中心（baidudev）。
如有侵權，請聯繫 [email protected] 刪除。
本文參與“OSC源創計劃”，歡迎正在閱讀的你也加入，一起分享。
