深度學習架構作為基礎軟體，不僅促進了深度學習技術的飛速進步，更為人工智能技術的廣泛應用鋪設了堅實的基礎。
深度學習架構為開發者提供了便捷易用的開發接口，這些接口對資料和操作進行了高度抽象，使得開發者能夠更專注于算法和模型的設計，而不必深陷底層資料的處理細節。通過這些接口，開發者無需直接感覺和應對複雜的硬體底層開發細節，進而極大地提升了開發效率和體驗。其次深度學習架構還提供了自動微分這一強大功能，開發者通常隻需要編寫前向傳播網絡的代碼，而繁瑣的反向傳播網絡則交由架構自動完成。
飛槳作為中國首個自主研發、功能豐富、開源開放的深度學習平台，從預設使用靜态圖的1.0版本，到預設采用動态圖并可實作動靜統一與訓推一體的2.0版本釋出，飛槳架構已經可以完美融合動态圖的靈活性與靜态圖的高效性，并支援模型的混合并行訓練；再到近日，為大模型時代而錘煉的3.0版本的正式出爐！飛槳正式開啟了新一代架構技術創新之路！
設計思想
深度學習架構的設計對于推動人工智能技術的發展至關重要，其核心設計目标是讓深度學習技術的創新與應用更簡單。
如何做到這一點呢？
架構需要充分考慮開發者和硬體廠商的需求。
從使用者角度出發，一個優秀的深度學習架構應當為開發者提供極緻的開發體驗。這不僅僅意味着提供一個使用者友好的開發環境，更重要的是要能夠大幅度減少開發者的學習成本和時間成本，同時顯著提升開發的便利性。為此，飛槳架構提出了“動靜統一、訓推一體、自動并行”的理念，極大地提高了開發效率。
從硬體适配角度出發，現代深度學習應用往往需要在多樣化的硬體平台上運作，是以，架構必須能夠相容并适配各種不同的硬體裝置。這要求架構能夠智能地隔離不同硬體接口之間的差異，實作廣泛的硬體适配性。同時，為了充分發揮硬體的性能，架構還需要具備軟硬體協同工作的能力，確定在利用硬體資源時能夠達到最優的性能表現。
與此同時，好的架構還需要考慮到 AI 技術發展的整體趨勢、産業的實際落地應用的需求。
技術發展層面，大語言模型（Large Language Model，簡稱LLM）、 MOE（Mixture of Experts）、多模态以及科學智能（AI for Science）等前沿技術逐漸成為新的研究熱點。随着模型複雜性的增加，計算瓶頸、存儲瓶頸、訪存瓶頸以及通信瓶頸等問題逐漸凸顯，對分布式訓練和通用性能優化的需求日益迫切。
在産業化層面，架構又需要具備支援訓練、壓縮、推理一體化的全流程能力。這意味着，從模型的訓練到優化，再到實際部署和推理，架構應當提供一套完整、高效的解決方案，才能滿足産業界對于深度學習技術的實際需求。
隻有跟得上趨勢、經得住打磨的架構，才能為産學研各界開發者提供持續穩定的支援。
飛槳架構3.0的設計理念和主要特色
綜上需求，飛槳将為開發者提供一個“動靜統一、訓推一體、自動并行、自動優化、廣泛硬體适配”的深度學習架構，開發者可以像寫單機代碼一樣寫分布式代碼，無需感覺複雜的通信和排程邏輯，即可實作大模型的開發；可以像寫數學公式一樣用 Python 語言寫神經網絡，無需使用硬體開發語言編寫複雜的算子核心代碼，即可實作高效運作。
飛槳架構 3.0 版本應運而生，延續了2.x 版本動靜統一、訓推一體的設計理念，其開發接口全面相容2.x 版本。這意味着，使用2.x 版本開發的代碼，在絕大多數情況下無需修改，即可直接在3.0版本上運作。着重推出了動靜統一自動并行、編譯器自動優化、大模型訓推一體、大模型多硬體适配四大新特性。這些特性在飛槳架構2.6版本或更早版本時就已經開始開發，目前已達到外部可試用的階段。這些新特性在使用體驗、性能、二次開發便利度以及硬體适配能力等方面帶來了顯著提升，飛槳正式釋出3.0版本。此版本包含了對架構2.x 版本部分已有功能的改進，并且在不使用新特性的情況下，表現成熟穩定。
架構架構一覽
為了實作深度學習架構的上述特性，必須對架構的架構進行精心設計，確定其能夠支援各種複雜的模型建構，同時與多樣化的晶片實作無縫對接。接下來，将通過直覺的架構圖，詳細展示飛槳新一代架構内所涵蓋的功能子產品，以及這些子產品之間的互相作用與聯系。以下為飛槳架構3.0的架構圖。
飛槳架構 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 性能和總線帶寬影響排程性能。基于編譯器的自動融合的通用優化技術，可以将多個算子融合成一個大算子，通過減少訪存量和算子數量，能夠大幅提升模型性能，編譯器技術會成為深度學習架構标配元件。
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完全相容，非常歡迎廣大的開發者去使用和回報。
— 完 —
量子位 QbitAI · 頭條号簽約
關注我們，第一時間獲知前沿科技動态