TVM優(yōu)化GPU機(jī)器翻譯
背景
神經(jīng)機(jī)器翻譯（NMT）是一種自動化的端到端方法，具有克服傳統(tǒng)基于短語的翻譯系統(tǒng)中的弱點的潛力。最近，阿里巴巴集團(tuán)正在為全球電子商務(wù)部署NMT服務(wù)。
將Transformer用作NMT系統(tǒng)的關(guān)鍵技術(shù)，相對于基于經(jīng)典RNN / LSTM的模型具有同等（甚至更高）的精度，對于高效的離線訓(xùn)練更為友好。盡管Transformer在離線訓(xùn)練階段很友好，打破了跨時間步長的依賴性，但在線推理效率不高。在生產(chǎn)環(huán)境中，已經(jīng)發(fā)現(xiàn)，初始版本的Transformer的推理速度約為1.5倍至2倍，比LSTM版本慢。已經(jīng)進(jìn)行了一些優(yōu)化來提高推理性能，例如圖形級op融合，循環(huán)不變節(jié)點運動。觀察到的一個特殊挑戰(zhàn)是，批處理matmul是Transformer中的主要性能熱點，而cuBLAS中的當(dāng)前實現(xiàn)并未得到很好的優(yōu)化。

下面的結(jié)果表明TVM生成內(nèi)核（與調(diào)度表優(yōu)化）帶來至少13X批量MATMUL計算加速，futher加快算子融合功能。

批處理Matmul
為什么批量Matmul
在Transformer中，批處理matmul廣泛用于multi-head attention計算。使用matmul批處理，關(guān)注層中multiple heads并行運行，幫助提高硬件的計算效率。

在推理階段對Transformer模型進(jìn)行了全面的分析，結(jié)果表明，批量matmul計算貢獻(xiàn)了約30％的GPU內(nèi)核執(zhí)行時間。使用nvprof，對cuBLAS批處理matmul內(nèi)核進(jìn)行一些第一性原理分析，可以清楚地表明，當(dāng)前的實現(xiàn)方式表現(xiàn)不佳，并且觀察到了一些有趣的現(xiàn)象。
什么是批量matmul
通常，批量矩陣計算對一批矩陣執(zhí)行乘法。批處理被認(rèn)為是“統(tǒng)一的”，即所有實例具有相同的維度（M，N，K），前導(dǎo)維度（lda，ldb，ldc），以及各自的A，B和C矩陣的轉(zhuǎn)置。
批量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)
}
}
批處理matmul形狀
在語言翻譯任務(wù)中，在其它工作負(fù)載中，批處理matmul的形狀顯著小于常規(guī)matmul計算。Transformer中的形狀與輸入語句的長度和解碼器步長有關(guān)。通常小于30。
對于批處理尺寸，給定一定的推斷批處理大小，這是一個固定的數(shù)字。例如，如果將16用作光束大小為4的批大小，則批大小為16 * 4 * #head（多頭注意中的頭數(shù)，通常為8）。矩陣M，K，N的形狀，在[1，最大解碼長度]或[1，最大編碼長度]的范圍內(nèi)。
cuBLAS batch matmul的性能問題
首先，對批處理matmul內(nèi)核進(jìn)行了理論上的FLOP分析。結(jié)果非常有趣：所有批處理Matmul的計算強(qiáng)度都有限（少于1個TFLOP）。
然后，通過nvprof剖析了具有多種形狀的matmul批處理cuBLAS性能。下表顯示了在帶有CUDA8.0的NVIDIA M40 GPU上獲得的一些指標(biāo)。
輸入形狀
[批，M，N，K] 核心 理論FLOP nvprof觀察到的FLOP 理論FLOPs /
觀察到的FLOPs
[512、17、17、128] maxwell_sgemmBatched_128x128_raggedMn_tn 18939904 2155872256 0.87％
[512、1、17、128] maxwell_sgemmBatched_128x128_raggedMn_tn 1114112 2155872256 0.052％
[512，17，1，128] maxwell_sgemmBatched_128x128_raggedMn_tn 1114112 2155872256 0.052％
[512，30，30，128] maxwell_sgemmBatched_128x128_raggedMn_tn 58982400 2155872256 2.74％
即使具有不同的形狀（M，N，K不同），所有maxwell_sgemmBatched_128x128_raggedMn_tn調(diào)用，也執(zhí)行相同數(shù)量的FLOP，這比理論值大得多。可以推斷出，所有這些不同的形狀都可以被填充成一定的形狀。在所有這些形狀中，即使在最佳情況下，理論上的FLOP仍然僅是實際執(zhí)行的FLOP的2.74％，因此，大多數(shù)計算都是相當(dāng)多余的。同樣，另一個cuBLAS內(nèi)核maxwell_sgemmBatched_64x64_raggedMn_tn的調(diào)用也顯示相同的現(xiàn)象。
顯然，cuBLAS的批量matmul實施遠(yuǎn)非效率。因此，使用TVM為NMT工作負(fù)載生成有效的批處理matmul內(nèi)核。
批量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’)
調(diào)度優(yōu)化
在宣布計算之后，需要精心設(shè)計自己的調(diào)度表以壓縮性能潛力。
塊/線程數(shù)的調(diào)整參數(shù)

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)
融合了批處理matmul的外部尺寸，即op尺寸的BB和FF，在批處理matmul計算中通常稱為“批處理”尺寸。然后將外部和內(nèi)部尺寸除以（number_thread * vthread）。
在批處理matmul中不需要交錯模式，因此虛擬線程號（vthread_y和vthread_x）都設(shè)置為1。
尋找number_thread的最佳組合
以下結(jié)果是在具有CUDA8.0的NVIDIA M40 GPU設(shè)備上獲得的。
輸入形狀[批處理，特征，M，N，K] num_thread_y，num_thread_x num_vthread_y，num_vthread_x 時間（us）
[64,8,1,17,128] 8,1 32,1 37.62
[64,8,1,17,128] 4,1 32,1 39.30
[64,8,1,17,128] 1,1 32,1 38.82
[64,8,1,17,128] 1,1 256,1 41.95
[64,8,1,17,128] 32,1 1,1 94.61
從以往的經(jīng)驗了解到，該方法找到的最佳組合num_thread_y，num_thread_x通過強(qiáng)力搜索。經(jīng)過蠻力搜索后，可以找到當(dāng)前形狀的最佳組合，在當(dāng)前計算中為num_thread_y= 8和num_thread_x= 32。
將批處理matmul 與其它算子融合
通常，現(xiàn)有的“黑盒” cuBLAS庫，調(diào)用扮演著通常使用的“ op Fusion”優(yōu)化策略的邊界的角色。但是，利用所生成的高效批處理matmul核，可以容易地打破融合邊界，不僅可以融合元素的操作方式，還可以進(jìn)一步提高性能。
從計算圖可以看出，批處理matmul總是跟在添加操作或轉(zhuǎn)置操作廣播之后。通過將“ add”或“ transpose”操作與批處理matmul融合，可以減少內(nèi)核啟動開銷和冗余內(nèi)存訪問時間。
批處理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和轉(zhuǎn)置融合計算可以聲明為：

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’)
融合內(nèi)核性能
選擇[batch = 64，heads = 8，M = 1，N = 17，K = 128]的形狀，詳細(xì)說明所生成代碼的性能。選擇17作為序列長度，這是生產(chǎn)場景中的平均輸入長度。
? TF-R1.4 BatchMatmul：513.9
? TF-R1.4 BatchMatmul+ Transpose（另購）：541.9
? TVM BatchMatmul：37.62美元
? TVM BatchMatmul+ Transpose（融合）：38.39美元
內(nèi)核融合優(yōu)化進(jìn)一步提高了1.7倍的速度。
與Tensorflow集成
批量matmul在工作量中的輸入形狀是有限的，可以很容易地預(yù)先枚舉。使用這些預(yù)定義的形狀，可以提前生成高度優(yōu)化的CUDA內(nèi)核（固定形狀計算可以帶來最佳的優(yōu)化潛力）。將生成適用于大多數(shù)形狀的通用批處理matmul內(nèi)核，為沒有相應(yīng)的提前生成的內(nèi)核的形狀提供回退機(jī)制。
針對特定形狀生成的高效內(nèi)核和后備內(nèi)核已集成到Tensorflow框架中。開發(fā)了融合操作，例如BatchMatMulTranspose或BatchMatMulAdd，使用TVM的runtime API生成特定生成的內(nèi)核，實現(xiàn)某些輸入形狀或調(diào)用后備內(nèi)核。進(jìn)行圖形優(yōu)化遍歷，用算子融合自動替換原始批處理matmul +添加/轉(zhuǎn)置模式。通過結(jié)合更積極的圖形優(yōu)化過程，試圖利用TVM為長尾算子模式生成更有效的融合內(nèi)核，以進(jìn)一步提高端到端性能。
概括
在阿里巴巴內(nèi)部，發(fā)現(xiàn)TVM是開發(fā)高性能GPU內(nèi)核，滿足內(nèi)部需求的非常有效的工具。以NMT Transformer模型為例來說明使用TVM的優(yōu)化策略。首先，通過第一性原理分析確定了Transformer模型的熱點。然后使用TVM生成高度優(yōu)化的CUDA內(nèi)核來取代CUBLAS版本（13X加速觀察）。接下來，利用TVM的內(nèi)核融合機(jī)制融合批處理matmul的先前/以下操作，以進(jìn)一步提高性能（進(jìn)一步提高1.7倍的性能）。端到端性能提高了1.4倍。基于這些生成的內(nèi)核，開發(fā)了圖優(yōu)化遍歷以自動用TVM融合內(nèi)核替換原始計算模式，確保優(yōu)化對最終用戶是透明的，作為AI基礎(chǔ)設(shè)施提供商，發(fā)現(xiàn)優(yōu)化策略的透明性對于推廣其優(yōu)化算法非常重要采用。最后，并非最不重要的一點是，所有這些優(yōu)化都以松散耦合的方式集成到TensorFlow中，展示了將TVM與不同的深度學(xué)習(xí)框架集成的潛在方法。此外，正在進(jìn)行一項將TVM集成為TensorFlow的代碼源后端的工作，希望將來能與社區(qū)共享更多結(jié)果。

總結(jié)

以上是生活随笔為你收集整理的TVM优化GPU机器翻译的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。