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

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

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

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

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

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

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

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

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

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

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

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

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

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

IREE的端到端流程

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

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

源代碼轉(zhuǎn)換

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

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

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

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

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

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

「Plane Strike」中的演示

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

此外，還有一個(gè)全新的用于多維陣列存儲(chǔ)的庫——TensorStore。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

在TPU/GPU上的平臺(tái)感知NAS概述

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

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

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

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

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

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

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

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

總結(jié)一下

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

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