在談及當(dāng)前的 AI 系統(tǒng)所面臨的問題時(shí)，低效是經(jīng)常被提及的一個(gè)。

谷歌人工智能主管 Jeff Dean 曾在一篇博文中寫道，「今天的人工智能系統(tǒng)總是從頭開始學(xué)習(xí)新問題 —— 數(shù)學(xué)模型的參數(shù)從隨機(jī)數(shù)開始。就像每次學(xué)習(xí)一項(xiàng)新技能（例如跳繩），你總會(huì)忘記之前所學(xué)的一切，包括如何平衡、如何跳躍、如何協(xié)調(diào)手的運(yùn)動(dòng)等，然后從無到有重新學(xué)習(xí)。這或多或少是我們今天訓(xùn)練大多數(shù)機(jī)器學(xué)習(xí)模型的方式：我們不是擴(kuò)展現(xiàn)有模型來學(xué)習(xí)新任務(wù)，而是從無到有訓(xùn)練新模型來做一件事（或者我們有時(shí)將通用模型專門用于特定任務(wù)）。結(jié)果是我們最終為數(shù)千個(gè)單獨(dú)的任務(wù)開發(fā)了數(shù)千個(gè)模型。以這種方式學(xué)習(xí)每項(xiàng)新任務(wù)不僅需要更長的時(shí)間，而且還需要更多的數(shù)據(jù)?！?/p>

為了改變這種局面，Jeff Dean 等人去年提出了一種名叫「Pathways」的通用 AI 架構(gòu)。他介紹說，Pathways 旨在用一個(gè)架構(gòu)同時(shí)處理多項(xiàng)任務(wù)，并且擁有快速學(xué)習(xí)新任務(wù)、更好地理解世界的能力。

該架構(gòu)的特點(diǎn)可以概括為：

• 能夠訓(xùn)練一個(gè)模型來做成千上萬件事情；
• 當(dāng)前模型只注重一種感官，Pathways 可做到多種；
• 當(dāng)前模型密集且效率低下，Pathways 會(huì)把模型變得稀疏而高效。

在發(fā)布想法大半年之后，Jeff Dean 終于公布了 Pathways 的論文，其中包含很多技術(shù)細(xì)節(jié)。

論文鏈接 https://arxiv.org/pdf/2203.12533.pdf

論文寫道，PATHWAYS 使用了異步算子的一個(gè)分片數(shù)據(jù)流圖（sharded dataflow graph），這些算子消耗并產(chǎn)生 futures，并在數(shù)千個(gè)加速器上高效地對(duì)異構(gòu)并行計(jì)算進(jìn)行 gang-schedule，同時(shí)在它們專用的 interconnect 上協(xié)調(diào)數(shù)據(jù)傳輸。PATHWAYS 使用了一種新的異步分布式數(shù)據(jù)流設(shè)計(jì)，它允許控制平面并行執(zhí)行，盡管數(shù)據(jù)平面中存在依賴關(guān)系。這種設(shè)計(jì)允許 PATHWAYS 采用單控制器模型，從而更容易表達(dá)復(fù)雜的新并行模式。

實(shí)驗(yàn)結(jié)果表明，當(dāng)在 2048 個(gè) TPU 上運(yùn)行 SPMD（single program multiple data）計(jì)算時(shí)，PATHWAYS 的性能（加速器利用率接近 100%）可以媲美 SOTA 系統(tǒng)，同時(shí)吞吐量可媲美跨越 16 個(gè) stage 或者被分割成兩個(gè)通過數(shù)據(jù)中心網(wǎng)絡(luò)連接的加速器島的 Transformer 模型的 SPMD 案例。

以下是論文細(xì)節(jié)：

研究背景

在過去的十年里，深度學(xué)習(xí)在圖像理解、自然語言處理等多個(gè)領(lǐng)域取得了顯著的進(jìn)展，這是 ML 模型、加速器硬件以及將兩者聯(lián)系在一起的軟件系統(tǒng)協(xié)同進(jìn)化的結(jié)果。這種協(xié)同進(jìn)化帶來的隱患是：深度學(xué)習(xí)系統(tǒng)可能過度專注于當(dāng)前的工作負(fù)載，無法預(yù)測未來的需求。

PATHWAYS 是一個(gè)為分布式 ML 構(gòu)建的新系統(tǒng)，劍指未來 ML 工作負(fù)載將需要的特定能力。當(dāng)前，這些工作負(fù)載缺乏 SOTA 系統(tǒng)的支持。

例如，當(dāng)今 SOTA ML 工作負(fù)載大多使用單程序多數(shù)據(jù)（SPMD）模型，該模型受到了 MPI 的啟發(fā)，其中所有加速器都在同步運(yùn)行相同的計(jì)算，加速器之間的通信由 AllReduce 等集體來描述。

但近年來，研究人員開始在 ML 計(jì)算中被 SPMD 掣肘。大型語言模型已經(jīng)使用流水線并行而不是純粹的數(shù)據(jù)并行來擴(kuò)展；混合專家（MoE）等模型已經(jīng)開始探索計(jì)算稀疏性，其最自然的表達(dá)方式是使用細(xì)粒度控制流和跨加速器的異構(gòu)計(jì)算；系統(tǒng)設(shè)計(jì)者們已經(jīng)開始采用巧妙的技術(shù)來在 MPI 風(fēng)格的系統(tǒng)上執(zhí)行流水線（pipelined）、同構(gòu) MoE 模型，但是，MPI 編程模型對(duì)于用戶和底層系統(tǒng)來說都太受限制了。

另一方面，隨著一代又一代新加速器的出現(xiàn)，ML 中的異構(gòu)環(huán)境變得越來越普遍。提供對(duì)（通過高帶寬 interconnect 連接的）同構(gòu)加速器的大「島」的獨(dú)占訪問是昂貴的，并且通常會(huì)造成浪費(fèi)，因?yàn)閱蝹€(gè)用戶程序必須努力保持所有加速器持續(xù)忙碌。這種限制進(jìn)一步推動(dòng)研究人員走向「多程序多數(shù)據(jù)」（MPMD）計(jì)算。MPMD 通過將整個(gè)計(jì)算的子部分映射到一組更容易獲得的小加速器島上來實(shí)現(xiàn)更大的靈活性。為了提高利用率，一些 ML 硬件資源管理研究人員以細(xì)粒度的方式在工作負(fù)載之間復(fù)用硬件，實(shí)現(xiàn)工作負(fù)載彈性，并提高容錯(cuò)能力。

最后，研究人員開始標(biāo)準(zhǔn)化一套基礎(chǔ)模型（foundation model），這些模型是在大數(shù)據(jù)上大規(guī)模訓(xùn)練的，可以適應(yīng)多種下游任務(wù)。通過在許多任務(wù)之間復(fù)用資源，并在它們之間有效地共享狀態(tài)，這種模型的訓(xùn)練和推理提供了提高集群利用率的機(jī)會(huì)。例如，幾個(gè)研究人員可能同時(shí)微調(diào)用于不同任務(wù)的一個(gè)基礎(chǔ)模型，使用相同的加速器來保持固定的基礎(chǔ)模型層。在共享的子模型上進(jìn)行的訓(xùn)練或推理可以受益于一些技術(shù)，這些技術(shù)允許來自不同任務(wù)的示例被組合在一個(gè) vectorized batch 中，以獲得更高的加速器利用率。

本文提出的 PATHWAYS 在功能和性能上可以媲美 SOTA ML 系統(tǒng)，同時(shí)提供了支持未來 ML 工作負(fù)載所需的能力。它使用了一個(gè) client-server 架構(gòu)，該架構(gòu)使得 PATHWAYS 的運(yùn)行時(shí)能夠代表許多 client 在系統(tǒng)管理計(jì)算島上執(zhí)行程序。

PATHWAYS 是第一個(gè)旨在透明、高效地執(zhí)行跨多個(gè) TPU pods 的程序的系統(tǒng)。通過采用新的數(shù)據(jù)流執(zhí)行模型，它可以擴(kuò)展到數(shù)千個(gè)加速器。PATHWAYS 的編程模型使得表達(dá)非 SPMD 計(jì)算變得很容易，并支持集中的資源管理和虛擬化，以提高加速器的利用率。

PATHWAYS 系統(tǒng)架構(gòu)

PATHWAYS 構(gòu)建在先前的系統(tǒng)的基礎(chǔ)上，包括用于表征和執(zhí)行 TPU 計(jì)算的 XLA (TensorFlow, 2019)、用于表征和執(zhí)行分布式 CPU 計(jì)算的 TensorFlow 圖和執(zhí)行器 (Abadi et al., 2016)，以及包括 JAX (Bradbury et al., 2016) 在內(nèi)的 Python 編程框架 (Bradbury et al., 2018) 和 TensorFlow API。利用這些構(gòu)建塊，PATHWAYS 在兼顧協(xié)調(diào)性的同時(shí)，僅用最少的代碼更改就能運(yùn)行現(xiàn)有的 ML 模型。

資源管理器

PATHWAYS 的后端由一組加速器組成，這些加速器組合成緊密耦合的 island，這些 island 又通過 DCN 相互連接，如上圖 3 所示。PATHWAYS 有一個(gè)「資源管理器」，負(fù)責(zé)集中管理所有 island 上的設(shè)備。client 可能會(huì)要求 island 的「虛擬 slice」具有適合其通信模式的特定 2D 或 3D 網(wǎng)格形狀。每個(gè)虛擬 slice 都包含「虛擬設(shè)備」，允許 client 表達(dá)計(jì)算在網(wǎng)格上的布局方式。資源管理器為滿足所需互連拓?fù)?、?nèi)存容量等的虛擬設(shè)備動(dòng)態(tài)分配物理設(shè)備。

最初的資源管理器使用一個(gè)簡單的啟發(fā)式方法來實(shí)現(xiàn)，嘗試通過在所有可用設(shè)備上傳播計(jì)算來靜態(tài)平衡負(fù)載，并在虛擬設(shè)備和物理設(shè)備之間保持一對(duì)一的映射。如果未來的工作負(fù)載需要，則可以采用更加復(fù)雜的分配算法，例如考慮所有 client 計(jì)算的資源需求和系統(tǒng)的當(dāng)前狀態(tài)，以近似計(jì)算物理設(shè)備的最佳分配。

PATHWAYS 允許動(dòng)態(tài)添加和移除后端計(jì)算資源，由資源管理器跟蹤可用設(shè)備。由單控制器設(shè)計(jì)啟用的虛擬設(shè)備和物理設(shè)備之間的間接層將允許未來支持透明的掛起 / 恢復(fù)和遷移等功能，其中 client 的虛擬設(shè)備可以臨時(shí)回收資源或重新分配而無需用戶程序的協(xié)助。

client

當(dāng)用戶想要運(yùn)行一個(gè)被跟蹤的程序時(shí)，可以調(diào)用 PATHWAYS client 庫，它首先將虛擬設(shè)備分配給之前沒有運(yùn)行過的任何計(jì)算，并用資源管理器注冊(cè)計(jì)算，觸發(fā) server 在后臺(tái)編譯計(jì)算。

然后，client 為程序構(gòu)建與設(shè)備位置無關(guān)的 PATHWAYS 中間表征 (IR)，表示為自定義 MLIR (Lattner et al., 2021) dialect。IR 通過一系列標(biāo)準(zhǔn)編譯器 pass 逐漸降低級(jí)別，最終輸出包含物理設(shè)備位置的低級(jí)表征。這種低級(jí)程序考慮了物理設(shè)備之間的網(wǎng)絡(luò)連接，并包含將輸出從源計(jì)算分片傳輸?shù)狡淠繕?biāo)分片（shard）位置的操作，包括需要數(shù)據(jù)交換時(shí)的分散和收集操作。在虛擬設(shè)備位置不變的通常情況下重復(fù)運(yùn)行低級(jí)程序是有效的，如果資源管理器改變了虛擬設(shè)備和物理設(shè)備之間的映射關(guān)系，可以 re-low 程序。

較舊的單控制器系統(tǒng)中的 client 可能很快成為性能瓶頸，因?yàn)樗?fù)責(zé)協(xié)調(diào)數(shù)千個(gè)單獨(dú)的計(jì)算，還要協(xié)調(diào)分布在數(shù)千個(gè)加速器中的計(jì)算分片相應(yīng)的數(shù)據(jù)緩沖區(qū)。PATHWAYS client 使用分片緩沖區(qū)抽象來表征可能分布在多個(gè)設(shè)備上的邏輯緩沖區(qū)。這種抽象通過以邏輯緩沖區(qū)而不是單個(gè)分片的粒度分?jǐn)?bookkeeping 任務(wù)（包括參考計(jì)數(shù)（reference counting））的成本來幫助 client 擴(kuò)展。

協(xié)調(diào)實(shí)現(xiàn)

 PATHWAYS 依賴 PLAQUE 完成所有使用 DCN 的跨主機(jī)協(xié)調(diào)。PLAQUE 是一種現(xiàn)有的（閉源）生產(chǎn)分片數(shù)據(jù)流系統(tǒng)，谷歌將它用于許多面向客戶的服務(wù)，這些服務(wù)需要高扇出或高扇入通信，并且可擴(kuò)展性和延遲都很重要。低級(jí) PATHWAYS IR 直接被轉(zhuǎn)換為 PLAQUE 程序，并表征為數(shù)據(jù)流圖。PATHWAYS 對(duì)其協(xié)調(diào) substrate 有嚴(yán)格的要求，而 PLAQUE 滿足所有要求。

首先，用于描述 PATHWAYS IR 的表征必須包含每個(gè)分片計(jì)算的單個(gè)節(jié)點(diǎn)，以確保能夠緊湊表征跨多個(gè)分片的計(jì)算，即帶有 N 個(gè)計(jì)算分片的 2 個(gè)計(jì)算 A 和 B 的鏈?zhǔn)綀?zhí)行，無論 N 是多少，每個(gè)計(jì)算分片在數(shù)據(jù)流表征中都有 4 個(gè)節(jié)點(diǎn)：Arg → Compute (A) → Compute (B) → Result。在 PLAQUE 運(yùn)行時(shí)實(shí)現(xiàn)中，每個(gè)節(jié)點(diǎn)都會(huì)生成帶有目標(biāo)分片標(biāo)記的輸出數(shù)據(jù)元組，因此在執(zhí)行數(shù)據(jù)并行執(zhí)行時(shí)，N 個(gè)數(shù)據(jù)元組將在每對(duì)相鄰的 IR 節(jié)點(diǎn)之間流動(dòng)。

協(xié)調(diào)運(yùn)行時(shí)還必須支持沿分片邊緣的稀疏數(shù)據(jù)交換，其中消息可以在動(dòng)態(tài)選擇的分片子集之間發(fā)送，使用標(biāo)準(zhǔn)的進(jìn)度跟蹤機(jī)制（Akidau et al., 2013; Murray et al., 2013）來檢測何時(shí)已收到分片的所有消息。高效的稀疏通信能夠避免 DCN 成為加速器上依賴于數(shù)據(jù)的控制流瓶頸，這是 PATHWAYS 啟用的關(guān)鍵功能之一。

如下圖 4 所示，協(xié)調(diào) substrate 用于發(fā)送傳輸調(diào)度消息和數(shù)據(jù) handle 的關(guān)鍵路徑中的 DCN 消息，因此它必須以低延遲發(fā)送關(guān)鍵消息，并在需要高吞吐量時(shí)將消息批量發(fā)送到同一個(gè) host。

使用可擴(kuò)展的通用數(shù)據(jù)流引擎來處理 DCN 通信也很方便，因?yàn)檫@意味著 PATHWAYS 還可以將其用于后臺(tái)管理任務(wù)，例如分發(fā)配置信息、監(jiān)控程序、清理程序、在出現(xiàn)故障時(shí)提示錯(cuò)誤等。

谷歌認(rèn)為，使用 Ray (Moritz et al., 2018) 等其他分布式框架而不是 PLAQUE 來重新實(shí)現(xiàn)完整的 PATHWAYS 設(shè)計(jì)以實(shí)現(xiàn)低級(jí)協(xié)調(diào)框架是可行的。在這種實(shí)現(xiàn)中，PATHWAYS 執(zhí)行器和調(diào)度器將被長期運(yùn)行的 Ray Actor 所取代，這些 Ray Actor 將在底層 Ray 集群調(diào)度之上實(shí)現(xiàn) PATHWAYS 調(diào)度，并且執(zhí)行器可以使用 PyTorch 進(jìn)行 GPU 計(jì)算和集合。

Gang-scheduled 動(dòng)態(tài)調(diào)度

如前所述，在一組共享加速器上支持 SPMD 計(jì)算的一個(gè)要求是支持高效的 gang-scheduling。

PATHWAYS 運(yùn)行時(shí)包括每個(gè) island 的集中式調(diào)度器，它對(duì) island 上所有計(jì)算進(jìn)行一致性排序。當(dāng) PATHWAYS 將一個(gè)程序加入隊(duì)列以執(zhí)行時(shí)，PLAQUE 數(shù)據(jù)流程序負(fù)責(zé)以下操作：

• 在每個(gè)加速器上將本地編譯函數(shù)執(zhí)行加入隊(duì)列，并將緩沖 future 作為輸入； 
• 將網(wǎng)絡(luò)發(fā)送（network sends）加入到遠(yuǎn)程加速器的隊(duì)列，以獲得函數(shù)執(zhí)行輸出的緩沖 future；
• 與調(diào)度器通信，以確定在 island 上運(yùn)行的所有程序中函數(shù)執(zhí)行的一致順序。 

調(diào)度器必須實(shí)施以毫秒為單位分配加速器的策略。不過，當(dāng)前的實(shí)現(xiàn)只是按照 FIFO 順序?qū)⒐ぷ骷尤腙?duì)列，但更復(fù)雜的調(diào)度器可能會(huì)根據(jù)估計(jì)的執(zhí)行時(shí)間重新排序計(jì)算。

并行異步調(diào)度

當(dāng)在加速器上運(yùn)行計(jì)算時(shí)，系統(tǒng)可以利用異步 API 將計(jì)算與協(xié)調(diào)重疊。如下圖 4a 中的三節(jié)點(diǎn)圖所示，正方形分別對(duì)應(yīng)三個(gè)節(jié)點(diǎn) A、B 和 C，它們?cè)谶B接到主機(jī) A、B 和 C 的加速器上運(yùn)行。所有節(jié)點(diǎn)計(jì)算都是常規(guī)編譯函數(shù)。主機(jī) A 將節(jié)點(diǎn) A 加入隊(duì)列，接收 A 輸出的 future 并將它傳輸給主機(jī) B。主機(jī) B 分類節(jié)點(diǎn) B 的輸入，將該輸入緩沖地址傳輸給節(jié)點(diǎn) A，并執(zhí)行大部分準(zhǔn)備工作以啟動(dòng)節(jié)點(diǎn) B 的功能。當(dāng)節(jié)點(diǎn) A 完成時(shí)，它的輸出直接通過加速器互聯(lián)發(fā)送至節(jié)點(diǎn) B 的輸入緩沖，然后主機(jī) B 啟動(dòng)節(jié)點(diǎn) B。一個(gè)節(jié)點(diǎn)完成和另一個(gè)節(jié)點(diǎn)啟動(dòng)之間的延遲時(shí)間要比數(shù)據(jù)傳輸時(shí)間更長。

當(dāng) predecessor 節(jié)點(diǎn)的計(jì)算時(shí)間超過主機(jī)之間調(diào)度、資源分類和協(xié)同所用時(shí)間時(shí)，上述設(shè)計(jì)運(yùn)行良好。但如果計(jì)算時(shí)間太短，異步 pipeline 就會(huì)停止，主機(jī)端的工作成為執(zhí)行整個(gè)計(jì)算序列過程中的關(guān)鍵瓶頸?？紤]到編譯的函數(shù)都是常規(guī)的，后續(xù)節(jié)點(diǎn)的輸入形狀實(shí)際上可以在 predecessor 計(jì)算加入隊(duì)列之前進(jìn)行計(jì)算。

因此，谷歌引入了一種全新的并行異步調(diào)度設(shè)計(jì)方案，具體如下圖 4 b 所示。該方案利用常規(guī)編譯函數(shù)的靜態(tài)已知資源來并行運(yùn)行計(jì)算節(jié)點(diǎn)的主機(jī)端工作，而不是在 predecessor 已經(jīng)加入隊(duì)列之后對(duì)節(jié)點(diǎn)工作進(jìn)行序列化處理?？紤]到常規(guī)函數(shù)下只能并行地調(diào)度工作，PATHWAYS 將并行調(diào)度作為一種優(yōu)化手段，并在節(jié)點(diǎn)資源需求在 predecessor 計(jì)算完成時(shí)才知道的情況下回退到傳統(tǒng)模型。

當(dāng)計(jì)算的子圖可以進(jìn)行靜態(tài)調(diào)度時(shí)，該程序會(huì)向調(diào)度器發(fā)送描述整個(gè)子圖的單條消息，該調(diào)度器能夠?qū)ψ訄D中所有活動(dòng)分片的執(zhí)行進(jìn)行背靠背排序。設(shè)計(jì)單條消息旨在最小化網(wǎng)絡(luò)流量，但不需要調(diào)度器將所有子圖的分片作為一個(gè)批次來加入隊(duì)列：計(jì)算仍可能與其他并發(fā)執(zhí)行程序提交的計(jì)算交錯(cuò)。

三節(jié)點(diǎn)程序的順序調(diào)度（a）與并行調(diào)度（b）比較。

實(shí)驗(yàn)結(jié)果

谷歌展示了 PATHWAYS 在訓(xùn)練真實(shí)機(jī)器學(xué)習(xí)模型（它們可以被表示為 SPMD 程序）中的性能。首先與使用編碼器 - 解碼器架構(gòu)運(yùn)行 Transformer 模型的 JAX 多控制器進(jìn)行比較。

下表 1 展示了在不同數(shù)量的加速器上訓(xùn)練時(shí)，不同大小的文本到文本 Transformer 模型的訓(xùn)練吞吐量（tokens / 秒）。正如所預(yù)期的一樣，由于模型代碼相同，在 JAX 和 PATHWAYS 上訓(xùn)練的模型在步數(shù)相同的情況下實(shí)現(xiàn)了相同的困惑度。

接著，谷歌比較了當(dāng)僅用解碼器架構(gòu)訓(xùn)練 Transformer 語言模型時(shí)，PATHWAYS 在不同配置上的性能。如表 2 所示，PATHWAYS 的訓(xùn)練吞吐量與每個(gè) pipeline 階段的 TPU 核心數(shù)量成比例增加，這與其他系統(tǒng)保持一致。

上述結(jié)果與下圖 5 一致， 表明 PATHWAYS 的吞吐量與主機(jī)數(shù)量呈線性縮放關(guān)系。增加 pipeline 階段的數(shù)量會(huì)提高最小開銷，當(dāng)階段數(shù)量從 4 增加到 16 時(shí)，吞吐量從 133.7k tokens / 秒減少到 131.4k tokens / 秒。谷歌將 pipelined 模型的性能與使用 SPMD 的等效模型進(jìn)行了比較，并觀察到至少在這種情況下，pipeline 的性能與 SPMD 相當(dāng)，這是因?yàn)?SPMD 計(jì)算內(nèi)部聚合通信產(chǎn)生的開銷比 pipeline 泡沫（bubble）開銷更高。

? ?