深挖Cerebras:世界上最大AI晶片的架構設計

語言: CN / TW / HK

作者|Sean Lie\ 翻譯|胡燕君、程浩源

近年來,神經網路模型規模呈指數級增長,從2018年擁有超1億引數的Bert到2020年擁有1750億個引數GPT-3,短短兩年模型的引數量增加了3個數量級,而且這種增長還看不到盡頭。

人們剛剛開始發掘神經網路的應用潛力,但傳統的訓練和推理方式已然無法跟上神經網路規模的飛速增長速度,無法滿足大規模機器學習所需的記憶體和算力需求。為此,國內外諸多創業公司尋求對軟硬體等進行實質性的底層技術革新來解決這一挑戰。

作為業內備受關注的AI加速器創業公司,成立於2016年的Cerebras希望通過構建全新AI加速器方案解決AI計算問題,以實現數量級計算效能:首先,需要改進計算核心架構,而不只是一味地提升每秒浮點運算次數;其次,需要以超越摩爾定律的速度提高晶片整合度;最後,還要簡化叢集連線,大幅度提升叢集計算效率。

為了實現上述目標,Cerebras設計了一種新的計算核心架構。它讓單臺裝置執行超大規模模型成為可能,此外,它開發出只需簡單資料並行的橫向擴充套件和本地非結構化稀疏加速技術,使大模型的應用門檻大幅降低。

圖1:近年來各SOTA神經網路模型的記憶體與算力需求

2021年,Cerebras曾推出全球最大AI晶片Wafer Scale Engine 2(WSE-2),面積是46225平方毫米,採用7nm工藝,擁有2.6萬億個電晶體和85萬個AI優化核,還推出了世界上第一個人類大腦規模的AI解決方案CS-2 AI計算機,可支援超過120萬億引數規模的訓練。今年6月,它又在基於單個WSE-2晶片的CS-2系統上訓練了世界上最大的擁有200億引數的NLP模型,顯著降低了原本需要數千個GPU訓練的成本。

在近期舉辦的Hot Chips大會上,Cerebras聯合創始人&首席硬體架構師Sean Lie深入介紹了Cerebras硬體,展示了他們在核心架構、縱向擴充套件和橫向擴充套件方面的創新方法。以下是他的演講內容,由OneFlow社群編譯。

1

Cerebras計算核心架構

計算核心(compute core)是所有計算機架構的“心臟”,而Cerebras針對神經網路的細粒度動態稀疏性重新設計了計算核心。

圖2:Cerebras計算核心

圖2是一款小型核心,它只有38,000平方微米,其中一半的矽面積用於48 KB記憶體,另一半是含110,000個標準單元(cell)的計算邏輯。整個計算核心以1.1 GHz的時鐘頻率高效執行,而峰值功率只有30毫瓦。

先從記憶體說起。GPU等傳統架構使用共享中央DRAM,但DRAM存取速度較慢,位置也較遠。即便使用中介層(interposer)和HBM等尖端技術,其記憶體頻寬也遠低於核心資料通路頻寬。例如,資料通路頻寬通常是記憶體頻寬的100倍。

這意味著每一個來自記憶體的運算元(operand)至少要在資料通路中被使用100次,才能實現高利用率。要做到這一點,傳統的方法是通過本地快取和本地暫存器實現資料複用。

然而,有一種方法可以讓資料通路以極致效能利用記憶體頻寬,就是將記憶體完全分佈在要使用記憶體的單元旁邊。這樣一來,記憶體頻寬就等於核心資料通路的運算元頻寬。

這是一個簡單的物理原理:將位元資料從本地記憶體移動到資料通路,中間只有幾十微米的距離,相比將它通過資料包移動到外部裝置要容易得多。

圖3:Cerebras計算核心的記憶體設計:每個核心配有獨立記憶體。

圖3展示了Cerebras計算核心的記憶體設計,每個核心配有48 KB本地SRAM,8個32位寬的單埠bank使其具備高密度,同時可保證充分發揮極致效能,這種級別的bank可提供超出資料通路所需的記憶體頻寬。

因此,我們可以從記憶體中提供極致資料通路效能,也就是每個迴圈只需2個64位讀取,一個64位寫入,因此它可以保證資料通路充分發揮效能。值得注意的是,每個核心的記憶體相互獨立,沒有傳統意義上的共享記憶體。

除了高效能的SRAM以外,Cerebras計算核心還具備一個256位元組的軟體管理快取,供頻繁訪問的資料結構使用,如累加器等。該快取離資料通路非常緊湊,所以消耗的功率極低。上述分散式記憶體架構造就了驚人的記憶體頻寬,相當於同等面積GPU記憶體頻寬的200倍。

2

所有BLAS級別的極致效能

圖4:稀疏GEMM即對每個非零權重執行一次AXPY操作。

有了極大的記憶體頻寬,就可以實現許多卓越的功能。比如,可以充分發揮所有BLAS級別(基礎線性代數程式集,BLAS levels)的極致效能。傳統的CPU和GPU架構的片上記憶體頻寬有限,因此只能實現GEMM(通用矩陣乘法)的極致效能,即矩陣-矩陣相乘。

從圖4可見,在低於矩陣-矩陣相乘的任何BLAS級別都需要比記憶體頻寬的大幅增加,這一點傳統架構無法滿足。

但有了足夠的記憶體頻寬後,就可以讓GEMV(矩陣-向量相乘)、DOT(向量-向量相乘)和AXPY(向量-標量相乘)均實現極致效能。高記憶體頻寬在神經網路計算中尤為重要,因為這可以實現非結構化稀疏的充分加速。一個稀疏GEMM操作可看作是多個AXPY操作的合集(對每個非零元素執行一次操作)。

Cerebras計算核心的基礎是一個完全可程式設計的處理器,以適應不斷變化的深度學習需求。與通用處理器一樣,Cerebras核心處理器支援算術、邏輯、載入/儲存、比較(compare)、分支等多種指令。這些指令和資料一樣儲存在每個核心的48 KB本地記憶體中,這意味著核心之間相互獨立,也意味著整個晶片可以進行細粒度動態計算。通用指令在16個通用暫存器上執行,其執行在緊湊的6級流水線中。

等式1,FMAC指令示例

除此之外,Cerebras核心還在硬體層面支援所有有關資料處理的張量指令。這些張量運算元在64-位資料通路中執行,資料通路由4個FP16 FMAC(融合乘積累加運算)單元組成。

為了提升效能與靈活性,Cerebras的指令集架構(ISA)將張量視為與通用暫存器和記憶體一樣的一等運算元(first-class operand)。上圖等式1是一個FMAC指令的例子,它將3D和2D張量視為運算元直接執行。

之所以可以做到這一點,是因為Cerebras核心使用資料結構暫存器(DSR)作為指令的運算元。Cerebras核心有44個DSR,每個DSR包含一個描述符,裡面有指標指向張量及其長度、形狀、大小等資訊。

有了DSR後,Cerebras核心的硬體架構更靈活,即可以在記憶體中支援記憶體中的4D張量,也可支援織構張量(fabric streaming tensors)、FIFO(先進先出演算法)和環形緩衝器。此外,Cerebras核心還配有硬體狀態機來管理整個張量在資料通路中的流動次序。

3

細粒度資料流排程

圖5:核心資料通路及核心資料流排程。細粒度動態計算核心可提升計算效能,稀疏利用率為GPU的10倍。

除了改進張量應用,Cerebras核心還可執行細粒度資料流排程。如圖5所示,所有計算都由資料觸發。Fabric直接在硬體中傳輸資料和關聯控制元件,一旦核心接收資料,就開始查詢執行指令,查詢工作完全基於接收到的資料。這一資料流機制使整個計算結構變成一個數據流引擎,可以支援稀疏加速——因為它只處理非零資料。傳送器會過濾所有零值,因此接收器只會接收到非零值,而所有計算都由非零資料觸發。

這樣做不但可以節省功率,還可以省略不必要的計算,加快運算效率。操作由單個數據元素觸發,使得Cerebras核心可以支援超細粒度、完全非結構化的稀疏性,同時不會造成效能損失。由於資料流具有動態性,所以Cerebras核心還支援8個張量操作同時執行,我們稱之為“微執行緒(micro-threads)”。

微執行緒之間相互獨立,每次迴圈時硬體可在其間切換。排程器持續為所有待處理張量監控輸入和輸出是否可用,還具有優先處理機制,保證關鍵任務得到優先處理。當不同任務間的切換產生大量動態行為時,微執行緒可以提升利用率,否則這些動態行為可能會導致流水線出現氣泡。

上述細粒度、動態、小型核心架構等特點使我們的架構具備前所未有的高效能,其非結構化稀疏計算的利用率是GPU的至少10倍。可見,通過對計算核心架構的改進,Cerebras可將效能進行數量級提升。

4

縱向擴充套件:超越摩爾定律

要縱向擴充套件晶片,傳統的方法都是從晶片製造方面入手,即提升晶片整合度。過去數十年,晶片行業的發展都符合摩爾定律,晶片整合度越來越高。如今,摩爾定律還在延續,但它的增量不夠大,每一代製程只能將整合度提升約兩倍,不足以滿足神經網路的計算需求。所以,Cerebras希望可以超越摩爾定律,實現數量級的效能提升。

為此,我們嘗試過傳統的方法——擴大芯片面積,並在這方面做到了極致,成果就是WSE-2(Wafer-Scale Engine,晶圓級引擎)。如今,WSE-2的應用已非常廣泛。它是全世界最大的晶片,尺寸超過46,000平方毫米,是目前最大的CPU的56倍。單塊WSE-2有2.6萬億個電晶體,核心數達850,000個。龐大的芯片面積可以實現極大的片上記憶體和極高的效能。

為了讓尺寸驚人的WSE-2也能在標準的資料中心環境中使用,我們還針對性地設計了Cerebras CS-2系統,做到了用單塊晶片實現叢集級計算。

圖6:從小型核心到大型晶圓級引擎

以下是我們從小型核心構建大型晶圓級引擎的過程:首先,我們在整片直徑約300毫米的晶圓上做出一個個傳統晶粒(Die),每個晶粒含有約10,000個核心;然後,不同於以往的是,我們不將單個晶粒切割出來做成傳統晶片,而是在整片晶圓內切割出一個邊長215毫米的方塊,方塊包含84個晶粒,共有850,000個計算核心(圖6)。

圖7:高頻寬、低延遲的晶片結構

實現這樣的超大晶片尺寸,離不開底層架構的配合,底層架構必須能使資料在整片晶圓上高效、高效能地傳輸(圖7)。Cerebras的晶片結構使用2D網格拓撲,這種結構非常適合擴充套件,而且只需消耗極低的開銷。

網格拓撲將所有核心連線起來,每個核心在網狀拓撲中有一個結構路由器(fabric router)。結構路由器有5個埠,4個方向各有1個,還有一個埠面向核心自身,各個埠都有32位的雙向介面。埠數量較少的好處是可以將節點間延時保持在一個時鐘週期以內,從而實現低成本、無損流控和非常低的緩衝。

晶片中的基本資料包是針對神經網路優化後的單個FP16資料元素,與之伴隨的是16位的控制資訊,它們共同組成32位的超細粒度資料包。

為了進一步優化晶片結構,我們使用了靜態路由(static routing),效率高,開銷低,而且可以充分利用神經網路的靜態連線。為了讓同一物理連線上可以有多條路由,我們提供24條相互獨立的靜態路由以供配置,路由之間無阻塞,且都可以通過時分複用(time-multiplexing)技術在同一物理連線上傳輸。

最後,由於神經網路傳輸需要高扇出(fan-out),因此Cerebras晶片的每個結構路由器都具有本地廣播(native broadcast)和多播(multi-cast)能力。

有了上述基礎後,我們就可以進行擴充套件。在單個晶粒內進行擴充套件比較簡單,但現在需要將晶粒與晶粒連線起來。為了跨越晶粒間不到一毫米寬的劃片槽(scribe line),我們使用了臺積電工藝中的高階金屬層。

我們將計算核心擴充套件為2D網格計算結構,然後又在整個晶圓上形成了完全同質的計算核心陣列。晶粒-晶粒介面是一種高效的源同步並行介面,但是,在如此大的晶圓規模上,總共有超過一百萬條線路,所以我們的底層協議必須採用冗餘度設計。我們通過訓練和自動校正狀態機來做到這一點。有了這些介面,即使在製造過程中存在瑕疵,整個晶圓的結構也能做到完全均質結構(uniform fabric)。

圖8:整個晶圓上的均質結構(uniform fabric)。

晶片上看似簡單的短線其實十分重要,它們在矽上的距離不到一毫米。這種線路設計與傳統的SERDES方法很不一樣。與前面提到的的記憶體設計相同,短線設計是出於簡單的物理原理:在晶片上將位元資料傳輸不到1毫米的距離,比通過封裝聯結器、PCB或者線纜傳輸都更容易。

與傳統IO相比,這種方法帶來了數量級的改進。從圖8資料可看出,WSE-2每單位面積的頻寬比GPU多出約一個數量級,並且每位元的功率效率提高了近兩個數量級。這些都表明整個晶圓結構具備了前所未有的高效能。

如果轉化為同等的GPU面積,WSE-2的頻寬是GPU的7倍,而功率僅約5瓦。正是這種級別的全域性結構效能,使晶圓能夠作為單個晶片執行。有了如此強大的單晶片,我們就可以解決一些極具挑戰性的問題。

5

通過權重流式技術支援超大模型

圖9:通過權重流式(Weight Streaming)技術可在單個晶片上支援所有模型大小。

高效能的晶片結構可以讓我們在單個晶片上執行大型神經網路。WSE-2具有足夠高的效能和容量來執行如今最大的模型,且無需分割槽或複雜的分散式處理,這是通過分解神經網路模型、權重和計算來完成的。

我們將所有模型權重儲存在名為MemoryX的外部裝置中,並將這些權重流式傳輸到CS-2系統。權重會在神經網路各層的計算中用到,而且一次只計算一層。權重不會儲存在CS-2系統上,哪怕是暫時儲存。CS-2接收到權重後,使用核心中的底層資料流機制執行計算(圖9)。

每個單獨的權重都會作為單獨的AXPY操作觸發計算。完成計算後,該權重就會被丟棄,硬體將繼續處理下一個元素。由於晶片不需要儲存權重,所以晶片的記憶體容量不會影響晶片可處理的模型大小。在反向傳播中,梯度以相反的方向流回到MemoryX單元,然後MemoryX單元進行權重更新。

圖10 完整的晶圓是MatMul陣列,可支援超大矩陣。

以下是晶片中執行計算的具體方法。神經網路各層的計算可歸結為矩陣乘法,由於CS-2的規模較大,我們能夠將晶圓的85萬個核心用作單個巨型矩陣乘法器。

它是這樣工作的:對於像GPT這樣的Transformer模型,啟用張量具有三個邏輯維度:批次(B)、序列(S)和隱藏(H)維度,我們將這些張量維度拆分到晶圓上的二維核心網格上。隱藏維度在晶片結構的x方向上劃分(split),而批次和序列維度在y方向上劃分。這樣可以實現高效的權重廣播以及序列和隱藏維度的高效歸約。

啟用函式儲存在負責執行計算工作的核心上,下一步是觸發這些啟用函式的計算,這是通過使用片上廣播結構來完成的。我們使用片上廣播結構來向每一列傳送權重、資料和命令的方法。

當然,在硬體資料流機制下,權重會直接觸發FMAC操作。這些是AXPY操作。由於廣播發生在列上,因此包含相同特徵子集的所有核心接收相同的權重。此外,我們傳送命令來觸發其他計算,例如歸約或非線性操作。

圖11:資料流排程以低開銷實現完全非結構化的稀疏MatMul運算。

舉個例子,我們首先在整個晶圓上廣播權重行(圖11)。每行的每個元素都是標量,當然,在單行中,有多個權重對映到單個列上,當存在稀疏性時,只有非零權重才會被廣播到列,觸發FMAC計算。我們跳過所有的零權重,並輸入下一個非零權重,這就是產生稀疏加速的原因。

圖12:稀疏輸入的GEMM:乘法和partial sum歸約。

如果我們現在放大一個核心,可以看到核心架構是如何進行此操作(圖12)。在資料流機制下,權重抵達後,就會觸發核心上的FMAC計算。權重值與每個啟用函式輸出相乘,然後累加到軟體管理快取中的本地累加器中。FMAC計算使用張量指令執行,將啟用函式輸出視為張量運算元。上述計算都不會對核心造成額外開銷。

此外,權重也不會產生記憶體開銷,因為一旦計算完成,核心就會轉而計算下一個權重,不需要儲存任何權重。若整行核心都接收到權重,每個核心就都會產生一個partial sum,然後該行核心的所有partial sum將進行歸約。

歸約計算由被廣播到每列所有核心的命令包觸發。同樣,在資料流排程機制下,一旦核心接收到命令包,它就會觸發partial sum歸約計算。實際的歸約計算本身是使用核心的張量指令完成,使用的是結構張量運算元。所有列都接收一個PSUM命令。但是其中一列會收到一個特殊的FSUM命令,它要求核心儲存final sum。這樣做是為了使用與輸入特徵相同的分佈來儲存輸出特徵,從而為下一層計算做好準備。

收到命令後,核心使用結構上的環形模式進行通訊,該模式使用結構靜態路由設定。使用微執行緒,所有歸約都與下一個權重行的FMAC計算重疊,該種FMAC計算並行開始。當所有的權重行都處理完畢,完整的GEMM操作就完成了,同時所有的啟用函式輸出都已完備,可以進行下一層計算。

上述設計能讓各種規模的神經網路都可以在單個晶片上高效能執行。獨特的核心記憶體和晶片架構使晶片可以無需分塊或分割槽即可支援超大矩陣,即使是具有多達100,000 x 100,000 MatMul層的超大模型也可以在不拆分矩陣的情況下執行。

若使用單個WSE-2晶片執行此模型,FP16稀疏效能可達75 PetaFLOPS(若稀疏性更高,效能還可更高),FP16密集效能可達7.5 PetaFLOPS。這就是我們應對機器學習硬體挑戰的第二個方面,通過擴充套件進一步帶來一個數量級的效能提升。

6

橫向擴充套件:為什麼這麼難

最後一個方面:叢集橫向擴充套件。如今已經存在叢集解決方案,但為什麼橫向擴充套件仍然如此困難?

圖13:分佈複雜性隨叢集規模顯著增加。

讓我們看看現有的橫向擴充套件技術(圖13)。最常見的是資料並行,這也是最簡單的方法,但它不適用於大型模型,因為它要求每個裝置都有足夠的容量容納整個模型。

為了解決這個問題,常見的方法是採用模型並行,即劃分模型,以流水線方式用不同的裝置執行模型的不同層。但隨著流水線變長,啟用值記憶體(activation memory)以二次方的速度增長。

為了避免這種情況,另一種常見的模型並行方法是跨裝置劃分層,但這會造成很大的通訊開銷,而且劃分單個層非常複雜。

由於上述種種限制,今天仍沒有一種萬能的方式來實現橫向擴充套件。在大多數情況下,訓練海量模型需要資料並行和模型並行混合的方法。現存的橫向擴充套件解決方案仍有許多不足,根本原因很簡單:在傳統的橫向擴充套件中,記憶體和計算是緊密聯絡的,如果在數千臺裝置上執行單個模型,擴充套件記憶體和計算就變成相互依賴的分散式約束問題。

圖14:GPU叢集在實踐中的複雜性。

這種複雜性導致的結果是:圖14顯示了過去幾年在GPU上訓練的最大模型及其使用的不同並行方法。從中可見,越大的模型需要的並行型別也越多,增加了複雜性。

例如,張量模型的並行級別始終限制為8,因為在單個伺服器中通常只有8個GPU。因此,大型模型大多采用流水式模型並行,這是最複雜的方法,原因就是之前提到的記憶體問題。在GPU叢集上訓練模型需要解決這些分散式系統問題。這種複雜性導致需要更長的開發時間,並且往往無法實現最佳擴充套件。

7

Cerebras架構使擴充套件變得容易

Cerebras架構能夠在單個晶片上執行所有模型,無需模型分割,因此擴充套件變得簡單而自然,可以僅通過資料並行進行擴充套件,不需要任何複雜的模型並行分割。

圖15:使用MemoryX和SwarmX進行擴充套件,只需近線性的資料並行。

我們為資料並行專門設計了SwarmX(圖15)互聯技術。它位於儲存權重的MemoryX單元和用於計算的CS-2系統之間,但又獨立於兩者。

SwarmX向所有CS-2系統廣播權重,並減少所有CS-2的梯度,它不僅僅是一個互聯,更是訓練過程中的一個活躍元件,專為資料並行橫向擴充套件而構建。

在內部,SwarmX使用樹形拓撲來實現模組化和低開銷擴充套件,因為它是模組化和可分解的,所以擴充套件到任意數量的具有與單個系統相同的執行模型的CS-2系統。要擴充套件到更多計算,只需在SwarmX拓撲中新增更多節點和更多CS-2系統。這就是我們應對機器學習硬體需求的最後一個方面:改進並大大簡化橫向擴充套件。

8

總結

在過去的幾年裡,機器學習工作負載的需求增加了三個數量級以上,而且沒有放緩的跡象。預計幾年後將增長到圖16的箭頭位置,我們問自己,可以滿足這種需求嗎?

圖16:各種最先進的神經網路的記憶體和計算要求。橫、縱座標每一格代表一個數量級的提升。

Cerebras相信,我們可以,但不是通過傳統技術做到這一點,而是通過非結構化稀疏加速、晶圓級晶片和叢集橫向擴充套件的結合將效能提升三個數量級。神經網路模型規模依然呈指數級增長,可以使用這些大模型的公司很少,而且未來只會更少。

然而,Cerebras架構支援用單個裝置執行超大模型,同時支援只需資料並行的橫向擴充套件,以及本地非結構化稀疏加速,將讓更多人都能使用大模型。

(原文:https://www.cerebras.net/blog/cerebras-architecture-deep-dive-first-look-inside-the-hw/sw-co-design-for-deep-learning; 影片:https://www.youtube.com/watch?v=8i1_Ru5siXc )

歡迎下載體驗 OneFlow v0.8.0 最新版本:\ https://github.com/Oneflow-Inc/oneflow/