勞倫斯利弗莫爾國家實(shí)驗(yàn)室、桑迪亞國家實(shí)驗(yàn)室和洛斯阿拉莫斯國家實(shí)驗(yàn)室在 HPC 社區(qū)中被簡稱為“三重實(shí)驗(yàn)室”，但這些 HPC 中心或許可以稱為“嘗試實(shí)驗(yàn)室”，因?yàn)樗鼈儦v史上嘗試過幾乎任何新的架構(gòu)，以了解它在推進(jìn)美國能源部的使命方面可能有何前景。

桑迪亞國家實(shí)驗(yàn)室是先鋒計(jì)劃測試新架構(gòu)的地方，現(xiàn)在它帶著 Cerebras Systems 的第三代晶圓級系統(tǒng)再次回歸，希望在一臺(tái)真正設(shè)計(jì)用于運(yùn)行 AI 訓(xùn)練和推理的機(jī)器上突破傳統(tǒng) HPC 代碼的性能障礙。

兩年前，桑迪亞從 Cerebras 收購了數(shù)量不詳?shù)?CS-2 系統(tǒng)，每個(gè)系統(tǒng)都有一個(gè) CPU 主機(jī)和一個(gè) WSE-2 晶圓級處理器，目的是將一些矩陣密集的 HPC 計(jì)算卸載到 WSE-2 引擎上的 16 位浮點(diǎn)核心上。

為什么桑迪亞甚至?xí)紤]將其 64 位或 32 位格式的計(jì)算精度降低四倍或兩倍？因?yàn)檫@些 WSE-2 引擎（正如我們在 2022 年 3 月詳細(xì)介紹的那樣）將 850,000 個(gè)內(nèi)核和 40 GB 的片上 SRAM 內(nèi)存（蝕刻在 2.6 萬億個(gè)晶體管中）塞進(jìn)一個(gè)餐盤大小的方形硅片中，具有 20 PB/秒的內(nèi)存帶寬和 6.25 千萬億次浮點(diǎn)運(yùn)算的密集矩陣和 62.5 千萬億次浮點(diǎn)運(yùn)算的稀疏矩陣。

我們的想法是，對于某些類型的工作負(fù)載，如果問題適合內(nèi)存，或者可以分解為跨越這些核心的部分，那么在單個(gè)大型設(shè)備上進(jìn)行計(jì)算可以使 HPC 模擬運(yùn)行得更快。

今年早些時(shí)候，桑迪亞國家實(shí)驗(yàn)室的一個(gè)研究小組證明了這一點(diǎn)，單個(gè) CS-2 系統(tǒng)在某種分子動(dòng)力學(xué)模擬中擊敗了橡樹嶺國家實(shí)驗(yàn)室的“Frontier”超級計(jì)算機(jī)。

更大的 Frontier 機(jī)器擁有 37,632 個(gè) AMD “Aldebaran” MI250X GPU 加速器，可以模擬晶格中的大量原子，但由于這些 GPU 集群的擴(kuò)展性較弱，它無法模擬長時(shí)間擺動(dòng)的原子。節(jié)點(diǎn)之間的延遲使這成為不可能。

但是，使用改進(jìn)的 LAMMPS 分子動(dòng)力學(xué)模擬，桑迪亞設(shè)立了一個(gè)測試，其中鎢、銅和鉭晶格由靜態(tài)原子數(shù)組成——801,792，足以讓一個(gè) WSE-2 核心保存一個(gè)原子的數(shù)據(jù)——然后模擬這些晶格受到輻射沖擊。與龐大的 Frontier 系統(tǒng)中的 GPU 相比，桑迪亞在一臺(tái) WSE-2 計(jì)算引擎上進(jìn)行的 LAMMPS 模擬中每秒可處理的時(shí)間步數(shù)對于銅高 109 倍，對于鎢高 96 倍，對于鉭高 179 倍。這為 Cerebras 鐵上的模擬提供了數(shù)十毫秒的時(shí)間，正如我們當(dāng)時(shí)指出的那樣，有足夠的時(shí)間來實(shí)際觀察當(dāng)你用能量戳晶格時(shí)晶格會(huì)發(fā)生什么。

在 Frontier 機(jī)器上，該應(yīng)用程序的擴(kuò)展在 32 個(gè) GPU 時(shí)逐漸減弱，這令人失望，并且表明對于某些類型的應(yīng)用程序來說，很難在單個(gè)計(jì)算設(shè)備之外進(jìn)行擴(kuò)展。

當(dāng)時(shí)，我們推測，如果 Sandia 升級到今年 3 月推出的 CS-3 系統(tǒng)中的 WSE-3 計(jì)算引擎，它就可以贏得更多時(shí)間（即模擬時(shí)間） 。借助 WSE-3 引擎，Cerebras 將晶體管縮小到 5 納米（WSE-2 為 7 納米），并將核心數(shù)量增加到 900,000 個(gè)，但轉(zhuǎn)移到 8 寬 FP16 SIMD 單元，是 WSE-2 和 WSE-1 引擎中使用的 SIMD 單元寬度的兩倍。我們認(rèn)為 WSE-3 的時(shí)鐘速度提高了約 5%，當(dāng)你將時(shí)鐘速度、核心跳躍和 SIMD 提升相乘時(shí)，這就是 WSE-3 的性能是 WSE-2 的 2 倍的原因。

我們猜測，通過從 WSE-2 移至 WSE-3，這種性能提升可能會(huì)將鉭晶格輻射的模擬窗口從 40 毫秒增加到 80 毫秒。相比之下，F(xiàn)rontier 機(jī)器上的這些節(jié)點(diǎn)模擬時(shí)間約為 200 納秒。

嗯，看起來桑迪亞國家實(shí)驗(yàn)室的人們想要獲得一些 WSE-3 計(jì)算引擎并一探究竟。我們還強(qiáng)烈懷疑他們想弄清楚是否可以將模擬擴(kuò)展到多個(gè)晶圓上并突破 1 秒模擬障礙。

也許是為了這個(gè)目的，也為了其他目的，桑迪亞國家實(shí)驗(yàn)室和 Cerebras 已經(jīng)開始構(gòu)建一個(gè)綽號為“Kingfisher”的系統(tǒng)，該系統(tǒng)將以四個(gè) CS-3 系統(tǒng)開始，并將在未來某個(gè)時(shí)候擴(kuò)展到八個(gè)系統(tǒng)。Kingfisher 集群將在傳統(tǒng)的 HPC 模擬工作和 AI 工作上加倍努力——當(dāng)然是生成式 AI，但不一定僅限于此——這可以增強(qiáng)三實(shí)驗(yàn)室在國家核安全局的支持下進(jìn)行的處理，該局資助三實(shí)驗(yàn)室管理美國軍方的核武器庫存。具體來說，Kingfisher 是由核威懾高級模擬和計(jì)算人工智能計(jì)劃資助的。

桑迪亞研究人員 Thuc Hoang、Ann Gentile、Andrew Younge、Si Hammond、James Laros 和 Kevin Stroup 站在 Kingfisher 旁邊。

在宣布 Kingfisher 系統(tǒng)的聲明中，一直領(lǐng)導(dǎo)桑迪亞 CS-2 系統(tǒng)工作的研究人員之一 James Laros 表示，實(shí)驗(yàn)室正在探索使用未來版本的 WSE 計(jì)算引擎“用于 Mod-Sim 和 AI 工作負(fù)載的組合”的可行性。我們過去曾與 Cerebras 聯(lián)合創(chuàng)始人兼首席執(zhí)行官 Andrew Feldman 開玩笑說，世界真正需要的是一個(gè)擁有 64 位 SIMD 引擎的 WSE，該引擎可以擴(kuò)展到 FP64 精度，甚至可以降低到 FP4 精度，并且可以動(dòng)態(tài)執(zhí)行，可能是在晶圓上的不同塊中，也可能是在代碼運(yùn)行時(shí)動(dòng)態(tài)執(zhí)行，因此理論上 HPC 中心的任何代碼都可以在 Cerebras 硬件上運(yùn)行。

如果有足夠多的人這么說，并且有足夠多的人為此提供資金，也許這就會(huì)實(shí)現(xiàn)。對我們來說，這些都不是笑翠鳥。Nvidia 不再專注于其 GPU 的 FP64 性能。

Kingfisher 系統(tǒng)的成本尚未公布，但我們知道，按照標(biāo)價(jià)（無論 HPC 領(lǐng)域的標(biāo)價(jià)是多少），去年 G42 開始與 Cerebras 合作，基于 CS-2 機(jī)器構(gòu)建“Condor Galaxy”集群時(shí)，配備單個(gè) WSE-2 的 CS-2 系統(tǒng)的成本約為 160 萬美元。也許這意味著要獲得 CS-3，價(jià)格要上漲 1.5 倍，或者可能上漲 2 倍。在一個(gè)要求每一代都物有所值的世界里，230 萬美元到 250 萬美元可能是合理的。在一個(gè)需要更高性能且選擇很少的世界里，320 萬美元的 CS-3 節(jié)點(diǎn)成本是合理的。無論如何，我們認(rèn)為桑迪亞不會(huì)為機(jī)器支付接近標(biāo)價(jià)的價(jià)格，但同時(shí)希望幫助資助那些可能幫助其更好地運(yùn)行模擬的公司。我們可以肯定地告訴你的是，F(xiàn)rontier 的 32 GPU 部分僅花費(fèi)約 425,000 美元，但它只能擴(kuò)展到這么大。

那么，增加原子數(shù)量和增加模擬時(shí)間對 NNSA 來說有什么價(jià)值呢？可能價(jià)值很大。

我們期待進(jìn)一步了解 Kingfisher 所做的事情以及它是如何做到的。

最后一件事：桑迪亞團(tuán)隊(duì)在 Cerbras 晶圓級系統(tǒng)上進(jìn)行的分子動(dòng)力學(xué)研究將角逐今年的戈登貝爾獎(jiǎng)。我們希望他們能獲獎(jiǎng)，這樣 GPU 供應(yīng)商們就得時(shí)刻保持警惕了。