TensorFlow+TVM優化NMT神經機器翻譯
背景
神經機器翻譯（NMT）是一種自動化的端到端方法，具有克服傳統基于短語的翻譯系統中的弱點的潛力。本文為全球電子商務部署NMT服務。
目前，將Transformer用作NMT系統的主要骨干，對基于經典RNN / LSTM模型的同等（甚至更高）精度進行高效的離線訓練更為友好。盡管Transformer在離線訓練階段很友好，打破了跨時間步長的依賴性，但在線推理效率不高。在生產環境中，已經發現，初始版本的Transformer的推理速度約為1.5倍至2倍比LSTM版本慢。進行一些優化來提高推理性能，例如圖形級op融合，循環不變節點運動。一個特殊挑戰是，批處理matmul是Transformer中的主要性能熱點，cuBLAS中的實現并未得到很好的優化。

下面的結果表明TVM生成內核（與調度表優化）帶來至少13X加速批量MATMUL計算和futher加快與運營商的融合功能。

批處理Matmul
為什么批量matmul
在Transformer中，批處理matmul廣泛用于計算多頭注意。使用批處理matmul，層中的多個頭部并行運行，幫助提高硬件的計算效率。

在推理階段對Transformer模型進行了全面的分析，結果表明，批量matmul計算貢獻了約30％的GPU內核執行時間。使用nvprof [2]對cuBLAS批處理matmul內核進行一些第一性原理分析，清楚地表明，當前的實現方式表現不佳，并且觀察到了一些有趣的現象。
什么是批量matmul
批量矩陣計算對一批矩陣執行矩陣矩陣乘法。批處理被認為是“統一的”，即所有實例的A，B和C矩陣具有相同的大小（M，N，K），前導大小（Ida，Idb，Idc）和置換。
批量matmul計算可以更具體地描述如下：
void BatchedGemm(input A, input B, output C, M, N, K, batch_dimension) {
for (int i = 0; i < batch_dimension; ++i) {
DoGemm(A[i],B[i],C[i],M,K,N)
}
}
批處理matmulshanpe
在語言翻譯任務中，批處理matmul的shanpe比其它工作負載中的常規matmul計算小得多。Transformer中的shanpe與輸入語句的長度和解碼器步長有關。通常小于30。
對于批處理大小，給定一定的推理批處理大小，它是一個固定的數字。例如，如果將16用作光束大小為4的批大小，則批大小為16 * 4 * #head（多頭注意中的頭數，通常為8）。矩陣M，K，N的shanpe在[1，最大解碼長度]或[1，最大編碼長度]的范圍內。
cuBLAS批次batch matmul的性能問題
首先，對批處理matmul內核進行了理論上的FLOP分析。結果非常有趣：所有批處理批量具有有限的計算強度（少于1個TFLOP）。
然后，通過nvprof剖析了具有多種shanpe的批處理matmul的cuBLAS性能。下表顯示了在帶有CUDA8.0的NVIDIA M40 GPU上獲得的一些指標。

即使具有不同的shanpe（M，N，K不同），所有maxwell_sgemmBatched_128x128_raggedMn_tn調用也執行相同數量的FLOP，這比理論值大得多。所有這些不同的shanpe都可以被填充成一定的shanpe。在所有這些shanpe中，即使在最佳情況下，理論上的FLOP仍僅是實際執行的FLOP的2.74％，大多數計算是相當多余的。另一個cuBLAS內核maxwell_sgemmBatched_64x64_raggedMn_tn的調用也顯示相同的現象。
cuBLAS的批量matmul實施遠非效率。因此，使用TVM為NMT工作負載生成有效的批處理matmul內核。
批量matmul計算
在TVM中，一般的批量Matmul計算可以聲明為：

computation representation

A = tvm.placeholder((batch, M, K), name=‘A’)
B = tvm.placeholder((batch, K, N), name=‘B’)
k = tvm.reduce_axis((0, K), ‘k’)
C = tvm.compute((batch, M, N),
lambda b, y, x: tvm.sum(A[b, y, k] * B[b, k, x], axis = k),
name = ‘C’)
日程優化
計算之后，需要精心設計自己的調度表以壓縮性能潛力。
塊/線程號的調整參數

thread indices

block_y = tvm.thread_axis(“blockIdx.y”)
block_x = tvm.thread_axis(“blockIdx.x”)
thread_y = tvm.thread_axis((0, num_thread_y), “threadIdx.y”)
thread_x = tvm.thread_axis((0, num_thread_x), “threadIdx.x”)
thread_yz = tvm.thread_axis((0, vthread_y), “vthread”, name=“vy”)
thread_xz = tvm.thread_axis((0, vthread_x), “vthread”, name=“vx”)

block partitioning

BB, FF, MM, PP = s[C].op.axis
BBFF = s[C].fuse(BB, FF)
MMPP = s[C].fuse(MM, PP)
by, ty_block = s[C].split(BBFF, factor = num_thread_y * vthread_y)
bx, tx_block = s[C].split(MMPP, factor = num_thread_x * vthread_x)
s[C].bind(by, block_y)
s[C].bind(bx, block_x)
vty, ty = s[C].split(ty_block, nparts = vthread_y)
vtx, tx = s[C].split(tx_block, nparts = vthread_x)
s[C].reorder(by, bx, vty, vtx, ty, tx)
s[C].reorder(by, bx, ty, tx)
s[C].bind(ty, thread_y)
s[C].bind(tx, thread_x)
s[C].bind(vty, thread_yz)
s[C].bind(vtx, thread_xz)
融合批處理mulmul的外部大小，即op大小的BB和FF，在批處理matmul計算中通常稱為“批處理”大小。然后將外部和內部大小除以（number_thread * vthread）。
在批處理matmul中不需要交錯模式，因此虛擬線程號（vthread_y和vthread_x）都設置為1。
尋找number_thread的最佳組合
以下結果是在具有CUDA8.0的NVIDIA M40 GPU設備上獲得的。

找到的最佳組合num_thread_y和num_thread_x，通過強力搜索。經過強力搜索后，可以找到當前shanpe的最佳組合，該組合為num_thread_y= 8和num_thread_x= 32。
將batch matmul與其它算子融合在一起
現有的“黑盒” cuBLAS庫調用充當通常使用的“ op Fusion”優化策略的邊界。但是，利用所生成的高效批處理matmul核，可以容易地打破融合邊界，不僅可以融合元素算子，還可以進一步提高性能。
從計算圖可以看出，批處理matmul總是跟在廣播添加算子或轉置算子之后。通過將“ add”或“ transpose”算子與批處理mult融合，可以減少內核啟動開銷和冗余內存訪問時間。
批處理matmul和廣播添加融合計算可以聲明如下：

computation representation

A = tvm.placeholder((batch_size, features, M, K), name=‘A’)

the shape of B is (N, K) other than (K, N) is because B is transposed is this fusion pattern

B = tvm.placeholder((batch_size, features, N, K), name=‘B’)
ENTER = tvm.placeholder((batch_size, 1, M, N), name = ‘ENTER’)
k = tvm.reduce_axis((0, K), ‘k’)
C = tvm.compute(
(batch_size, features, M, N),
lambda yb, yf, m, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, x, k], axis = k),
name = ‘C’)
D = topi.broadcast_add(C, ENTER)
批處理matmul和轉置融合計算可以聲明為：

computation representation

A = tvm.placeholder((batch_size, features, M, K), name=‘A’)
B = tvm.placeholder((batch_size, features, K, N), name=‘B’)
k = tvm.reduce_axis((0, K), ‘k’)
C = tvm.compute(
(batch_size, M, features, N),
lambda yb, m, yf, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, k, x], axis = k),
name = ‘C’)
融合內核性能
選擇[batch = 64，heads = 8，M = 1，N = 17，K = 128]的shanpe以詳細說明所生成代碼的性能。選擇17作為序列長度，輸出場景中的平均輸入長度。
? TF-R1.4 BatchMatmul：513.9
? TF-R1.4 BatchMatmul+ Transpose（另購）：541.9
? TVM BatchMatmul：37.62美元
? TVM BatchMatmul+ Transpose（融合）：38.39美元
內核融合優化進一步提高了1.7倍的速度。
與Tensorflow集成
批量matmul在工作量中的輸入shanpe是有限的，可以很容易地預先枚舉。使用這些預定義的shanpe，可以提前生成高度優化的CUDA內核（固定shanpe計算可以帶來最佳的優化潛力）。同時，還將生成適用于大多數shanpe的通用批處理matmul內核，為沒有相應的提前生成的內核的shanpe提供回退機制。
針對特定shanpe生成的高效內核和后備內核已集成到Tensorflow框架中。開發融合算子，例如BatchMatMulTranspose或BatchMatMulAdd，以使用TVM的運行時API針對特定輸入shanpe啟動特定生成的內核，或調用后備內核。進行圖形優化遍歷，用融合的算子自動替換原始批處理matmul +添加/轉置模式。結合更積極的圖形優化過程，試圖利用TVM為長尾算子模式生成更有效的融合內核，以進一步提高端到端性能。
總結
在阿里巴巴內部，發現TVM是開發高性能GPU內核以滿足內部需求的非常有效的工具。以NMT變壓器模型為例來說明使用TVM的優化策略。首先，通過第一性原理分析確定了Transformer模型的熱點。然后使用TVM生成高度優化的CUDA內核來取代CUBLAS版本（13X加速觀察）。接下來，利用TVM的內核融合機制融合批處理matmul的先前/以下算子，以進一步提高性能（進一步提高1.7倍的性能）。端到端性能提高了1.4倍。基于這些生成的內核，開發了圖優化遍歷，自動用TVM融合內核替換原始計算模式，確保優化對最終用戶是透明的，因為作為AI基礎架構提供商，發現優化策略的透明性對于推廣其優化算法非常重要。最后，但并非最不重要的一點是，所有這些優化都以松散耦合的方式集成到TensorFlow中，展示了將TVM與不同的深度學習框架集成的潛在方法。此外，將TVM集成為TensorFlow的代碼源后端的工作，希望將來能共享更多結果。

總結

以上是生活随笔為你收集整理的TensorFlow+TVM优化NMT神经机器翻译的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。