AI 幾小時設計芯片超越人類,谷歌 AlphaChip 已設計出三代旗艦 TPU
谷歌的 AlphaChip,幾小時內就能設計出芯片布局,直接碾壓人類專家!這種超人芯片布局,已經應用在 TPU、CPU 在內的全球硬件中。人類設計芯片的方式,已被 AI 徹底改變。

能設計芯片的 AI 黑科技來了!就在剛剛,谷歌 DeepMind 推出名為 AlphaChip 的 AI 系統。

無論是設計最先進的用于構建 AI 模型的 TPU,還是數據中心的 CPU,它在相關的眾多領域,都產生了廣泛影響。
在谷歌的許多款芯片設計中,它都取得了出色的效果,比如 Axion 芯片(一種基于 Arm 的通用數據中心 CPU)。

AlphaChip 設計芯片,用的是強化學習的原理。
也就是說,芯片布局設計對它來說是一種游戲,就像 AlphaGo 一樣,它在游戲中,學習如何設計出最好的芯片布局。
幾小時內,它就能生成超出人類水平,或是與人類專家相當的芯片布局了。
現在,它已經用于設計多代 TPU 芯片(TPU v5e、TPU v5p 和 Trillium)。而且跟人類專家相比,AlphaChip 放置的塊數越來越多,線長也減少了許多。
01 布局五年,谷歌多代 TPU 全由 AI 設計
其實谷歌對于這個 AI,已經布局多年了。早在 2020 年,團隊就發表了一篇預印本論文,介紹了谷歌的全新強化學習方法,用于設計芯片布局。
后來在 2021 年,這項工作發表在了 Nature 上,并且進行了開源。
而自從首次發表這項工作以來,谷歌內部一直在對它進行改進。今天,團隊發表了 Nature 附錄,詳細描述了具體方法,及其對芯片設計領域的影響。同時,他們還發布了一個預訓練的檢查點,分享了模型權重,公布模型名稱為 AlphaChip。
谷歌表示,AlphaChip 是最早用于解決現實世界工問題的強化學習方法之一。

在數小時內,它就可以生成超人或類似的芯片布局,而不需要花費數周或數月的人類。它設計的芯片布局,已經被用于世界各地的芯片中,包括數據中心和移動電話。
為了設計 TPU 布局,AlphaChip 首先在來自前幾代的各種芯片模塊上進行實踐,例如片上和片間網絡模塊、內存控制器和數據傳輸緩沖區。這一過程被稱為預訓練。
然后,團隊在當前的 TPU 模塊上運行 AlphaChip,以生成高質量的布局。
與之前的方法不同,AlphaChip 在解決更多芯片布局任務時變得更好、更快,類似于人類專家的工作方式。
對于每一代新的 TPU,包括谷歌最新的 Trillium (第 6 代),AlphaChip 都設計了更好的芯片布局,并提供了更多的總體布局圖,從而加快了設計周期,產生了更高性能的芯片。

條形圖顯示了谷歌三代 TPU 上 AlphaChip 設計的芯片塊的數量,包括 v5e、v5p 和 Trillium。跟 TPU 物理設計團隊生成的布局相比,AlphaChip 在 TPU 三代產品中的平均有線長度減少。
02 工作原理:一邊設計,一邊獎勵
其實,設計芯片布局并不是一項簡單的任務。
一般來說,計算機芯片有許多相互連接的模塊、多層電路元件組成,所有這些部件都由纖細無比的導線連接起來。
此外,還有許多復雜且相互交織的設計約束,必須同時滿足。
由于設計的復雜性,60 多年來,芯片設計工程師一直在努力自動化芯片布局規劃過程。
谷歌表示,AlphaChip 的研發,從 AlphaGo 和 AlphaZero 中汲取了經驗。

眾所周知,通過深度學習和博弈論,AlphaGo 和 AlphaZero 逐漸從 0 掌握了圍棋、國際象棋和將棋的潛在規則。
AlphaChip 同樣是采用了,將芯片底層規劃視為一種游戲的策略。
從空白柵格開始,AlphaChip 每次放置一個電路元件,直至放置完所有元件。
然后,根據最終布局的質量,給予模型獎勵。
一種全新的「基于邊」的圖神經網絡讓 AlphaChip,能夠學習相互連接的芯片元件之間的關系,并在芯片之間進行泛化,讓 AlphaChip 在設計的每種布局中都有所改進。
03 AI 大牛帶隊,2 頁濃縮版力作
讓我們從最新論文中深扒一下,AlphaChip 的整個訓練過程。
值得一提的是,這項研究依舊是由 Jeff Dean 帶隊,所有核心要素全都濃縮在了這兩頁論文中。

預訓練
與以往方法不同的是,AlphaChip 是基于一種「強化學習」的方法。
這意味著,當它解決了更多芯片布局問題的實例時,會變得更好、更快。
正如 Nature 論文(2021 年),以及 ISPD 2022 后續研究中所討論的那樣,這種預訓練過程顯著提升了 AI 的速度、可靠性、布局的性能。
順便提一句,預訓練也能培養出 LLM,像 Gemini、ChatGPT 令人印象深刻的能力。
自此前研究發表以來,谷歌便開源了一個軟件庫,以重現論文中描述的方法。
開發者們可以使用這個庫,對各種芯片進行預訓練,然后將預訓練的模型應用到新的塊。
基于最新的 AlphaChip 訓練過程,研究人員在庫中添加了預訓練的 20 個 TPU 塊模型檢查點(checkpoint)。
顯然,如果不進行任何預訓練,AlphaChip 就無法從先前的經驗中學習,從而規避了學習方面的問題。
訓練和計算資源
隨著 RL 智能體(任何 ML 模型)的投入訓練,它的損失通常會逐漸減少。
最終會趨于平穩,這代表著模型對其正在執行的任務有了盡可能多的了解,對外表現就是「收斂」。
從訓練到收斂,是機器學習的標準做法。如果不按照這個路徑來,可能會損害模型的性能。
AlphaChip 的性能隨應用的計算資源而擴展,在 ISPD 2022 論文中,谷歌團隊曾進一步探討了這一特性。
正如 Nature 論文中所描述的,在對特定塊進行微調時,使用了 16 個工作單元,每個單元由 1 個 GPU 和 32 個 RL 環境組成,通過多進程處理共享 10 個 CPU。
總言之,用較少的計算資源可能會損害性能,或者需要運行相當長的時間,才能實現相同(或更差)性能。
初始布局
在運行 Nature 論文中評估方法之前,團隊使用了來自物理綜合的近似初始布局,以解決 hMETIS 標準單元集群大小不平衡的問題。
RL 智能體無權訪問初始布局,并且不復雜放置標準單元。
盡管如此,谷歌作者還是進行了一項消融研究,排除了任何初始布局的使用,并且也沒有觀察到 AlphaChip 性能下降。
如下表 1 所示。
具體來說,他們跳過了單元集群重新平衡的一步,而是將 hMETIS 集群不平衡參數降低到最低設置(UBfactor = 1)。
由此,這使得 hMETIS 生成更平衡的集群。

基準
在 Nature 論文中,研究人員采用了 10nm 以下制程的 TPU 塊進行實驗得出的結果。
這個技術節點的大小,正是現代芯片的典型尺寸。之前許多論文報告中,采用較早的 45nm、12nm。
從物理設計角度來看,這種較老的技術節點尺寸的芯片,有著顯著的不同。
比如,在 10nm 以下的芯片中,通常使用多重圖案設計,這會在較低密度下導致布線擁堵的問題。

因此,對于較早的技術節點尺寸,AlphaChip 可能需要調整其獎勵函數,以便更好地適應技術。
04 展望未來:AI 將改變整個芯片設計流程
自從 2020 年發布以來,AlphaChip 已經生成了每一代谷歌 TPU 使用的超人芯片布局。
可以說,正是因為它,才能使大規模放大基于 Transformer 架構的 AI 模型成為可能。
無論是在 Gemini 這樣的 LLM,還是 Imagen 和 Veo 這樣的圖像和視頻生成器中,TPU 都位于谷歌強大的生成式 AI 系統的核心。
另外,這些 AI 加速器也處于谷歌 AI 服務的核心,外部用戶可以通過谷歌云獲得服務。

谷歌數據中心的一排Cloud TPU v5p AI加速器超算
如今,谷歌的三代旗艦 TPU 芯片,已經在世界各地的數據中心中制造、部署。
隨著每一代 TPU 的發展,AlphaChip 和人類專家之間的性能差距不斷擴大。
從 TPU v5e 中的 10 個 RL 放置模塊和 3.2% 的布線長度減少,到 TPU v5p 中的 15 個模塊和 4.5% 的減少,再到 Trillium 中的 25 個模塊和 6.2% 的減少。
AlphaChip 還為數據中心 CPU(Axion)和谷歌尚未公布的其他芯片,生成了超越人類的布局設計。
而其他公司,也在谷歌研究的基礎上進行了改進。
比如聯發科就不僅用 AlphaChip 加速了最先進芯片的開發,還在功耗、性能和面積上對芯片做了優化。
如今,AlphaChip 僅僅是一個開始。
谷歌對未來做出了大膽暢想:AI 將實現芯片設計全流程的自動化,通過超人算法以及硬件、軟件和機器學習模型的端到端協同優化,芯片設計的周期會顯著加快,還會解鎖性能的新領域。
谷歌表示,非常期待和社區合作,實現 AI 芯片以及芯片 AI 之間的閉環。



