簡(jiǎn)介

TensorFlow是Google研發(fā)的第二代人工智能學(xué)習(xí)系統(tǒng)，能夠處理多種深度學(xué)習(xí)算法模型，以功能強(qiáng)大和高可擴(kuò)展性而著稱(chēng)。TensorFlow完全開(kāi)源，所以很多公司都在使用，但是美團(tuán)點(diǎn)評(píng)在使用分布式TensorFlow訓(xùn)練WDL模型時(shí)，發(fā)現(xiàn)訓(xùn)練速度很慢，難以滿(mǎn)足業(yè)務(wù)需求。
經(jīng)過(guò)對(duì)TensorFlow框架和Hadoop的分析定位，發(fā)現(xiàn)在數(shù)據(jù)輸入、集群網(wǎng)絡(luò)和計(jì)算內(nèi)存分配等層面出現(xiàn)性能瓶頸。主要原因包括TensorFlow數(shù)據(jù)輸入接口效率低、PS/Worker算子分配策略不佳以及Hadoop參數(shù)配置不合理。我們?cè)谡{(diào)整對(duì)TensorFlow接口調(diào)用、并且優(yōu)化系統(tǒng)配置后，WDL模型訓(xùn)練性能提高了10倍，分布式線性加速可達(dá)32個(gè)Worker，基本滿(mǎn)足了美團(tuán)點(diǎn)評(píng)廣告和推薦等業(yè)務(wù)的需求。

術(shù)語(yǔ)

TensorFlow - Google發(fā)布的開(kāi)源深度學(xué)習(xí)框架
OP - Operation縮寫(xiě)，TensorFlow算子
PS - Parameter Server 參數(shù)服務(wù)器
WDL - Wide & Deep Learning，Google發(fā)布的用于推薦場(chǎng)景的深度學(xué)習(xí)算法模型
AFO - AI Framework on YARN的簡(jiǎn)稱(chēng) - 基于YARN開(kāi)發(fā)的深度學(xué)習(xí)調(diào)度框架，支持TensorFlow，MXNet等深度學(xué)習(xí)框架

TensorFlow分布式架構(gòu)簡(jiǎn)介

為了解決海量參數(shù)的模型計(jì)算和參數(shù)更新問(wèn)題，TensorFlow支持分布式計(jì)算。和其他深度學(xué)習(xí)框架的做法類(lèi)似，分布式TensorFlow也引入了參數(shù)服務(wù)器（Parameter Server，PS），用于保存和更新訓(xùn)練參數(shù)，而模型訓(xùn)練放在Worker節(jié)點(diǎn)完成。

TensorFlow分布式架構(gòu)

TensorFlow支持圖并行（in-graph）和數(shù)據(jù)并行（between-graph）模式，也支持同步更新和異步更新。因?yàn)閕n-graph只在一個(gè)節(jié)點(diǎn)輸入并分發(fā)數(shù)據(jù)，嚴(yán)重影響并行訓(xùn)練速度，實(shí)際生產(chǎn)環(huán)境中一般使用between-graph。
同步更新時(shí)，需要一個(gè)Woker節(jié)點(diǎn)為Chief，來(lái)控制所有的Worker是否進(jìn)入下一輪迭代，并且負(fù)責(zé)輸出checkpoint。異步更新時(shí)所有Worker都是對(duì)等的，迭代過(guò)程不受同步barrier控制，訓(xùn)練過(guò)程更快。

AFO架構(gòu)設(shè)計(jì)

TensorFlow只是一個(gè)計(jì)算框架，沒(méi)有集群資源管理和調(diào)度的功能，分布式訓(xùn)練也欠缺集群容錯(cuò)方面的能力。為了解決這些問(wèn)題，我們?cè)赮ARN基礎(chǔ)上自研了AFO框架解決這個(gè)問(wèn)題。
AFO架構(gòu)特點(diǎn)：

• 高可擴(kuò)展，PS、Worker都是任務(wù)（Task），角色可配置
• 基于狀態(tài)機(jī)的容錯(cuò)設(shè)計(jì)
• 提供了日志服務(wù)和Tensorboard服務(wù)，方便用戶(hù)定位問(wèn)題和模型調(diào)試
  
  AFO 架構(gòu)
  

AFO模塊說(shuō)明：

• Application Master：用來(lái)管理整個(gè)TensorFlow集群的資源申請(qǐng)，對(duì)任務(wù)進(jìn)行狀態(tài)監(jiān)控
• AFO Child：TensorFlow執(zhí)行引擎，負(fù)責(zé)PS、Worker運(yùn)行時(shí)管理和狀態(tài)同步
• History Server：管理TensorFlow訓(xùn)練生成的日志
• AFO Client：用戶(hù)客戶(hù)端

WDL模型

在推薦系統(tǒng)、CTR預(yù)估場(chǎng)景中，訓(xùn)練的樣本數(shù)據(jù)一般是查詢(xún)、用戶(hù)和上下文信息，系統(tǒng)返回一個(gè)排序好的候選列表。推薦系統(tǒng)面臨的主要問(wèn)題是，如何同時(shí)可以做到模型的記憶能力和泛化能力，WDL提出的思想是結(jié)合線性模型（Wide，用于記憶）和深度神經(jīng)網(wǎng)絡(luò)（Deep，用于泛化）。
以論文中用于Google Play Store推薦系統(tǒng)的WDL模型為例，該模型輸入用戶(hù)訪問(wèn)應(yīng)用商店的日志，用戶(hù)和設(shè)備的信息，給應(yīng)用App打分，輸出一個(gè)用戶(hù)“感興趣”App列表。

WDL 模型網(wǎng)絡(luò)

其中，installed apps和impression apps這類(lèi)特征具有稀疏性（在海量大小的App空間中，用戶(hù)感興趣的只有很少一部分），對(duì)應(yīng)模型“寬的部分”，適合使用線性模型；在模型“深的部分”，稀疏特征由于維度太高不適合神經(jīng)網(wǎng)絡(luò)處理，需要embedding降維轉(zhuǎn)成稠密特征，再和其他稠密特征串聯(lián)起來(lái)，輸入到一個(gè)3層ReLU的深度網(wǎng)絡(luò)。最后Wide和Deep的預(yù)估結(jié)果加權(quán)輸入給一個(gè)Logistic損失函數(shù)（例如Sigmoid）。
WDL模型中包含對(duì)稀疏特征的embedding計(jì)算，在TensorFlow中對(duì)應(yīng)的接口是tf.embedding_lookup_sparse，但該接口所包含的OP無(wú)法使用GPU加速，只能在CPU上計(jì)算，因此TensorFlow在處理稀疏特征性能不佳。不僅如此，我們發(fā)現(xiàn)分布式TensorFlow在進(jìn)行embedding計(jì)算時(shí)會(huì)引發(fā)大量的網(wǎng)絡(luò)傳輸流量，嚴(yán)重影響訓(xùn)練性能。

性能瓶頸分析與調(diào)優(yōu)

在使用TensorFlow訓(xùn)練WDL模型時(shí)，我們主要發(fā)現(xiàn)3個(gè)性能問(wèn)題：

每輪訓(xùn)練時(shí)，輸入數(shù)據(jù)環(huán)節(jié)耗時(shí)過(guò)多，超過(guò)60%的時(shí)間用于讀取數(shù)據(jù)。

訓(xùn)練時(shí)產(chǎn)生的網(wǎng)絡(luò)流量高，占用大量集群網(wǎng)絡(luò)帶寬資源，難以實(shí)現(xiàn)分布式性能線性加速。

Hadoop的默認(rèn)參數(shù)配置導(dǎo)致glibc malloc變慢，一個(gè)保護(hù)malloc內(nèi)存池的內(nèi)核自旋鎖成為性能瓶頸。

TensorFlow輸入數(shù)據(jù)瓶頸

TensorFlow支持以流水線（Pipeline）的方式輸入訓(xùn)練數(shù)據(jù)。如下圖所示，典型的輸入數(shù)據(jù)流水線包含兩個(gè)隊(duì)列：Filename Queue對(duì)一組文件做shuffle，多個(gè)Reader線程從此隊(duì)列中拿到文件名，讀取訓(xùn)練數(shù)據(jù)，再經(jīng)過(guò)Decode過(guò)程，將數(shù)據(jù)放入Example Queue，以備訓(xùn)練線程從中讀取數(shù)據(jù)。Pipeline這種多線程、多隊(duì)列的設(shè)計(jì)可以使訓(xùn)練線程和讀數(shù)據(jù)線程并行。理想情況下，隊(duì)列Example Queue總是充滿(mǎn)數(shù)據(jù)的，訓(xùn)練線程完成一輪訓(xùn)練后可以立即讀取下一批的數(shù)據(jù)。如果Example Queue總是處于“饑餓”狀態(tài)，訓(xùn)練線程將不得不阻塞，等待Reader線程將Example Queue插入足夠的數(shù)據(jù)。使用TensorFlow Timeline工具，可以直觀地看到其中的OP調(diào)用過(guò)程。

TensorFlow輸入數(shù)據(jù)流水線

使用Timeline，需要對(duì)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 timelinefetched_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)

這樣訓(xùn)練到global step在1000輪左右時(shí)，會(huì)將該輪訓(xùn)練的Timeline信息保存到timeline_01.json文件中，在Chrome瀏覽器的地址欄中輸入chrome://tracing，然后load該文件，可以看到圖像化的Profiling結(jié)果。
業(yè)務(wù)模型的Timeline如圖所示：

Timeline顯示數(shù)據(jù)輸入是性能瓶頸

可以看到QueueDequeueManyV2這個(gè)OP耗時(shí)最久，約占整體時(shí)延的60%以上。通過(guò)分析TensorFlow源碼，我們判斷有兩方面的原因：
（1）Reader線程是Python線程，受制于Python的全局解釋鎖（GIL），Reader線程在訓(xùn)練時(shí)沒(méi)有獲得足夠的調(diào)度執(zhí)行；
（2）Reader默認(rèn)的接口函數(shù)TFRecordReader.read函數(shù)每次只讀入一條數(shù)據(jù)，如果Batch Size比較大，讀入一個(gè)Batch的數(shù)據(jù)需要頻繁調(diào)用該接口，系統(tǒng)開(kāi)銷(xiāo)很大；
針對(duì)第一個(gè)問(wèn)題，解決辦法是使用TensorFlow Dataset接口，該接口不再使用Python線程讀數(shù)據(jù)，而是用C++線程實(shí)現(xiàn)，避免了Python GIL問(wèn)題。
針對(duì)第二個(gè)問(wèn)題，社區(qū)提供了批量讀數(shù)據(jù)接口TFRecordReader.read_up_to，能夠指定每次讀數(shù)據(jù)的數(shù)量。我們?cè)O(shè)置每次讀入1000條數(shù)據(jù)，使讀數(shù)句接口被調(diào)用的頻次從10000次降低到10次，每輪訓(xùn)練時(shí)延降低2-3倍。

優(yōu)化數(shù)據(jù)輸入使性能提升2-3倍

可以看到經(jīng)過(guò)調(diào)優(yōu)后，QueueDequeueManyV2耗時(shí)只有十幾毫秒，每輪訓(xùn)練時(shí)延從原來(lái)的800多毫秒降低至不到300毫秒。

集群網(wǎng)絡(luò)瓶頸

雖然使用了Mellanox的25G網(wǎng)卡，但是在WDL訓(xùn)練過(guò)程中，我們觀察到Worker上的上行和下行網(wǎng)絡(luò)流量抖動(dòng)劇烈，幅度2-10Gbps，這是由于打滿(mǎn)了PS網(wǎng)絡(luò)帶寬導(dǎo)致丟包。因?yàn)榉植际接?xùn)練參數(shù)都是保存和更新都是在PS上的，參數(shù)過(guò)多，加之模型網(wǎng)絡(luò)較淺，計(jì)算很快，很容易形成多個(gè)Worker打一個(gè)PS的情況，導(dǎo)致PS的網(wǎng)絡(luò)接口帶寬被打滿(mǎn)。
在推薦業(yè)務(wù)的WDL模型中，embedding張量的參數(shù)規(guī)模是千萬(wàn)級(jí)，TensorFlow的tf.embedding_lookup_sparse接口包含了幾個(gè)OP，默認(rèn)是分別擺放在PS和Worker上的。如圖所示，顏色代表設(shè)備，embedding lookup需要在不同設(shè)備之前傳輸整個(gè)embedding變量，這意味著每輪Embedding的迭代更新需要將海量的參數(shù)在PS和Worker之間來(lái)回傳輸。

embedding_lookup_sparse的OP拓?fù)鋱D

有效降低網(wǎng)絡(luò)流量的方法是盡量讓參數(shù)更新在一個(gè)設(shè)備上完成，即

with tf.device(PS):do embedding computing

社區(qū)提供了一個(gè)接口方法正是按照這個(gè)思想實(shí)現(xiàn)的：embedding_lookup_sparse_with_distributed_aggregation接口，該接口可以將embedding計(jì)算的所使用的OP都放在變量所在的PS上，計(jì)算后轉(zhuǎn)成稠密張量再傳送到Worker上繼續(xù)網(wǎng)絡(luò)模型的計(jì)算。
從下圖可以看到，embedding計(jì)算所涉及的OP都是在PS上，測(cè)試Worker的上行和下行網(wǎng)絡(luò)流量也穩(wěn)定在2-3Gpbs這一正常數(shù)值。

embedding_lookup_sparse_with_distributed_aggregation的OP拓?fù)鋱D

PS上的UniqueOP性能瓶頸

在使用分布式TensorFlow 跑廣告推薦的WDL算法時(shí)，發(fā)現(xiàn)一個(gè)奇怪的現(xiàn)象：WDL算法在AFO上的性能只有手動(dòng)分布式的1/4。手動(dòng)分布式是指：不依賴(lài)YARN調(diào)度，用命令行方式在集群上分別啟動(dòng)PS和Worker作業(yè)。
使用Perf診斷PS進(jìn)程熱點(diǎn)，發(fā)現(xiàn)PS多線程在競(jìng)爭(zhēng)一個(gè)內(nèi)核自旋鎖，PS整體上有30%-50%的CPU時(shí)間耗在malloc的在內(nèi)核的spin_lock上。

Perf診斷PS計(jì)算瓶頸

進(jìn)一步查看PS進(jìn)程棧，發(fā)現(xiàn)競(jìng)爭(zhēng)內(nèi)核自旋鎖來(lái)自于malloc相關(guān)的系統(tǒng)調(diào)用。WDL的embedding_lookup_sparse會(huì)使用UniqueOp算子，TensorFlow支持OP多線程，UniqueOp計(jì)算時(shí)會(huì)開(kāi)多線程，線程執(zhí)行時(shí)會(huì)調(diào)用glibc的malloc申請(qǐng)內(nèi)存。
經(jīng)測(cè)試排查，發(fā)現(xiàn)Hadoop有一項(xiàng)默認(rèn)的環(huán)境變量配置:

export MALLOC_ARENA_MAX="4"

該配置意思是限制進(jìn)程所能使用的glibc內(nèi)存池個(gè)數(shù)為4個(gè)。這意味著當(dāng)進(jìn)程開(kāi)啟多線程調(diào)用malloc時(shí)，最多從4個(gè)內(nèi)存池中競(jìng)爭(zhēng)申請(qǐng)，這限制了調(diào)用malloc的線程并行執(zhí)行數(shù)量最多為4個(gè)。
翻查Hadoop社區(qū)相關(guān)討論，當(dāng)初增加這一配置的主要原因是：glibc的升級(jí)帶來(lái)多線程ARENA的特性，可以提高malloc的并發(fā)性能，但同時(shí)也增加進(jìn)程的虛擬內(nèi)存（即top結(jié)果中的VIRT）。YARN管理進(jìn)程樹(shù)的虛擬內(nèi)存和物理內(nèi)存使用量，超過(guò)限制的進(jìn)程樹(shù)將被殺死。將MALLOC_ARENA_MAX的默認(rèn)設(shè)置改為4之后，可以不至于VIRT增加很多，而且一般作業(yè)性能沒(méi)有明顯影響。
但這個(gè)默認(rèn)配置對(duì)于WDL深度學(xué)習(xí)作業(yè)影響很大，我們?nèi)サ袅诉@個(gè)環(huán)境配置，malloc并發(fā)性能極大提升。經(jīng)過(guò)測(cè)試，WDL模型的平均訓(xùn)練時(shí)間性能減少至原來(lái)的1/4。

調(diào)優(yōu)結(jié)果

注意：以下測(cè)試都去掉了Hadoop MALLOC_ARENA_MAX的默認(rèn)配置

我們?cè)贏FO上針對(duì)業(yè)務(wù)的WDL模型做了性能調(diào)優(yōu)前后的比對(duì)測(cè)試，測(cè)試環(huán)境參數(shù)如下：
模型：推薦廣告模型WDL
OS：CentOS 7.1
CPU: Xeon E5 2.2G, 40 Cores
GPU：Nvidia P40
磁盤(pán): Local Rotational Disk
網(wǎng)卡：Mellanox 25G（未使用RoCE）
TensorFlow版本：Release 1.4
CUDA/cuDNN: 8.0/5.1

分布式線性加速效果

可以看到調(diào)優(yōu)后，訓(xùn)練性能提高2-3倍，性能可以達(dá)到32個(gè)GPU線性加速。這意味著如果使用同樣的資源，業(yè)務(wù)訓(xùn)練時(shí)間會(huì)更快，或者說(shuō)在一定的性能要求下，資源節(jié)省更多。如果考慮優(yōu)化MALLOC_ARENA_MAX的因素，調(diào)優(yōu)后的訓(xùn)練性能提升約為10倍左右。

總結(jié)

我們使用TensorFlow訓(xùn)練WDL模型發(fā)現(xiàn)一些系統(tǒng)上的性能瓶頸點(diǎn)，通過(guò)針對(duì)性的調(diào)優(yōu)不僅可以大大加速訓(xùn)練過(guò)程，而且可以提高GPU、帶寬等資源的利用率。在深入挖掘系統(tǒng)熱點(diǎn)瓶頸的過(guò)程中，我們也加深了對(duì)業(yè)務(wù)算法模型、TensorFlow框架的理解，具有技術(shù)儲(chǔ)備的意義，有助于我們后續(xù)進(jìn)一步優(yōu)化深度學(xué)習(xí)平臺(tái)性能，更好地為業(yè)務(wù)提供工程技術(shù)支持。

https://tech.meituan.com/tensorflow-performance-bottleneck-analysis-on-hadoop.html

總結(jié)

以上是生活随笔為你收集整理的使用TensorFlow训练WDL模型性能问题定位与调优的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。