能設計芯片的AI黑科技來了！

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

無論是設計最先進的用于構建AI模型的TPU，還是數(shù)據(jù)中心的CPU，它在相關的眾多領域，都產生了廣泛影響。

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

AlphaChip設計芯片，用的是強化學習的原理。

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

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

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

布局五年，谷歌多代TPU全由AI設計

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

工作原理：一邊設計，一邊獎勵

其實，設計芯片布局并不是一項簡單的任務。

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

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

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

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

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

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

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

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

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

左圖：動畫顯示AlphaChip在沒有任何經(jīng)驗的情況下，將開源的Ariane RISC-V CPU置入。右圖：動畫顯示AlphaChip在對20個TPU相關設計進行練習后，放置相同的塊。

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

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

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

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

預訓練

與以往方法不同的是，AlphaChip是基于一種「強化學習」的方法。

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

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

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

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

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

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

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

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

訓練和計算資源

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

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

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

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

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

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

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

初始布局

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

RL智能體無權訪問初始布局，并且不復雜放置標準單元。

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

如下表1所示。

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

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

基準

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

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

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

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

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

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

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

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

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

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

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

谷歌數(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）和谷歌尚未公布的其他芯片，生成了超越人類的布局設計。

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

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

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

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

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

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