摘要
本文描述如何結合專家組合、流式dataflow和三層記憶體系統來克服AI記憶體壁障。我們介紹Samba-CoE，一個擁有150個專家和總計一兆參數的專家組合系統。我們將Samba-CoE部署在SambaNova SN40L可重構dataflow單元(RDU)上——一種為企業推論和訓練應用程式而共同設計的商用dataflow加速器架構。該晶片推出一個新的三層記憶體系統，具有晶片內分散式SRAM、封裝內HBM和封裝外DDR DRAM。一個專用的RDU間網路能夠跨多個插槽進行向上和向外擴展。我們展示了在8個RDU插槽上執行各種基準測試相較於非融合基線的2倍到13倍加速。我們表明，對於CoE推理部署，8插槽RDU節點相較於DGX H100和DGX A100，將機器規模縮小高達19倍、加速模型切換時間15倍到31倍，並實現3.7倍和6.6倍的整體加速。
參考論文: https://arxiv.org/pdf/2405.07518
I. 導言
近期大型語言模型(LLM)的訓練和推理進展令世人驚艷。ChatGPT和Gemini等先進的生成AI/ML應用程式建立在整體式LLM之上，可能擁有數十億或兆級參數。它們使用從網路抓取的數兆級token所組成的精心策劃資料集進行訓練。然而，訓練和服務一個先進整體式LLM既非常昂貴，也是一項複雜的系統工程挑戰。訓練需要建造和操作一台由數千台主機、專用網路、電力和冷卻基礎設施以及數千個加速器（通常是GPU或TPU）組成的超級電腦。訓練和服務數千億參數所需的高昂成本和專業知識使得許多學術研究人員和小型組織難以獲得先進AI能力，尤其是在需要內部部署時。例如，估計訓練OpenAI的GPT-4的計算成本為7800萬美元，而谷歌的Gemini Ultra為1.91億美元。為超大規模公司或任何需要持續在其資料上訓練和更新強大AI模型的組織建造和部署大型整體模型可能是不可持續的。最後，適應整體模型的系統已經將計算TFLOP的增長速度遠大於記憶體頻寬和容量，產生了記憶體壁障，記憶體系統無法再有效地供給計算。
ML研究社群已經做出回應，提供了一個由更小、更模組化的模型組成的生態系統，這些模型同樣強大，但訓練和服務成本更低、更容易。像7B參數的Llama 3, Llama 2和Mistral 7B等小型模型通常已足夠。它們可能無法在一般任務套件上匹配大型模型的性能，但在針對它們進行微調的一組較窄專業任務上，小型模型可以提供優越的精確度。例如，Flan-T5-XL只有3B參數，但在MMLU分數上超過175B參數的GPT-3近10%。像這樣的案例加強了社群在透過將基礎模型專業化到一個領域、透過將基礎模型微調到特定任務或一組任務，以及透過將大型模型蒸餾或壓縮到小型模型中來構建和訓練小型模型的活動。此外，這些較小模型的組合已被證明展現出匹配大型整體模型的新興行為。它們讓更廣泛的社群能夠使用AI。
我們相信，未來成功的AI系統將有效地託管和執行許多小型模型。這反映在學術界成功追求的方向以及業界正在採用的新產品中。類似CoE的複合AI系統在推進AI前沿方面發揮了關鍵作用。在本文中，我們將這些由專業化小型模型組成的模組化系統稱為專家組合(CoE)。
CoE由幾個小型專家模型組成，協同完成一項任務。一個專家的輸出決定下一步要執行哪個專家。運行一個專家涉及將模型參數權重載入到加速器的主記憶體中，然後執行模型。因此，執行CoE涉及一系列模型切換和模型執行。如圖1所示，目前先進的AI加速器無法有效地處理這一系列操作。
有效加速CoE歸結為在最小化模型切換成本的同時有效執行專家模型。我們將其細分為三個關鍵要求：
1) 積極的運算子融合和流水線平行化，以有效執行專家模型。較小的模型具有較低的運算強度和運算子之間複雜的存取模式。傳統的運算子融合技術在任意存取模式下取得的成功有限。
2) 高頻寬記憶體，以在生成式推理期間利用權重和中間結果中的時間和空間局部性，以及
3) 高容量記憶體，以最小化切換成本並儲存許多專家模型的參數
在本文中，我們描述了一種硬體/軟體解決方案，透過解決上述挑戰來克服記憶體壁障。
首先，我們描述Samba-CoE，一個擁有150個7B專家的兆級參數CoE系統，以及如何有效運行它需要硬體支援積極的運算子融合和新穎的記憶體系統。我們介紹SambaNova SN40L可重構dataflow單元(RDU)，這是一種結合流式dataflow平行化與新穎的三層記憶體系統的商用dataflow加速器，包含大型晶片內SRAM、HBM和直接連接到加速器的DDR DRAM。
RDU的流式dataflow架構允許我們將數百個具有任意存取模式的複雜運算子融合到單一核心呼叫中——無需程式設計師手動編寫該核心。這透過利用晶片內硬體支援流水線、資料和張量平行化的混合來提供巨大的加速。我們的積極融合技術遠超過傳統架構使用的先進技術的能力。
使用TSMC 5nm技術製造的SN40L RDU是一種採用2.5D晶片上晶圓封裝(CoWoS)的晶片級設計，包含兩個SN40L可重構dataflow晶片(RDD)和HBM。每個SN40L RDU插槽使用1040個分散式模式計算單元(PCU)提供638 BF16 TFLOP的峰值計算性能。這些由1040個分散式模式記憶體單元(PMU)補充，合計提供數百TBps的晶片內記憶體頻寬以及PMU內和PMU之間的高銀行級平行度。靈活的晶片內位址產生邏輯為任意張量記憶體存取模式提供高頻寬。SN40L中的三層記憶體是：520 MiB晶片內PMU SRAM、64 GiB封裝內HBM和高達1.5 TiB的DDR DRAM（使用可插拔DIMM）。在單個SN40L節點中，模型以超過1 TB/s的速度從DDR載入到HBM。
我們在幾個真實世界基準測試上量化並討論流式dataflow平行化的影響，顯示相較於優化的基線加速2倍到13倍。我們在包含八個SN40L RDU插槽和一個主機的單個SN40L節點上部署Samba-CoE。我們討論了SN40L節點上Samba-CoE的性能，並與DGX A100和DGX H100進行比較。我們表明，對於CoE推理部署，SN40L將機器規模減少高達19倍，將模型切換時間加速15倍到31倍，相較於DGX H100和DGX A100實現3.7倍到6.6倍的整體加速。
本文的組織如下：第二節描述我們的兆級參數CoE Samba-CoE。第三節描述流式dataflow及其轉化為關鍵硬體需求的挑戰。第四節詳細描述SN40L硬體架構，並列出與先前RDU相比的主要變化。第五節描述管理DDR和HBM的軟體支援。第六節量化流式dataflow的好處以及Samba-CoE的性能。第七節討論從硬體/軟體協同設計過程中學到的關鍵經驗。第八節涵蓋相關工作。我們在第九節作結。
II. 背景：專家組合
在本節中，我們描述在SN40L上構建和部署的一個CoE實例，名為Samba-CoE。圖2顯示了Samba-CoE從提示到回應的流程。
Samba-CoE由多個專家模型和一個路由器模型組成。每個專家都在特定領域中進行了微調。我們利用了開源社群中在編程、數學和語言翻譯等領域微調的幾個優秀專家模型。路由器是另一個專業模型，可動態地將每個輸入提示分配給最相關的專家。例如，與數學相關的查詢將被路由到數學專家，而編程問題將去到編程專家。
Samba-CoE的靈感來自混合專家(MoE)架構，但有一些關鍵區別。雖然MoE和CoE都比傳統的密集型整體模型更稀疏地啟動，但MoE比CoE更不靈活。MoE需要像整體模型一樣作為單一模型進行訓練/微調，而CoE是由獨立且異構的專家模型組成，這些模型彼此獨立地訓練/微調。CoE也更強大：先前的工作已經表明CoE可以優於MoE以及GPT-3.5和GPT-4等大型整體模型, 。我們注意到CoE和MoE是正交的技術，可以輕鬆結合：CoE可以利用內部實現為MoE的專家模型。
為簡單起見，在本文中，路由器模型和專家模型都來自Llama2-7B。請注意，路由器和專家模型不需要是同質的——它們可以是具有不同參數數量的不同架構。選擇Llama2-7B作為本工作的基礎是因為其方便的大小、令人印象深刻的能力以及強大的社群支援。CoE概念和Samba-CoE系統不限於Llama2。
III. CoE的硬體需求
CoE執行時間細分為模型執行時間和模型切換時間，如圖1所示。最小化CoE執行時間可用於在每個用戶基礎上減少機器佔用空間或在給定佔用空間下增加支持的用戶數量。為了減少模型執行時間，我們展示了流式dataflow相對於傳統運算子融合的優勢。為了最小化模型切換時間，我們說明了同時需要高容量的加速器本地DDR介面和HBM的必要性。
A. 流式Dataflow
傳統運算子融合不足：運算子融合是一種常見的優化技術，用於增加運算強度並提高硬體利用率。融合還減少了執行模型所需的核心呼叫數量，並攤銷了核心啟動開銷。然而，專家模型通常包含低運算強度的運算子，加上涉及洗牌和轉置的複雜存取模式。複雜的存取模式嚴重限制了GPU上融合的功效。像PyTorch2和TensorRT等框架對融合明確不支持的模式有文檔記錄的限制。因此，許多複雜的融合核心仍然是為GPU手動編寫的。
圖3描述了Monarch FFT分解的簡化版本，張量形狀標註在邊上。表I顯示了融合對運算強度的影響。較高的運算強度允許應用程式針對給定的目標加速器實現屋頂線性能。例如，A100 GPU的TFLOPS/TBps比率約為300/2=150，這意味著運算強度小於150 FLOP/byte的核心在A100上受記憶體限制。在表I中，前兩行在A100上受記憶體限制，最後一行受計算限制。
然而，由於以下原因，GPU無法融合圖3的所有部分：
1) 僵化的記憶體層次結構和程式設計模型造成資料移動瓶頸：GPU核心使用線程塊網格啟動。網格結構在核心的持續時間內是固定的。融合Gemm0和Mul將是微不足道的。然而，Transpose迫使線程存取其他SM中線程的值，觸發通過共享快取和HBM在SM之間交換資料。由於沒有其他方式在SM之間傳輸資料，這種缺乏靈活性在共享快取和HBM上造成瓶頸。
2) 晶片內SRAM容量不足迫使Transpose的輸出具體化到HBM，防止融合機會。
3) 運算子之間沒有利用流水線平行化：更高階的Monarch FFT分解（在第VI節中研究）創建許多小的矩陣乘法，其大小為32×32×32或更小，這無法有效利用所有SM。然而，在所有矩陣乘法和元素級運算子之間存在大量的流水線級平行度。GPU SIMT程式設計模型沒有提供直接的方式將圖3中的運算子作為流水線執行。
流式Dataflow實現流水線和任意存取模式的自動融合：
與傳統融合不同，流式dataflow將運算子作為粗粒度流水線執行。張量被切塊並通過此流水線流式傳輸。切塊之間可以具有任意的讀寫存取模式。
圖4描繪了空間融合的實現。藍色框表示晶片內緩衝區單元，灰色框表示晶片內計算單元。運算子Gemm0、Mul和Gemm1作為粗粒度流水線中的階段執行。中間的藍色記憶體單元用作解耦的階段緩衝區，保存中間結果。由於Gemm0和Gemm1佔總操作的較大部分，因此分配更多計算單元給它們。通過使用適當數量的記憶體單元來匹配各自階段所需的輸入和輸出頻寬。例如，邏輯階段緩衝區I0被分割為兩個記憶體單元I00和I01，以匹配Gemm0所需的輸入頻寬。出於容量原因，緩衝區S0-S3被分割為四個記憶體單元。轉置操作被融合到緩衝區T0*和T1*中作為存取模式。
我們從上面的觀察中提煉出以下晶片內架構特性列表，以實現流式dataflow：
1) 可組合的記憶體單元：單個記憶體單元提供固定的容量和頻寬。由於晶片內張量的容量和頻寬需求各不相同，硬體應支援跨記憶體單元對邏輯位址進行可程式設計的交錯處理。
2) 位址產生頻寬和靈活性：高資料頻寬需要高位址產生頻寬。此外，每個記憶體單元應支援非阻塞的並發讀寫，以有效實現階段緩衝區。簡而言之，位址產生硬體應允許以高輸送量為任意複雜的位址表達式產生多個並發位址。
3) 全身和流式計算：ML加速器架構通常實現全身陣列以提高GEMM等運算的計算密度。然而，在許多ML模型中，GEMM經常伴隨著需要高輸送量流式計算能力的元素級運算子和約簡。
4) 一對多、多對一和資料重排：生產者和消費者單元數量的差異會在RDN上創建一生產者對多消費者和多生產者對一消費者的流量流。對於一對多，需要硬體支援在互連中從來源創建分支路徑到軟體決定的一組目的地。對於多對一流量，來自不同路徑的資料可能以亂序方式到達目的地。這個亂序序列通常必須在目的地重新排序，以匹配程式的預期。換句話說，硬體必須提供一個協議來重新排序資料流。最後，需要軟體控制的頻寬管理（例如節流）和路由來滿足具有不同頻寬需求的流。所有這些硬體支援需要能夠被編譯器內的自動佈局和布線演算法使用。
B. 模型託管和切換成本
HBM的有限容量限制了當託管在GPU或TPU上時CoE中的專家數量。僅使用HBM，運行大型CoE需要(a)使用更多機器來提供HBM容量，這增加了成本、複雜了部署並引入了負載平衡挑戰，或者(b)使用主機的記憶體，如圖1所示，這增加了切換延遲。直接連接到加速器的更高容量DDR記憶體可以同時減少模型託管和模型切換成本。此外，專家參數在CoE中表現出時間局部性，因為它們被多次使用（例如，在自迴歸解碼期間）。HBM通過充當DDR和SRAM之間的軟體管理快取層，在利用這種時間資料局部性方面發揮關鍵作用。
因此，我們得出結論，執行較小模型組合的系統需要兩種類型的晶片外記憶體：(1)高頻寬記憶體，以利用專家參數的時間局部性，以及(2) 高容量記憶體，以在小佔用空間中儲存專家參數。
在下一節中，我們將描述基於上述原則構建的SN40L可重構dataflow單元。
IV. SN40L硬體架構
SN40L dataflow加速器使用TSMC的5FF工藝製造，並使用晶片上晶圓封裝(CoWoS)多晶片封裝技術封裝為雙晶片插槽。圖5顯示了SN40L的顯著組件，描述如下。
RDU Tile：粗粒度可重構dataflow核心陣列。包括模式計算單元(PCU)、模式記憶體單元(PMU)以及連接在稱為可重構dataflow網路(RDN)的二維網格互連中的位址產生和合併單元(AGCU)。
記憶體介面：SN40L與兩層晶片外記憶體介面——HBM和DDR。這兩個記憶體空間都由軟體管理。DDR層可以有高達1.5 TiB的峰值記憶體容量，峰值頻寬超過200 GB/s。HBM層每個插槽有64 GiB的容量，峰值頻寬約為2 TB/s。
晶片間(D2D)介面：SN40L tile組件可以直接在兩個晶片之間流式傳輸資料，而無需通過晶片外記憶體。
主機介面：SN40L通過PCIe連結與主機x86 CPU介面。此介面支援主機和設備晶片外記憶體之間的DMA以及主機和tile之間的直接通信。
對等(P2P)介面：將SN40L連接到其他SN40L RDU。第IV-D節中描述的對等協議提供了實現集體通信原語的基本體。
頂級網路(TLN)：此網路將SN40L tile連接到主機、記憶體和對等介面。
圖6展示了一個具有關鍵dataflow組件的SN40L tile：PCU、PMU、RDN交換機和AGCU。以下小節將更詳細地描述它們。
A. 模式計算單元(PCU)
SN40L中的PCU提供SN40L中的全身和流式計算能力。PCU的資料路徑由頭部、體部和尾部組成。頭部消耗輸入dataflow並驅動體部。PCU的體部可配置為輸出固定的全身陣列或具有多個向量計算級的流水線SIMD核心。尾部執行特殊的元素級函數並驅動輸出FIFO。這實現了GEMM等運算、元素級運算或約簡的高效執行。
圖7展示了PCU作為2D全身陣列和SIMD核心的示意圖。全身陣列加速矩陣乘法，如圖3中的Gemm0和Gemm1。輸入到全身陣列通過稱為廣播緩衝區的結構從左到右和從上到下流式傳輸。累積結果通過尾部單元從左到右排出到輸出FIFO。矩陣乘法可以在多個PCU上進一步並行化，類似於圖4中的描述。作為SIMD核心，PCU以完全流水線方式執行並行的多維張量運算，如圖4中的Mul。每個SIMD階段支持FP32、BF16和INT32格式的常見算術、邏輯和位元運算。PCU可以配置為實現可選的跨通道約簡網路，如圖中的藍色三角形所示。還支援純SIMD方式的通道內約簡。跟蹤迴圈反覆運算的計數器在達到程式設計的最大值時產生控制事件，表明迴圈已完成執行。
尾部支援超越函數、隨機數產生、隨機取整和格式轉換。尾部操作可以與體部的計算融合和流水線化。
一個操作可以以資料並行、張量並行或流水線並行的方式在多個PCU上並行化。資料並行通過分割輸入和輸出來創建由不同PCU處理的多個獨立資料流來實現。張量並行通過分叉到資料並行流，然後將它們連接來實現。流水線並行通過將多個PCU鏈接在一起以融合運算並增加運算強度來實現。
B. 模式記憶體單元(PMU)
SN40L中的PMU提供晶片內記憶體容量、頻寬和運算子融合所需的位址靈活性。圖8顯示了PMU的高級塊圖。PMU用於儲存晶片內張量，如輸入、參數、元資料和中間結果。例如，圖4中的所有藍色塊對應於PMU。PMU的關鍵組件如下所述。
Scratchpad banks：每個PMU都包含一個程式設計員管理的scratchpad記憶體，它被組織為SRAM實例陣列。SRAM陣列整體支援併發寫入和讀取。
標量ALU管道：PMU包含幾個整數ALU級，可配置為同時產生讀取和寫入scratchpad的位址。PMU ALU實現了一組經常用於位址計算的特殊複雜指令，如位域提取和移位設置。這使得能夠以更少的ALU級更有效地生成複雜的位址，因此延遲更少。ALU管道還有一條路徑，可以從標量RDN攝取標量作為操作數，並將計算值作為標量輸出到標量RDN中。此路徑有助於位址可組合性：如果需要，可以將複雜的整數計算分解並映射到多個PMU上。
如第III節所述並在圖4中所示（緩衝區T00-T03和T10-T13），空間融合核心中的階段緩衝區需要併發讀寫，它們可能具有不同的存取模式。雖然不是普遍正確的，但我們已經偶然觀察到張量的寫入和讀取存取模式經常相互影響對方的複雜性；複雜的寫入存取模式通常實現更簡單的讀取存取模式，反之亦然。標量ALU管道允許軟體利用這一見解。它可以分割為獨立的讀取和寫入位址產生管道，具有軟體配置的分配給每個管道的級數。
位址謂詞和Banking：圖4顯示，由於容量需求（S0-S3）、頻寬需求（W00-W01和I00-I01）或兩者兼而有之（T00-T03和T10-T13），單個邏輯張量可以跨越多個PMU。PMU透過提供鉤子以程式設計方式控制跨PMU的張量位址交錯來實現這一點。具體而言，每個PMU可以程式設計為該PMU的一系列有效位址，或為每個產生的位址提供一個可程式設計的謂詞位元。如果在程式設計的範圍內或有效謂詞，則處理位址，否則丟棄。此外，使用軟體可程式設計的銀行位元位置將位址映射到scratchpad銀行。
資料對齊單元：資料對齊單元支援常見的張量轉換運算，如轉置、跨通道向量置換、向量非對齊存取、查找表(LUT)、資料格式和資料佈局轉換。要轉置的張量以特殊的對角線條紋格式跨scratchpad bank寫入，這使得能夠以完整頻寬讀取相同張量的常規和轉置格式。此功能實現了透過將轉置運算子融合到圖4中緩衝區T00-T03和T10-T13之間的讀寫存取模式優化中來實現轉置運算子。
C. 可重構Dataflow網路（RDN）
RDN是晶片內可程式設計互連，可促進PCU、PMU和AGCU之間的通信。RDN由三個物理網狀架構組成——向量、標量和控制。向量和標量網狀是以封包為基礎的。控制網狀是電路交換的，由一束可以單獨路由的單位元線組成。向量網狀是張量資料的主要管道。標量網狀主要用於傳輸元資料，如位址，但在某些情況下也可以用於攜帶資料或控制。控制網狀用於攜帶分佈式粗粒度流控制的token以及集體編排圖的執行。控制token通常對應於計數器的「完成」事件，表示迴圈的結束，如第IV-A節中簡要討論的。
RDN使用非阻塞交換機網格實現。如圖7和8所示，從交換機到單元的輸入標量和向量封包落在輸入FIFO中，透過輸出FIFO退出。向量和標量網狀上的傳輸在每一跳都受到基於信用的流控制。封包流在RDN上的通信單元之間也透過軟體token、細粒度硬體信用和硬體中的前向進度保證的組合進行端到端流控制。所有三個RDN網狀的路由表都由軟體使用編譯器中的佈局和布線（PnR）層配置。
現在我們描述支援第III節中確定的關鍵通信模式的機制。
- 多播和可程式設計路由：標量和向量網狀上的封包路由是透過2-D維度順序路由或軟體控制的靜態流路由動態完成的。在靜態流路由中，軟體為封包流分配流ID域，該域隨封包一起攜帶。在每個交換機埠對流ID域進行解碼，並在將封包轉發到下一個目的地之前重新分配。靜態流路由機制透過交換機支援封包多播。
- 多對一和資料重排：向量封包包含一個稱為序列ID的中繼資料域，這是在晶片內支援任意多對一流的主要機制。單元的向量輸出埠配備了可程式設計邏輯，用於為每個輸出向量產生序列ID。這樣，軟體就可以程式設計序列ID，以表示跨多個來源的給定運算的邏輯向量順序。序列ID域用作PMU中的輸入運算元，以計算寫入位址來重新排序封包。
D. 位址產生和合併單元（AGCU）
AGCU是RDU tile到本地設備記憶體（HBM/DDR）、主機記憶體、遠程RDU設備記憶體和透過TLN到遠程RDU tile的可重構dataflow橋。在tile側，它像dataflow核心一樣暴露RDN向量、標量和控制埠。在TLN側，它產生讀寫請求並合併回應。它配備了標量位址產生管道和計數器，與PMU邏輯有一些相似之處（不包括SRAM）。位址產生管道由計數器組成，並提供記憶體管理的位址轉換層。
對等通信：AGCU支援對等（P2P）通信協議，以直接在不同插槽上的RDU tile之間流式傳輸資料，而無需涉及DDR或HBM。P2P協議提供了在RDU之間構建集體通信原語所需的原語，如AllReduce。
核心啟動編排：AGCU實現了一個核心啟動機制，由三個命令的序列組成：程式載入、參數載入和核心執行。運行模型涉及執行核心啟動的時間表，可以是軟體編排的或硬體編排的。軟體編排允許更靈活地調度核心，並提供更多主機軟體對模型執行的可見性。然而，軟體編排會產生影響性能的開銷。硬體編排將靜態核心時間表卸載到AGCU中的專用硬體，這顯著降低了開銷。在第VI節中，我們量化並討論了軟體與硬體編排執行對各種基準的影響。
E. SN40L中的新功能
先前的工作涵蓋了從學術原型到第一代工業產品的經驗教訓。本文首次深入探討了生產RDU的內部工作原理。
SN40L對於更廣泛的工作負載相較於先前的RDU是一個顯著的改進，它還改善了軟體可用性。SN10主要專注於訓練。我們討論一些關鍵的增強功能。
HBM：SN40L是第一個包含HBM的RDU。這是由需要繼續將大型模型放入緊湊的佔用空間（使用DDR）但也需要改善更多受記憶體限制的推理工作負載的加速（使用HBM）所驅動的。添加HBM記憶體層對CoE的可行性至關重要，但它也能夠在整體模型上獲得更好的性能。
GEMM和小張量：PCU的全身陣列得到改進，可以更有效地處理較小的矩陣乘法。PCU和PMU增加了對具有整合偏置（可以節省硬體資源）的GEMM以及對有效的晶片內pad產生而無需具體化pad值的支援。
動態dataflow：SN40L包括幾項改進，以更好地支援動態dataflow。原始SN10 RDU要求所有封包目的地都按每個dataflow核心進行靜態配置。這在執行將封包在多個目的地之間解多工的序列化一對多運算（如張量「拆分」運算）時被證明是有局限性的。SN40L在PCU中增加了對每個封包的動態目的地的支援。另外，SN40L PMU可以透過增強資料對齊單元，使其具有高輸送量可程式設計通道改組和遮罩，更好地支援FFT和排序等運算子。
晶片內頻寬利用率：在SN10上，必須全域分配有限的RDN流量流ID池。這意味著，如果具有流ID X和Y的兩個流經過相同的交換機，則這些ID在這兩個流使用的任何交換機中都不可用。這為編譯器創建了一個複雜的全域最佳化問題，並浪費了晶片內頻寬。在SN40L中，流ID分配被改為類似於MPLS[15]的靈活方案。每個流在編譯時分配一個流ID序列，該序列被程式設計到交換機流表中。當封包流過交換機時，封包中的流ID域使用交換機流表重寫。這種靈活性導致更好的晶片內頻寬利用率和可用性。
增強的電源管理：SN10需要軟體幫助緩解電壓下垂（L di/dt），但實現過於保守，導致性能損失高達25%。SN40L透過先進的時脈技術和更複雜的硬體電源管理更好地緩解了電壓下垂。
V. 軟體支援
在本節中，我們描述如何在SN40L上部署Samba-CoE。
Samba-CoE由150個Llama-7B專家組成，總共超過1T參數。它部署在單個具有8個RDU插槽的SN40L節點上。圖9顯示了Samba-CoE如何利用DDR和HBM，以及單個提示的簡化事件流。所有150個專家的權重保存在高容量DDR中，而路由器的權重保存在HBM中。單個Samba-CoE推理有三個高級步驟：(1)運行路由器以確定專家，(2)將專家從DDR複製到HBM，以及(3)運行專家。
在模型切換之後，專家模型的權重在自迴歸解碼迴圈中被多次讀取以產生多個token。Samba-CoE中這種固有的模型層次時間局部性透過將權重移動到HBM來利用。
然而，這確實給軟體堆疊帶來了管理多個非一致記憶體空間的額外負擔。在本節中，我們討論在應用層和SN40L的低層次執行時期驅動程式之間構建的軟體基礎設施，以實現這一願景。
A. 記憶體分配
我們這個系統設計的目標之一是使設備記憶體管理對應用開發人員透明。為此，我們在SN40L編譯器中加入了自動異構設備記憶體管理。基本假設是只要一切都適合，預設就使用HBM。因此，我們將DDR用於兩個主要目的：
1) 當給定模型的常駐記憶體太大而無法容納在HBM中時，將資料從HBM溢出到DDR，以及 2) 保存CoE的所有其他當前非活動的模型。請注意，雖然在HBM和DDR之間移動資料會產生顯著的性能成本，但它仍然比將所有內容一直溢出到主機的記憶體要便宜得多，如圖1所量化的那樣。在本節中，我們將重點關注第一種使用情況，並在第V-B節中討論第二種情況。
我們發現需要積極的垃圾回收才能將表II中的Llama 7B等模型放入每個插槽64 GiB的HBM容量中。然而，典型的動態垃圾回收方案對於這種使用情況而言開銷太大，因為它會要求我們經常在應用程式中間將控制權返回給CPU以重新組織SN40L的記憶體。相反，我們利用了SN40L的程式設計模型既沒有動態記憶體分配也沒有指標別名的事實，因此我們可以靜態執行符號生命週期分析，並透過在其生命週期不重疊時將多個邏輯符號分配給相同的設備虛擬位址來實現垃圾回收。
最後的問題是當記憶體仍然不適合時要將哪些符號溢出到DDR。因為這兩種記憶體之間存在一個數量級的頻寬差異，所以我們分析每個符號的時間局部性及其傳輸佔用空間，以估計該符號在整個應用程式中的總頻寬需求。然後按總傳輸大小對符號進行排序，以便我們首先溢出具有最小頻寬需求的符號。
在實踐中，我們發現對於表II中的LLM模型，權重獲得最高優先級以保留在HBM中，而激活符號和其他中間結果可以在必要時溢出。我們正在研究對該演算法的進一步改進。
B. CoE執行時期
我們為創建CoE模型架構的主要設計目標之一是每個單獨的模型專家的生命週期應該是獨立的：開發、編譯、訓練、微調、量化、維護、服務和共享。這為軟體模組化提供了顯著的飛躍，優於整體模型。模型精確度和性能的進展非常快，CoE使得逐步利用最新創新變得容易。
為了實現這一目標，我們需要在執行時期將任意數量的獨立編譯的模型鏈接在一起，並在新請求從應用層進入時在它們之間動態切換。該方法類似於動態鏈結器/載入器在傳統軟體應用程式中的工作方式。
我們在先前描述的靜態管理器之上添加了一個輕量級動態記憶體管理器。它可以使用主機執行時期有效實現，因為它只需要在整個專家模型的邊界運行。每個編譯的模型二進位檔案提前告訴我們該模型需要多少HBM和DDR空間。然後，CoE執行時期與低層次設備驅動程式介面，以動態分配DDR中每個模型的記憶體塊。此分配包括編譯器打算用於HBM的記憶體段；這些記憶體段最初也儲存在DDR中。當請求特定模型時，CoE執行時期透過將打算用於HBM的記憶體段「激活」，然後將控制權轉移到編譯的應用程式二進位檔案以正常運行。一旦應用程式退出，控制權返回到CoE執行時期，等待下一個請求。
CoE運行時嘗試在HBM中同時保持盡可能多的模型處於活動狀態，並在達到容量限制時使用LRU驅逐策略。如果下一個請求是針對同一模型的，則可以立即恢復，無需額外開銷。如果請求需要一個新模型，而新模型無法適應剩餘的HBM空間，則從HBM中驅逐最舊模型的HBM記憶體，並換入新模型的HBM記憶體。為避免不必要地複製回大型唯讀權重，編譯器在適當的地方將符號標註為唯讀，然後運行時跳過將任何此類符號複製回DDR。
VI. 案例研究
在本節中，我們使用幾個真實世界的訓練和推理基準來量化並討論運算子融合的影響。然後，我們在SN40L節點上評估Samba-CoE，並與DGX A100和DGX H100進行比較。
A. 運算子融合的影響
1) 基準和設置：表II描述了用於量化運算子融合影響的一組語言模型基準。基準包括各種大小的幾個訓練和推理工作負載。在推理工作負載中，我們將「預填充」階段——生成第一個token——與「解碼」階段——透過自迴歸解碼生成第二個和後續token分開。自迴歸解碼步驟利用KV緩存，與構建KV緩存的預填充階段相比，計算量和運算強度要低得多。FlashFFTConv的計算圖結構是第III節中描述的激勵示例的複雜版本。
除了FlashFFTConv之外的所有實驗都在包含8個SN40L插槽和一個主機的系統上進行評估，FlashFFTConv是一個較小的核心，我們在單個SN40L上進行評估。我們在三種配置中測量並比較每個基準的性能：
- 非融合：模型中的每個PyTorch運算子在SN40L上作為一個或多個核心執行，中間結果具體化到DDR或HBM。核心由軟體安排執行。每個核心仍然被並行化以在SN40L上有效運行。
- 融合+軟體編排（SO）：使用自動編譯器優化和程式設計員提示的組合將運算子融合到更少的核心中。核心調度從主機上運行的軟體執行。
- 融合+硬體編排（HO）：與上面相同的融合核心，但核心調度使用第IV-D節中描述的功能卸載到硬體。
2) 基準結果：圖10顯示了運算子融合對所有基準的影響。藍色條形圖量化了運算子融合的影響。橙色條形圖量化了硬體編排核心啟動的影響。
運算子融合加速 如果沒有融合，FlashFFTConv的運算強度非常低，性能受到影響。在SN40L上進行融合時，整個FlashFFTConv基準使用單個核心啟動執行，類似於圖4中所示的簡化圖。增加的運算強度提供了13倍於非融合基線的加速。
推理預填充、訓練和sparseGPT基準實現了1.5倍到3倍的加速範圍。使用空間融合，在具有足夠粗粒度流水線平行性的情況下實現更高的運算強度，以保持晶片內單元被佔用。對於這些基準，批次大小和序列大小提供了所需級別的流水線平行性，以利用融合的好處。
儘管自迴歸解碼推理基準的運算強度較低，但它們透過消除核心啟動和不必要的HBM流量的開銷而從融合中獲益。我們觀察到從1倍到13倍的加速，其中Mistral實現了最高的加速。
圖11比較了以融合和非融合配置執行各種基準所涉及的核心呼叫啟動數量。對於llama7B-4k-inf-prefill，11倍的比率意味著使用運算子融合，llama7B-4k-inf-prefill執行的核心啟動數量是其非融合基線的11倍。此比率量化了在給定基準中執行的融合級別。對於某些基準，較高的數字意味著更積極的融合，如FlashFFTConv和sparseGPT的情況。對於其他基準如llama70B，較高的數字也是由於模型的較大規模。
硬體編排核心啟動加速：現在我們討論由於硬體編排核心啟動獲得的改進。在這裡，我們看到與前面討論相反的趨勢：自迴歸解碼推理基準實現了顯著的1.4倍到8倍的加速。在這些基準中，核心具有非常短的執行時間，其中大部分時間都由載入權重和其他輸入所主導。因此，核心啟動開銷開始佔總時間的更大部分。將核心調度卸載到SN40L可以消除軟體調度的開銷，從而提供加速。
另一方面，推理預填充和訓練基準僅看到最多1.1倍的改進。在這些基準中，每個核心執行的時間要長得多，因此核心啟動開銷被攤銷了。融合的FlashFFTConv基準只有一個核心呼叫，在兩種核心調度方法下花費相同的持續時間。
B. 專家組合(Composition of Experts, COE)
我們量化在SN40L節點上部署專家數量不斷增加的Samba-CoE的延遲和系統佔用空間，並與DGX A100和DGX H100進行比較。表III總結了結果。我們研究了兩種專家數量不斷增加的情況：單個節點上的延遲影響，以及系統佔用空間影響以維持相同的TP8延遲。
延遲影響：我們模擬兩個用例：聊天機器人對話用例，每個輸入提示產生20個輸出token，以及翻譯用例，產生200個輸出token。模型和路由器在所有平台上以張量並行的8插槽（TP8）方式映射。路由器和KV緩存總是在HBM中。SN40L節點數是在真實硬體上測量的。由於Samba-CoE未部署在DGX上，我們使用已發佈的模型延遲數和基於DGX規格的樂觀模型切換估計值來估計延遲。總延遲包括運行路由器、複製所需的專家權重以及運行專家的時間。研究了兩個專家情況：生成20個輸出token和生成200個輸出token。我們樂觀地假設在DGX上，HBM和主機記憶體的整個容量都可用於保存權重和KV緩存（僅HBM）。
我們報告批次大小（BS）= 1和BS=8情況下的延遲。請注意，「批次」適用於整個Samba-CoE模型，而不是單個專家。BS=8意味著Samba-CoE模型收到8個批次提示。首先使用BS=8運行路由器以獲得每個提示的專家。樣本中的每個提示可能需要不同的專家，因為批次中的樣本彼此無關。然後將所需的專家複製到HBM。然後按順序運行每個（提示，專家）對。
圖12比較了三個平台的延遲。我們從兩個廣泛類別分析這些結果：
50個專家以下：所有專家都適合HBM，因此性能受專家執行時間的限制。SN40L節點生成20個token的速度比DGX A100快2倍，比DGX H100快1.5倍。對於200個token，加速數字分別為3.2倍和2.3倍。請注意，生成式推理受記憶體限制，SN40L節點具有與A100相當的HBM頻寬（低於H100）。加速效果清楚地展示了流式dataflow的好處：整個解碼器層被融合到單個核心呼叫中，使用近90%的PCU和PMU，並飽和接近85%的HBM頻寬。此外，由於模型主要包含多個相同的解碼器層，SN40L幾乎看不到核心啟動開銷。
50個專家以上：當專家溢出到主機DRAM時，DGX上的延遲會激增（大約50個7B專家）。對於BS=8，SN40L節點分別實現了6.6倍和3.7倍於DGX A100和DGX H100的加速。對於BS=1，SN40L節點分別實現了4.8倍和2.8倍的加速。BS=8需要複製更多的專家，因此佔總時間的更大部分。圖1a和1b顯示了專家切換與模型執行的時間細分。在SN40L節點上，借助超過1 TB/s的DDR到HBM總頻寬，複製時間比DGX A100快31倍（提供32 GB/s的主機到GPU頻寬），比H100快16倍（提供64 GB/s的主機到GPU頻寬）。DGX在150個專家時耗盡記憶體。
系統佔用空間影響：接下來，我們量化隨著專家數量增加以維持每個平台上相同TP8延遲的系統佔用空間。實現這一點需要消除GPU上的專家複製。因此，所有專家都應位於GPU HBM中。SN40L DDR到HBM的切換成本計入SN40L的延遲。
單個SN40L節點可以以TP8延遲保存和服務多達850個專家的CoE。使用DGX實現這一點需要19個DGX節點才能將所有專家保存在HBM中。
VII. 經驗教訓
SN40L是一個複雜的系統，是多年來許多軟硬體工程師集體智慧和辛勤工作的結晶。在本節中，我們討論了在此過程中學到的三個經驗教訓。
軟體管理頻寬：軟體必須管理來自各種實體的頻寬：Tile級單元通信、HBM、DDR、晶片間、對等和主機頻寬。頻寬轉化為TLN和RDN中的一個或多個並發資料流。為了從HBM等單元利用更多頻寬，需要創建更多的載入和儲存資料流。相反，需要更少頻寬的單元應分配更少的資源，以避免過度提供和浪費。在編譯器中構建靜態頻寬模型對於啟用適當的頻寬分配和流量管理至關重要。對良好的靜態頻寬模型的投資在其他方面得到了回報：應用程式可以靜態地分析和調整性能到第一個數量級。
性能偵錯：積極的運算子融合和流水線化會在晶片內創建大量並發流量流，從而產生頻寬瓶頸。我們注意到，頻寬問題通常歸結為兩件事之一：網路擁塞或記憶體銀行衝突。SN40L交換機和PMU中的性能計數器統計停頓次數，有助於識別SN40L tile中的熱點。在RDN擁塞方面，
我們觀察到，如果不加管理，突發流量很容易減慢整個核心的速度。硬體中可程式設計的封包節流功能使軟體能夠減少突發行為並緩解許多RDN擁塞問題。為了處理PMU銀行衝突，我們觀察到PMU經常被程式設計為任意張量形狀的雙緩衝區，如果這些緩衝區靜態映射到不同的銀行，則可以避免銀行衝突。第IV-B節中描述的可程式設計銀行位元幫助根據這一見解採取行動，消除了這種多緩衝區配置的銀行衝突。
流水線計算和集體通信：對於我們的編譯器的更高層而言，跨插槽將資料/張量/流水線並行dataflow的映射問題類似於將它們映射到插槽內的問題。借助第IV-D節中描述的硬體對等協議，集體通信運算子可以與其他計算融合並流水線化到同一個核心中，就像任何其他運算子組一樣。此外，SN40L之間的流式對等協議在許多情況下避免了到HBM的躍點，這有助於節省一些通信資源。
VIII. 相關工作
1) 商用AI加速器：在本節中，我們廣泛討論其他商用AI加速系統。由於這是一個競爭激烈的市場領域，關於當代AI硬體的許多細節通常是未知的或模糊的。
NVIDIA A100、H100、AMD MI300X、Google TPUv4、和Intel Gaudi等架構都是帶有HBM的AI加速器。雖然這些架構在其程式設計模型、記憶體系統和擴展能力方面差異很大，但它們沒有有效執行大型CoE和巨型模型所需的三層記憶體系統。因此，如第VI-B節所述，在它們上部署CoE會產生低效率。此外，SN40L中的流式dataflow提供了相較於上述加速器的獨特優勢，如第VI節所量化的。先前的研究量化並利用了TPU和GPU上的運算子融合，但沒有執行本文中描述的積極融合級別。最後，SN40L提供了每個插槽約2.5倍於最近宣佈的NVIDIA GH200的總記憶體容量，這使得在SN40L上支援更大的CoE模型成為可能。
Graphcore、Cerebras、Groq和SambaNova的早期RDU提供了替代的AI加速器。然而，它們都缺乏有效執行CoE所需的三層記憶體系統。據我們所知，SN40L是唯一已經成功演示在單個節點部署萬億參數CoE和其他巨型模型並達到所報告性能的系統。
2) 研究dataflow架構：關於dataflow架構的主題有幾篇先前的出版物，涵蓋了計算、記憶體、互連和程式設計模型的各個方面，如綜述論文所述。據我們所知，SN40L是第一個結合流式dataflow與三層記憶體系統的dataflow架構，並量化其對真實世界基準的影響。
3) 運算子融合：傳統的運算子融合是一個廣泛研究的主題。然而，映射在PCU和PMU上的流式dataflow流水線通常包含20多個運算子（見圖11），這些運算子是由編譯器從Python框架層自動生成的。相比之下，傳統的運算子融合針對1-5個運算子，這些運算子通常是手寫的，並且具有存取模式限制。
4) 參數高效微調技術（PEFT）：像LoRA這樣的技術通常用於將專家權重縮小到應用於基礎模型的小型低秩適配器。然而，在幾種情況下，PEFT技術無法達到與監督微調（SFT）相同的品質。因此，較小的專家模型通常是使用額外訓練或SFT專業化的整個模型（在撰寫本文時，HuggingFace上有超過9000個Llama 2變體）。
IX. 結論
在本文中，我們將專家組合（CoE）描述為大型整體式LLM的一種模組化和成本效益高的替代方案。我們描述了具有150個專家的Samba-CoE，並說明了CoE的硬體需求。然後，我們介紹了SN40L dataflow加速器和SN40L節點，它們旨在使用流式dataflow和新穎的三層記憶體系統來解決記憶體牆問題。SN40L的記憶體系統包括晶片內分散式SRAM、晶片外HBM和高容量DDR DRAM。我們討論了跨DDR和HBM管理地址空間的軟體影響，以及在SN40L上部署Samba-CoE的執行時期複雜性。我們證明了流式dataflow相較於非融合基線提供了2倍到13倍的好處。我們表明，在SN40L節點上部署Samba-CoE將機器佔用空間減少高達19倍，將專家複製時間加速15倍到31倍，並相較於DGX H100和DGX A100分別實現3.7倍到6.6倍的整體加速。
我的觀點
這篇論文展示了SambaNova在解決大型語言模型(LLM)訓練和推理面臨的挑戰方面取得的重大進展。作者提出了一種創新的硬體/軟體協同設計方法,通過結合專家組合(CoE)、流式dataflow和新穎的三層記憶體系統來克服AI記憶體壁障。我認為CoE是一個非常有前景的研究方向。相較於訓練和服務單一的大型整體式模型,CoE提供了更高的模組化、靈活性和成本效益。它允許針對特定任務或領域專業化和優化較小的專家模型,同時仍能透過路由器實現整體系統的強大功能。這種方法不僅使學術界和較小組織更容易獲得尖端AI能力,而且還為持續改進和更新复合AI系統開闢了道路。
SN40L硬體架構在支援CoE方面給我留下了深刻的印象。流式dataflow設計實現了積極的運算子融合和自動流水線化,大大提高了加速器的效率。三層記憶體系統在提供所需的容量和頻寬來儲存和切換專家參數方面發揮了關鍵作用。晶片內SRAM、HBM和大容量的直連DDR DRAM的獨特組合使SN40L能夠在單個節點上高效地執行具有數兆參數的CoE。
此外,我欣賞論文強調了硬體/軟體協同設計在實現卓越系統性能方面的重要性。記憶體管理、頻寬分配、運算子融合和集體通信等方面的軟體創新與硬體功能緊密結合,以充分利用底層架構。這凸顯出AI系統設計中跨學科協作的價值。
SN40L與DGX A100和DGX H100的性能比較結果令人印象深刻。在CoE推理部署中,SN40L在延遲、機器規模和模型切換時間方面的顯著優勢說明了其架構創新的有效性。隨著AI工作負載變得越來越複雜和多樣化,這種專為CoE等新興範例量身定制的加速器將變得越來越重要。
我認為這項工作為探索複合AI系統和專用硬體架構開闢了令人興奮的道路。它不僅推進了LLM訓練和推理的技術發展,而且還強調了系統層面和跨層協同設計對於驅動AI創新的重要性。我期待看到未來在CoE、記憶體系統和dataflow加速器方面的更多研究。
沒有留言:
發佈留言