天池 DeepRec CTR 模型效能優化大賽 - 奪冠技術分享

語言: CN / TW / HK

作者:niceperf 團隊 (李揚, 郭琳)

大家好,我們是 niceperf 團隊,在天池 DeepRec CTR 模型效能優化大賽中,很榮幸取得了冠軍的成績 (Top 1/3802)。這篇文章覆盤一下我們的參賽經驗,希望對大家有所啟發。

1.背景介紹

我們團隊包括兩名成員:李揚、郭琳,現就職於國內知名網際網路公司,擔任廣告演算法工程師。本次比賽的賽題,是在給定的深度學習框架 DeepRec 下,優化 WDL、DeepFM、DLRM、DIN、DIEN、MMoE 六大經典模型的單機 CPU 訓練速度。

賽題具有一定的挑戰性,我們在日常工作中經常使用的訓練效能優化手段主要是分散式訓練和資料IO 優化,而本次比賽限定了是單機條件,而且在資料 IO 方面的效能提升空間很有限。這就要求我們結合模型結構與訓練框架,做出更細緻深入的優化。

好在賽題涉及的技術棧與我們是較為契合的,比賽初期短暫適應後就可以上手優化了。首先,DeepRec的底層框架是 TensorFlow 1.15,我們在日常工作中已對其核心程式碼非常熟悉;此外,比賽所涉及的模型,團隊成員在實際業務中已使用多年,能夠結合對模型結構的理解,更快定位訓練效能瓶頸所在之處,便於合理分配比賽精力,在優勢方向重點發力。

2.優化內容

在做優化之前,要先確定優化的切入點,先把 low-hanging fruit 拿到,再做更深入的優化,有利於比賽的推進。我們逐模型分析了單步耗時,統計結果如下圖所示:

image.png

從資料來看,DeepFM 耗時如此之高是不符合預期的,因此我們將該模型作為切入點,展開後續的優化。

(1) IndicatorColumn 運算元優選 (DeepFM)

DeepFM 的單步訓練timeline 如下圖所示:

image.png

對上圖資料做 op 粒度的下鑽分析後發現,OneHot 為耗時頭部運算元,而該運算元來自於 IndicatorColumn,這裡可能存在較大優化空間。

圖片

進一步,我們精讀了 IndicatorColumn 的原始碼。為了相容定長 & 非定長的特徵,其輸入型別為 SparseTensor,輸出型別為 Tensor,為該特徵的 multi-hot 表達。具體處理過程分為三步:

Step 1:呼叫sparse_tensor_to_dense,將 SparseTensor 轉為 Tensor;Step 2:呼叫 one_hot,得到每個子特徵的 one-hot 表達;Step 3:呼叫reduce_sum,將屬於同樣本的子特徵 one-hot 求和,得到 multi-hot 表達。

針對部分高度稀疏的特徵,sparse_tensor_to_dense 帶來了額外的 padding,引入了無效計算。針對這一問題,我們建立了子類IndicatorColumnV2,改變原有的 multi-hot 生成方式,採用 scatter_nd 代替 sparse_tensor_to_dense + one_hot + reduce_sum。程式碼示意如下:

tf.scatter_nd(indices=multi-hot 非零值座標,updates=全 1 向量,shape=原輸入 sparse tensor 的 dense shape)

優化後單步訓練耗時從 500 ms 降低至 75 ms,前後效能對比如下兩表所示,其中紅框中的 ConcatV2 運算元是用來將多個特徵向量拼接送入後續 MLP 使用的。

圖片

圖片

ConcatV2 運算元是 "並行特徵處理" 與 "序列MLP 多層計算" 的分界點,它的耗時排名從第 9 名變為了第 1 名,可見特徵處理階段的 op 並行效率得到大幅提升。這一提升一方面來自於 IndicatorColumn 本身的效能優化,另一方面來自於有限硬體資源下,優化後的IndicatorColumn 讓出了更多執行緒供其他運算元使用,從而帶來了更顯著的效能提升。

(2) RNN cell 運算元融合 (DIEN)

針對 DeepFM 的優化告一段落,我們將關注點轉向單步耗時第二高的 DIEN。其整體結構如下圖所示:

圖片

在 DIEN 中有兩個 RNN 層,即 GRU 與 VecAttGRU,分別對應上圖中兩個紅框內的結構。由於 RNN 採用迴圈結構,所以很自然的想法是針對單步 cell 做優化,提升單步執行效能從而帶動整體效能提升。

比賽中我們分別針對 GRU cell 與 VecAttGRU cell 編寫了前向計算與梯度計算的 c++ 運算元,以替代原有的多個運算元構成的子圖,DIEN 整體耗時降低 ~67.96 秒。以 VecAttGRUCell 為例,前向 op、梯度 op 執行邏輯如下圖所示:

圖片

(3) Attention layer 運算元優選 (DIN & DIEN)

DIN 與 DIEN 中都用到了 Attention layer,在它們的原始實現中,存在大量無效 padding。這是因為不同使用者的歷史行為序列長度不一,當多個樣本 batching 在一起時,容易出現部分樣本序列短、部分樣本序列長的現象。

在原始實現中,會將各個樣本的 query padding 到統一長度,與 facts 互動後生成三維張量送入後續 MLP 計算。由於引入了無效 padding,導致 MLP 計算過程存在額外計算量,效率較低。

圖片

優化方法就是通過組合合適的運算元,去除 padding 部分,具體過程如下圖所示。

圖片

對比兩張圖的紅框部分,進入 MLP 的張量 shape 從 [B,T,4C] 變為[N,4C]。訓練過程中 BT >= N 恆成立,通常一個 batch 內的序列長度不均衡,BT/N 越大,效能收益越大,最終 DIN 與 DIEN 合計耗時降低約 141.38 秒。

(4) 序列特徵解析運算元融合

在模型輸入特徵解析環節,存在諸多瑣碎的小運算元構成的子圖,比如求序列均值特徵的子圖,典型案例如 history_price,原始資料為數值序列通過某個分隔符拼接構成的字串,構建模型輸入時需要將該字串 split 後轉成數值再通過多個步驟求均值,本次比賽中我們編寫了兩個 c++ 運算元 (SparseSequenceLength、StringSplitToNumberAndMean)替換掉了整個繁雜的過程。其他被替換的還有:求序列特徵 hash 分桶的子圖、序列特徵截斷的子圖,累積耗時降低約 88.47 秒。

圖片

(5) 工作流排程優化

對工作流的優化也是我們的一項重點工作。

第一個優化點是 "非同步檢查點",思路是將 checkpoint (即檢查點) 的儲存過程非同步化,減少對模型訓練過程的干擾。具體的優化方法是開發了 AsynchronousCheckpointSaverHook 替代原有的CheckpointSaverHook,開闢新的執行緒專門用來儲存 checkpoint,合計耗時降低約 31.7 秒。

圖片

另一項優化是多執行緒配置超參優化。如:針對 elm 資料集的相關模型開啟 smart stage;stage prefetch 使用獨立執行緒池,運算元並行度 intra_op 與 inter_op 均設為 8 (主執行緒池的配置也一致);stage prefetch threads 設為 1,capacity 設為 2;dataset map threads 設為 1 等等。

配置相關的超參優化暫無放之四海而皆準的規則,是我們結合給定的硬體資源、資料集,在多個模型整體效果上的粗調結果,調參思路是嘗試降低資料 parsing 佔用的執行緒資源,觀察是否有效能提升。最終多個版本優化後,整體耗時降低約 33.4 秒。

3.優化效果與總結

各優化點的效果彙總如下圖所示,整體訓練耗時降低 36.65%  (2063.25 秒 → 1307.17 秒)。

圖片

主要優化方法分為運算元優選與運算元融合。兩者在思路上是一致的,都是選擇效能更好的運算元 (組合) 來替換原來的運算元 (組合),在本文中的區別只在於是否開發了新的 macro op 替換原有的子圖。

一些其他的優化措施包括:非同步檢查點、多執行緒超參優化,也對效能提升有較大幫助。此外,其他嘗試過但沒有取得全域性穩定收益的措施還有很多,本文不做詳盡探討。

最後感謝主辦方提供的這次比賽機會,過程中主辦方技術團隊與我們多次溝通協作,不斷升級比賽評測系統,使評測結果更加穩定,為我們提供了卓越的競技環境。通過參加 DeepRec CTR 模型效能優化大賽,對深度學習框架效能優化加深了理解,期待DeepRec 框架繼續推出更多效能優化新特性。

DeepRec開源地址: https://github.com/alibaba/DeepRec