Layout工程師危,谷歌自動晶片設計AlphaChip問世,開放權重可外部定製

机器之心發表於2024-09-27
2020 年,谷歌發表了預印本論文《Chip Placement with Deep Reinforcement Learning》,介紹了其設計晶片佈局的新型強化學習方法。後來在 2021 年,谷歌又發表在 Nature 上並開源了出來。

圖片

今天,谷歌發表了這篇 Nature 文章的附錄,更詳細介紹了該方法及其對晶片設計領域的影響。同時,谷歌也開放了一個在 20 個 TPU 模組上預訓練的檢查點,分享模型權重並命名為「AlphaChip」

圖片

  • Nature 附錄地址:https://www.nature.com/articles/s41586-024-08032-5
  • 預訓練檢查點地址:https://github.com/google-research/circuit_training/?tab=readme-ov-file#PreTrainedModelCheckpoint

谷歌首席科學家 Jeff Dean 表示,開放預訓練 AlphaChip 模型檢查點以後,外部使用者可以更容易地使用 AlphaChip 來啟動自己的晶片設計。

圖片

x@JeffDean

計算機晶片推動了 AI 的顯著進步,AlphaChip 利用 AI 來加速和最佳化晶片設計。該方法已被用於設計谷歌自定義 AI 加速器(TPU)最近三代的「超人」晶片佈局。

作為首批用於解決現實世界工程問題的強化學習方法,AlphaChip 只需要數小時便能完成媲美或者超越人類的晶片佈局,而無需再花費數週或數月人工努力。並且該方法設計的佈局已應用於世界各地的晶片,覆蓋場景包括資料中心到手機。

谷歌 DeepMind 表示,AlphaChip 已經徹底改變了我們設計微晶片的方式,從幫助設計用於構建 AI 模型的 SOTA TPU 到資料中心 CPU,它的廣泛影響已經擴充套件到了 Alphabet 內外。

圖片

谷歌 DeepMind 聯合創始人兼 CEO Demis Hassabis 表示,如今我們形成了這樣一種反饋迴路:訓練 SOTA 晶片設計模型(AlphaChip)→使用 AlphaChip 來設計更好的 AI 晶片→使用這些 AI 晶片來訓練更好的模型→再設計更好的晶片,這正是谷歌 TPU 堆疊表現如此好的部分原因。

圖片

各路網友對谷歌的 AlphaChip 寄予厚望,稱「晶片設計晶片的時代來了」,還預言谷歌將贏得未來 AGI 之爭。

圖片

圖片

圖片

圖片

AlphaChip 是如何工作的?

晶片設計並非易事,部分原因在於計算機晶片由許多相互連線的塊組成,這些塊具有多層電路元件,所有元件都透過極細的導線連線。此外,晶片還有很多複雜且相互交織的設計約束,設計時必須同時滿足所有約束。由於這些複雜性,晶片設計師們在 60 多年來一直在努力實現晶片佈局規劃過程的自動化。

與 AlphaGo 和 AlphaZero 類似,谷歌構建了 AlphaChip,將晶片佈局規劃視為一種博弈。

AlphaChip 從空白網格開始,一次放置一個電路元件,直到完成所有元件的放置。然後根據最終佈局的質量給予獎勵。谷歌提出了一種新穎的「基於邊」的圖神經網路使 AlphaChip 能夠學習互連晶片元件之間的關係,並在整個晶片中進行推廣,讓 AlphaChip 在其設計的每一個佈局中不斷進步。Layout工程師危,谷歌自動晶片設計AlphaChip問世,開放權重可外部定製
左圖:動畫展示了 AlphaChip 在沒有任何經驗的情況下放置開源處理器 Ariane RISC-V CPU。右圖:動畫展示 AlphaChip 在對 20 個 TPU 相關設計進行練習後放置相同的塊。

谷歌藉助 AI 設計 AI 加速器晶片

自 2020 年釋出以來,谷歌已經採用 AlphaChip 為每一代 Google TPU 生成超級晶片佈局。這些晶片使得大規模擴充套件基於 Google Transformer 架構的 AI 模型成為可能。

TPU 作為谷歌強大的生成式 AI 系統的核心,應用範圍從大語言模型(如 Gemini)到影像和影片生成器(Imagen 和 Veo)。這些 TPU 是 Google AI 服務的核心,可透過 Google Cloud 供外部使用者使用。

圖片

谷歌資料中心存放的基於 Cloud TPU v5p 的超級計算機。

為了設計 TPU 佈局,AlphaChip 首先在前幾代的各種晶片塊上進行練習,例如片上和晶片間網路塊、記憶體控制器和資料傳輸緩衝區。這個過程稱為預訓練。然後谷歌在當前的 TPU 塊上執行 AlphaChip 以生成高質量的佈局。與之前的方法不同,AlphaChip 解決了更多晶片佈局任務例項,因此變得更好、更快,就像人類專家所做的那樣。

隨著每一代新 TPU(包括谷歌最新的 Trillium(第 6 代))的推出,AlphaChip 設計出了更好的晶片佈局並提供了更多的整體平面圖,從而加快了設計週期併產生了效能更高的晶片。

圖片

圖中顯示了 Google 三代張量處理單元 (TPU)(包括 v5e、v5p 和 Trillium)中 AlphaChip 設計的晶片塊的數量。

圖片

圖中顯示了 AlphaChip 在三代 Google 張量處理單元 (TPU) 中的平均線長(wirelength)減少量,並與 TPU 物理設計團隊生成的位置進行了比較。

AlphaChip 帶來的更廣泛影響

AlphaChip 的影響力體現在 Alphabet、研究界和晶片設計行業的應用上。除了設計 TPU 等專用 AI 加速器外,AlphaChip 還為 Alphabet 的其他晶片設計佈局,例如 Google Axion 處理器,這是谷歌首款基於 Arm 的通用資料中心 CPU。

外部組織也在採用和構建 AlphaChip。例如,全球頂級晶片設計公司之一聯發科擴充套件了 AlphaChip,以加速其最先進晶片(如三星手機使用的 Dimensity Flagship 5G)的開發,同時提高了功耗、效能和晶片面積。

AlphaChip 引發了晶片設計 AI 工作的爆炸式增長,並已擴充套件到晶片設計的其他關鍵階段,例如邏輯綜合和宏選擇。

開創晶片新未來

谷歌堅信,AlphaChip 有潛力最佳化從計算架構到製造的晶片設計週期的每個階段,並改變智慧手機、醫療裝置、農業感測器等日常裝置中定製硬體的晶片設計。

目前,AlphaChip 的未來版本正在開發中。谷歌期待與社群合作,繼續變革自動晶片設計領域,從而在未來迎來速度更快、價格更低、能效更高的晶片。

參考連結:https://deepmind.google/discover/blog/how-alphachip-transformed-computer-chip-design/

相關文章