作者｜Sean Lie

　　翻譯｜胡燕君、程浩源

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

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

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

　　為了實現上述目標，Cerebras設計了一種新的計算核心架構。它讓單臺設備運行超大規模模型成為可能，此外，它開發出只需簡單數據并行的橫向擴展和本地非結構化稀疏加速技術，使大模型的應用門檻大幅降低。

　　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是一款小型核心，它只有38,000平方微米，其中一半的硅面積用于48 KB內存，另一半是含110,000個標準單元(cell)的計算邏輯。整個計算核心以1.1 GHz的時鐘頻率高效運行，而峰值功率只有30毫瓦。

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

　　這意味著每一個來自內存的操作數(operand)至少要在數據通路中被使用100次，才能實現高利用率。要做到這一點，傳統的方法是通過本地緩存和本地寄存器實現數據復用。

　　然而，有一種方法可以讓數據通路以極致性能利用內存帶寬，就是將內存完全分布在要使用內存的單元旁邊。這樣一來，內存帶寬就等于核心數據通路的操作數帶寬。

　　這是一個簡單的物理原理：將比特數據從本地內存移動到數據通路，中間只有幾十微米的距離，相比將它通過數據包移動到外部設備要容易得多。

　　圖3展示了Cerebras計算核心的內存設計，每個核心配有48 KB本地SRAM，8個32位寬的單端口bank使其具備高密度，同時可保證充分發揮極致性能，這種級別的bank可提供超出數據通路所需的內存帶寬。

　　因此，我們可以從內存中提供極致數據通路性能，也就是每個循環只需2個64位讀取，一個64位寫入，因此它可以保證數據通路充分發揮性能。值得注意的是，每個核心的內存相互獨立，沒有傳統意義上的共享內存。

　　除了高性能的SRAM以外，Cerebras計算核心還具備一個256字節的軟件管理緩存，供頻繁訪問的數據結構使用，如累加器等。該緩存離數據通路非常緊湊，所以消耗的功率極低。上述分布式內存架構造就了驚人的內存帶寬，相當于同等面積GPU內存帶寬的200倍。

　　2

　　所有BLAS級別的極致性能

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

　　從圖4可見，在低于矩陣-矩陣相乘的任何BLAS級別都需要比內存帶寬的大幅增加，這一點傳統架構無法滿足。

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

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

　　除此之外，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

　　細粒度數據流調度

　　除了改進張量應用，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系統，做到了用單塊芯片實現集群級計算。

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

　　實現這樣的超大芯片尺寸，離不開底層架構的配合，底層架構必須能使數據在整片晶圓上高效、高性能地傳輸(圖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)。

　　芯片上看似簡單的短線其實十分重要，它們在硅上的距離不到一毫米。這種線路設計與傳統的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)。每行的每個元素都是標量，當然，在單行中，有多個權重映射到單個列上，當存在稀疏性時，只有非零權重才會被廣播到列，觸發FMAC計算。我們跳過所有的零權重，并輸入下一個非零權重，這就是產生稀疏加速的原因。

　　如果我們現在放大一個核心，可以看到核心架構是如何進行此操作(圖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)。最常見的是數據并行，這也是最簡單的方法，但它不適用于大型模型，因為它要求每個設備都有足夠的容量容納整個模型。

　　為了解決這個問題，常見的方法是采用模型并行，即劃分模型，以流水線方式用不同的設備運行模型的不同層。但隨著流水線變長，激活值內存(activation memory)以二次方的速度增長。

　　為了避免這種情況，另一種常見的模型并行方法是跨設備劃分層，但這會造成很大的通信開銷，而且劃分單個層非常復雜。

　　由于上述種種限制，今天仍沒有一種萬能的方式來實現橫向擴展。在大多數情況下，訓練海量模型需要數據并行和模型并行混合的方法。現存的橫向擴展解決方案仍有許多不足，根本原因很簡單：在傳統的橫向擴展中，內存和計算是緊密聯系的，如果在數千臺設備上運行單個模型，擴展內存和計算就變成相互依賴的分布式約束問題。

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

　　例如，張量模型的并行級別始終限制為8，因為在單個服務器中通常只有8個GPU。因此，大型模型大多采用流水式模型并行，這是最復雜的方法，原因就是之前提到的內存問題。在GPU集群上訓練模型需要解決這些分布式系統問題。這種復雜性導致需要更長的開發時間，并且往往無法實現最佳擴展。

　　7

　　Cerebras架構使擴展變得容易

　　Cerebras架構能夠在單個芯片上運行所有模型，無需模型分割，因此擴展變得簡單而自然，可以僅通過數據并行進行擴展，不需要任何復雜的模型并行分割。

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

　　SwarmX向所有CS-2系統廣播權重，并減少所有CS-2的梯度，它不僅僅是一個互聯，更是訓練過程中的一個活躍組件，專為數據并行橫向擴展而構建。

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

　　8

　　總結

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

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

　　然而，Cerebras架構支持用單個設備運行超大模型，同時支持只需數據并行的橫向擴展，以及本地非結構化稀疏加速，將讓更多人都能使用大模型。