簡介
TensorFlow是Google研發的第二代人工智能學習系統，能夠處理多種深度學習算法模型，以功能強大和高可擴展性而著稱。TensorFlow完全開源，所以很多公司都在使用，但是美團點評在使用分佈式TensorFlow訓練WDL模型時，發現訓練速度很慢，難以滿足業務需求。
經過對TensorFlow框架和Hadoop的分析定位，發現在數據輸入、集群網絡和計算內存分配等層面出現性能瓶頸。主要原因包括TensorFlow數據輸入接口效率低、PS/Worker算子分配策略不佳以及Hadoop參數配置不合理。我們在調整對TensorFlow接口調用、並且優化系統配置後，WDL模型訓練性能提高了10倍，分佈式線性加速可達32個Worker，基本滿足了美團點評廣告和推薦等業務的需求。
術語
TensorFlow - Google發佈的開源深度學習框架
OP - Operation縮寫，TensorFlow算子
PS - Parameter Server 參數服務器
WDL - Wide & Deep Learning，Google發佈的用於推薦場景的深度學習算法模型
AFO - AI Framework on YARN的簡稱 - 基於YARN開發的深度學習調度框架，支持Tensorflow，MXNet等深度學習框架
TensorFlow分佈式架構簡介
為了解決海量參數的模型計算和參數更新問題，TensorFlow支持分佈式計算。和其他深度學習框架的做法類似，分佈式TensorFlow也引入了參數服務器（Parameter Server，PS），用於保存和更新訓練參數，而模型訓練放在Worker節點完成。
圖1 TensorFlow分佈式架構
TensorFlow支持圖並行（in-graph）和數據並行（between-graph）模式，也支持同步更新和異步更新。因為in-graph只在一個節點輸入並分發數據，嚴重影響並行訓練速度，實際生產環境中一般使用between-graph。
同步更新時，需要一個Woker節點為Chief，來控制所有的Worker是否進入下一輪迭代，並且負責輸出checkpoint。異步更新時所有Worker都是對等的，迭代過程不受同步barrier控制，訓練過程更快。
AFO架構設計
TensorFlow只是一個計算框架，沒有集群資源管理和調度的功能，分佈式訓練也欠缺集群容錯方面的能力。為了解決這些問題，我們在YARN基礎上自研了AFO框架解決這個問題。
AFO架構特點：
高可擴展，PS、Worker都是任務（Task），角色可配置
基於狀態機的容錯設計
提供了日誌服務和Tensorboard服務，方便用戶定位問題和模型調試
圖2 AFO 架構
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列表。
圖3 WDL 模型網絡
其中，installed apps和impression apps這類特徵具有稀疏性（在海量大小的App空間中，用戶感興趣的只有很少一部分），對應模型“寬的部分”，適合使用線性模型；在模型“深的部分”，稀疏特徵由於維度太高不適合神經網絡處理，需要embedding降維轉成稠密特徵，再和其他稠密特徵串聯起來，輸入到一個3層ReLU的深度網絡。最後Wide和Deep的預估結果加權輸入給一個Logistic損失函數（例如Sigmoid）。
WDL模型中包含對稀疏特徵的embedding計算，在TensorFlow中對應的接口是tf.embedding_lookup_sparse，但該接口所包含的OP（例如tf._gather）無法使用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調用過程。
圖4 TensorFlow輸入數據流水線
使用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如圖所示：
圖5 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倍。
圖6 優化數據輸入使性能提升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之間來回傳輸。
圖7 embedding_lookup_sparse的OP拓撲圖
有效降低網絡流量的方法是儘量讓參數更新在一個設備上完成，即
with tf.device(PS): do embedding computing
社區提供了一個接口方法正是按照這個思想實現的：embedding_lookup_sparse_with_distributed_aggregation接口，該接口可以將embedding計算的所使用的OP都放在變量所在的PS上，計算後轉成稠密張量再傳送到Worker上繼續網絡模型的計算。
從下圖可以看到，embedding計算所涉及的OP都是在PS上，測試Worker的上行和下行網絡流量也穩定在2-3Gpbs這一正常數值。
圖8 embedding_lookup_sparse_with_distributed_aggregation的OP拓撲圖
PS上的UniqueOP性能瓶頸
在使用分佈式TensorFlow 跑廣告推薦的WDL算法時，發現一個奇怪的現象：WDL算法在AFO上的性能只有手動分佈式的1/4。手動分佈式是指：不依賴YARN調度，用命令行方式在集群上分別啟動PS和Worker作業。
使用Perf診斷PS進程熱點，發現PS多線程在競爭一個內核自旋鎖，PS整體上有30%-50%的CPU時間耗在malloc的在內核的spin_lock上。
圖9 Perf診斷PS計算瓶頸
進一步查看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
圖10 分佈式線性加速效果
可以看到調優後，訓練性能提高2-3倍，性能可以達到32個GPU線性加速。這意味著如果使用同樣的資源，業務訓練時間會更快，或者說在一定的性能要求下，資源節省更多。如果考慮優化MALLOC_ARENA_MAX的因素，調優後的訓練性能提升約為10倍左右。
總結
我們使用TensorFlow訓練WDL模型發現一些系統上的性能瓶頸點，通過針對性的調優不僅可以大大加速訓練過程，而且可以提高GPU、帶寬等資源的利用率。在深入挖掘系統熱點瓶頸的過程中，我們也加深了對業務算法模型、TensorFlow框架的理解，具有技術儲備的意義，有助於我們後續進一步優化深度學習平臺性能，更好地為業務提供工程技術支持。
鄭坤，美團點評技術專家，2015年加入美團點評，負責深度學習平臺、Docker平臺的研發工作。
招聘信息
美團點評GPU計算團隊，致力於打造公司一體化的深度學習基礎設施平臺，涉及到的技術包括：資源調度、高性能存儲、高性能網絡、深度學習框架等。目前平臺還在建設中期，不論在系統底層、分佈式架構、算法工程優化上都有很大的挑戰。誠邀對這個領域感興趣的同學加盟，不論是工程背景，還是算法背景我們都非常歡迎。有興趣的同學可以發送簡歷到 zhengkun#meituan.com。
閱讀更多 美團技術團隊 的文章
