阿里巴巴最新實踐：TVM+TensorFlow提升神經機器翻譯性能

策劃 & 編譯｜Natalie

AI 前線導讀： 本文是阿里巴巴 PAI-Blade 團隊發表於 TVM 的最新博文，文中闡述瞭如何將 TVM 引入 TensorFlow，使 TensorFlow 中的 batchmul 速度提升 13 倍，同時將端到端神經機器翻譯性能提升 1.7 倍。AI 前線對原文進行了編譯。

更多幹貨內容請關注微信公衆號「AI 前線」，（ID：ai-front）

TVM 是由華盛頓大學在讀博士陳天奇等人提出的深度學習自動代碼生成方法，該技術能自動爲大多數計算硬件生成可佈署優化代碼，其性能可與當前最優的供應商提供的優化計算庫相比，且能夠適應新型專用加速器後端。

背景

神經機器翻譯（NMT）是一種端到端的自動翻譯方法，有可能克服傳統短語翻譯系統的不足。最近，阿里巴巴集團正在嘗試將 NMT 服務部署到全球電子商務業務場景中。

目前，咱們正在嘗試開發 Transformer[1] 做爲咱們 NMT 系統的核心組件，由於它比基於 RNN/LSTM 的經典模型更有利於高效的離線訓練，具備同等（甚至更高）的精確度。雖然 Transformer 打破了時間步長之間的依賴關係，所以對於離線訓練階段很是友好，可是對於在線推理來講，就沒有那麼高效了。在咱們的生產環境中，初始版本的 Transformer 推理速度比 LSTM 版本慢大約 1.5 倍到 2 倍。爲了提升推理性能，咱們已經嘗試進行了多種優化，如圖級別 op 融合、循環不變式節點運動 [3]。咱們發現批量 matmul 是 Transformer 中的一個性能關鍵點，而當前 cuBLAS 的實現並未進行很好的優化。

咱們的實踐結果代表，經過 TVM 生成的內核（優化了 schedule ）可使批量 matmul 計算速度提升至少 13 倍，而且在啓用了算子融合的狀況下還將進一步提升速度。

批量 Matmul

爲何要使用批量 matmul

在 Transformer 中，批量 matmul 被普遍應用於多頭 attention（multi-head attention）的計算。利用批量 matmul，attention 層中的多頭能夠並行運行，有助於提升硬件的計算效率。

咱們在推理階段對 Transformer 模型進行了深刻的剖析，結果代表批量 matmul 計算佔 GPU 內核執行時間的 30% 左右。咱們利用 nvprof[2] 對 cuBLAS 的批量 matmul 內核進行了第一性原理分析，結果代表當前的實現效果不好，同時咱們還發現了一些有趣的現象。

什麼是批量 matmul

一般，批量 matmul 計算是指在一批矩陣上執行矩陣 - 矩陣乘法。其中 batch 被認爲是「均勻的」，即全部實例對於它們各自的 A、B 和 C 矩陣具備相同的維度（M、N、K）、前導維度（lda、ldb、ldc）和轉換率。

批量 matmul 計算能夠更具體地描述以下：

批量 matmul 的 shape

在語言翻譯任務中，批量 matmul 的 shape 明顯小於其餘工做負載下的正常 matmul 計算。Transformer 的 shape 與輸入句子的長度和解碼器步驟的長度有關。一般會小於 30。

當給定推理的 batch size 時，批量維度（batch dimension）是一個固定值。例如，若是將 batch size 設定爲 16，beam 大小爲 4，則批量維度爲 16*4*#head（多頭 attention 中的頭數，一般爲 8），矩陣 M、K、N 的 shape 在 [1, 最大解碼長度] 或 [1, 最大編碼長度] 的範圍內。

cuBLAS 批量 matmul 的性能問題

首先，咱們對批量 matmul 內核進行了 FLOPs 理論分析。結果至關有趣：全部批量 matmul 的計算強度都很是有限（小於 1TFLOPs）。

接着咱們利用 nvprof 軟件對不一樣 shape 的批量 matmul 的 cuBLAS 進行了分析。下表顯示了使用 CUDA 8.0 在 NVIDIA M40 GPU 上獲得的一些指標。

即便使用不一樣的 shape（改變 M , N , K），全部的 Maxwell_sgembatched_128x128_raggedMn_tn 調用都會執行相同數量的 FLOP，並且實際值比理論值大得多。能夠推斷，全部這些不一樣的 shape 可能都被填充成了特定的 shape。在全部這些 shape 中，即便在最好的狀況下，理論 FLOP 仍僅佔實際執行 FLOP 的 2.74%，所以大部分計算至關冗餘。相似地，調用另外一個 cuBLAS 內核的 Maxwell_sgembatched_64x64_raggedMn_tn 出現了一樣的現象。

顯而易見，cuBLAS 的批量 matmul 實現效率很低。所以，咱們使用 TVM 爲 NMT 工做負載生成高效的批量 matmul 內核。

批量 matmul 計算

在 TVM 中，一般批量 matmul 計算能夠聲明爲：

Schedule優化

在對批量 matmul 計算進行聲明以後，咱們須要仔細設計Schedule，以充分挖掘其性能潛力。

調參：塊 / 線程數

咱們融合了批量 matmul 的外部維度，即 op 維度的 BB 和 FF，在批量 matmul 計算中一般稱爲「批量」維度。而後咱們用因子（number _ thread * vthread）分割外部維度和內部維度。

批量 matmul 中不須要 stridged 模式，所以虛擬線程數（vthread_y 和 vthread_x）都設置爲 1。

尋找 number_thread 的最佳組合

下圖是使用 CUDA8.0 在 NVIDIA M40 GPU 上獲得的結果。

根據以往的經驗，尋找 num_thread_y 和 num_thread_x 最佳組合的方法是暴力搜索。經過暴力搜索，能夠找到當前形狀（shape）的最佳組合，咱們在當前計算中獲得的最佳組合爲 num_thread_y=8，num_thread_x=32。

將批量 matmul 與其餘運算融合

當前的「黑盒」cuBLAS 庫調用一般會被做爲經常使用「op 融合」優化策略的邊界。然而，利用生成的高效批量 matmul 內核，能夠很容易地打破這一融合邊界，不只能夠融合 element-wise 運算，從而進一步提升性能。

從計算圖中能夠看出，批量 matmul 以後老是會進行廣播加法運算或轉置運算。經過將「加法」或「轉置」運算與批量 matmul 融合，能夠減小內核啓動開銷和冗餘內存訪問時間。

批量 matmul 和廣播加法融合計算可聲明以下：

批量 matmul 和轉置融合計算可聲明以下：

融合後的內核性能

咱們選擇 [batch=64, heads=8, M=1, N=17, K=128] 的形狀來測試前文生成的代碼的性能。咱們將序列長度設置爲 17，由於這個值是咱們生產場景中的平均輸入長度。

測試結果以下：

• tf-r1.4 BatchMatmul：513.9 us

• tf-r1.4 BatchMatmul+Transpose (separate)：541.9 us

• TVM BatchMatmul：37.62 us

• TVM BatchMatmul+Transpose (fused)：38.39 us

內核融合優化進一步將速度提升了 1.7 倍。

與 TensorFlow 集成

在咱們的工做負載中，批量 matmul 的輸入形狀是有限的，而且易於預先枚舉。利用這些預約義的形狀，咱們能夠提早生成高度優化的 CUDA 內核（固定形狀計算能夠帶來最佳的優化潛力）。同時，咱們還將生成適用於大多數形狀的通用批量 matmul 內核，爲沒有提早生成內核的形狀提供回退機制。

咱們將生成的針對特定形狀的高效內核和回退機制集成到了 TensorFlow 框架中。咱們開發了融合op，如 BatchMatMulTranspose 或 BatchMatMulAdd，以使用 TVM 的運行時 API 爲特定輸入形狀啓動特定的生成內核，或調用回退內核。經過執行圖形優化通道，能夠用融合 op 自動替換原始的批量 matmul+ 加法 / 轉置模式。同時，經過結合更激進的圖形優化通道，咱們嘗試利用 TVM 爲長尾操做模式生成更高效的融合內核，以進一步提升端到端性能。

總結

在阿里巴巴內部，咱們發現 TVM 是一個很是高效的工具，咱們能夠用它開發高性能的 GPU 內核以知足內部需求。

本文以 NMT Transformer 模型爲例，對 TVM 的優化策略進行了詳細說明。首先，經過第一性原理分析找出 Transformer 模型的關鍵問題。而後使用 TVM 生成高度優化的 CUDA 內核來替換 cuBLAS 版本（此時已經獲得了 13 倍的加速效果）。接下來，咱們利用 TVM 的內核融合機制來融合批量 matmul 的前 / 後操做，從而帶來進一步的性能改善（性能進一步提升 1.7 倍）。在此基礎上，咱們開發了一個圖形優化通道，自動用 TVM 融合內核替換原有的計算模式，保證優化過程對最終用戶透明。做爲 AI 基礎設施提供商，咱們發現透明度對於推廣優化策略的應用很是重要。

最後，全部這些優化都以鬆散耦合的方式集成到了 TensorFlow 中，展現了將 TVM 與不一樣深度學習框架集成的一種可能方式。此外，咱們目前還在進行將 TVM 整合爲 TensorFlow 代碼生成後端的工做，但願從此可以與社區分享更多成果。

資源

TVM implementation of fused batch matmul + transpose computation項目代碼

https://github.com/Orion34C/tvm-batch-matmul-example/blob/master/tvm_batch_matmul_transpose_m1_kX.py

PAI 團隊中文介紹

https://zhuanlan.zhihu.com/p/33513874

參考文獻

[1] Attention is All You Need

https://arxiv.org/pdf/1706.03762.pdf

[2] nvprof is Your Handy Universal GPU Profiler

https://devblogs.nvidia.com/cuda-pro-tip-nvprof-your-handy-universal-gpu-profiler/

[3] Add Loop Invariant Node Motion Optimization in GraphOptimizer

https://github.com/tensorflow/tensorflow/pull/16306

原文連接：

http://www.tvmlang.org/2018/03/23/nmt-transformer-optimize.html

更多幹貨內容請關注微信公衆號「AI 前線」，（ID：ai-front）