使用TensorFlow訓練WDL模型性能問題定位與調優

時間 2019-11-17

標籤使用 tensorflow 訓練 wdl 模型性能問題定位欄目系統性能简体版

原文原文鏈接

簡介

TensorFlow是Google研發的第二代人工智能學習系統，可以處理多種深度學習算法模型，以功能強大和高可擴展性而著稱。TensorFlow徹底開源，因此不少公司都在使用，可是美團點評在使用分佈式TensorFlow訓練WDL模型時，發現訓練速度很慢，難以知足業務需求。python

通過對TensorFlow框架和Hadoop的分析定位，發如今數據輸入、集羣網絡和計算內存分配等層面出現性能瓶頸。主要緣由包括TensorFlow數據輸入接口效率低、PS/Worker算子分配策略不佳以及Hadoop參數配置不合理。咱們在調整對TensorFlow接口調用、而且優化系統配置後，WDL模型訓練性能提升了10倍，分佈式線性加速可達32個Worker，基本知足了美團點評廣告和推薦等業務的需求。git

術語

TensorFlow - Google發佈的開源深度學習框架github

OP - Operation縮寫，TensorFlow算子算法

PS - Parameter Server 參數服務器chrome

WDL - Wide & Deep Learning，Google發佈的用於推薦場景的深度學習算法模型apache

AFO - AI Framework on YARN的簡稱 - 基於YARN開發的深度學習調度框架，支持TensorFlow，MXNet等深度學習框架json

TensorFlow分佈式架構簡介

爲了解決海量參數的模型計算和參數更新問題，TensorFlow支持分佈式計算。和其餘深度學習框架的作法相似，分佈式TensorFlow也引入了參數服務器（Parameter Server，PS），用於保存和更新訓練參數，而模型訓練放在Worker節點完成。api

TensorFlow支持圖並行（in-graph）和數據並行（between-graph）模式，也支持同步更新和異步更新。由於in-graph只在一個節點輸入並分發數據，嚴重影響並行訓練速度，實際生產環境中通常使用between-graph。同步更新時，須要一個Woker節點爲Chief，來控制全部的Worker是否進入下一輪迭代，而且負責輸出checkpoint。異步更新時全部Worker都是對等的，迭代過程不受同步barrier控制，訓練過程更快。瀏覽器

AFO架構設計

TensorFlow只是一個計算框架，沒有集羣資源管理和調度的功能，分佈式訓練也欠缺集羣容錯方面的能力。爲了解決這些問題，咱們在YARN基礎上自研了AFO框架解決這個問題。 AFO架構特色：bash

高可擴展，PS、Worker都是任務（Task），角色可配置
基於狀態機的容錯設計
提供了日誌服務和Tensorboard服務，方便用戶定位問題和模型調試

AFO模塊說明：

Application Master：用來管理整個TensorFlow集羣的資源申請，對任務進行狀態監控
AFO Child：TensorFlow執行引擎，負責PS、Worker運行時管理和狀態同步
History Server：管理TensorFlow訓練生成的日誌
AFO Client：用戶客戶端

WDL模型

在推薦系統、CTR預估場景中，訓練的樣本數據通常是查詢、用戶和上下文信息，系統返回一個排序好的候選列表。推薦系統面臨的主要問題是，如何同時能夠作到模型的記憶能力和泛化能力，WDL提出的思想是結合線性模型（Wide，用於記憶）和深度神經網絡（Deep，用於泛化）。以論文中用於Google Play Store推薦系統的WDL模型爲例，該模型輸入用戶訪問應用商店的日誌，用戶和設備的信息，給應用App打分，輸出一個用戶「感興趣」App列表。

其中，installed apps和impression apps這類特徵具備稀疏性（在海量大小的App空間中，用戶感興趣的只有不多一部分），對應模型「寬的部分」，適合使用線性模型；在模型「深的部分」，稀疏特徵因爲維度過高不適合神經網絡處理，須要embedding降維轉成稠密特徵，再和其餘稠密特徵串聯起來，輸入到一個3層ReLU的深度網絡。最後Wide和Deep的預估結果加權輸入給一個Logistic損失函數（例如Sigmoid）。 WDL模型中包含對稀疏特徵的embedding計算，在TensorFlow中對應的接口是tf.embedding_lookup_sparse，但該接口所包含的OP沒法使用GPU加速，只能在CPU上計算，所以TensorFlow在處理稀疏特徵性能不佳。不只如此，咱們發現分佈式TensorFlow在進行embedding計算時會引起大量的網絡傳輸流量，嚴重影響訓練性能。

性能瓶頸分析與調優

在使用TensorFlow訓練WDL模型時，咱們主要發現3個性能問題：

每輪訓練時，輸入數據環節耗時過多，超過60%的時間用於讀取數據。
訓練時產生的網絡流量高，佔用大量集羣網絡帶寬資源，難以實現分佈式性能線性加速。
Hadoop的默認參數配置致使glibc malloc變慢，一個保護malloc內存池的內核自旋鎖成爲性能瓶頸。

TensorFlow輸入數據瓶頸

TensorFlow支持以流水線（Pipeline）的方式輸入訓練數據。以下圖所示，典型的輸入數據流水線包含兩個隊列：Filename Queue對一組文件作shuffle，多個Reader線程今後隊列中拿到文件名，讀取訓練數據，再通過Decode過程，將數據放入Example Queue，以備訓練線程從中讀取數據。Pipeline這種多線程、多隊列的設計可使訓練線程和讀數據線程並行。理想狀況下，隊列Example Queue老是充滿數據的，訓練線程完成一輪訓練後能夠當即讀取下一批的數據。若是Example Queue老是處於「飢餓」狀態，訓練線程將不得不阻塞，等待Reader線程將Example Queue插入足夠的數據。使用TensorFlow Timeline工具，能夠直觀地看到其中的OP調用過程。

使用Timeline，須要對tf.Session.run()增長以下幾行代碼：

with tf.Session as sess:
    ptions = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    _ = sess.run([train_op, global_step], options=run_options, run_metadata=run_metadata)
    if global_step > 1000 && global_step < 1010:
        from tensorflow.python.client import timeline
        fetched_timeline = timeline.Timeline(run_metadata.step_stats)
        chrome_trace = fetched_timeline.generate_chrome_trace_format()
        with open('/tmp/timeline_01.json', 'w') as f:
            f.write(chrome_trace)
複製代碼

這樣訓練到global step在1000輪左右時，會將該輪訓練的Timeline信息保存到timeline_01.json文件中，在Chrome瀏覽器的地址欄中輸入chrome://tracing，而後load該文件，能夠看到圖像化的Profiling結果。業務模型的Timeline如圖所示：

能夠看到QueueDequeueManyV2這個OP耗時最久，約佔總體時延的60%以上。經過分析TensorFlow源碼，咱們判斷有兩方面的緣由：（1）Reader線程是Python線程，受制於Python的全局解釋鎖（GIL），Reader線程在訓練時沒有得到足夠的調度執行；（2）Reader默認的接口函數TFRecordReader.read函數每次只讀入一條數據，若是Batch Size比較大，讀入一個Batch的數據須要頻繁調用該接口，系統開銷很大；針對第一個問題，解決辦法是使用TensorFlow Dataset接口，該接口再也不使用Python線程讀數據，而是用C++線程實現，避免了Python GIL問題。針對第二個問題，社區提供了批量讀數據接口TFRecordReader.read_up_to，可以指定每次讀數據的數量。咱們設置每次讀入1000條數據，使讀數句接口被調用的頻次從10000次下降到10次，每輪訓練時延下降2-3倍。

能夠看到通過調優後，QueueDequeueManyV2耗時只有十幾毫秒，每輪訓練時延從原來的800多毫秒下降至不到300毫秒。

集羣網絡瓶頸

雖然使用了Mellanox的25G網卡，可是在WDL訓練過程當中，咱們觀察到Worker上的上行和下行網絡流量抖動劇烈，幅度2-10Gbps，這是因爲打滿了PS網絡帶寬致使丟包。由於分佈式訓練參數都是保存和更新都是在PS上的，參數過多，加之模型網絡較淺，計算很快，很容易造成多個Worker打一個PS的狀況，致使PS的網絡接口帶寬被打滿。在推薦業務的WDL模型中，embedding張量的參數規模是千萬級，TensorFlow的tf.embedding_lookup_sparse接口包含了幾個OP，默認是分別擺放在PS和Worker上的。如圖所示，顏色表明設備，embedding lookup須要在不一樣設備以前傳輸整個embedding變量，這意味着每輪Embedding的迭代更新須要將海量的參數在PS和Worker之間來回傳輸。

有效下降網絡流量的方法是儘可能讓參數更新在一個設備上完成，即

with tf.device(PS):
    do embedding computing
複製代碼

社區提供了一個接口方法正是按照這個思想實現的：embedding_lookup_sparse_with_distributed_aggregation接口，該接口能夠將embedding計算的所使用的OP都放在變量所在的PS上，計算後轉成稠密張量再傳送到Worker上繼續網絡模型的計算。從下圖能夠看到，embedding計算所涉及的OP都是在PS上，測試Worker的上行和下行網絡流量也穩定在2-3Gpbs這一正常數值。

PS上的UniqueOP性能瓶頸

在使用分佈式TensorFlow 跑廣告推薦的WDL算法時，發現一個奇怪的現象：WDL算法在AFO上的性能只有手動分佈式的1/4。手動分佈式是指：不依賴YARN調度，用命令行方式在集羣上分別啓動PS和Worker做業。使用Perf診斷PS進程熱點，發現PS多線程在競爭一個內核自旋鎖，PS總體上有30%-50%的CPU時間耗在malloc的在內核的spin_lock上。

進一步查看PS進程棧，發現競爭內核自旋鎖來自於malloc相關的系統調用。WDL的embedding_lookup_sparse會使用UniqueOp算子，TensorFlow支持OP多線程，UniqueOp計算時會開多線程，線程執行時會調用glibc的malloc申請內存。經測試排查，發現Hadoop有一項默認的環境變量配置:

export MALLOC_ARENA_MAX="4"
複製代碼

該配置意思是限制進程所能使用的glibc內存池個數爲4個。這意味着當進程開啓多線程調用malloc時，最多從4個內存池中競爭申請，這限制了調用malloc的線程並行執行數量最多爲4個。翻查Hadoop社區相關討論，當初增長這一配置的主要緣由是：glibc的升級帶來多線程ARENA的特性，能夠提升malloc的併發性能，但同時也增長進程的虛擬內存（即top結果中的VIRT）。YARN管理進程樹的虛擬內存和物理內存使用量，超過限制的進程樹將被殺死。將MALLOC_ARENA_MAX的默認設置改成4以後，能夠不至於VIRT增長不少，並且通常做業性能沒有明顯影響。但這個默認配置對於WDL深度學習做業影響很大，咱們去掉了這個環境配置，malloc併發性能極大提高。通過測試，WDL模型的平均訓練時間性能減小至原來的1/4。

調優結果

注意：如下測試都去掉了Hadoop MALLOC_ARENA_MAX的默認配置

咱們在AFO上針對業務的WDL模型作了性能調優先後的比對測試，測試環境參數以下：模型：推薦廣告模型WDL OS：CentOS 7.1 CPU: Xeon E5 2.2G, 40 Cores GPU：Nvidia P40 磁盤: Local Rotational Disk 網卡：Mellanox 25G（未使用RoCE） TensorFlow版本：Release 1.4 CUDA/cuDNN: 8.0/5.1

能夠看到調優後，訓練性能提升2-3倍，性能能夠達到32個GPU線性加速。這意味着若是使用一樣的資源，業務訓練時間會更快，或者說在必定的性能要求下，資源節省更多。若是考慮優化MALLOC_ARENA_MAX的因素，調優後的訓練性能提高約爲10倍左右。

總結

咱們使用TensorFlow訓練WDL模型發現一些系統上的性能瓶頸點，經過針對性的調優不只能夠大大加速訓練過程，並且能夠提升GPU、帶寬等資源的利用率。在深刻挖掘系統熱點瓶頸的過程當中，咱們也加深了對業務算法模型、TensorFlow框架的理解，具備技術儲備的意義，有助於咱們後續進一步優化深度學習平臺性能，更好地爲業務提供工程技術支持。