Julia加入TPU，這是一個靠自己也要融入機器學習的編程語言

選自arxiv

作者：Keno Fischer,Elliot Saba

機器之心編輯部

Julia 語言發展非常迅速，它可以視為同時具備了 Python 的靈活性與 C 的速度，但目前 TensorFlow 和 PyTorch 等框架官方都不支持 Julia 語言。因此近日有研究者藉助 XLA 底層編譯器為 Julia 構建 TPU 支持，他們表示該方法能夠將 Julia 程序編寫的 VGG19 模型融合到 TPU 可執行文件中，並調用 TPU 實現高效計算。而 Google.ai 的負責人 Jeff Dean 在推特上也表示「Julia + TPUs = fast and easily expressible ML computations!」

1. 引言

過去的幾年裡推動機器學習技術穩步發展的根本性改變之一是訓練和優化機器學習模型的巨大計算力。許多技術都是很年前就已經提出，唯有近幾年提升的計算力可以為現實世界的問題提供足夠優質的解決方案。這些計算能力的很大一部分是通過 GPU 獲取的，其針對向量的計算能力最初是為圖形而設計的，但機器學習模型通常需要執行複雜的矩陣運算，因此 GPU 同樣表現出了非常好的性能。

這些方法及 GPU 在現實世界，尤其是在機器學習領域的成功引發了硬體設計者的一系列創新，他們致力於為機器學習工作負載研發新的加速器。然而，儘管 GPU 很長一段時間都在 CUDA 等軟體系統發力，但這些庫通常不會擴展到新的非 GPU 加速器，為這些加速器開發軟體仍然是一大挑戰。

2017 年，谷歌宣布他們將通過雲服務向大眾提供他們專有的 TPU 機器學習加速器。最初，TPU 的使用局限於根據谷歌 TensorFlow 機器學習框架編寫的應用。幸運的是，2018 年 9 月，谷歌通過底層 XLA（Accelerated Linear Algebra）編譯器的 IR 開放了 TPU 的訪問許可權。這個 IR 是一個通用的優化編譯器，用於表達線性代數基元的任意計算，因此為使用 TPU 的非 TensorFlow 用戶以及非機器學習工作負載提供了良好的基礎。

在本文中，我們介紹了使用這個介面編譯通用 Julia 代碼的初步工作，它們可以進一步訪問谷歌雲的 TPU。這一方法與 TensorFlow（Abadi et al., 2016）採用的方法形成對比，後者沒有編譯 python 代碼，而是先用 Python 構建一個計算圖，然後再對這個計算圖進行編譯。它在美學上類似於 JAX（Frostig et al., 2018），JAX 的目標是通過跟蹤和 Offload 高級數組運算來 Offload Python 本身編寫的計算。然而重要的是，我們不依賴於追蹤，而是利用 Julia 的靜態分析和編譯能力來編譯整個程序，包括傳遞到設備端的所有控制流。

值得一提的是，我們的方法允許用戶在編寫模型時充分利用 Julia 語言的表現力。這些表現力主要體現在一些高級特徵上，如多重派發、高階函數和現有庫，如微分方程求解器（Rackauckas & Nie，2017）和通用線性代數常式等。由於只在純 Julia 代碼上運行，所以它也與 Zygote.jl（Innes, 2018）自動微分工具兼容，該工具能執行自動微分作為高級編譯過程。總的來說，我們能夠編譯使用 Flux 機器學習框架編寫的完整機器學習模型，將模型的前向、反向傳播及訓練迴路融合成一個可執行文件，並 Offload 到 TPU 中。

論文：Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

論文鏈接：https://arxiv.org/abs/1810.09868

摘要：谷歌的雲 TPU 是一種前景廣闊的新型機器學習工作負載硬體架構，近年來已經成就了谷歌很多里程碑式的機器學習突破。如今，谷歌已經在其雲平台上為大眾提供 TPU，最近又進一步開放，允許非 TensorFlow 前端使用。我們描述了一種通過這一新 API 及谷歌 XLA 編譯器將 Julia 程序的適當部分 Offload 到 TPU 的方法和實現。我們的方法能夠將 Julia 程序編寫的 VGG19 模型及其正向傳播完全融合到單個 TPU 可執行文件中，以便 Offload 到設備上。我們的方法與 Julia 代碼上現有的基於編譯器的自動微分技術很好地結合在一起，因此也能夠自動獲得 VGG19 反向傳播並採用類似的方法將其 Offload 到 TPU。使用我們的編譯器訪問 TPU，我們能夠在 0.23 秒內完成批量為 100 張圖像的 VGG19 前向傳播，而 CPU 上的原始模型則需要 52.4s。我們的實現僅需不到 1000 行的 Julia 代碼，無需根據 TPU 對核心 Julia 編譯器或任何其他 Julia 包進行特有的更改。

5. 將 Julia 語義映射到 XLA

只要 Julia 程序是按照 XLA 基元來編寫的，我們就能將其編譯到 XLA。然而，Julia 程序不是根據晦澀難懂的 HLO 操作來編寫的，而是根據由 Julia 基本庫提供的函數和抽象來編寫的。幸運的是，Julia 使用了多重派發，使得根據 HLO 操作來表達標準庫的抽象變得容易。下面展示了幾個簡單的例子：

除了這些簡單的操作以外，我們還提供了高級數組抽象的實現，尤其是 mapreduce 和 broadcast。依據 HLO 操作實現的 broadcast 大約有 20 行代碼，為節省空間起見，此處不予展開，但「mapreduce」的實現非常簡單：

從上圖可以看到將任意 Julia 函數作為靜態計算運算的效果。由於 Julia 對泛型抽象的依賴，它只需指定極少數定義，就能覆蓋大量 API。具體來說，從 mapreduce 的定義中，我們可以自動得到在 base 中所定義運算（如 sum 和 prod）的降維。事實上，獲取足夠的 API 覆蓋來編譯 VGG19 模型的前向傳播和反向傳播需要不到 200 行定義。

5.1 結構映射

我們做了一個額外的識別。embedded IR 中的任意元組或 immutable 結構被映射至一個 XLA 元組，即 julia 值 1 + 2im（由兩個整數結構組成的複雜數字）將被映射至 XLA 元組 (s64[], s64[])。我們在 XLA IR 的 Julia 嵌入中保存該結構類型，但很顯然 XLA 不了解 julia 類型，因此在最終的轉換步驟中這些類型被轉換成適當的元組。類似地，（julia）元組構造函數（以及 immutable 結構的構造函數）變成了 XLA 的元組構件。元組引用（immutable 結構的欄位引用）變成了 XLA 的元組引用。

5.2 處理控制流

有一個額外的複雜問題我們還沒討論：Julia 提供的命令式控制流和 XLA 提供的函數式控制流之間的語義不匹配。為了解決 if/else 控制流模塊，我們在 Julia 編譯器的 SSA IR 中查看 φ 節點，然後將這些節點作為 XLA 函數式控制流的結果（如果在同一個合併點存在多個 φ 節點，則我們構造這些節點的元組）。導致計算流分化的條件變成了函數式控制流的條件，二者之間的任意計算都可作為函數調用。循環控制流類似條件控制流的構建，我們識別控制流圖的強連接區域，將其作為循環的主體。

7 結果

7.2 VGG19 前向傳播

我們的第一個複雜示例是完整 VGG19 前向傳播。我們使用 Metalhead 包中的 VGG19 實現 (Mike Innes & Contributors, 2018)，它利用 Flux (Innes & Contributors, 2017) 框架將熟悉的機器學習層（卷積層、全連接層）轉換成線性代數運算。但重要的是，Flux 框架中的每個層只是一般的函數，它們可以反過來調用一般的線性代數運算。因此，Flux 中表達的機器學習模型（包括 VGG19）只是一般的 Julia 函數，因此能夠使用本論文介紹的方法。

我們的編譯器能夠完全推斷、offload 和融合（fuse）VGG19 的全部前向傳播。在 Julia 級別的優化之後，頂層函數的最終 IR 包括 181 個指令（每個 HloOp 都是具備適當推斷的常數靜態參數和適當形態推斷的動態參數）。每個級別計算的 HLO operands 總數是 183（多出的兩個用於嵌入中隱藏的參數指令），29 個計算一共有 361 個 HLO operands，指令數詳情見圖 3。由於我們能夠 offload 全部前向傳播計算，因此 Julia 不參與任何評估步驟，從而可以同步執行其他任務（如為下一批準備數據）。此外，得到代碼的性能僅受限於 XLA 生成的代碼質量，不受限於前端（性能評估見 7.4）。我們在 ImageNet 驗證集上評估了 VGG19 模型，並驗證了得到結果與原版 Metalhead 的結果相匹配，從而驗證了生成的 XLA 代碼準確性。

7.3 VGG19 反向傳播

為了獲取反向傳播，我們利用基於 Zygote.jl 編譯器的 AD 框架 (Innes, 2018)。Zygote 在 Julia 代碼上運行，其輸出也是 Julia 函數（適合重新導入 Zygote 以獲取更高階的導數，也適合編譯成針對 TPU 的模型）。如下是一個具體示例：

即模型當前值和特定訓練樣本（或者訓練樣本批）所對應的導數。我們使用 sum 作為損失函數的簡單替代。意外的是，第 6 章介紹的類型推斷修改也能夠提高所有 VGG19 反向傳播的類型推斷精度。至於前向傳播，優化和未優化的指令總數如圖 1 所示。反向傳播生成的 XLA 指令明顯多於前向傳播，其最大貢獻者之一便是 Zygote 的混合模式廣播融合（mixed mode broadcast fusion）——在一個映射內核（map kernel）中同時計算前向傳播和反向傳播。由於 XLA 目前不支持來自一個映射指令的多個輸出，該函數在多個映射指令上重複運行，因此後續需要清洗 XLA 的 DCE。一般，我們的編譯過程解決了 XLA 對映射指令的處理，因為在泛型代碼中調用 Julia 映射和 broadcast 函數非常普遍。

7.4 在 TPU 上進行評估

圖 2：不同批大小對應的 VGG19 前向傳播時長。Flux CPU 是 Flux master/Julia master，但不使用 XLA 編譯器。PyTorch CPU 是同一 CPU 上的相同 PyTorch 模型。FluXLA CPU 是我們的研究在 CPU 上的 xrt 實現；FluXLA TPU (total) 是端到端時間，和客戶端報告的時間一致（包括 kernel launch 開銷和從谷歌雲把數據遷移回來，注意由於額外的網路遷移，該測量結果會出現極大的變動）；FluXLA TPU (compute) 是 TPU 上的總計算時間，和雲分析器報告的時間一致（與 FluXLA TPU (total) 不同，該測量很穩定）。所有 CPU 測量基於支持 AVX512 的 Intel(R) Xeon(R) Silver 4114 CPU @ 2.20GHz CPU。可獲取高達 20 個內核，且 CPU 基準不限於單個內核（即使在實踐中，也不是所有 CPU 基準都使用並行化）。TPU 基準僅限單個 TPU 內核。所有時間至少經過 4 次運行（除了 FluXLA CPU for N=100，因為它無法在 10 分鐘內完成一次運行）。

圖 3：被編譯為 XLA 後，Metalhead.jl VGG19 前向傳播和反向傳播的指令數分解，上圖展示了未優化（Julia 前端之後）和優化指令數（XLA 優化流程之後，與 CPU 後端所用流程類似，但沒有 HLO fusion）。每個指令數被進一步拆分為實體計算中的指令（E）和所有計算中的指令數（T）。

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

------------------------------------------------

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！