剛剛，Jeff Dean發推表示，谷歌重磅打造的超級硬核年終大總結，出第三彈了！

第一彈：「超詳超硬Jeff Dean萬字總結火熱出爐！圖解谷歌2022年AIGC、LLM、CV三大領域成就」

第二彈：「谷歌2022年度回顧：讓AI更負責任，主要做了4點微小的工作」

偉大的機器學習研究需要偉大的系統。

隨著算法和硬件越來越復雜，以及運行規模越來越大，執行日常任務所需的軟件的復雜性也在不斷增加。

在這篇文章中，研究人員概述了過去一年整個谷歌在ML系統方面取得的眾多進展，這些進展使谷歌能夠支持復雜模型的服務和訓練，同時減輕了終端用戶的實施復雜性。

同時，這篇文章還提到了谷歌如何利用ML本身來改進和設計下一代系統堆棧的研究。

機器學習編程語言

對于機器學習的工作，基礎架構的穩健性和正確性至關重要。

谷歌一直在努力，確保基礎架構建立在可靠的技術和理論基礎之上。并且，作為后盾，谷歌一直在做編程語言和構建編譯器方面的前沿研究。

谷歌會繼續對開源MLIR編譯器的基礎架構投資，構建更加可控、可組合和模塊化的編譯器堆棧。

論文地址：https://research.google/pubs/pub49988/

此外，谷歌在稀疏線性代數的代碼生成方面也取得了很大進展，現在可以從幾乎相同的MLIR程序中生成密集和稀疏的代碼。

最后，谷歌還繼續開發了IREE編譯器，這個編譯器既可以在位于數據中心的強大計算機上使用，在可以在智能手機之類的移動設備上使用。

IREE的端到端流程

在更理論的層面，谷歌探索了哪些方法可以形式化（formalize）和驗證自己使用的代碼生成技術。

谷歌還發布了一種新穎的方法，用于執行和形式化一套自動微分（AD）系統，它正是ML庫的核心。

源代碼轉換

谷歌將反向模式的AD算法分解成三個獨立的程序轉換，這就變得更簡單，更容易驗證，從而突出了JAX實現的獨特性。

反向模式自動微分作為正向微分、解壓縮和轉置

利用抽象解釋和程序合成等編程語言技術，谷歌成功地減少了進行神經結構搜索（NAS）所需的資源數量。這項??NAS成果，可以讓我們在不降低準確性的前提下，發現了更有效的模型。

在用于圖像分類的視覺Transformer架構演化過程中由??NAS合成的突變

在過去的一年里，谷歌在JAX生態系統中發布了許多新的開源庫，比如Rax和T5X。

隨著圍繞jax2tf的持續努力，JAX模型現在可以使用TensorFlow Lite部署在移動設備上，并使用TensorFlow.js部署在網絡上。

「Plane Strike」中的演示

用于機器學習的分布式系統

2022年，谷歌在更好地支持ML和通用科學計算進行大規模計算方面取得了重大進展。

不僅為大型模型設計了SOTA的服務技術，改進了張量程序的自動分區，而且還重新設計了庫的API，以確保所有這些發展能夠被廣大用戶所接受。

其中最大的改進之一，便是用于評估大規模矩陣乘法運算的CollectiveEinsum策略，這是神經網絡的核心。

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

與之前流行的SPMD分區策略不同，CollectiveEinsum會將通信與設備本地計算分開，并通過快速的TPU ICI鏈接進行疊加，進而使性能提高了1.38倍。

同時，CollectiveEinsum算法也是谷歌擴展Transformer推理工作的一個關鍵組成部分。比如，在吞吐量優化的配置中達到SOTA模型76%的FLOPs利用率（MFU）。

此外，谷歌還將SPMD風格的分區概念整合進了TensorFlow（通過DTensor擴展）和JAX（通過重新設計的數組類型）。

在這兩個庫中，那些程序員看來是完整的張量，可以通過附加聲明性的布局注釋，在一些設備上透明地進行分片。

事實上，這兩種方法不僅和為單設備計算編寫的現有代碼兼容，并且還可以擴展到多設備程序中，而不需要修改任何代碼！

論文地址：https://arxiv.org/abs/2105.04663

然而，GSPMD在很大程度上依賴于啟發式方法，也就是有時仍然需要手動做出決定，而這通常會讓性能無法達到最優。

為了使分區推理完全自動化，谷歌開發了Alpa——一個它探索了運算器級（模型）并行和較大子計算之間管線并行策略的自動化系統。

Alpa不僅實現了在Transformer等主流模型上與「人工微調」相媲美的性能，同時也能夠擴展到其他模型之中，如卷積網絡和專家混合模型（MOE）。

與之類似，谷歌最近提出的Pathways系統，在TPU運行時間之上增加了一個額外的虛擬化層——加速器由長期存在的進程管理，而不是直接分配給用戶。

然后，單個終端用戶可以連接到任意數量的Pathways控制的設備，并編寫他們的程序。就像所有的設備都直接連接到他們的進程一樣，即使現實中的情況是跨越多個數據中心的。

論文地址：https://arxiv.org/abs/2203.12533

由于Pathways：（1）作業啟動時間減少，（2）更容易實現容錯，以及（3）使多租戶成為一個可行的選擇，從而讓多個作業可以同時執行，更有效地利用硬件。

更重要的是，Pathways能夠輕松實現跨越多個TPU pods的計算，而這可以有效避免未來的擴展瓶頸。

左上：用有向無環圖表征的分布式計算；右上：資源管理器為每個編譯的函數（如A、B和C）分配虛擬的加速器網格片；下：集中的調度器對計算進行分組調度，然后由每個分片的執行器進行調度

此外，還有一個全新的用于多維陣列存儲的庫——TensorStore。

TensorStore在訓練具有多控制器運行時間的大型語言模型（LLM）時非常實用，其中每個進程只用管理參數的一個子集，而所有的參數則需要被整理成一個一致的檢查點。

TensorStore為高效和并發的多維數組序列化提供了數據庫級的保證（ACID），并已成功用于計算密集型工作負載，如PaLM和人類皮層和果蠅大腦的重建。

一個蒼蠅大腦的重建，其基礎數據可以使用TensorStore輕松訪問和操作

硬件加速器和機器學習

用于ML的硬件設計

使用定制的硬件（如TPU和GPU），在性能提升和能源效率上會有巨大的優勢，還能減少碳足跡。

在最近的MLPerf競賽中，谷歌在TPU v4上的五項基準測試中創造了新的性能記錄，實現了比第二名平均高1.42倍的速度。

不過，為了跟上最近的進展，谷歌也在為特定的流行模型開發定制的硬件架構。

在已公布的五個基準測試（MLPerf 2.0）中，谷歌的TPU都比競品（NVIDIA on-premises）速度更快。（條形圖內的數字代表使用的芯片/加速器的數量）

然而，構建新的硬件加速器會產生很高的初始成本，并且需要大量的開發和部署時間。

為了使單工作負載加速器（single-workload accelerators）可行，必須減少設計周期時間。

全棧加速器搜索技術

而全棧搜索技術（FAST）通過引入一個硬件加速器搜索框架，就解決了這個問題。

這個框架同時優化了數據路徑、調度和重要的編譯器決策。

FAST引入了一個近似的模板，能夠描述不同類型的架構和多功能的內存層次，從而使加速器的單位熱設計功率（與單位總成本的性能高度相關）的單工作負載性能比TPU v3提高3.7倍。

這表明，單工作負載加速器對于中等規模的數據中心部署是實用的。

用于硬件設計的機器學習

為了盡可能地實現芯片設計過程的自動化，谷歌在硬件設計的各個階段，都在推動機器學習的功能，包括高級架構探索、驗證以及布局和布線。

方法和訓練方案概述

谷歌最近開源了一個名為Circuit Training的分布式RL基礎設施，以及一個電路環境，后者谷歌在發于Nature的論文中詳細做了介紹。

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

谷歌在生產中使用了這個基礎設施，為最新一代的TPU芯片生成了宏觀布局。

在解決架構探索問題時，PRIME引入了一種基于ML的方法來搜索硬件設計空間，只利用現有的數據（比如來自傳統加速器設計工作的數據），而不需要進一步的硬件模擬。

這種方法減輕了運行耗時的模擬的需要，即使在目標應用程序集發生變化時。

PRIME比最先進的模擬驅動方法提高了約1.2-1.5倍的性能，同時減少了93%-99%的模擬時間。

AutoApprox通過將每個神經網絡層映射到適當的近似級別，自動生成近似的低功耗深度學習加速器，而沒有任何精度損失。

PRIME使用記錄的加速器數據（包括可行的和不可行的加速器）來訓練模型，其設計的加速器的延遲小了1.5倍，同時減少了99%的硬件模擬時間

依賴于硬件的模型設計

雖然神經架構搜索（NAS）在SOTA模型的發現方面展示出了巨大的能力，但它仍然受到缺乏硬件知識的限制。

而基于平臺感知（Platform-aware）的NAS，則可以通過將硬件結構的知識納入NAS搜索空間的設計中，來解決這一問題。

由此產生的EfficientNet-X模型在TPU v3和GPU v100上的速度分別是EfficientNet的1.5倍-2倍，而精度卻相差無幾。

目前，平臺感知的NAS和EfficientNet-X都已在生產中部署。實踐證明，對于各種生產型視覺模型來說，都有明顯的精度提升和高達40%的效率提升。

論文地址：https://arxiv.org/abs/2102.05610

NaaS通過共同搜索神經網絡架構和硬件架構，更進一步。

測試結果顯示，NaaS可以在Edge TPU上發現同等精度但能效提高了2倍的視覺模型。

在TPU/GPU上的平臺感知NAS概述

用于大規模生產系統的機器學習

在生產中運行的各種大規模系統上，谷歌也利用機器學習實現了效率的提升。

比如，最近發布的第一個在LLVM基礎設施中系統地整合ML技術的工業級通用框架——MLGO，可以用RL策略取代LLVM中的啟發式方法來做出優化決策。

測試發現，在優化內聯決策時，經過訓練的策略可以減少3%-7%的二進制大小，而在優化寄存器分配決策時，可以提高0.3%~1.5%的吞吐量。

論文地址：https://arxiv.org/abs/2101.04808

在生產型ML編譯器中，幾年前發布的學習成本模型XLA，也被用于指導頂級ML工作負載的TPU內核的最佳瓦片大小的選擇，進而在數據中心上節省了2%的TPU總計算時間。

論文地址：https://arxiv.org/abs/2008.01040

此外，谷歌還用新的混合算法取代了YouTube緩存替換算法中現有的啟發式算法，該算法結合了簡單的啟發式算法和學習模型，在峰值時將byte miss提高了9%。

總結一下

谷歌表示，隨著機器學習領域的發展，自己將在開發高性能、高能效和易于使用的系統和基礎設施上持續投入，進而實現對新想法的快速探索。

同時，谷歌也會繼續探索機器學習的能力、提高復雜系統的性能，并使系統設計中的勞動密集型任務自動化。