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

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

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

偉大的機(jī)器學(xué)習(xí)研究需要偉大的系統(tǒng)。

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

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

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

機(jī)器學(xué)習(xí)編程語言

對于機(jī)器學(xué)習(xí)的工作，基礎(chǔ)架構(gòu)的穩(wěn)健性和正確性至關(guān)重要。

谷歌一直在努力，確?；A(chǔ)架構(gòu)建立在可靠的技術(shù)和理論基礎(chǔ)之上。并且，作為后盾，谷歌一直在做編程語言和構(gòu)建編譯器方面的前沿研究。

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

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

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

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

IREE的端到端流程

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

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

源代碼轉(zhuǎn)換

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

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

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

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

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

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

「Plane Strike」中的演示

用于機(jī)器學(xué)習(xí)的分布式系統(tǒng)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

硬件加速器和機(jī)器學(xué)習(xí)

用于ML的硬件設(shè)計

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

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

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

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

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

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

全棧加速器搜索技術(shù)

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

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

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

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

用于硬件設(shè)計的機(jī)器學(xué)習(xí)

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

方法和訓(xùn)練方案概述

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

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

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

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

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

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

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

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

依賴于硬件的模型設(shè)計

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

總結(jié)一下

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

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