文章轉(zhuǎn)載于:知乎
作者:楊軍

本文是AI編譯優(yōu)化系列連載的第四篇，總綱請移步：

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

在之前的文章中，我們介紹了PAI團(tuán)隊(duì)在計(jì)算密集算子方面的優(yōu)化工作。

在這篇文章里，我們會介紹一個(gè)AI編譯優(yōu)化技術(shù)在阿里內(nèi)部實(shí)際業(yè)務(wù)場景的落地case。AI編譯技術(shù)在業(yè)務(wù)落地的過程中，不僅完善了自身的完備性和優(yōu)化效果，同時(shí)在和業(yè)務(wù)結(jié)合的過程中，加深了我們對于線上環(huán)境復(fù)雜度的認(rèn)知，理解了分布式訓(xùn)練中帶來的各種挑戰(zhàn)和問題，在過程中也積累了一些獲取最優(yōu)模型迭代性能的best practice。

機(jī)器翻譯模型訓(xùn)練是我們的一個(gè)典型的落地場景，在這個(gè)場景中，我們將訪存密集算子優(yōu)化+計(jì)算密集算子優(yōu)化+分布式+IO優(yōu)化多項(xiàng)技術(shù)進(jìn)行了結(jié)合，形成了組合優(yōu)化的效果，幫助業(yè)務(wù)方提升了模型迭代的效率，模型收斂時(shí)間由最初的三天半(85h)訓(xùn)練時(shí)間下降到了一天(25h)左右，同時(shí)結(jié)合落地過程中業(yè)務(wù)方同學(xué)的使用反饋我們也進(jìn)行了性能和可用性的改進(jìn)。

模型介紹

業(yè)務(wù)方使用的Alitranx-Transformer2.0翻譯質(zhì)量大幅提升，追平了目前世界上最好的開源NMT系統(tǒng)Marian。是業(yè)務(wù)方結(jié)合Transformer模型和阿里特有的電商場景，進(jìn)行了大量模型創(chuàng)新之后的模型結(jié)構(gòu)(部分工作也被AI頂會工作接受)。使用這個(gè)模型，在WMT2018 國際機(jī)器翻譯大賽上，阿里機(jī)器翻譯團(tuán)隊(duì)共提交5項(xiàng)結(jié)果，并全數(shù)獲得冠軍，Transformer-2.0模型在其中功不可沒。

AI編譯技術(shù)應(yīng)用介紹

針對業(yè)務(wù)方的Transformer 2.0模型的訓(xùn)練優(yōu)化需求，我們進(jìn)行了訪存密集算子優(yōu)化，計(jì)算密集算子優(yōu)化，以及和TF社區(qū)的Distribution Strategy自動(dòng)分布式工作配合的聯(lián)合優(yōu)化，在下面會逐一進(jìn)行介紹。

訪存密集算子優(yōu)化

宏觀來看，AI編譯器與其他編譯器架構(gòu)沒有本質(zhì)區(qū)別，分別由前端、中端、后端三個(gè)部分構(gòu)成。前端負(fù)責(zé)完成TensorFlow Op計(jì)算子圖到DL IR子圖的翻譯轉(zhuǎn)換工作，中端會在DL IR子圖層面執(zhí)行一系列圖優(yōu)化動(dòng)作，包括邏輯化簡、以及IR融合等等，后端會進(jìn)行Codegen以及可執(zhí)行碼的構(gòu)建。

下圖中我們以業(yè)務(wù)方使用的Transformer模型中的基本組成結(jié)構(gòu)，_DotProductAttention_為例，具體說明AI編譯器的整個(gè)工作流程。在我們獲取了用戶的構(gòu)圖代碼之后，前端會將原始Op構(gòu)圖描述(_DotProductAttention_主要為 _matmul_，_bias add_和_softmax_op)，轉(zhuǎn)換成更細(xì)粒度的DL IR表示(下圖中可以看到包含了細(xì)粒度的_dot_，_add_，_reduce_等指令)，經(jīng)過中端將這些DL IR表示進(jìn)行融合，可以看到融合之后DL IR指令數(shù)目大大減少，最后由后端生成優(yōu)化后的高效fusion kernel，由TensorFlow的Runtime調(diào)用執(zhí)行。

首先介紹我們研發(fā)的大尺度fusion codegen框架Fusion Stitching。

Fusion Stitching

我們針對NVIDIA GPU硬件平臺，原創(chuàng)性地提出了通過片上共享內(nèi)存進(jìn)行數(shù)據(jù)中轉(zhuǎn)從而大幅提升編譯圖融合力度的Fusion Stitching codegen框架，在編譯器中進(jìn)行了大量的精細(xì)優(yōu)化設(shè)計(jì)，詳細(xì)技術(shù)細(xì)節(jié)參見這里。

這個(gè)框架會同時(shí)涉及到編譯器中端和后端的動(dòng)作。中端負(fù)責(zé)激進(jìn)力度的計(jì)算子圖融合，后端基于融合后的子圖完成對應(yīng)的代碼生成動(dòng)作。

編譯器中端改進(jìn)

Fusion Stitching不再受限于OpFusion CodeGen的模版的設(shè)計(jì)理念，在Op Fusion的顆粒度上采用了一種比較激進(jìn)的策略，_只要是有生產(chǎn)消費(fèi)關(guān)系的且類型被支持的節(jié)點(diǎn)即可fuse在一起_。

我們依然以業(yè)務(wù)方機(jī)器翻譯模型使用的 DotProductionAttention_為例，下圖為其中的 _softmax 前向計(jì)算部分的DL IR表示，使用社區(qū)XLA的編譯優(yōu)化，這部分將產(chǎn)生 4 個(gè)Kernel，而基于我們的Stitching Fusion Pass，這個(gè)計(jì)算圖可以全部合并為 1 個(gè)Kernel,從而可以顯著節(jié)省包括kernel launch以及顯存訪問的性能開銷。

編譯器后端Codegen

社區(qū)XLA CodeGen對每個(gè)Codegen出的Kernel基于一個(gè)單一的Parallel Loop模板，即每個(gè)cuda thread處理tensor中的一個(gè)element，這種簡單的模板在處理線性連接的elementwise計(jì)算圖時(shí)能夠解決問題，但當(dāng)計(jì)算圖的連接關(guān)系變復(fù)雜，同時(shí)計(jì)算圖中出現(xiàn)Reduce／Broadcast等復(fù)雜計(jì)算節(jié)點(diǎn)時(shí)，在實(shí)際業(yè)務(wù)中效果并不好，特別是對于包含了后向處理部分的訓(xùn)練過程計(jì)算圖更是如此。

而我們的工作將一個(gè)Kernel的計(jì)算圖通過shared memory做隔離劃分為多個(gè)子圖，每個(gè)部分為一個(gè)獨(dú)立計(jì)算的Parallel Loop，可以根據(jù)實(shí)際shape需要決定自己的schedule（即如何為每個(gè)cuda thread劃分工作量）。

下面的兩張圖分別為社區(qū)XLA Codegen的情況和FusionStitching Codegen的情況，可以看到社區(qū)的fusion做法中沒有涉及到shared memory的使用，而我們的codegen將社區(qū)做法中難以fuse在一起的計(jì)算 _通過低訪存開銷的shared memory作為橋接，將多個(gè)Kernel縫合在一起_，這就是FusionStitching的含義來源。

社區(qū)XLA Codegen的情況

Fusion Stitching CodeGen的情況如下：

其他優(yōu)化

在整個(gè)研發(fā)過程中，我們還進(jìn)行了對前端的優(yōu)化，解決了不合理的編譯子圖劃分引發(fā)不必要的memcpy，同時(shí)我們根據(jù)實(shí)際業(yè)務(wù)場景中遇到的需求，還提供了更完備的Op支持。在中端，我們完善了由于過于激進(jìn)的fusion策略可能帶來的潛在的可用性和性能問題。在后端，我們通過對于codegen指令優(yōu)化以及根據(jù)硬件和kernel計(jì)算自適應(yīng)的調(diào)節(jié)launch dimension的改進(jìn)，都在真實(shí)的workload中拿到了可觀的收益。

自動(dòng)混合精度

自動(dòng)混合精度可以使得用戶無需經(jīng)過繁瑣的模型改寫即可實(shí)現(xiàn)混合精度訓(xùn)練、充分利用NVIDIA V100 GPU TensorCore來優(yōu)化計(jì)算密集型kernel（比如Transformer模型中大量存在的矩陣乘計(jì)算）的性能、提升模型迭代速度。而靈活易用的loss scale接口，使得用戶可以在使用混合精度進(jìn)行訓(xùn)練的過程中方便的對可能存在的精度問題進(jìn)行控制，可以做到和fp32訓(xùn)練出的模型精度無損。同時(shí)，自動(dòng)混合精度訓(xùn)練出來的模型和全精度之前可以無縫的進(jìn)行restore加載，進(jìn)行finetune。

混合精度使用Best Practice

業(yè)務(wù)方在新的Transformer 2.0的訓(xùn)練中已經(jīng)使用了自動(dòng)混合精度的功能。但是取得的加速比并不明顯，經(jīng)過profiling發(fā)現(xiàn)，并不是所有的fp16的GEMM計(jì)算都使用了TensorCore的高效實(shí)現(xiàn)，這是由于TensorCore計(jì)算kernel實(shí)現(xiàn)的一個(gè)限制，要求在一個(gè)矩陣計(jì)算中，其兩個(gè)輸入矩陣的尺寸M、N、K都是8的整數(shù)倍時(shí)才會啟用TensorCore的kernel。

在用戶的模型里面中一般隱層的數(shù)目都是類似于 512/1024 這樣是8的倍數(shù)的magic number，已經(jīng)滿足是8的整數(shù)倍的要求。根據(jù)Transformer模型的特性，矩陣乘計(jì)算中尺寸的MNK有一個(gè)值或者兩個(gè)值就是隱層數(shù)目的大小，而另外的值的大小都是和計(jì)算中輸入的batch size相關(guān)的。基于模型的這個(gè)特性，我們提供了一個(gè)在用戶側(cè)輕巧修改的best practice方法，控制輸入的batch size是8的整數(shù)倍（如下代碼所示，用戶使用的TensorFlow中的_GroupByWindowDataset_，我們修改了其中控制batch size的window/_size/_func），從而保證了矩陣乘計(jì)算能夠使用TensorCore充分加速?；旌暇鹊牟糠旨铀俦扔芍暗?strong>1.18X提升到了1.48X。

def window_size_func(key):
  key += 1  # For bucket_width == 1, key 0 is unassigned.
  size = (num_gpus * batch_size_words // (key * bucket_width))
  size = size - size % (8 * num_gpus)
  return tf.to_int64(size)

混合精度和訪存密集算子優(yōu)化的結(jié)合

在上圖中，我們將GPU計(jì)算分成計(jì)算密集型的部分和訪存密集部分，在Transformer模型中，計(jì)算密集型部分占比約為50%，主要是矩陣乘計(jì)算，這個(gè)計(jì)算部分可以通過自動(dòng)混合精度的使用進(jìn)行加速，剩余的訪存密集型部分則由Fusion Stitching進(jìn)行優(yōu)化。我們可以看到，單獨(dú)自動(dòng)混合精度的性能提升是1.48倍，單獨(dú)Fusion Stitching優(yōu)化的提升是1.33倍，而我們同時(shí)使用兩種優(yōu)化手段，得到了2.4倍左右的性能提升，取得了1+1>2的組合優(yōu)化效果。一方面是由于使用了混合精度之后，提升了訪存密集計(jì)算的比例，擴(kuò)展了編譯優(yōu)化的可優(yōu)化空間；同時(shí)在進(jìn)行自動(dòng)混合精度過程中，為了控制精度而插入的cast節(jié)點(diǎn)和loss scale等計(jì)算，也可以被fusion優(yōu)化，消除了這些節(jié)點(diǎn)帶來的額外開銷。

自動(dòng)分布式

業(yè)務(wù)方的Transfromer-2.0模型使用了TF社區(qū)提供的MirroredStrategy，便捷高效的實(shí)現(xiàn)了單機(jī)8卡的訓(xùn)練(除了在社區(qū)已有的Distribution Strategy進(jìn)行完善強(qiáng)化之外，PAI團(tuán)隊(duì)在分布式方面也開展了大量的自研性質(zhì)的工作，后續(xù)也會在不同渠道有相關(guān)的文章進(jìn)行分享，也敬請大家期待)。在8卡情況下，由于我們充分利用了NVLink提供的高速卡間互聯(lián)進(jìn)行數(shù)據(jù)通信，同時(shí)隨著自動(dòng)混合精度和Fusion Stitching優(yōu)化的加入隨著計(jì)算的粒度變大和計(jì)算效率提高，通信和IO方面的壓力突顯出來，對我們提出了更多的挑戰(zhàn)。

這里著重講一下對于跨device control dependency依賴的聚合簡化和對數(shù)據(jù)IO的優(yōu)化，在這兩個(gè)優(yōu)化的共同作用下，我們消除了數(shù)據(jù)IO的bottleneck，使得我們計(jì)算優(yōu)化能夠充分發(fā)揮其威力，進(jìn)而提升業(yè)務(wù)方同學(xué)的模型迭代效率。

Dependency Optimizer

在profiling過程中發(fā)現(xiàn)，單機(jī)8卡情況下，存在著大量的_Send/Recv/Identity/Const Op_的launch阻塞了GPU設(shè)備側(cè)計(jì)算kernel的發(fā)射時(shí)機(jī)，這就影響到了GPU計(jì)算性能的表現(xiàn)。首先我們分析一下這些Op的來源。

ControlDependency來源

業(yè)務(wù)方模型訓(xùn)練中使用的是AdamOptimizer。在TensorFlow的AdamOptimizer實(shí)現(xiàn)中可以看到，AdamOptimizer需要在等到所有的update操作完成之后才會進(jìn)行最后的 乃至于global step的更新。

由于單機(jī)8卡情況下，我們使用MirroredStrategy作為同步更新的機(jī)制，因此需要等到所有卡上的Variable更新都完成之后才會進(jìn)行最終的global step等的更新。而Transformer2.0模型一共有198個(gè)Variable，因此每張卡上的global/_step更新都要依賴于當(dāng)前卡以及其他卡上的Variable更新完成，因此產(chǎn)生了7 * 198個(gè)跨device的依賴邊 ，8張卡每張卡上Variable都有這樣的跨device依賴，因此一共產(chǎn)生了8 * 7 * 198個(gè)跨device依賴邊。

TensorFlow實(shí)現(xiàn)機(jī)制中會將1個(gè)跨device control dependency依賴邊轉(zhuǎn)換為1個(gè)Send，1個(gè)Recv，1個(gè)Const以及1個(gè)Identity 4個(gè)Op，這部分引入了一共44,352個(gè)Op，這些op的launch不能被overlap隱藏起來，因此其造成了后續(xù)計(jì)算的阻塞。

ControlDependency聚合

假設(shè)我們存在A->D, B->D,C->D這樣的依賴邊，同時(shí)A、B、C在同一個(gè)device X上，而C在另外一個(gè)device Y上，則我們可以在device X上引入一個(gè)相同device上的D‘節(jié)點(diǎn)，使得A、B和C都依賴于D’，將原始的依賴關(guān)系等價(jià)轉(zhuǎn)換為了

A->D', B->D', C'->D‘，D'->D，將原有的3個(gè)跨device的依賴邊，聚合之后只存在一個(gè)跨device依賴邊，如下圖所示，圖中的紅色虛線就是跨device依賴邊。

效果和性能提升

經(jīng)過Dependency Optimizer對Control Dependency的合并簡化，減少了44,128個(gè)冗余op的使用。我們通過跨device依賴邊的聚合優(yōu)化，獲得了約10%的性能提升，同時(shí)由于op launch的大大減少，也減輕了CPU的開銷。

數(shù)據(jù)預(yù)處理

業(yè)務(wù)方的模型訓(xùn)練中的數(shù)據(jù)預(yù)處理流程是經(jīng)典的文本數(shù)據(jù)bucketing的做法。輸入是原始文本數(shù)據(jù)，在預(yù)處理流程中轉(zhuǎn)換為數(shù)字id，并且按照數(shù)據(jù)長度組成各個(gè)bucket，每次讀取數(shù)據(jù)時(shí)，會從一個(gè)有數(shù)據(jù)的bucket中獲取一個(gè)batch的數(shù)據(jù)。

數(shù)據(jù)預(yù)處理阻塞計(jì)算

圖中的step time上可以看出300步之前，每個(gè)step時(shí)間比較平均，但是300步之后，平均時(shí)間開始增加，且伴隨著比較大的時(shí)間抖動(dòng)。我們profiling也發(fā)現(xiàn)了，經(jīng)過約300個(gè)step計(jì)算的之后，數(shù)據(jù)的預(yù)取隊(duì)列就變?yōu)榭?，意味著后續(xù)的計(jì)算中，每次取數(shù)據(jù)的時(shí)候，因?yàn)轭A(yù)取隊(duì)列為空，獲取數(shù)據(jù)時(shí)都要等數(shù)據(jù)處理好，再進(jìn)行計(jì)算，即數(shù)據(jù)供給部分阻塞了計(jì)算的進(jìn)行。

經(jīng)過分析發(fā)現(xiàn)，計(jì)算數(shù)據(jù)吞吐速率高于數(shù)據(jù)預(yù)處理時(shí)候的吞吐速率，因此出現(xiàn)的現(xiàn)象是計(jì)算等待數(shù)據(jù)，根據(jù)木桶效應(yīng)，此時(shí)的短板就在于數(shù)據(jù)預(yù)處理的效率，因?yàn)橹挥薪鉀Q掉解決掉數(shù)據(jù)預(yù)處理的短板，才能充分發(fā)揮計(jì)算的性能。

數(shù)據(jù)預(yù)處理中的性能熱點(diǎn)及優(yōu)化

經(jīng)過profiling我們發(fā)現(xiàn)，IO預(yù)處理部分pipeline中的瓶頸一是在vocabulary lookup和異常樣本過濾部分，另外一個(gè)是在最后bucketing的組batch部分，這兩個(gè)部分中，第一個(gè)部分是可以轉(zhuǎn)到離線過程中的，這樣就可以消除數(shù)據(jù)預(yù)處理對于計(jì)算的阻塞。

將數(shù)據(jù)IO性能熱點(diǎn)中的，vocabulary lookup和異常樣本過濾轉(zhuǎn)移到離線環(huán)節(jié)中，業(yè)務(wù)方同學(xué)使用UDF和SQL完成了這兩個(gè)步驟，并加入到了訓(xùn)練的pipeline當(dāng)中。

最終優(yōu)化效果

我們將訪存密集算子優(yōu)化，計(jì)算密集算子優(yōu)化以及分布式優(yōu)化相結(jié)合，改造了IO通路之后，目前的訓(xùn)練過程已經(jīng)集成到了業(yè)務(wù)生產(chǎn)訓(xùn)練平臺中，端到端加速比為3.65X，訓(xùn)練收斂時(shí)間由之前的85小時(shí)縮短到了25小時(shí)，顯著提升了業(yè)務(wù)模型訓(xùn)練的迭代效率。

值得一提的是，這里的加速結(jié)果，是在使用相同模型配置，數(shù)據(jù)以及相同的資源情況下所得到的，相當(dāng)于大幅提升了現(xiàn)有硬件資源的使用效率，從而達(dá)到PAI作為AI基礎(chǔ)設(shè)施提供方所一直倡導(dǎo)的“使用更少的硬件，支持更多業(yè)務(wù)更快完成迭代“的宗旨。

下面的曲線為訓(xùn)練過程中bleu值隨著迭代的step的變化曲線，可以看出上述所有優(yōu)化打開之后，訓(xùn)練過程中收斂趨勢和baseline基本一致，做到了精度的無損。

總結(jié)

在對機(jī)器翻譯團(tuán)隊(duì)模型訓(xùn)練進(jìn)行加速的過程中，我們將訪存密集算子優(yōu)化，計(jì)算密集算子優(yōu)化以及分布式/IO通用優(yōu)化手段相結(jié)合，形成優(yōu)化的組合拳，獲得了理想的加速效果。在業(yè)務(wù)落地過程中解決的可用性和性能問題，都已經(jīng)沉淀到了PAI的訓(xùn)練工具中，作為通用優(yōu)化手段的一部分，在其他的workload上面也得到了充分的檢驗(yàn)，助力我們不斷完善AI編譯的技術(shù)建設(shè)。

在機(jī)器翻譯模型使用編譯技術(shù)加速訓(xùn)練過程中，我們幾乎遇到了所有的訓(xùn)練任務(wù)的常見問題，編譯優(yōu)化的可用性問題，性能問題，通信問題，IO問題，甚至還遇到了一個(gè)不合理的padding帶來的精度損失問題，在這個(gè)過程中，不僅僅錘煉了AI編譯技術(shù)的成熟度和可用性，我們也對未來的發(fā)力點(diǎn)有了進(jìn)一步的認(rèn)知。希望通過AI編譯優(yōu)化工作的持續(xù)推進(jìn)，可以讓業(yè)務(wù)方算法團(tuán)隊(duì)的同學(xué)能夠更加專注于模型和算法本身，真正的實(shí)現(xiàn)讓天下沒有難train的model。

這篇文章分享的其實(shí)是18年底團(tuán)隊(duì)的一個(gè)業(yè)務(wù)落地工作，在過去的近兩年時(shí)間，我們在AI編譯技術(shù)以及落地上有了更進(jìn)一步的發(fā)展，也在不斷啟動(dòng)一個(gè)又一個(gè)更exciting的項(xiàng)目，歡迎對我們的工作感興趣的同學(xué)聯(lián)系我們，一起來推進(jìn)AI編譯技術(shù)的建設(shè)打造。

推薦閱讀

• 如何在OpenCV DNN模塊中添加Tengine后端？
• Android手機(jī)身體檢測＆關(guān)鍵點(diǎn)

更多嵌入式AI相關(guān)內(nèi)容請關(guān)注嵌入式AI專欄。

審核編輯：符乾江