首頁 > 科技要聞 > 科技> 正文

AI幾小時設(shè)計芯片超越人類！谷歌AlphaChip登Nature，已設(shè)計出三代旗艦TPU

新智元 整合編輯：太平洋科技 發(fā)布于：2024-09-28 00:15

能設(shè)計芯片的AI黑科技來了！

就在剛剛，谷歌DeepMind推出名為AlphaChip的AI系統(tǒng)。

無論是設(shè)計最先進的用于構(gòu)建AI模型的TPU，還是數(shù)據(jù)中心的CPU，它在相關(guān)的眾多領(lǐng)域，都產(chǎn)生了廣泛影響。

在谷歌的許多款芯片設(shè)計中，它都取得了出色的效果，比如Axion芯片（一種基于Arm 的通用數(shù)據(jù)中心CPU）。

AlphaChip設(shè)計芯片，用的是強化學(xué)習(xí)的原理。

也就是說，芯片布局設(shè)計對它來說是一種游戲，就像AlphaGo一樣，它在游戲中，學(xué)習(xí)如何設(shè)計出最好的芯片布局。

幾小時內(nèi)，它就能生成超出人類水平，或是與人類專家相當?shù)男酒季至恕?/span>

現(xiàn)在，它已經(jīng)用于設(shè)計多代TPU芯片（TPU v5e、TPU v5p和Trillium）。而且跟人類專家相比，AlphaChip放置的塊數(shù)越來越多，線長也減少了許多。

布局五年，谷歌多代TPU全由AI設(shè)計

其實谷歌對于這個AI，已經(jīng)布局多年了。

早在2020年，團隊就發(fā)表了一篇預(yù)印本論文，介紹了谷歌的全新強化學(xué)習(xí)方法，用于設(shè)計芯片布局。

后來在2021年，這項工作發(fā)表在了Nature上，并且進行了開源。

而自從首次發(fā)表這項工作以來，谷歌內(nèi)部一直在對它進行改進。

今天，團隊發(fā)表了Nature附錄，詳細描述了具體方法，及其對芯片設(shè)計領(lǐng)域的影響。

同時，他們還發(fā)布了一個預(yù)訓(xùn)練的檢查點，分享了模型權(quán)重，公布模型名稱為AlphaChip。

谷歌表示，AlphaChip是最早用于解決現(xiàn)實世界工問題的強化學(xué)習(xí)方法之一。

在數(shù)小時內(nèi)，它就可以生成超人或類似的芯片布局，而不需要花費數(shù)周或數(shù)月的人類。它設(shè)計的芯片布局，已經(jīng)被用于世界各地的芯片中，包括數(shù)據(jù)中心和移動電話。

為了設(shè)計TPU布局，AlphaChip首先在來自前幾代的各種芯片模塊上進行實踐，例如片上和片間網(wǎng)絡(luò)模塊、內(nèi)存控制器和數(shù)據(jù)傳輸緩沖區(qū)。這一過程被稱為預(yù)訓(xùn)練。

然后，團隊在當前的TPU模塊上運行AlphaChip，以生成高質(zhì)量的布局。

與之前的方法不同，AlphaChip在解決更多芯片布局任務(wù)時變得更好、更快，類似于人類專家的工作方式。

對于每一代新的TPU，包括谷歌最新的Trillium(第6代)，AlphaChip都設(shè)計了更好的芯片布局，并提供了更多的總體布局圖，從而加快了設(shè)計周期，產(chǎn)生了更高性能的芯片。

條形圖顯示了谷歌三代TPU上AlphaChip設(shè)計的芯片塊的數(shù)量，包括v5e、v5p和Trillium

條形圖顯示，跟TPU物理設(shè)計團隊生成的布局相比，AlphaChip在TPU三代產(chǎn)品中的平均有線長度減少

工作原理：一邊設(shè)計，一邊獎勵
其實，設(shè)計芯片布局并不是一項簡單的任務(wù)。

一般來說，計算機芯片有許多相互連接的模塊、多層電路元件組成，所有這些部件都由纖細無比的導(dǎo)線連接起來。

此外，還有許多復(fù)雜且相互交織的設(shè)計約束，必須同時滿足。

由于設(shè)計的復(fù)雜性，60多年來，芯片設(shè)計工程師一直在努力自動化芯片布局規(guī)劃過程。

谷歌表示，AlphaChip的研發(fā)，從AlphaGo和AlphaZero中汲取了經(jīng)驗。

眾所周知，通過深度學(xué)習(xí)和博弈論，AlphaGo和AlphaZero逐漸從0掌握了圍棋、國際象棋和將棋的潛在規(guī)則。

AlphaChip同樣是采用了，將芯片底層規(guī)劃視為一種游戲的策略。

從空白柵格開始，AlphaChip每次放置一個電路元件，直至放置完所有元件。

然后，根據(jù)最終布局的質(zhì)量，給予模型獎勵。

一種全新的「基于邊」的圖神經(jīng)網(wǎng)絡(luò)讓AlphaChip，能夠?qū)W習(xí)相互連接的芯片元件之間的關(guān)系，并在芯片之間進行泛化，讓AlphaChip在設(shè)計的每種布局中都有所改進。

左圖：動畫顯示AlphaChip在沒有任何經(jīng)驗的情況下，將開源的Ariane RISC-V CPU置入。

右圖：動畫顯示AlphaChip在對20個TPU相關(guān)設(shè)計進行練習(xí)后，放置相同的塊。

AI大牛帶隊，2頁濃縮版力作

讓我們從最新論文中深扒一下，AlphaChip的整個訓(xùn)練過程。

值得一提的是，這項研究依舊是由Jeff Dean帶隊，所有核心要素全都濃縮在了這兩頁論文中。

與以往方法不同的是，AlphaChip是基于一種「強化學(xué)習(xí)」的方法。

這意味著，當它解決了更多芯片布局問題的實例時，會變得更好、更快。

正如Nature論文（2021年），以及ISPD 2022后續(xù)研究中所討論的那樣，這種預(yù)訓(xùn)練過程顯著提升了AI的速度、可靠性、布局的性能。

順便提一句，預(yù)訓(xùn)練也能培養(yǎng)出LLM，像Gemini、ChatGPT令人印象深刻的能力。

自此前研究發(fā)表以來，谷歌便開源了一個軟件庫，以重現(xiàn)論文中描述的方法。

開發(fā)者們可以使用這個庫，對各種芯片進行預(yù)訓(xùn)練，然后將預(yù)訓(xùn)練的模型應(yīng)用到新的塊。

基于最新的AlphaChip訓(xùn)練過程，研究人員在庫中添加了預(yù)訓(xùn)練的20個TPU塊模型檢查點（checkpoint）。

顯然，如果不進行任何預(yù)訓(xùn)練，AlphaChip就無法從先前的經(jīng)驗中學(xué)習(xí)，從而規(guī)避了學(xué)習(xí)方面的問題。

訓(xùn)練和計算資源

隨著RL智能體（任何ML模型）的投入訓(xùn)練，它的損失通常會逐漸減少。

最終會趨于平穩(wěn)，這代表著模型對其正在執(zhí)行的任務(wù)有了盡可能多的了解，對外表現(xiàn)就是「收斂」。

從訓(xùn)練到收斂，是機器學(xué)習(xí)的標準做法。如果不按照這個路徑來，可能會損害模型的性能。

AlphaChip的性能隨應(yīng)用的計算資源而擴展，在ISPD 2022論文中，谷歌團隊曾進一步探討了這一特性。

論文地址：https://dl.acm.org/doi/10.1145/3505170.3511478

正如Nature論文中所描述的，在對特定塊進行微調(diào)時，使用了16個工作單元，每個單元由1個GPU和32個RL環(huán)境組成，通過多進程處理共享10個CPU。

總言之，用較少的計算資源可能會損害性能，或者需要運行相當長的時間，才能實現(xiàn)相同（或更差）性能。

初始布局

在運行Nature論文中評估方法之前，團隊使用了來自物理綜合的近似初始布局，以解決hMETIS標準單元集群大小不平衡的問題。

RL智能體無權(quán)訪問初始布局，并且不復(fù)雜放置標準單元。

盡管如此，谷歌作者還是進行了一項消融研究，排除了任何初始布局的使用，并且也沒有觀察到AlphaChip性能下降。

如下表1所示。

具體來說，他們跳過了單元集群重新平衡的一步，而是將hMETIS集群不平衡參數(shù)降低到最低設(shè)置（UBfactor = 1）。

由此，這使得hMETIS生成更平衡的集群。

基準

在Nature論文中，研究人員采用了10nm以下制程的TPU塊進行實驗得出的結(jié)果。

這個技術(shù)節(jié)點的大小，正是現(xiàn)代芯片的典型尺寸。之前許多論文報告中，采用較早的45nm、12nm。

從物理設(shè)計角度來看，這種較老的技術(shù)節(jié)點尺寸的芯片，有著顯著的不同。

比如，在10nm以下的芯片中，通常使用多重圖案設(shè)計，這會在較低密度下導(dǎo)致布線擁堵的問題。

來源：Pushing Multiple Patterning in Sub-10nm: Are We Ready?

因此，對于較早的技術(shù)節(jié)點尺寸，AlphaChip可能需要調(diào)整其獎勵函數(shù)，以便更好地適應(yīng)技術(shù)。

展望未來：AI將改變整個芯片設(shè)計流程

自從2020年發(fā)布以來，AlphaChip已經(jīng)生成了每一代谷歌TPU使用的超人芯片布局。

可以說，正是因為它，才能使大規(guī)模放大基于Transformer架構(gòu)的AI模型成為可能。

無論是在Gemini這樣的LLM，還是Imagen和Veo這樣的圖像和視頻生成器中，TPU都位于谷歌強大的生成式AI系統(tǒng)的核心。

另外，這些AI加速器也處于谷歌AI服務(wù)的核心，外部用戶可以通過谷歌云獲得服務(wù)。

谷歌數(shù)據(jù)中心的一排Cloud TPU v5p AI加速器超算

如今，谷歌的三代旗艦TPU芯片，已經(jīng)在世界各地的數(shù)據(jù)中心中制造、部署。

隨著每一代TPU的發(fā)展，AlphaChip和人類專家之間的性能差距不斷擴大。

從TPU v5e中的10個RL放置模塊和3.2%的布線長度減少，到TPU v5p中的15個模塊和4.5%的減少，再到Trillium中的25個模塊和6.2%的減少。

AlphaChip還為數(shù)據(jù)中心CPU（Axion）和谷歌尚未公布的其他芯片，生成了超越人類的布局設(shè)計。

而其他公司，也在谷歌研究的基礎(chǔ)上進行了改進。

比如聯(lián)發(fā)科就不僅用AlphaChip加速了最先進芯片的開發(fā)，還在功耗、性能和面積上對芯片做了優(yōu)化。

如今，AlphaChip僅僅是一個開始。

谷歌對未來做出了大膽暢想：AI將實現(xiàn)芯片設(shè)計全流程的自動化，

通過超人算法以及硬件、軟件和機器學(xué)習(xí)模型的端到端協(xié)同優(yōu)化，芯片設(shè)計的周期會顯著加快，還會解鎖性能的新領(lǐng)域。

谷歌表示，非常期待和社區(qū)合作，實現(xiàn)AI芯片以及芯片AI之間的閉環(huán)。

本文來源：新智元

網(wǎng)友評論

聚超值•精選

推薦 手機 筆記本 影像 硬件 家居 商用 企業(yè) 出行 未來
    • 二維碼 回到頂部