能設(shè)計(jì)芯片的AI黑科技來(lái)了！

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

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

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

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

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

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

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

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

其實(shí)谷歌對(duì)于這個(gè)AI，已經(jīng)布局多年了。

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

論文地址：https://arxiv.org/pdf/2004.10746

后來(lái)在2021年，這項(xiàng)工作發(fā)表在了Nature上，并且進(jìn)行了開(kāi)源。

論文地址：https://www.nature.com/articles/s41586-021-03544-w

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

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

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

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

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

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

然后，團(tuán)隊(duì)在當(dāng)前的TPU模塊上運(yùn)行AlphaChip，以生成高質(zhì)量的布局。

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

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

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

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

工作原理：一邊設(shè)計(jì)，一邊獎(jiǎng)勵(lì)

其實(shí)，設(shè)計(jì)芯片布局并不是一項(xiàng)簡(jiǎn)單的任務(wù)。

一般來(lái)說(shuō)，計(jì)算機(jī)芯片有許多相互連接的模塊、多層電路元件組成，所有這些部件都由纖細(xì)無(wú)比的導(dǎo)線(xiàn)連接起來(lái)。

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

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

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

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

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

從空白柵格開(kāi)始，AlphaChip每次放置一個(gè)電路元件，直至放置完所有元件。

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

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

左圖：動(dòng)畫(huà)顯示AlphaChip在沒(méi)有任何經(jīng)驗(yàn)的情況下，將開(kāi)源的Ariane RISC-V CPU置入。右圖：動(dòng)畫(huà)顯示AlphaChip在對(duì)20個(gè)TPU相關(guān)設(shè)計(jì)進(jìn)行練習(xí)后，放置相同的塊。

AI大牛帶隊(duì)，2頁(yè)濃縮版力作

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

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

論文地址：https://www.nature.com/articles/s41586-024-08032-5

預(yù)訓(xùn)練

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

這意味著，當(dāng)它解決了更多芯片布局問(wèn)題的實(shí)例時(shí)，會(huì)變得更好、更快。

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

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

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

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

GitHub地址：https://github.com/google-research/circuit_training

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

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

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

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

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

從訓(xùn)練到收斂，是機(jī)器學(xué)習(xí)的標(biāo)準(zhǔn)做法。如果不按照這個(gè)路徑來(lái)，可能會(huì)損害模型的性能。

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

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

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

總言之，用較少的計(jì)算資源可能會(huì)損害性能，或者需要運(yùn)行相當(dāng)長(zhǎng)的時(shí)間，才能實(shí)現(xiàn)相同（或更差）性能。

初始布局

在運(yùn)行Nature論文中評(píng)估方法之前，團(tuán)隊(duì)使用了來(lái)自物理綜合的近似初始布局，以解決hMETIS標(biāo)準(zhǔn)單元集群大小不平衡的問(wèn)題。

RL智能體無(wú)權(quán)訪(fǎng)問(wèn)初始布局，并且不復(fù)雜放置標(biāo)準(zhǔn)單元。

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

如下表1所示。

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

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

基準(zhǔn)

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

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

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

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

來(lái)源：Pushing Multiple Patterning in Sub-10nm: Are We Ready?

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

展望未來(lái)：AI將改變整個(gè)芯片設(shè)計(jì)流程

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

可以說(shuō)，正是因?yàn)樗?，才能使大?guī)模放大基于Transformer架構(gòu)的AI模型成為可能。

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

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

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

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

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

從TPU v5e中的10個(gè)RL放置模塊和3.2%的布線(xiàn)長(zhǎng)度減少，到TPU v5p中的15個(gè)模塊和4.5%的減少，再到Trillium中的25個(gè)模塊和6.2%的減少。

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

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

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

如今，AlphaChip僅僅是一個(gè)開(kāi)始。

谷歌對(duì)未來(lái)做出了大膽暢想：AI將實(shí)現(xiàn)芯片設(shè)計(jì)全流程的自動(dòng)化，

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

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