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

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

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

偉大的機器學習研究需要偉大的系統(tǒng)。

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

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

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

機器學習編程語言

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

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

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

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

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

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

IREE的端到端流程

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

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

源代碼轉(zhuǎn)換

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

反向模式自動微分作為正向微分、解壓縮和轉(zhuǎn)置

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

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

在過去的一年里，谷歌在JAX生態(tài)系統(tǒng)中發(fā)布了許多新的開源庫，比如Rax和T5X。

隨著圍繞jax2tf的持續(xù)努力，JAX模型現(xiàn)在可以使用TensorFlow Lite部署在移動設備上，并使用TensorFlow.js部署在網(wǎng)絡上。

「Plane Strike」中的演示

用于機器學習的分布式系統(tǒng)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

硬件加速器和機器學習

用于ML的硬件設計

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

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

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

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

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

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

全棧加速器搜索技術

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

這個框架同時優(yōu)化了數(shù)據(jù)路徑、調(diào)度和重要的編譯器決策。

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

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

用于硬件設計的機器學習

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

方法和訓練方案概述

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

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

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

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

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

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

AutoApprox通過將每個神經(jīng)網(wǎng)絡層映射到適當?shù)慕萍墑e，自動生成近似的低功耗深度學習加速器，而沒有任何精度損失。

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

依賴于硬件的模型設計

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

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

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

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

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

NaaS通過共同搜索神經(jīng)網(wǎng)絡架構和硬件架構，更進一步。

測試結(jié)果顯示，NaaS可以在Edge TPU上發(fā)現(xiàn)同等精度但能效提高了2倍的視覺模型。

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

用于大規(guī)模生產(chǎn)系統(tǒng)的機器學習

在生產(chǎn)中運行的各種大規(guī)模系統(tǒng)上，谷歌也利用機器學習實現(xiàn)了效率的提升。

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

測試發(fā)現(xiàn)，在優(yōu)化內(nèi)聯(lián)決策時，經(jīng)過訓練的策略可以減少3%-7%的二進制大小，而在優(yōu)化寄存器分配決策時，可以提高0.3%~1.5%的吞吐量。

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

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

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

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

總結(jié)一下

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

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