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

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

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

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

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

設計思想

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

如何做到這一點呢？

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

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

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

與此同時，好的框架還需要考慮到 AI 技術發展的整體趨勢、產業的實際落地應用的需求。

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

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

只有跟得上趨勢、經得住打磨的框架，才能為產學研各界開發者提供持續穩定的支持。

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

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

框架架構一覽

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

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

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

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

調度層：負責對代碼或計算圖進行智能編排與高效調度，并且能夠根據實際需求進行顯存和內存的管理優化，支持動態圖和靜態圖高效執行。無論開發者選擇使用動態圖還是靜態圖進行模型開發，飛槳框架都能提供高效的執行環境，同時確保資源利用的最優化。

算子層：由神經網絡編譯器 CINN 和算子庫 PHI 共同構成，涵蓋了張量定義、算子定義、算子自動融合和算子內核實現等關鍵功能。

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

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

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

2）神經網絡編譯器自動優化，通過自動融合和策略調優，大幅提升模型端到端表現；

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

高擴展中間表示PIR

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

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

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

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

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

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

神經網絡編譯器自動優化

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

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

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

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

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

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

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

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

為此，飛槳把神經網絡編譯器技術作為一個重要的研發方向，下面是飛槳編譯器的整體架構圖。

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

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

動靜統一自動并行

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

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

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

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

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

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

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

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

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

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

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

產業優勢

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

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

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

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

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

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