Google以機(jī)器學(xué)習(xí)技術(shù)優(yōu)化編譯器配置，強(qiáng)化模型執(zhí)行效率

Google發(fā)布機(jī)器學(xué)習(xí)最新研究，目的是通過改善機(jī)器學(xué)習(xí)編譯器，來提升機(jī)器學(xué)習(xí)工作負(fù)載的執(zhí)行效率，借由采用機(jī)器學(xué)習(xí)技術(shù)優(yōu)化編譯器決策，進(jìn)一步增加機(jī)器學(xué)習(xí)模型的性能。作為研究的一部分，Google還推出了新數(shù)據(jù)集TpuGraphs，用以促進(jìn)機(jī)器學(xué)習(xí)程序優(yōu)化研究。

隨著機(jī)器學(xué)習(xí)技術(shù)不斷發(fā)展，現(xiàn)在的機(jī)器學(xué)習(xí)系統(tǒng)已能理解自然語言，生成對(duì)話、圖像和視頻等。而這些先進(jìn)的機(jī)器學(xué)習(xí)模型，主要使用TensorFlow、JAX與PyTorch等機(jī)器學(xué)習(xí)程序開發(fā)框架開發(fā)和訓(xùn)練而成，這些框架提供了諸如線性代數(shù)運(yùn)算、神經(jīng)網(wǎng)絡(luò)等高端功能。

由于機(jī)器學(xué)習(xí)框架通過底層編譯器，自動(dòng)優(yōu)化用戶的模型，因此機(jī)器學(xué)習(xí)工作負(fù)載的性能，很大程度取決于編譯器的性能，研究人員指出，由于編譯器通常依賴啟發(fā)式方法來解決復(fù)雜的優(yōu)化問題，而這也常造成性能不佳的原因。

機(jī)器學(xué)習(xí)編譯器可以將用戶編寫的程序轉(zhuǎn)換成為可在硬件上運(yùn)行的可執(zhí)行指令，機(jī)器學(xué)習(xí)編譯器必需要解決許多復(fù)雜的優(yōu)化問題。而強(qiáng)化機(jī)器學(xué)習(xí)編譯器的性能，也就能夠提升機(jī)器學(xué)習(xí)模型的性能。

在編譯器優(yōu)化策略中，分配內(nèi)存布局給所有中介張量是一個(gè)關(guān)鍵的環(huán)節(jié)，內(nèi)存布局可以有不同的排列方式，常見的包括列優(yōu)先與行優(yōu)先布局，在機(jī)器學(xué)習(xí)的程序中，編譯器需要決定每個(gè)中介張量的內(nèi)存布局，而布局的優(yōu)化會(huì)影響計(jì)算效率。

編譯器在分配內(nèi)存布局時(shí)，需要權(quán)衡局部運(yùn)算效率和內(nèi)存布局轉(zhuǎn)換成本，而這個(gè)布局分配優(yōu)化的目的，便是在計(jì)算效率和內(nèi)存管理間，找到一個(gè)最佳平衡點(diǎn)。機(jī)器學(xué)習(xí)技術(shù)便可被用來找到該平衡點(diǎn)，通過分析程序的不同特性，像是操作種類、數(shù)據(jù)大小和結(jié)構(gòu)，預(yù)測(cè)可以提供最佳性能的內(nèi)存布局。

在編譯器中運(yùn)用學(xué)習(xí)成本模型，就能有效協(xié)助改進(jìn)經(jīng)編譯的機(jī)器學(xué)習(xí)模型效率。輸入程序和編譯器配置到學(xué)習(xí)成本模型中，模型便會(huì)輸出程序執(zhí)行時(shí)間預(yù)測(cè)。

雖然原本就存在一些用于程序性能預(yù)測(cè)數(shù)據(jù)集，但主要是針對(duì)小型子程序，為此，Google創(chuàng)建了TpuGraphs數(shù)據(jù)集來訓(xùn)練學(xué)習(xí)成本模型，特別是針對(duì)大型的計(jì)圖像運(yùn)算（Computational Graph）。由于這是一個(gè)用于Google機(jī)器學(xué)習(xí)加速器TPU執(zhí)行程序的學(xué)習(xí)成本模型數(shù)據(jù)集，因此稱作TpuGraphs，主要針對(duì)兩種XLA編譯器配置，涵蓋了布局和切割（Tiling）。

數(shù)據(jù)集中的樣本，包含了機(jī)器學(xué)習(xí)工作負(fù)載的計(jì)圖像運(yùn)算、編譯配置，還有使用該配置的執(zhí)行時(shí)間。數(shù)據(jù)集的內(nèi)容來自于開源機(jī)器學(xué)習(xí)程序，涵蓋多種熱門模型架構(gòu)，像是ResNet、EfficientNet、Mask R-CNN和Transformer等。

與之前的圖屬性預(yù)測(cè)數(shù)據(jù)集相比，TpuGraphs數(shù)據(jù)集圖的數(shù)量是25倍，平均圖大小比現(xiàn)有的機(jī)器學(xué)習(xí)性能預(yù)測(cè)數(shù)據(jù)集大770倍。官方指出，隨著數(shù)據(jù)集的擴(kuò)大，研究人員第一次可以探索大圖上針對(duì)圖層級(jí)的預(yù)測(cè)任務(wù)，解決包括可擴(kuò)展性、訓(xùn)練效率和模型品質(zhì)等問題。

Google還采用了一種稱為GST（Graph Segment Training）的訓(xùn)練方法，解決在內(nèi)存容量有限的設(shè)備上，訓(xùn)練大型圖神經(jīng)網(wǎng)絡(luò)的問題。該方法使端到端訓(xùn)練時(shí)間加速達(dá)3倍，有效提高訓(xùn)練的可行性和效率。

該篇研究的具體貢獻(xiàn)是將機(jī)器學(xué)習(xí)技術(shù)，用于優(yōu)化XLA編譯器配置，加速編譯后模型的執(zhí)行效率，另外，Google發(fā)展TpuGraphs數(shù)據(jù)集和GST訓(xùn)練方法高效訓(xùn)練學(xué)習(xí)成本模型，預(yù)測(cè)機(jī)器學(xué)習(xí)程序在特定編譯器配置下的執(zhí)行時(shí)間，以此作為優(yōu)化模型的基礎(chǔ)。