1313久久国产午夜精品理论片,91爱逼,爱爱91传媒

文章轉載于:知乎
作者:楊軍

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

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

在之前的文章中，我們介紹了PAI團隊在計算密集算子方面的優化工作。

在這篇文章里，我們會介紹一個AI編譯優化技術在阿里內部實際業務場景的落地case。AI編譯技術在業務落地的過程中，不僅完善了自身的完備性和優化效果，同時在和業務結合的過程中，加深了我們對于線上環境復雜度的認知，理解了分布式訓練中帶來的各種挑戰和問題，在過程中也積累了一些獲取最優模型迭代性能的best prac tice。

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

模型介紹

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

AI編譯技術應用介紹

針對業務方的Transformer 2.0模型的訓練優化需求，我們進行了訪存密集算子優化，計算密集算子優化，以及和TF社區的Distribution Strategy自動分布式工作配合的聯合優化，在下面會逐一進行介紹。

訪存密集算子優化

宏觀來看，AI編譯器與其他編譯器架構沒有本質區別，分別由前端、中端、后端三個部分構成。前端負責完成TensorFlow Op計算子圖到DL IR子圖的翻譯轉換工作，中端會在DL IR子圖層面執行一系列圖優化動作，包括邏輯化簡、以及IR融合等等，后端會進行Codegen以及可執行碼的構建。

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

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

Fusion Stitching

我們針對NVIDIA GPU硬件平臺，原創性地提出了通過片上共享內存進行數據中轉從而大幅提升編譯圖融合力度的Fusion Stitching codegen框架，在編譯器中進行了大量的精細優化設計，詳細技術細節參見這里。

這個框架會同時涉及到編譯器中端和后端的動作。中端負責激進力度的計算子圖融合，后端基于融合后的子圖完成對應的代碼生成動作。

編譯器中端改進

Fusion Stitching不再受限于OpFusion CodeGen的模版的設計理念，在Op Fusion的顆粒度上采用了一種比較激進的策略，_只要是有生產消費關系的且類型被支持的節點即可fuse在一起_。

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

編譯器后端Codegen

社區XLA CodeGen對每個Codegen出的Kernel基于一個單一的Parallel Loop模板，即每個cuda thread處理tensor中的一個element，這種簡單的模板在處理線性連接的elementwise計算圖時能夠解決問題，但當計算圖的連接關系變復雜，同時計算圖中出現Reduce／Broadcast等復雜計算節點時，在實際業務中效果并不好，特別是對于包含了后向處理部分的訓練過程計算圖更是如此。

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

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

社區XLA Codegen的情況

Fusion Stitching CodeGen的情況如下：

其他優化

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

自動混合精度

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

混合精度使用Best Practice

業務方在新的Transformer 2.0的訓練中已經使用了自動混合精度的功能。但是取得的加速比并不明顯，經過profiling發現，并不是所有的fp16的GEMM計算都使用了TensorCore的高效實現，這是由于TensorCore計算kernel實現的一個限制，要求在一個矩陣計算中，其兩個輸入矩陣的尺寸M、N、K都是8的整數倍時才會啟用TensorCore的kernel。

在用戶的模型里面中一般隱層的數目都是類似于 512/1024 這樣是8的倍數的magic number，已經滿足是8的整數倍的要求。根據Transformer模型的特性，矩陣乘計算中尺寸的MNK有一個值或者兩個值就是隱層數目的大小，而另外的值的大小都是和計算中輸入的batch size相關的。基于模型的這個特性，我們提供了一個在用戶側輕巧修改的best practice方法，控制輸入的batch size是8的整數倍（如下代碼所示，用戶使用的TensorFlow中的_GroupByWindowDataset_，我們修改了其中控制batch size的window/_size/_func），從而保證了矩陣乘計算能夠使用TensorCore充分加速。混合精度的部分加速比由之前的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)

混合精度和訪存密集算子優化的結合

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

自動分布式

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

這里著重講一下對于跨device control dependency依賴的聚合簡化和對數據IO的優化，在這兩個優化的共同作用下，我們消除了數據IO的bottleneck，使得我們計算優化能夠充分發揮其威力，進而提升業務方同學的模型迭代效率。

Dependency Optimizer

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

ControlDependency來源

業務方模型訓練中使用的是AdamOptimizer。在TensorFlow的AdamOptimizer實現中可以看到，AdamOptimizer需要在等到所有的update操作完成之后才會進行最后的 乃至于global step的更新。

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

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

ControlDependency聚合

假設我們存在A->D, B->D,C->D這樣的依賴邊，同時A、B、C在同一個device X上，而C在另外一個device Y上，則我們可以在device X上引入一個相同device上的D‘節點，使得A、B和C都依賴于D’，將原始的依賴關系等價轉換為了

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

效果和性能提升

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

數據預處理

業務方的模型訓練中的數據預處理流程是經典的文本數據bucketing的做法。輸入是原始文本數據，在預處理流程中轉換為數字id，并且按照數據長度組成各個bucket，每次讀取數據時，會從一個有數據的bucket中獲取一個batch的數據。

數據預處理阻塞計算

圖中的step time上可以看出300步之前，每個step時間比較平均，但是300步之后，平均時間開始增加，且伴隨著比較大的時間抖動。我們profiling也發現了，經過約300個step計算的之后，數據的預取隊列就變為空，意味著后續的計算中，每次取數據的時候，因為預取隊列為空，獲取數據時都要等數據處理好，再進行計算，即數據供給部分阻塞了計算的進行。

經過分析發現，計算數據吞吐速率高于數據預處理時候的吞吐速率，因此出現的現象是計算等待數據，根據木桶效應，此時的短板就在于數據預處理的效率，因為只有解決掉解決掉數據預處理的短板，才能充分發揮計算的性能。

數據預處理中的性能熱點及優化

經過profiling我們發現，IO預處理部分pipeline中的瓶頸一是在vocabulary lookup和異常樣本過濾部分，另外一個是在最后bucketing的組batch部分，這兩個部分中，第一個部分是可以轉到離線過程中的，這樣就可以消除數據預處理對于計算的阻塞。

將數據IO性能熱點中的，vocabulary lookup和異常樣本過濾轉移到離線環節中，業務方同學使用UDF和SQL完成了這兩個步驟，并加入到了訓練的pipeline當中。

最終優化效果

我們將訪存密集算子優化，計算密集算子優化以及分布式優化相結合，改造了IO通路之后，目前的訓練過程已經集成到了業務生產訓練平臺中，端到端加速比為3.65X，訓練收斂時間由之前的85小時縮短到了25小時，顯著提升了業務模型訓練的迭代效率。

值得一提的是，這里的加速結果，是在使用相同模型配置，數據以及相同的資源情況下所得到的，相當于大幅提升了現有硬件資源的使用效率，從而達到PAI作為AI基礎設施提供方所一直倡導的“使用更少的硬件，支持更多業務更快完成迭代“的宗旨。

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

總結

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

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

這篇文章分享的其實是18年底團隊的一個業務落地工作，在過去的近兩年時間，我們在AI編譯技術以及落地上有了更進一步的發展，也在不斷啟動一個又一個更exciting的項目，歡迎對我們的工作感興趣的同學聯系我們，一起來推進AI編譯技術的建設打造。

推薦閱讀