一、背景介紹

在介紹自動并行之前，我們思考一下為什么需要自動并行？一方面現在有著不同的模型結構，另一方面還有各種各樣的并行策略，兩者之間一般是多對多的映射關系。假設我們能實現一個統一的模型結構滿足各種任務需求，那么我們的并行策略是不是在這種統一的模型結構上實現收斂？

答案是否定的，因為并行策略不僅僅跟模型結構相關，還跟模型的規模以及實際使用的機器資源息息相關。這就體現出自動并行的價值，它的目標是：用戶給定一個模型和所使用的機器資源后，能夠自動地幫用戶選擇一個比較好或者最優的并行策略來高效執行。

這里羅列了個人感興趣的一些工作，不一定完整，想跟大家討論一下自動并行的現狀和歷史。大概分了幾個維度：第一個維度是自動并行的程度，分為全自動和半自動；第二個維度是并行粒度，分別是針對每個 Layer 來提供并行策略，或者是針對每一個算子或者張量來提供并行策略；第三個是表示能力，這里簡化為 SPMD（Single Program Multiple Data）并行和 Pipeline 并行兩大類；第四個是特色，這里列出了個人覺得相關工作比較有特色的地方；第五個是支持硬件，主要寫出相關工作所支持最大規模的硬件類型和數量。其中，標紅部分主要是對飛槳自動并行研發有啟發性的點。

對于全自動并行來說，我們可以看到并行粒度，是由粗粒度到細粒度的發展過程；表示能力是從比較簡單的 SPMD 到非常通用的 SPMD 與 Pipeline 的方式；支持的模型是從簡單的 CNN 到 RNN 再到比較復雜的 GPT；雖然支持多機多卡，但整體規模不是特別大。

對于半自動并行來說，并行粒度基本上都是以算子為粒度的，而表示能力從簡單的  SPMD 到完備的 SPMD 加上 Pipeline 的并行策略，模型支持規模達到千億和萬億量級，所使用的硬件數量達到千卡量級。

再從框架角度來看，我們可以看到現有的框架基本上已經支持或計劃支持半自動這種模式，而并行粒度也發展到算子粒度，表示能力基本上都采用 SPMD 加上 Pipeline 的完備表示，都面向各種各樣模型和各種各樣硬件。

這里總結一下個人的一些思考：

① 第一點，分布式策略在底層表示上逐漸統一。

② 第二點，半自動會逐漸成為框架的一種分布式編程范式，而全自動會結合特定的場景和經驗規則去探索落地。

③ 第三點，實現一個極致端到端性能，需要采用并行策略與優化策略聯合調優來實現。

二、架構設計

一般完整分布式訓練包括 4 個具體的流程。首先是模型切分，無論是手動并行還是自動并行都需要將模型切分為多個可以并行的任務；其次是資源獲取，可以通過自己搭建或者從平臺申請來準備好我們訓練所需要的設備資源；然后是任務放置（或者任務映射），也就是將切分后的任務放置到對應資源上；最后是分布式執行，就是各個設備上的任務并行執行，并通過消息通信來進行同步和交互。

現在一些主流的解決方案存在一些問題：一方面可能只考慮分布式訓練中的部分流程，或者只側重部分流程；第二個就是過于依賴專家的經驗規則，比如模型切分和資源分配；最后是在整個訓練的過程中，缺乏對任務和資源的感知能力。

而飛槳所設計的端到端自適應分布式訓練架構，在全面考慮 4 個流程基礎上，又加入第五個流程，即彈性調度。我們核心設計理念主要包括這 3 點：

第一，計算和資源統一表示，并且計算和資源同等重要。往往大家比較關心怎么切分模型，但是對資源關注度比較少。我們一方面用統一的分布式計算圖來表示各種各樣的并行策略；另一方面，我們用統一的分布式資源圖來建模各種各樣的機器資源，既能表示同構的，又能表示異構的資源連接關系，還包括資源本身的計算和存儲能力。

第二，最大化解耦，除了模塊之間解耦外，我們還將邏輯切分跟物理放置以及分布式執行等進行解耦，這樣能夠更好地實現不同模型在不同集群資源高效執行。

第三，端到端自適應，涵蓋分布式訓練所涉及的全面流程，并采用一個全局的代表模型來驅動并行策略或者資源放置的自適應決策，來盡可能代替人工定制化決策。上圖淺藍色框住的部分就是本次報告所介紹的自動并行相關工作。

1、統一分布式計算圖

首先是統一的分布式計算圖。統一目的是便于我們采用統一方式來表示現有的各種各樣的并行策略，這樣利于做自動化處理。眾所周知，串行計算圖能表示各種各樣的模型，類似地，我們在串行計算圖的基礎上，對每個算子和張量加上分布式屬性來作為分布式計算圖，這種細粒度方式能表示現有并行策略，而且語義會更豐富和通用，還能表示新的并行策略。分布式計算圖中的分布式屬性主要包括三個方面信息：1）需要表示張量怎么切分或者算子怎么切分；2）需要表示在哪些資源進行分布式計算；3）如何將切分后的張量或算子映射到資源上。對比串行計算圖，分布式計算圖有 3 個基礎組成概念：分布式張量，類似于串行的張量；分布式算子，類似于串行的算子；分布式重排，分布式計算圖獨有。

（1）分布式張量

首先，介紹分布式張量所包括的三個方面信息：

① 串行張量信息：主要包含張量 shape、dtype 等一些元信息，一般實際計算不需要對串行張量進行實例化。

② ProcessMesh：進程的 cartesion topology 表示，有別于 DeviceMesh，我們之所以采用 ProcessMesh，主要希望邏輯的進程跟物理設備進行一個解耦，這樣便于做更高效任務映射。

③ ShardSpec：用來表示串行張量每個維度用 ProcessMesh 哪個維度進行切分，具體可看下圖示例。

假如有一個二維的 6*12 張量和一個 3*2 的 ProcessMesh（第一維是 x，第二維是 y，元素是進程 ID）。如果 ShardSpec 是 [None，None]，就表示張量第 0 維和第 2 維都不切分，每個進程上都有一個全量張量。如果 ShardSpec是 ['x’, ‘y’]，表示用 ProcessMesh 的 x 軸去切張量第 0 維，用 ProcessMesh 的 y 軸去切張量第 1 維，這樣每個進程都有一個 2*6 大小的 Local 張量。總之，通過 ProcessMesh 和 ShardSpec 以及張量未切分前的串行信息，就能夠表示一個張量在相關進程上切分情況。

（2）分布式算子

分布式算子表示是基于分布式張量的，包括串行算子信息，輸入和輸出張量的分布式屬性。類似一個分布式張量可能會對應多種切分方式，分布式算子里面的分布式屬性不一樣，對應著不同切分。以矩形乘 Y=X*W 算子為例，如果輸入和輸出分布式屬性不同，就對應不同的分布式算子實現（分布屬性包括 ProcessMesh 和 ShardSpec）。對于分布式算子來說，其輸入和輸出張量的 ProcessMesh 相同。

（3）分布式重排

最后一個是分布式重排，這是分布式計算圖所必須具有的概念，用來處理源張量和目的張量分布式屬性不同的情況。比如有 2 個算子的計算，上一個算子產生 y，跟下一個算子使用 y 分布式屬性不同（圖中用不同顏色表示），這時我們需要插入額外的一個 Reshard 操作來通過通信進行張量分布式重排，本質就是處理生產和消費不匹配問題。

導致不匹配的原因主要有三個方面：1）支持支持數據和計算分離，所以對張量和使用它的算子有不同的分布式屬性；2）支持用戶自定義標記分布式屬性，用戶可能對張量和使用它的算子標記不同的分布式屬性；3）分布式算子底層實現有限，如果出現了輸入或者輸出分布式屬性不支持的情況，也需要通過分布式重排。

2、統一分布式資源圖

介紹完統一的分布計算圖三個基本概念后，再看統一的分布式資源圖，主要設計的考量：1）支持異構集群，異構集群就是集群中可能有 CPU、GPU、XPU 資源；2）表示拓撲連接，這里面涵蓋了集群的層次結構連接關系，包括對連接能力的量化，比如帶寬或者延遲；3）設備本身建模，包括一個設備的存儲和計算能力。為了滿足上面設計需求，我們用Cluster來表示分布式資源，它包含多個同構 DeviceMesh。每個 DeviceMesh 內會隱含一個由 Device 鏈接組成的 Graph。

這里舉個例子，上圖可以看到有 4 臺機器，包括 2 個 GPU 機器和 2 個 XPU 機器。對于 2 臺 GPU 機器，會用一個同構的 DeviceMesh 表示，而對于 2 臺 XPU 機器，會用另一個同構的 DeviceMesh 表示。對于一個固定集群來說，它的 DeviceMesh 是固定不變的，而用戶操作的是 ProcessMesh，可以理解是 DeviceMesh 的抽象，用戶可以隨意 Reshape 和 Slice，最后會統一地將 ProcessMesh 進程映射到 DeviceMesh 設備上。

采用前面基于張量和算子細粒度的分布式計算圖表示，能涵蓋現有并行策略以及未來可能會出現新的并行策略。數據并行就是對數據張量的 Batch 維度進行切分。模型并行對權重相關維度進行切分。流水線并行使用不同 ProcessMesh 來表示，它可以表示為更靈活 Pipeline 并行，比如一個 Pipeline Stage 可以連接多個 Pipeline Stage，而且不同 Stage 使用 ProcessMesh 的 shape 可以不同。其他有些框架的流水線并行是通過 Stage Number 或者 Placement 實現，不夠靈活通用。混合并行就是數據并行，張量模型并行和流水線并行三者混合。

三、關鍵實現

前面是飛槳自動并行架構設計和一些抽象概念介紹。基于前面的基礎，下面我們通過 2 層 FC 網絡例子，來介紹飛槳自動并行內部實現流程。

上圖是飛槳整個自動并行的流程圖。首先我們會基于一個串行前向計算圖，進行反向生成，獲得包括前向、后向和更新子圖的完整計算圖。然后，需要明確組網中每個張量和每個算子的分布式屬性。既可以采用半自動的推導方式，也可以采用全自動搜索方式。本報告主要講解半自動推導方式，即基于用戶少量標記來推導其他未標記張量和算子的分布式屬性。通過分布式屬性推導后，串行計算圖中每個張量和每個算子都有自己的分布式屬性。基于分布式屬性，先通過自動切分模塊，將串行計算圖變成支持 SPMD 并行的邏輯分布式計算圖，再通過分布式重排，實現支持 Pipeline 并行的邏輯分布式計算圖。生成的邏輯分布式計算圖會通過物理映射，變成物理分布式計算圖，目前只支持一個進程和一個設備的一一映射。最后，將物理分布式計算圖變成一個實際任務依賴圖交給異步執行器進行實際執行。

1、分布式屬性推導

分布式屬性推導就是給定計算圖中部分張量和算子的分布式屬性，自動補全其他所有的張量和算子的分布式屬性。例子是兩個 Matmul 計算，用戶只標記了兩個參數分布式屬性，表示 W1 在 0，1 進程上進行列切，W2 是在 2，3 進程上進行行切，這里有兩個不同 ProcessMesh，用不同的顏色表示。

分布式屬性推導分為兩個步驟：1）先進行 ProcessMesh 傳導，實現 Pipeline 切分；2）再進行 ShardSpec 傳導，實現一個 Stage 內的 SPMD 切分。ProcessMesh 推導利用了飛槳線性 Program lR， 按靜態 Program Order 采用就近選擇策略進行推導，支持包含計算，即如果兩個 ProcessMesh，一個大一個小，就選較大作為最終 ProcessMesh。而 ShardSpec 推導利用飛槳 SSA Graph IR 進行前向和后向數據流分析進行推導，之所以可以用數據流分析，是因為 ShardSpec 語義，滿足數據流分析的 Semilattice 性質。數流分析理論上能保證收斂，通過結合前向和后向分析，能夠將計算圖任何一個位置標記信息傳播到整張計算圖，而不是只能單方向傳播。

基于分布式屬性推導，串行計算圖中的每個張量和算子都擁有自己的分布式屬性，這樣就可以基于分布式屬性進行計算圖的自動切分。按照例子來說，就是把單機串行計算圖，變成 Rank0、Rank1，Rank2、Rank3 四個計算圖。 

簡單來說，會遍歷每個算子，先對算子輸入和輸出進行張量切分，然后再對每個算子進行計算切分。張量切分會通過 Distributed Tensor 對象來構造 Local Tensor 對象，而算子切分會通過 Distributed Operator 對象來基于實際輸入和輸出的分布屬性來選擇對應分布式實現，類似于單機框架的算子到 Kernel 的分發過程。

通過前面自動切分，只能獲得支持 SPMD 并行的分布式計算圖。為了支持 Pipeline 并行，還需要通過分布式重排來處理，這樣通過插入一個合適的 Reshard 操作，例子中每個 Rank 擁有自己真正獨立計算圖。雖然左圖 Rank0 的 Y 跟 Rank2 的 Y，切分一樣，但是由于他們在不同 ProcessMesh 上，導致了生產消費分布式屬性不匹配，所以也需要插入 Reshard。

飛槳目前支持兩類分布式重排。第一類，是比較常見源張量分布跟目標張量分布都在同一個 ProcessMesh 上，但是源張量分布和目標張量分布所使用切分方式不一樣（即 ShardSpec 不一樣）。第二類，是源張量分布和目標張量分布在不同的 ProcessMesh 上，而且 ProcessMesh 大小可以不一樣，比如圖中情形 2 中的 0-5 進程和 6-9 進程。為了盡可能減少通信，飛槳也對 Reshard 操作進行相關優化。

經過分布式重排后，得到一個邏輯上的分布式計算圖，這個時候還沒有決定進程和具體設備映射。基于邏輯的分布式計算圖和前面統一的資源表示圖，會進行物理映射操作，也就是 Rank Mapping，就是從多種映射方案（一個進程具體跟哪個設備進行映射）中找到一個最優的映射方案。 

這里介紹一個比較簡單基于貪心規則的實現。先構建進程和進程間通信量的鄰接表，邊表示通信量，節點表示設備需求，再構建設備與設備之間的鄰接表，邊表示通信帶寬，節點表示設備計算和內存。我們會任選一個進程 R 放置到滿足需求的設備 D 上，放置后，選擇與 R 通信量最大的進程放到 D 所在機器其他設備上，按照這種方式直到完成所有進程映射。在映射過程中，需要判斷所選擇的設備與進程圖所需求的設備類型以及所需要計算量和內存是否匹配。

通過物理映射后，我們會根據所獲得的物理分布式組網構建實際任務依賴圖。圖中示例是基于計算圖的前向、后向和更新角色來構建任務依賴圖，相同角色的算子會組成一個任務。為了支持 Micro-batching 優化，一個任務依賴圖會生成多個任務實例依賴圖，每個示例雖然計算邏輯一樣，但是使用不同內存。目前飛槳會自動地根據計算圖角色去構建任務圖，但是用戶可以根據合適粒度自定義任務構建。每個進程有了任務多實例依賴圖后就會基于 Actor 模式進行異步執行，通過消息驅動方式就可以自動地實現 1F1B 執行調度。

基于上面整個流程，我們已經實現一個功能比較完備的自動并行。但只有并行策略還無法獲得一個比較好的端到端性能，所以我們還需要加入相應的優化策略。對飛槳自動并行，我們會在一個自動切分之前和組網切分之后加一些優化策略，這是因為一些優化在串行邏輯上實現比較自然，有一些優化在切分后比較容易實現，通過統一的優化 Pass 管理機制，我們能夠保證飛槳自動并行中并行策略與優化策略自由結合。

四、應用實踐

下面介紹應用實踐。

首先是接口，不管如何實現，用戶最終還是通過接口來使用我們所提供的自動并行能力。如果對用戶分布式需求進行拆解，包括模型組網切分，資源表示，數據分布式加載，分布式執行過程控制和分布式保存和恢復等。為了滿足這些需求，我們提供了一個 Engine 類，它同時兼顧易用性和靈活性。易用性方面，它提供了高階 API，能支持自定義 Callback，分布式過程對用戶透明。靈活性方面，它提供了低階 API，包括分布式 dataloader 構建，自動并行切圖和執行等接口，用戶可以進行更細粒度控制。兩者會共享 shard_tensor、shard_op 以及 save 和 load 等接口。

這里有兩個標記接口 shard_op 和 shard_tensor。其中，shard_op 既可以標記單個算子，也可以對整個 Module 進行標記，是函數式。上圖是一個非常簡單的使用示例。首先，使用飛槳已有 API 進行一個串行組網，在其中，我們會使用 shard_tensor 或者  shard_op 進行非侵入式分布式屬性標記。然后，構建一個自動并行 Engine，傳入模型相關信息和配置。這個時候用戶有兩個選擇，一個選擇是直接使用 fit /evaluate/predict 高階接口，另一種選擇是使用 dataloader+prepare+run 接口。如果選擇 fit 接口，用戶只需要傳 Dataset，框架會自動進行分布式數據集加載，自動并行過程編譯和分布式訓練執行。如果選擇 dataloader+prepare+run 接口，用戶可以將分布式數據加載、自動并行編譯和分布式執行進行解耦，能更好進行單步調試。

PaddleFleetX 是具備易用性和高性能的一個端到端的一站式大模型套件，支持自動并行功能。用戶要想調用飛槳自動并行端到端功能，只需要提供串行動態圖模型組網即可。在獲得用戶動態圖串行組網后，內部實現會利用飛槳動轉靜模塊，將動態圖單卡組網變成靜態圖單卡組網，然后通過自動并行編譯，最后再進行分布式訓練。在推理生成時候，使用的機器資源可能與訓練時使用的資源不太一樣，內部實現還會進行參數和組網的自適應參數切分。目前 PaddleFleetX 中的自動并行涵蓋了大家常用的并行策略和優化策略，并支持兩者任意組合，而且對生成任務來說，還支持 While 控制流的自動切分。

五、總結展望

飛槳自動并行還有很多工作在開展中，目前的特色可以總結為以下幾方面：

首先，統一的分布式計算圖，能夠支持完備的 SPMD 和 Pipeline 的分布式策略，能夠支持存儲和計算的分離式表示；

第二，統一的分布式資源圖，能夠支持異構資源的建模和表示；

第三，支持并行策略和優化策略的有機結合；

第四，提供了比較完備的接口體系；

最后，作為關鍵組成，支撐飛槳端到端的自適應分布式架構。

并行一般可以分為兩個領域（沒有明確分界），一個是傳統的分布式計算，一個是傳統高性能計算，兩者各有優劣。基于傳統分布式計算的框架代表是 TensorFlow，它側重 MPMD（Multiple Program-Multiple Data）并行模式，能夠很好支持彈性和容錯，分布式計算的用戶體驗會好一些，編程較簡單，用戶一般是以一個串行全局視角進行編程；基于傳統高性能的計算的框架代表是 PyTorch，更側重 SPMD（Single Program-Multiple Data）模式，追求極致性能，用戶需要直接面臨物理集群進行編程，自己負責切分模型和插入合適通信，對用戶要求較高。而自動并行或者自適應分布式計算，可以看作兩者結合。當然不同架構設計側重點不一樣，需要根據實際需求進行權衡，我們希望飛槳自適應架構能夠兼顧兩個領域的優勢。