晶圓級(jí)AI芯片WSE-3性能公布：80億參數(shù)模型上每秒生成1800個(gè)Token

今年3月，新創(chuàng)AI芯片公司Cerebras Systems推出了其第三代的晶圓級(jí)AI芯片WSE-3，性能達(dá)到了上一代WSE-2的兩倍，可用于訓(xùn)練業(yè)內(nèi)一些最大的人工智能模型。在近日的Hot Chips 2024大會(huì)上，Cerebras Systems詳細(xì)介紹了這款芯片在AI推理方面的性能。

根據(jù)官方資料顯示，WSE-3依然是采用了一整張12英寸晶圓來制作，基于臺(tái)積電5nm制程，芯片面積為46225平方毫米，擁有的晶體管數(shù)量達(dá)到了4萬億個(gè)，擁有90萬個(gè)AI核心，44GB片上SRAM，整體的內(nèi)存帶寬為21PB/s，結(jié)構(gòu)帶寬高達(dá)214PB/s。使得WSE-3具有125 FP16 PetaFLOPS的峰值性能，相比上一代的WSE-2提升了1倍。

作為對(duì)比，WSE-2芯片面積同樣是46225平方毫米，基于臺(tái)積電7nm制程，晶體管數(shù)量為2.6萬億個(gè)，AI內(nèi)核數(shù)量為85萬個(gè)，片上內(nèi)存SRAM為40GB，內(nèi)存帶寬為20PB/s，結(jié)構(gòu)帶寬高達(dá)220PB/s。

如果將其與英偉達(dá)的H100相比，WSE-3面積將是H100的57倍，內(nèi)核數(shù)量是H100的52倍，片上內(nèi)存是H100的880倍，內(nèi)存帶寬是H100的7000倍，結(jié)構(gòu)帶寬是H100的3715倍。（ H200 的 HBM3e 僅擁有 4.8TBps 的帶寬。）

在此次的Hot Chips 2024大會(huì)上，Cerebras公布了更多關(guān)于WSE-3在運(yùn)行AI大模型上的性能表現(xiàn)。

Cerebras表示，它在Llama3.1-8B上的推理速度比微軟Azure等公司使用NVIDIA H100快了20倍。需要指出的是，在許多現(xiàn)代生成式 AI 工作負(fù)載中，推理性能通常更得益于內(nèi)存帶寬的大小，而不單單是計(jì)算能力。也就是說，擁有更大的內(nèi)存帶寬，模型的推理速度就越快。

Cerebra Systems首席執(zhí)行官 Andrew Feldman 稱，WSE-3通過使用 44GB片上SRAM，使得其能夠以 16 位精度運(yùn)行 Llama 3.1 8B 時(shí)，每秒能夠生成超過 1,800 個(gè)Token，而性能最好的基于英偉達(dá)H100的實(shí)例每秒只能生成超過 242 個(gè)Token。

與此同時(shí)，Cerebras還推出了基于WSE-3的CS-3超級(jí)計(jì)算機(jī)，可用于訓(xùn)練參數(shù)高達(dá)24萬億的人工智能模型，這比相比基于WSE-2和其他現(xiàn)代人工智能處理器的超級(jí)計(jì)算機(jī)有了重大飛躍。該超級(jí)計(jì)算機(jī)可以支持1.5TB、12TB或1.2PB的外部內(nèi)存，這使它能夠在單個(gè)邏輯空間中存儲(chǔ)大量模型，而無需分區(qū)或重構(gòu)，從而簡化了訓(xùn)練過程，提高了開發(fā)人員的效率。

最新的 Cerebras 軟件框架可以為PyTorch 2.0 和最新的 AI 模型和技術(shù)（如多模態(tài)模型、視覺轉(zhuǎn)換器、專家混合和擴(kuò)散）提供原生支持。Cerebras 仍然是唯一為動(dòng)態(tài)和非結(jié)構(gòu)化稀疏性提供本機(jī)硬件加速的平臺(tái)，相比英偉達(dá)的DGX-100計(jì)算機(jī)系統(tǒng)，將訓(xùn)練速度提高了8 倍。

在運(yùn)行分布在四個(gè)CS-3加速器上的 700 億參數(shù)版本的 Llama 3.1 大模型時(shí)，也能夠?qū)崿F(xiàn)每秒 450 個(gè)Token。相比之下，H100 可以管理的最佳狀態(tài)是每秒 128 個(gè)Token。

Feldman 認(rèn)為，這種性能水平，就像寬帶的興起一樣，將為 AI 的采用開辟新的機(jī)會(huì)?！敖裉?，我認(rèn)為我們正處于 Gen AI 的撥號(hào)時(shí)代，”他說，并指出了生成式 AI 的早期應(yīng)用，其中提示的響應(yīng)會(huì)有明顯的延遲。

他認(rèn)為，如果能夠足夠快地處理請(qǐng)求，就可以基于多個(gè)模型構(gòu)建代理應(yīng)用程序，而不會(huì)因?yàn)檠舆t變得難以為繼。Feldman 認(rèn)為這種性能有益的另一個(gè)應(yīng)用是允許 LLM 在多個(gè)步驟中迭代他們的答案，而不僅僅是吐出他們的第一個(gè)響應(yīng)。如果您可以足夠快地處理Token，則可以在幕后做更多的處理。

雖然WSE-3能夠以 16 位精度運(yùn)行 Llama 3.1 8B 時(shí)，每秒能夠生成超過 1,800 個(gè)Token，但是如果不是因?yàn)橄到y(tǒng)受計(jì)算限制，WSE-3的速度應(yīng)該能夠更快。

該產(chǎn)品代表了 Cerebras 的一些轉(zhuǎn)變，因?yàn)榇饲?，Cerebras 主要專注于 AI 訓(xùn)練。雖然現(xiàn)在也開始應(yīng)用于AI推理，但是硬件本身實(shí)際上并沒有改變。Feldman表示，他們正在使用相同的 WSE-3 芯片和 CS-3 系統(tǒng)進(jìn)行推理和訓(xùn)練。

“我們所做的是擴(kuò)展了編譯器的功能，可以同時(shí)在芯片上放置多個(gè)層，”Feldman解釋說。

SRAM 速度很快，但使 HBM 容量更大

雖然 SRAM 在性能方面比 HBM 具有明顯的優(yōu)勢(shì)，但它的不足之處在于容量。對(duì)于大型語言模型 （LLM）來說，44GB的容量并不多，因?yàn)楸仨毧紤]到鍵值緩存在WSE-3所針對(duì)的高批處理大小下占用了相當(dāng)多的空間。

Meta 的 Llama 3 8B 模型是 WSE-3 的理想化場景，因?yàn)榇笮?16GB（FP16），整個(gè)模型可以安裝在芯片的 SRAM 中，為鍵值緩存留下大約 28GB 的空間。

Feldman 聲稱，除了極高的吞吐量外，WSE-3 還可以擴(kuò)展到更高的批量大小。盡管它究竟可以擴(kuò)展到多大程度并保持每個(gè)用戶Token的生成率，這家初創(chuàng)公司不愿透露。“我們目前的批次大小經(jīng)常變化。我們預(yù)計(jì)第四季度的批量規(guī)模將達(dá)到兩位數(shù)，“Cerebras說道。

當(dāng)被追問更多細(xì)節(jié)時(shí)，Feldman補(bǔ)充說：“我們目前的批量大小還不成熟，因此我們寧愿不提供它。系統(tǒng)架構(gòu)旨在以高批量運(yùn)行，我們預(yù)計(jì)在未來幾周內(nèi)實(shí)現(xiàn)這一目標(biāo)。”

與現(xiàn)代 GPU 非常相似，Cerebras 通過跨多個(gè) CS-3 系統(tǒng)并行化模型來應(yīng)對(duì)這一挑戰(zhàn)。具體來說，Cerebras 正在使用管道并行性將模型的層分布到多個(gè)系統(tǒng)。

對(duì)于需要 140GB 內(nèi)存的 Llama 3 70B，該模型的 80 層分布在四個(gè)通過以太網(wǎng)互連的 CS-3 系統(tǒng)中。這確實(shí)會(huì)帶來性能損失，因?yàn)閿?shù)據(jù)必須通過這些鏈接。

△由于單個(gè)WSE-3只有 44GB SRAM，因此需要將多個(gè)加速器連接在一起以支持更大的模型

然而，根據(jù) Feldman 的說法，節(jié)點(diǎn)到節(jié)點(diǎn)的延遲并不像您想象的那么大?！斑@里的延遲是真實(shí)的，但很小，并且它與通過芯片上所有其他層的Token分?jǐn)?，”他解釋說?！白詈螅琓oken上的晶圓到晶圓延遲約占總數(shù)的 5%。”

對(duì)于更大的模型，例如最近宣布的 4050 億參數(shù)變體的 Llama 3，Cerebras 估計(jì)它將能夠使用 12 個(gè) CS-3 系統(tǒng)實(shí)現(xiàn)每秒約 350 個(gè)Token。

利用更高速的片上SRAM來替代HBM并不是一個(gè)新鮮事，Cerebra的競爭對(duì)手 Groq也是這么做的。

Groq 的語言處理單元 （LPU） 實(shí)際上使用了片上SRAM。不同之處在于，由于單個(gè) Groq LPU SRAM 容量較低，因此需要通過光纖連接更多的加速器來支持更大的模型。

Cerebras 需要四個(gè) CS-3 系統(tǒng)才能以每秒 450 個(gè)令牌的速度運(yùn)行 Llama 3 70B，Groq 此前曾表示，它需要 576 個(gè) LPU 才能實(shí)現(xiàn)每秒300個(gè)Token。而Cerebras 引用的人工智能分析 Groq 基準(zhǔn)測(cè)試略低，為每秒 250 個(gè)Token。

Feldman 還指出，Cerebras 能夠在不求助于量化的情況下做到這一點(diǎn)。認(rèn)為Groq 正在使用 8bit量化來達(dá)到他們的性能目標(biāo)，這減少了模型大小、計(jì)算開銷和內(nèi)存壓力，但代價(jià)是準(zhǔn)確性有所損失。

不過，僅比較性能，而忽略整體的成本是不公平的對(duì)比。畢竟單個(gè)晶圓級(jí)的WSE-3芯片的成本也遠(yuǎn)遠(yuǎn)高于Groq LPU的成本。

編輯：芯智訊-浪客劍