Google以機器學習技術最佳化編譯器配置，強化模型執行效率

Google發表機器學習最新研究，目的是透過改善機器學習編譯器，來提升機器學習工作負載的執行效率，藉由採用機器學習技術最佳化編譯器決策，進一步增加機器學習模型的效能。作為研究的一部分，Google還推出了新資料集TpuGraphs，用以促進機器學習程式最佳化研究。

隨著機器學習技術不斷發展，現在的機器學習系統已能理解自然語言，生成對話、圖像和影片等。而這些先進的機器學習模型，主要使用TensorFlow、JAX與PyTorch等機器學習程式開發框架開發和訓練而成，這些框架提供了諸如線性代數運算、神經網路等高階功能。

由於機器學習框架透過底層編譯器，自動最佳化使用者的模型，因此機器學習工作負載的效能，很大程度取決於編譯器的效能，研究人員指出，由於編譯器通常仰賴啟發式方法來解決複雜的最佳化問題，而這也常造成效能不佳的原因。

機器學習編譯器可以將用戶編寫的程式轉換成為可在硬體上運作的可執行指令，機器學習編譯器必須要解決許多複雜的最佳化問題。而強化機器學習編譯器的效能，也就能夠提升機器學習模型的效能。

在編譯器最佳化策略中，分配記憶體布局給所有中介張量是一個關鍵的環節，記憶體布局可以有不同的排列方式，常見的包括列優先與行優先布局，在機器學習的程式中，編譯器需要決定每個中介張量的記憶體布局，而布局的最佳化會影響計算效率。

編譯器在分配記憶體布局時，需要權衡局部運算效率和記憶體布局轉換成本，而這個布局分配最佳化的目的，便是在計算效率和記憶體管理間，找到一個最佳平衡點。機器學習技術便可被用來找到該平衡點，透過分析程式的不同特性，像是操作種類、資料大小和結構，預測可以提供最佳效能的記憶體布局。

在編譯器中運用學習成本模型，就能有效協助改進經編譯的機器學習模型效率。輸入程式和編譯器配置到學習成本模型中，模型便會輸出程式執行時間預測。

雖然原本就存在一些用於程式效能預測資料集，但主要是針對小型子程式，為此，Google建立了TpuGraphs資料集來訓練學習成本模型，特別是針對大型的計算圖（Computational Graph）。由於這是一個用於Google機器學習加速器TPU執行程式的學習成本模型資料集，因此稱作TpuGraphs，主要針對兩種XLA編譯器配置，涵蓋了布局和切割（Tiling）。

資料集中的樣本，包含了機器學習工作負載的計算圖、編譯配置，還有使用該配置的執行時間。資料集的內容來自於開源機器學習程式，涵蓋多種熱門模型架構，像是ResNet、EfficientNet、Mask R-CNN和Transformer等。

與之前的圖屬性預測資料集相比，TpuGraphs資料集圖的數量是25倍，平均圖大小比現有的機器學習效能預測資料集大770倍。官方指出，隨著資料集的擴大，研究人員第一次可以探索大圖上針對圖層級的預測任務，解決包括可擴展性、訓練效率和模型品質等問題。

Google還採用了一種稱為GST（Graph Segment Training）的訓練方法，解決在記憶體容量有限的裝置上，訓練大型圖神經網路的問題。該方法使端到端訓練時間加速達3倍，有效提高訓練的可行性和效率。

該篇研究的具體貢獻是將機器學習技術，用於最佳化XLA編譯器配置，加速編譯後模型的執行效率，另外，Google發展TpuGraphs資料集和GST訓練方法高效訓練學習成本模型，預測機器學習程式在特定編譯器配置下的執行時間，以此作為最佳化模型的基礎。