飛槳框架進階3.0！一文講透“大模型訓推一體”等五大新特性

十三 2024-08-01 14:10:23 來源：量子位

深度學習框架作為基礎軟件，不僅促進了深度學習技術的飛速進步，更為人工智能技術的廣泛應用鋪設了堅實的基礎。

深度學習框架為開發(fā)者提供了便捷易用的開發(fā)接口，這些接口對數(shù)據(jù)和操作進行了高度抽象，使得開發(fā)者能夠更專注于算法和模型的設計，而不必深陷底層數(shù)據(jù)的處理細節(jié)。通過這些接口，開發(fā)者無需直接感知和應對復雜的硬件底層開發(fā)細節(jié)，從而極大地提升了開發(fā)效率和體驗。其次深度學習框架還提供了自動微分這一強大功能，開發(fā)者通常只需要編寫前向傳播網(wǎng)絡的代碼，而繁瑣的反向傳播網(wǎng)絡則交由框架自動完成。

飛槳作為中國首個自主研發(fā)、功能豐富、開源開放的深度學習平臺，從默認使用靜態(tài)圖的1.0版本，到默認采用動態(tài)圖并可實現(xiàn)動靜統(tǒng)一與訓推一體的2.0版本發(fā)布，飛槳框架已經(jīng)可以完美融合動態(tài)圖的靈活性與靜態(tài)圖的高效性，并支持模型的混合并行訓練；再到近日，為大模型時代而錘煉的3.0版本的正式出爐！飛槳正式開啟了新一代框架技術創(chuàng)新之路！

設計思想

深度學習框架的設計對于推動人工智能技術的發(fā)展至關重要，其核心設計目標是讓深度學習技術的創(chuàng)新與應用更簡單。

如何做到這一點呢？

框架需要充分考慮開發(fā)者和硬件廠商的需求。

從用戶角度出發(fā)，一個優(yōu)秀的深度學習框架應當為開發(fā)者提供極致的開發(fā)體驗。這不僅僅意味著提供一個用戶友好的開發(fā)環(huán)境，更重要的是要能夠大幅度減少開發(fā)者的學習成本和時間成本，同時顯著提升開發(fā)的便利性。為此，飛槳框架提出了“動靜統(tǒng)一、訓推一體、自動并行”的理念，極大地提高了開發(fā)效率。

從硬件適配角度出發(fā)，現(xiàn)代深度學習應用往往需要在多樣化的硬件平臺上運行，因此，框架必須能夠兼容并適配各種不同的硬件設備。這要求框架能夠智能地隔離不同硬件接口之間的差異，實現(xiàn)廣泛的硬件適配性。同時，為了充分發(fā)揮硬件的性能，框架還需要具備軟硬件協(xié)同工作的能力，確保在利用硬件資源時能夠達到最優(yōu)的性能表現(xiàn)。

與此同時，好的框架還需要考慮到 AI 技術發(fā)展的整體趨勢、產(chǎn)業(yè)的實際落地應用的需求。

技術發(fā)展層面，大語言模型（Large Language Model，簡稱LLM）、 MOE（Mixture of Experts）、多模態(tài)以及科學智能（AI for Science）等前沿技術逐漸成為新的研究熱點。隨著模型復雜性的增加，計算瓶頸、存儲瓶頸、訪存瓶頸以及通信瓶頸等問題逐漸凸顯，對分布式訓練和通用性能優(yōu)化的需求日益迫切。

在產(chǎn)業(yè)化層面，框架又需要具備支持訓練、壓縮、推理一體化的全流程能力。這意味著，從模型的訓練到優(yōu)化，再到實際部署和推理，框架應當提供一套完整、高效的解決方案，才能滿足產(chǎn)業(yè)界對于深度學習技術的實際需求。

只有跟得上趨勢、經(jīng)得住打磨的框架，才能為產(chǎn)學研各界開發(fā)者提供持續(xù)穩(wěn)定的支持。

綜上需求，飛槳將為開發(fā)者提供一個“動靜統(tǒng)一、訓推一體、自動并行、自動優(yōu)化、廣泛硬件適配”的深度學習框架，開發(fā)者可以像寫單機代碼一樣寫分布式代碼，無需感知復雜的通信和調(diào)度邏輯，即可實現(xiàn)大模型的開發(fā)；可以像寫數(shù)學公式一樣用 Python 語言寫神經(jīng)網(wǎng)絡，無需使用硬件開發(fā)語言編寫復雜的算子內(nèi)核代碼，即可實現(xiàn)高效運行。

飛槳框架 3.0 版本應運而生，延續(xù)了2.x 版本動靜統(tǒng)一、訓推一體的設計理念，其開發(fā)接口全面兼容2.x 版本。這意味著，使用2.x 版本開發(fā)的代碼，在絕大多數(shù)情況下無需修改，即可直接在3.0版本上運行。著重推出了動靜統(tǒng)一自動并行、編譯器自動優(yōu)化、大模型訓推一體、大模型多硬件適配四大新特性。這些特性在飛槳框架2.6版本或更早版本時就已經(jīng)開始開發(fā)，目前已達到外部可試用的階段。這些新特性在使用體驗、性能、二次開發(fā)便利度以及硬件適配能力等方面帶來了顯著提升，飛槳正式發(fā)布3.0版本。此版本包含了對框架2.x 版本部分已有功能的改進，并且在不使用新特性的情況下，表現(xiàn)成熟穩(wěn)定。

框架架構一覽

為了實現(xiàn)深度學習框架的上述特性，必須對框架的架構進行精心設計，確保其能夠支持各種復雜的模型構建，同時與多樣化的芯片實現(xiàn)無縫對接。接下來，將通過直觀的架構圖，詳細展示飛槳新一代框架內(nèi)所涵蓋的功能模塊，以及這些模塊之間的相互作用與聯(lián)系。以下為飛槳框架3.0的架構圖。

豐富接口：飛槳框架對外提供了豐富的深度學習相關的各種開發(fā)接口，如張量表示、數(shù)學計算、模型組網(wǎng)、優(yōu)化策略等。通過這些接口，開發(fā)者能夠便捷地構建和訓練自己的深度學習模型，無需深入到底層的技術細節(jié)中去。

在開發(fā)接口之下，飛槳框架可以劃分為4個層次：表示層、調(diào)度層、算子層和適配層。

表示層：專注于計算圖的表達與轉(zhuǎn)換，通過高可擴展中間表示 PIR，為動轉(zhuǎn)靜（動態(tài)圖轉(zhuǎn)為靜態(tài)圖）、自動微分、自動并行、算子組合以及計算圖優(yōu)化等核心功能提供堅實支撐。

調(diào)度層：負責對代碼或計算圖進行智能編排與高效調(diào)度，并且能夠根據(jù)實際需求進行顯存和內(nèi)存的管理優(yōu)化，支持動態(tài)圖和靜態(tài)圖高效執(zhí)行。無論開發(fā)者選擇使用動態(tài)圖還是靜態(tài)圖進行模型開發(fā)，飛槳框架都能提供高效的執(zhí)行環(huán)境，同時確保資源利用的最優(yōu)化。

算子層：由神經(jīng)網(wǎng)絡編譯器 CINN 和算子庫 PHI 共同構成，涵蓋了張量定義、算子定義、算子自動融合和算子內(nèi)核實現(xiàn)等關鍵功能。

適配層：則用于實現(xiàn)與底層芯片適配，包括設備管理、算子適配、通信適配以及編譯接入等功能。

下面將重點介紹飛槳3.0版本架構全新重大升級，這次升級主要包含以下模塊：

1）高擴展中間表示 PIR，通過打造全架構統(tǒng)一的中間表示，突破框架層各模塊壁壘，提升飛槳在科學計算、編譯優(yōu)化、大模型領域的潛力；

2）神經(jīng)網(wǎng)絡編譯器自動優(yōu)化，通過自動融合和策略調(diào)優(yōu)，大幅提升模型端到端表現(xiàn)；

3）自動并行，降低大模型場景模型開發(fā)和性能優(yōu)化的成本，大幅提升大模型場景的用戶體驗。

高擴展中間表示PIR

計算圖中間表示（Intermediate Representation，即 IR）是深度學習框架性能優(yōu)化、推理部署、編譯器等方向的重要基石。近些年來，越來越多的框架和研究者將編譯器技術引入到深度學習的神經(jīng)網(wǎng)絡模型優(yōu)化中，并在此基礎上借助編譯器的理念、技術和工具對神經(jīng)網(wǎng)絡進行自動優(yōu)化和代碼生成。在大模型時代，對 IR 在靈活性、擴展性、完備性有了更高的要求。

因此在3.0版本下，飛槳在基礎架構層面規(guī)范了中間表示 IR 定義，實現(xiàn)全架構統(tǒng)一表示，實現(xiàn)上下游各個方向共享開發(fā)成果。飛槳的新一代 IR 架構聚焦于高度靈活和高擴展性兩個重要維度，通過更加完備且魯棒的語義表達能力、訓推全架構統(tǒng)一表示和高效可插拔的性能優(yōu)化策略（Pass）開發(fā)機制，實現(xiàn)復雜語義支持，更便捷地支撐大模型自動并行下豐富的切分策略，無縫對接神經(jīng)網(wǎng)絡編譯器實現(xiàn)自動性能優(yōu)化和多硬件適配。

飛槳中間表示（PIR）在底層抽象了一套高度可擴展的基礎組件，涵蓋 Type、Attribute、Op、Trait 和 Interface，并引入了 Dialect 的概念，賦予開發(fā)者靈活擴展與自由定制的能力，從而提供了全面且穩(wěn)健的語義表達能力。在模型表示層，通過多 Dialect 的模塊化管理和統(tǒng)一多端表示，實現(xiàn)了訓練與推理一體化的全架構統(tǒng)一表示，實現(xiàn)了算子和編譯器的無縫銜接，支持自動優(yōu)化和多硬件適配。在圖變換層，通過統(tǒng)一底層模塊并簡化基礎概念，向用戶提供了低成本、易用且高性能的開發(fā)體驗，以及豐富且可插拔的 Pass 優(yōu)化機制。飛槳 PIR 堅守靜態(tài)單賦值（SSA）原則，確保模型等價于一個有向無環(huán)圖，并采用 Value 和 Operation 對計算圖進行抽象，其中 Operation 代表節(jié)點，Value 代表邊。

Operation 表示計算圖中的節(jié)點：每個 Operation 表示一個算子，并包含零個或多個 Region。Region 表示一個閉包，它內(nèi)部可以包含零個或多個 Block。而 Block 則代表一個符合靜態(tài)單賦值（SSA）原則的基本塊，其中包含零個或多個 Operation。這三者之間通過循環(huán)嵌套的方式，能夠構建出任意復雜的語法結構。

Value 表示計算圖中的有向邊：它用于連接兩個 Operation，從而描述了程序中的 Use-Define 鏈（即 UD 鏈）。其中，OpResult 作為定義端，用于定義一個 Value；而 OpOperand 則作為使用端，描述了對某個 Value 的使用情況。

飛槳提供了 PatternRewriter 和 Declarative Rewrite Rule（簡稱 DRR）這兩種 Pass 開發(fā)機制，兼顧了自定義的靈活性與開發(fā)的易用性。采用三段式的 Pass 開發(fā)方式，使開發(fā)者能夠更加專注于 Pass 邏輯的處理，而無需關注底層IR的細節(jié)。利用 PIR 的 Pass 開發(fā)機制，實現(xiàn)了 Pass 開發(fā)成本降低58%；應用于推理場景，超過84%的模型推理加速超10%。

神經(jīng)網(wǎng)絡編譯器自動優(yōu)化

為什么我們要研發(fā)編譯器技術，有3個維度的原因：

1）硬件發(fā)展趨勢：結合硬件發(fā)展歷史和技術演進特點，算力發(fā)展速度遠大于訪存性能、CPU 性能和總線帶寬；其中訪存性能影響訪存密集型算子(norm 類，activation 等)性能，CPU 性能和總線帶寬影響調(diào)度性能?；诰幾g器的自動融合的通用優(yōu)化技術，可以將多個算子融合成一個大算子，通過減少訪存量和算子數(shù)量，能夠大幅提升模型性能，編譯器技術會成為深度學習框架標配組件。

2）模型發(fā)展趨勢：模型結構存在多樣性的特點，多樣性的需求非常依賴編譯器的通用優(yōu)化。

3）多硬件優(yōu)化：當前市面存在有多款硬件，不同的硬件平臺有不同的特性和優(yōu)化需求，每個硬件均需要投入大量的人力進行優(yōu)化，借助編譯器技術，能夠大幅降低這類優(yōu)化技術成本。

讓我們通過一個實例來闡釋這一點。我們以 Llama 模型中經(jīng)常使用的 RMS Normalization （Root Mean Square Layer Normalization）為例，其計算公式相對簡單明了。

假設我們需要是實現(xiàn) RMS Normalization 的計算，最簡單的辦法是，我們可以使用飛槳框架提供的張量運算開發(fā)接口，調(diào)用平方、求和、除法、開根號等操作來完成，代碼如下：

上述代碼開發(fā)簡單，但是性能較差，且顯存占比較多；開發(fā)者可以進行 FusedRMSNorm 的實現(xiàn)，但是對于開發(fā)者要求更高，成本也更高。

借助神經(jīng)網(wǎng)絡編譯器技術，我們能夠在維持高度靈活性和易用性的基礎上，實現(xiàn)性能的顯著提升。以下 A100平臺上 RMSNorm 算子的性能測試結果便是一個明證：相較于采用 Python 開發(fā)接口組合實現(xiàn)的方式，經(jīng)過編譯優(yōu)化后的算子運行速度提升了 4 倍；即便與手動算子融合的方式相比，也實現(xiàn)了 14%的性能提升。這一成果充分展示了飛槳框架在靈活性與性能之間尋找到的理想平衡點。

為此，飛槳把神經(jīng)網(wǎng)絡編譯器技術作為一個重要的研發(fā)方向，下面是飛槳編譯器的整體架構圖。

在表示層，借助 PIR 的擴展能力，實現(xiàn) CINN 前端模塊，處理圖層相關變換，包含算子拆分、重計算、子圖劃分、維度推導模塊等模塊，最終得到多個可被編譯器后端生成優(yōu)化的子圖。在編譯器后端，對于這些可融合的子圖，編譯器會進一步調(diào)用Compute 函數(shù)，將它們轉(zhuǎn)換為由抽象語法樹（AST）構成的低層中間表示（IR），并在此基礎上進行循環(huán)的融合，確保能融合成一個 kernel；在 CINN 底層 IR 上，會進行性能的調(diào)優(yōu)分析，得到最優(yōu)配置；最后底層 IR 會被進一步精心轉(zhuǎn)換成具體的代碼實現(xiàn)。

在生成式大語言模型 Llama 和文生圖模型 Stable Diffusion 上的實驗結果顯示，通過使用編譯器的優(yōu)化技術，相較于未采用手動性能優(yōu)化的基礎版本，推理速度分別實現(xiàn)了 36%和 30%的提升。

動靜統(tǒng)一自動并行

為什么我們要做自動并行？

當前大模型主流訓練方式，會用到多種并行策略，這些并行策略基于動態(tài)圖模式實現(xiàn)的“手動”并行方式，即在單卡的基礎上，手工處理切分（切分 Tensor、計算圖）、通信（添加通信算子）、顯存優(yōu)化（顯存共享、Re-Compute）、調(diào)度優(yōu)化（流水線編排、計算和通信異步）等策略，開發(fā)者既要熟知模型結構，也要深入了解并行策略和框架調(diào)度邏輯， 使得大模型的開發(fā)和性能優(yōu)化門檻非常高。除了要有專門算法團隊負責模型算法創(chuàng)新，還必須有專門負責模型并行優(yōu)化的團隊配合，這給大模型的創(chuàng)新和迭代帶來了諸多障礙。

我們舉一個簡單的例子，來闡釋下大模型開發(fā)和單卡邏輯的差異，由于并行策略會引起 Tensor 運行時 shape 發(fā)生變化，所以跟 shape 處理相關算子均需考慮是否會受到并行策略的影響。如下面 reshape 的處理，切分策略導致輸入 shape 發(fā)生了變換，所以輸出的 shape 需要根據(jù)切分策略進行合理的調(diào)整：

為此，我們提出了動靜統(tǒng)一的自動并行方案。開發(fā)者僅需少量的張量切分標注，框架便能自動推導出所有張量和算子的分布式切分狀態(tài)，并添加合適的通信算子，保證結果正確性；最后會根據(jù)模型結構和集群信息，結合顯存、調(diào)度層優(yōu)化，自動尋找最高效的分布式并行策略。

在自動并行設計中，開發(fā)者僅需少量的張量切分標注，我們將切分方式進行抽象，共需兩類切分方式：切分張量（參數(shù)，輸入）和切分計算圖（流水線）。為實現(xiàn)這兩類切分方式，框架需要一種機制來描述分布式張量和計算設備之前的映射關系，為此我們引入 ProcessMesh 和 Placements 兩個分布式概念，其中 ProcessMesh 將一塊 GPU 卡映射為一個進程，將多個設備映射為多個進程組成的一維或多維數(shù)組，下圖展示了由 8 個設備構成的兩種不同 ProcessMesh 抽象表示。

Placements 是由 Replicate、Shard、Partial 三種分布式標記組成的列表，長度和 ProcessMesh 的維度一致，用于表示分布式張量在對應計算設備的維度上，按照哪種分布式標記做切分，這三種分布式標記的詳細描述如下：

如下圖所示，Replicate 表示張量在不同設備上會以復制的形式存在；Shard 表示按照特定的維度在不同設備上進行切分；Partial 表示設備上的張量不完整，需要進行 Reduce Sum 或者 Reduce Mean 等不同方式的操作后，才能得到完整的狀態(tài)。

在完成分布式標記抽象后，我們通過調(diào)用
paddle.distributed.shard_tensor()接口，實現(xiàn)對張量切分的標記。通過張量切分的標記和自動推導，我們可以表示復雜的分布式混合并行，下圖展示了一個具體的數(shù)據(jù)并行、張量模型并行、流水線并行組成的混合并行的例子。

以下代碼展示了混合并行的具體例子。

通過采用自動并行的開發(fā)方式，開發(fā)者無需再考慮復雜的通信邏輯。以 Llama 任務為例，分布式訓練核心代碼量減少了50%，從而大大降低了開發(fā)的難度；從我們的一些實驗可知，借助全局的分析等優(yōu)化，性能也優(yōu)于動態(tài)圖手動并行的性能。

未來，我們將進一步探索無需使用張量切分標記的全自動并行，讓開發(fā)者可以像寫單機代碼一樣寫分布式代碼，進一步提升大模型的開發(fā)體驗。

產(chǎn)業(yè)優(yōu)勢

總的來說，飛槳新一代框架——飛槳框架3.0-Beta 是面向大模型、異構多芯進行專屬設計，向下適配異構多芯，充分釋放硬件潛能；向上一體化支撐大模型的訓練、推理。同時具有動靜統(tǒng)一自動并行、編譯器自動優(yōu)化、大模型訓推一體、大模型多硬件適配四大能力，全面地提升了服務產(chǎn)業(yè)的能力。

動靜統(tǒng)一自動并行：這一功能大幅度降低了產(chǎn)業(yè)開發(fā)和訓練的成本。用戶只需在單卡基礎上進行少量的張量切分標記，飛槳框架便會自動完成分布式切分信息的推導，并添加通信算子以確保邏輯的正確性。同時，根據(jù)模型結構和集群信息，結合顯存和調(diào)度層的優(yōu)化，飛槳能自動尋找最高效的分布式并行策略，從而大幅降低混合并行訓練的開發(fā)成本，使開發(fā)者能夠更專注于模型和算法的創(chuàng)新。

編譯器自動優(yōu)化：這一功能顯著降低了性能優(yōu)化的成本。飛槳的編譯器采用與框架一體化的設計，能夠支持生成式模型、科學計算模型等多種模型的高效訓練與可變形狀推理，為計算靈活性與高性能之間提供了良好的平衡點。通過算子的自動融合和代碼生成技術，Llama2和 Stable Diffusion 等生成式模型的推理性能得到了超過30%的提升。

大模型訓推一體：這一特性為產(chǎn)業(yè)提供了極致的開發(fā)體驗。它使訓練和推理的能力能夠相互復用，為大模型的全流程提供了統(tǒng)一的開發(fā)體驗和極致的訓練效率。通過動轉(zhuǎn)靜的工作，訓練和推理的工作得以無縫銜接。在 RLHF（人類反饋強化學習）訓練過程中的生成計算可以復用推理優(yōu)化，實現(xiàn)2.1倍的加速。同時，推理量化場景復用訓練的分布式自動并行策略，效率提升了3.8倍。

大模型多硬件適配：飛槳的重要特色之一是適配異構多芯并充分釋放硬件潛能。在接入機制上，飛槳提供了簡潔高效的抽象接口和基礎算子體系，降低了適配成本。在運行機制上，它優(yōu)化了調(diào)度編排和存儲共享等機制，提升了調(diào)度效率。從算子內(nèi)核角度，飛槳提供了編譯器自動融合調(diào)優(yōu)方案，以提升端到端的性能。同時，飛槳還為新硬件廠商建設了代碼合入、持續(xù)集成、模型回歸測試等研發(fā)基礎設施。這些機制保障了新硬件被納入飛槳的正常發(fā)版體系中，用戶無需編譯即可直接安裝試用。飛槳這種功能完善、低成本接入的機制吸引了硬件廠商共同為飛槳貢獻了3,456個PR，共包含25,000多個 commits。

這就是飛槳的新一代框架 3.0 ，目前3.0-Beta 版本已面向開發(fā)者開放，并且所有的開發(fā)接口跟2.0完全兼容，非常歡迎廣大的開發(fā)者去使用和反饋。