轉載：AI深度研究員
在 NeurIPS 2024 大會上，谷歌首席科學家 Jeff Dean 帶來了一個令人震撼的訊息：AI 技術正在徹底改變晶片設計的遊戲規則。傳統晶片設計需要數百位工程師投入 18 個月時間，而藉助 AI 技術，這一過程有望縮短到驚人的 1 秒。
這個突破性的宣佈來自 12 月 15 日在加拿大溫哥華舉辦的 NeurIPS 2024 大會。在這場重要演講中，Dean 不僅展示了 AI 在晶片設計領域的革命性突破，更深入探討了機器學習如何全面革新計算機系統。他特別引用了 Rich Sutton 的觀點：在解決複雜問題時，搜尋和學習往往能幫助我們超越傳統方法，找到最優解決方案。
這不僅僅是速度的提升，而是整個計算機系統設計正規化的革命。"Dean 指出，傳統的計算機系統，無論是作業系統、編譯器還是記憶體分配器，都尚未充分利用機器學習的力量。但這種情況正在發生根本性的改變。早在 2018 年，他就預見了這場變革，並提出了三個核心挑戰：如何將機器學習無縫整合到傳統系統、如何在複雜環境中應用機器學習，以及如何確保機器學習決策的可靠性和安全性。
雖然這些挑戰曾經看似難以逾越，但近年來的技術進展證明，革命已經悄然展開。從記憶體管理到編譯器最佳化，從系統排程到晶片設計，機器學習正在重寫計算機系統的每一個環節。這不僅帶來了效率的提升，更預示著計算機工業即將迎來一個全新的時代。
演講文稿
編譯器設計最佳化
在編譯器和自動調優系統領域，機器學習為效能最佳化提供了前所未有的機會。這種突破首先體現在用機器學習替代傳統的啟發式編譯器選擇上。研究者孟婆提出了一種自動調優多遍編譯器的創新方法，包含兩個核心元件：學習成本模型的評估器和程式碼最佳化器。
學習成本模型的評估器能夠評估程式碼在特定編譯策略下的執行速度，其獨特之處在於有時甚至無需實際執行程式碼就能做出準確判斷。程式碼最佳化器和學習策略則負責在高階程式碼與低階程式碼的對映過程中，選擇最優的編譯器引數，以獲得最佳效能。
運算子融合是這種方法的典型應用案例。在編譯器中，當處理不同形狀的陣列時，融合操作可能會提升效率，但這種提升並非必然。它取決於多個因素，包括裝置的記憶體頻寬和陣列的具體形狀。傳統方法難以權衡這些因素，而機器學習系統則能夠根據執行時的硬體特性，智慧地決定何時進行融合。
佈局分配是另一個重要應用。系統需要決定如何將抽象張量轉化為實際張量，並最佳化記憶體佈局。這個過程涉及複雜的權衡，傳統的規則基礎方法難以應對。透過機器學習，系統能夠根據實際執行環境和需求，自動找到最優的佈局策略。這些最佳化在實踐中取得了令人矚目的成效。在谷歌的機器學習工作負載中，生產模型的效能提升了5%到25%。正如Dean幽默地說："這相當於每個人免費獲得了一塊TPU。"這種提升不僅顯著，而且具有廣泛的實用價值。
突破性的記憶體管理創新
記憶體管理領域的創新同樣令人矚目。谷歌開發的系統採用了革命性的方法,透過機器學習預測和管理物件生命週期。系統的核心創新在於將呼叫堆疊視為自然語言,使用 LSTM 網路進行分析,從而準確預測物件的生命週期。具體實現機制包括:
即時堆疊雜湊處理和快速快取查詢
動態更新預測模型以適應程式行為變化
智慧分配策略最佳化記憶體使用
系統能夠精確預測物件生命週期長短,並據此最佳化分配策略。短生命週期物件被放入執行緒本地快取,長生命週期物件則使用中心化記憶體管理。系統還會將相似生命週期的物件分組儲存,顯著減少記憶體碎片。更重要的是,系統具備自適應能力。如果預測出現錯誤,比如將物件生命週期預測為"短暫"而實際為"中等",系統會自動調整並更新預測模型。這種方法在實踐中取得了顯著成效,將記憶體碎片減少了19%到78%。
SmartChoice：革命性的輕量級學習決策系統
SmartChoice是谷歌開發的另一項重要創新，旨在將輕量級學習決策無縫整合到現有的複雜系統程式碼中。這個系統提供了簡潔而強大的API，包括上下文型別（用於決策的資訊）和手臂型別（待選擇的決策集），透過持續的選擇和反饋過程來最佳化決策質量。
系統最令人印象深刻的特點是其極快的決策速度。得益於基於Bandit演算法（多臂老虎機問題）的實現，SmartChoice能在幾十微秒內完成決策。這種高效性使它能夠應用於即時性要求極高的場景。
系統不僅能處理即時反饋，還能巧妙地處理來自分散式系統不同部分的延遲反饋。例如，當用戶在網路伺服器97上的操作與伺服器4上的早期決策相關時，系統能夠正確關聯這些資訊，並在離線處理中更新模型。
新的模型會被推送到所有相關服務系統，實現全域性最佳化。
在YouTube的實踐應用中，SmartChoice成功優化了影片快取策略。系統學習如何在全球各地的快取位置儲存最受歡迎的影片，顯著降低了使用者透過昂貴網路連結訪問遠端資料中心的需求。在相對適中的計算成本下，系統將快取未命中率降低了9.1%，這個成果在影片流務領域具有重要意義。除了快取最佳化，SmartChoice還被廣泛應用於其他場景。線上程計數最佳化中，系統能夠根據請求型別（如航班搜尋）智慧地決定最佳執行緒數，在延遲和吞吐量之間取得平衡。在工作分割槽最佳化中，系統能夠最佳化廣告資料的重新整理頻率。在使用者介面最佳化方面，系統能根據使用者歷史行為預測並定製介面元素。
推理系統最佳化：降低成本與延遲
在機器學習推理領域，降低成本和延遲是一個核心挑戰。正如Jeff Dean強調的，這直接關係到先進模型能否惠及更多使用者。最近釋出的Gemini 2.0獲得積極反饋，很大程度上就得益於其極低的延遲。
谷歌開發了多項創新技術來最佳化推理系統。首要的是過度訓練小型模型：透過精心調整模型大小，在固定的訓練計算資源下進行更多資料遍歷（epoch）。實踐表明，將模型縮小到原來的三分之一甚至五分之一，雖然顯著降低了推理成本，但在適當的訓練策略下，模型質量幾乎不受影響。
知識蒸餾是另一個重要技術，它能夠將大模型的知識有效轉移到更小的模型中。這種方法不僅能保持模型的核心能力，還能顯著降低推理成本。此外，選擇性啟用技術允許模型根據不同型別的上下文和輸入，智慧地啟用其特定部分，從而在不同任務中實現更高的效率。
推測性解碼是一項特別值得關注的創新，它能使自迴歸模型的解碼速度提升2到3倍。這種方法的獨特之處在於，它無需改變模型架構或重新訓練，完全透過最佳化解碼演算法來實現效能提升。其核心思想是引入一個小型模型快速生成候選標記，再由大模型並行驗證這些預測，從而提高整體效率。
AlphaChip：重新定義晶片設計
在晶片設計領域,谷歌的創新最為引人注目。傳統晶片設計需要數百人投入18個月時間,成本高達數億美元。AlphaChip 系統試圖將這個過程縮短到幾周。
AlphaChip 採用了多項創新技術:
智慧晶片分塊和元件聚類
強化學習最佳化佈局佈線
預訓練模型加速設計過程系統最令人印象深刻的是其預訓練能力。如果從零開始,系統需要20小時才能達到較好狀態。但透過預訓練,系統能在1秒內完成高質量的佈局評估,效率超越人類專家。
AlphaChip 不僅在速度上實現了突破，更重要的是開創了晶片設計的新正規化。傳統的晶片設計流程主要依賴人工經驗，從高階架構規範到低階 RTL 設計的轉換過程高度依賴硬體工程師的手動工作。而 AlphaChip 透過端到端學習方法，使整個設計過程更加自動化和智慧化。一個關鍵創新是引入了多層次的反饋迴圈。系統使用快速評估機制，透過低成本代理快速評估設計決策的下游影響，同時保留複雜評估作為整體驗證手段。這種方法讓設計團隊能夠快速迭代，同時確保最終設計的質量。
Dean 提出了一個大膽的設想：如果我們在21天內投入1000萬美元的計算資源，使用16,000塊 TPU 晶片（約15 exaflops的算力），可能徹底改變晶片設計的方式。這種計算密集型方法雖然前期投入較大，但能顯著縮短設計週期，降低總體成本。按照他的預測，晶片設計成本有望降低20到100倍。
強化學習在系統最佳化中的應用
強化學習在系統最佳化中展現出獨特優勢。以晶片佈局為例，AlphaChip 透過強化學習演算法確定元件位置，並透過輕量級代理評估匯流排長度、面積利用率和潛在擁塞等關鍵指標。更重要的是，系統能夠在幾秒鐘內評估決策影響，使設計過程變得更加敏捷。
在架構探索方面，強化學習同樣發揮著重要作用。系統需要在龐大的設計空間中做出多項決策，包括：
快取層次結構設計
記憶體頻寬配置
計算單元規劃
透過結合高階模擬器和強化學習，系統能夠快速探索這個複雜的決策空間，並根據下游反饋不斷最佳化設計方案。這種方法在針對特定領域的晶片設計中特別有效，因為它可以充分利用領域特定的知識和約束。在學習綜合(learning synthesis)領域，儘管相關研究尚未正式發表，但谷歌 DeepMind 團隊已經在多個競賽中展示了這項技術的潛力。與傳統方法相比，機器學習綜合能夠自動處理特定的邏輯設計問題，比如五位加法器的最優實現。這種方法的優勢在於它能夠自動發現和最佳化關鍵子系統，而不是試圖一次性解決整個晶片設計問題。這種分而治之的策略不僅降低了問題的複雜度，也提高了解決方案的質量和可靠性。
面向未來的系統設計
機器學習如何從根本上改變計算機系統的設計和最佳化方式。從編譯器最佳化到晶片設計，從記憶體管理到快取策略，機器學習正在重新定義計算機系統的各個方面。特別值得注意的是端到端學習的重要性。在晶片設計中，早期決策會對最終的佈局和佈線產生深遠影響。透過端到端學習，系統能夠在架構探索和綜合階段做出更優決策，從而實現真正的全域性最佳化。未來的發展方向包括：
深化機器學習在傳統系統中的應用
提升端到端學習的效率和可靠性
擴大自動化決策的範圍和準確性
谷歌在系統機器學習領域的領導地位，更描繪了計算機系統發展的新願景。透過將學習和搜尋這兩個強大工具應用到系統的各個層面，我們正在見證並參與計算機科學史上的一次重大變革。
更多精彩：
相關閱讀：
轉載：AI深度研究員。本文版權歸屬作者和原載媒體所有。