天池 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開源地址: http://github.com/alibaba/DeepRec