美團(tuán)內(nèi)部深度定制的TensorFlow版本,基于原生TensorFlow 1.x架構(gòu)與接口,從大規(guī)模稀疏參數(shù)的支持、訓(xùn)練模式、分布式通信優(yōu)化、流水線優(yōu)化、算子優(yōu)化融合等多維度進(jìn)行了深度優(yōu)化。在推薦系統(tǒng)場(chǎng)景中,分布式擴(kuò)展性提升10倍以上,單位算力性能也有顯著提升,并在美團(tuán)內(nèi)部業(yè)務(wù)中大量使用,本文介紹了相關(guān)的優(yōu)化與實(shí)踐工作。1 背景TensorFlow(下文簡(jiǎn)稱TF)是谷歌推出的一個(gè)開源深度學(xué)習(xí)框架,在美團(tuán)推薦系統(tǒng)場(chǎng)景中得到了廣泛的使用。但TensorFlow官方版本對(duì)工業(yè)級(jí)場(chǎng)景的支持,目前做得并不是特別的完善。美團(tuán)在大規(guī)模生產(chǎn)落地的過(guò)程中,遇到了以下幾方面的挑戰(zhàn):- 所有參數(shù)都是用Variable表達(dá), 對(duì)于百億以上的稀疏參數(shù)開辟了大量的內(nèi)存,造成了資源的浪費(fèi);
- 只支持百級(jí)別Worker的分布式擴(kuò)展,對(duì)上千Worker的擴(kuò)展性較差;
- 由于不支持大規(guī)模稀疏參數(shù)動(dòng)態(tài)添加、刪除,增量導(dǎo)出,導(dǎo)致無(wú)法支持Online Learning;
- 大規(guī)模集群運(yùn)行時(shí),會(huì)遇到慢機(jī)和宕機(jī);由于框架層不能處理,導(dǎo)會(huì)致任務(wù)運(yùn)行異常。
以上這些問(wèn)題,并不是TensorFlow設(shè)計(jì)的問(wèn)題,更多是底層實(shí)現(xiàn)的問(wèn)題??紤]到美團(tuán)大量業(yè)務(wù)的使用習(xí)慣以及社區(qū)的兼容性,我們基于原生TensorFlow 1.x架構(gòu)與接口,從大規(guī)模稀疏參數(shù)的支持、訓(xùn)練模式、分布式通信優(yōu)化、流水線優(yōu)化、算子優(yōu)化融合等多維度進(jìn)行了深度定制,從而解決了該場(chǎng)景的核心痛點(diǎn)問(wèn)題。首先新系統(tǒng)在支持能力層面,目前可以做到千億參數(shù)模型,上千Worker分布式訓(xùn)練的近線性加速,全年樣本數(shù)據(jù)能夠1天內(nèi)完成訓(xùn)練,并支持Online Learning的能力。同時(shí),新系統(tǒng)的各種架構(gòu)和接口更加友好,美團(tuán)內(nèi)部包括美團(tuán)外賣、美團(tuán)優(yōu)選、美團(tuán)搜索、廣告平臺(tái)、大眾點(diǎn)評(píng)Feeds等業(yè)務(wù)部門都在使用。本文將重點(diǎn)介紹大規(guī)模分布式訓(xùn)練優(yōu)化的工作,希望對(duì)大家能夠有所幫助或啟發(fā)。2 大規(guī)模訓(xùn)練優(yōu)化挑戰(zhàn)2.1 業(yè)務(wù)迭代帶來(lái)的挑戰(zhàn)隨著美團(tuán)業(yè)務(wù)的發(fā)展,推薦系統(tǒng)模型的規(guī)模和復(fù)雜度也在快速增長(zhǎng),具體表現(xiàn)如下:- 訓(xùn)練數(shù)據(jù):訓(xùn)練樣本從到百億增長(zhǎng)到千億,增長(zhǎng)了近10倍。
- 稀疏參數(shù):個(gè)數(shù)從幾百到幾千,也增長(zhǎng)了近10倍;總參數(shù)量從幾億增長(zhǎng)到百億,增長(zhǎng)了10~20倍。
- 模型復(fù)雜度:越來(lái)越復(fù)雜,模型單步計(jì)算時(shí)間增長(zhǎng)10倍以上。
對(duì)于大流量業(yè)務(wù),一次訓(xùn)練實(shí)驗(yàn),從幾個(gè)小時(shí)增長(zhǎng)到了幾天,而此場(chǎng)景一次實(shí)驗(yàn)保持在1天之內(nèi)是基本的需求。2.2 系統(tǒng)負(fù)載分析2.2.1 問(wèn)題分析工具鏈TensorFlow是一個(gè)非常龐大的開源項(xiàng)目,代碼有幾百萬(wàn)行之多,原生系統(tǒng)的監(jiān)控指標(biāo)太粗,且不支持全局的監(jiān)控,如果要定位一些復(fù)雜的性能瓶頸點(diǎn),就比較困難。我們基于美團(tuán)已經(jīng)開源的監(jiān)控系統(tǒng)CAT[2],構(gòu)建了TensorFlow的細(xì)粒度監(jiān)控鏈路(如下圖1所示),可以精準(zhǔn)定位到性能的瓶頸問(wèn)題。 圖1 TensorFlow PS架構(gòu)全鏈路監(jiān)控同時(shí),在性能優(yōu)化的過(guò)程中,會(huì)涉及到大量的性能測(cè)試和結(jié)果分析,這也是一個(gè)非常耗費(fèi)人力的工作。我們抽象了一套自動(dòng)化的實(shí)驗(yàn)框架(如下圖2所示),可以自動(dòng)化、多輪次地進(jìn)行實(shí)驗(yàn),并自動(dòng)采集各類監(jiān)控指標(biāo),然后生成報(bào)告。 圖2 自動(dòng)化實(shí)驗(yàn)框架2.2.2 業(yè)務(wù)視角的負(fù)載分析在推薦系統(tǒng)場(chǎng)景中,我們使用了TensorFlow Parameter Server[3](簡(jiǎn)稱PS)異步訓(xùn)練模式來(lái)支持業(yè)務(wù)分布式訓(xùn)練需求。對(duì)于這套架構(gòu),上述的業(yè)務(wù)變化會(huì)帶來(lái)什么樣的負(fù)載變化?如下圖3所示: 圖3 TensorFlow PS架構(gòu)大規(guī)模訓(xùn)練負(fù)載分析總結(jié)來(lái)看,主要包括通信壓力、PS并發(fā)壓力、Worker計(jì)算壓力。對(duì)于分布式系統(tǒng)來(lái)說(shuō),通常是通過(guò)橫向擴(kuò)展來(lái)解決負(fù)載問(wèn)題。雖然看來(lái)起可以解決問(wèn)題,但從實(shí)驗(yàn)結(jié)果來(lái)看,當(dāng)PS擴(kuò)展到一定數(shù)量后,單步訓(xùn)練時(shí)間反而會(huì)增加,如下圖4所示: 圖4 擴(kuò)展PS提升訓(xùn)練性能實(shí)驗(yàn) 導(dǎo)致這種結(jié)果的核心原因是:Worker單步訓(xùn)練需要和所有的PS通信同步完成,每增加1個(gè)PS要增加N條通信鏈路,這大大增加了鏈路延遲(如下圖5所示)。而一次訓(xùn)練要執(zhí)行上百萬(wàn)、上千萬(wàn)步訓(xùn)練。最終導(dǎo)致鏈路延遲超過(guò)了加PS算力并發(fā)的收益。 圖5 增加PS帶來(lái)的鏈路開銷而對(duì)于這個(gè)系統(tǒng),優(yōu)化的核心難點(diǎn)在于:如何在有限的PS實(shí)例下,進(jìn)行分布式計(jì)算的優(yōu)化。3 優(yōu)化實(shí)踐3.1 大規(guī)模稀疏參數(shù)介紹對(duì)于推薦系統(tǒng)模型,絕大多數(shù)參數(shù)都是稀疏參數(shù),而對(duì)稀疏參數(shù)來(lái)說(shuō)有一個(gè)非常重要的操作是Embedding,這個(gè)操作通常也是負(fù)載最重的,也是后續(xù)優(yōu)化的重點(diǎn)。由于我們對(duì)稀疏參數(shù)進(jìn)行了重新定義,后續(xù)的優(yōu)化也基于此之上,所以我們先介紹一下這部分的工作。在原生的TensorFlow中構(gòu)建Embedding模塊,用戶需要首先創(chuàng)建一個(gè)足夠裝得下所有稀疏參數(shù)的Variable,然后在這個(gè)Variable上進(jìn)行Embedding的學(xué)習(xí)。然而,使用Variable來(lái)進(jìn)行Embedding訓(xùn)練存在很多弊端:- Variable的大小必須提前設(shè)定好,對(duì)于百億千億的場(chǎng)景,該設(shè)定會(huì)帶來(lái)巨大的空間浪費(fèi);
- 訓(xùn)練速度慢,無(wú)法針對(duì)稀疏模型進(jìn)行定制優(yōu)化。
我們首先解決了有無(wú)的問(wèn)題,使用HashTable來(lái)替代Variable,將稀疏特征ID作為Key,Embedding向量作為Value。相比原生使用Variable進(jìn)行Embedding的方式,具備以下的優(yōu)勢(shì):- HashTable的大小可以在訓(xùn)練過(guò)程中自動(dòng)伸縮,避免了開辟冗余的存儲(chǔ)空間,同時(shí)用戶無(wú)需關(guān)注申請(qǐng)大小,從而降低了使用成本。
- 針對(duì)HashTable方案實(shí)施了一系列定制優(yōu)化,訓(xùn)練速度相比Variable有了很大的提高,可以進(jìn)行千億規(guī)模模型的訓(xùn)練,擴(kuò)展性較好。
- 得益于稀疏參數(shù)的動(dòng)態(tài)伸縮,我們?cè)诖嘶A(chǔ)上支持了Online Learning。
- API設(shè)計(jì)上保持與社區(qū)版本兼容,在使用上幾乎與原生Variable一致,對(duì)接成本極低。
簡(jiǎn)化版的基于PS架構(gòu)的實(shí)現(xiàn)示意如下圖6所示: 圖6 支撐大規(guī)模稀疏參數(shù)的HashTable方案- 稀疏特征ID(通常我們會(huì)提前完成統(tǒng)一編碼的工作)進(jìn)入Embedding模塊,借助TensorFlow搭建的Send-Recv機(jī)制,這些稀疏特征ID被拉取到PS端,PS端上的Lookup等算子會(huì)實(shí)際從底層HashTable中查詢并組裝Embedding向量。
- 上述Embedding向量被Worker拉回進(jìn)行后續(xù)訓(xùn)練,并通過(guò)反向傳播計(jì)算出這部分參數(shù)的梯度,這些梯度進(jìn)一步被位于PS端的優(yōu)化器拉回。
- PS端的優(yōu)化器首先調(diào)用Find算子,從HashTable獲取到梯度對(duì)應(yīng)的原始稀疏參數(shù)向量和相應(yīng)的優(yōu)化器參數(shù),最終通過(guò)優(yōu)化算法,完成對(duì)Embedding向量和優(yōu)化器參數(shù)的更新計(jì)算,再通過(guò)Insert算子插入HashTable中。
3.2 分布式負(fù)載均衡優(yōu)化這部分優(yōu)化,是分布式計(jì)算的經(jīng)典優(yōu)化方向。PS架構(gòu)是一個(gè)典型的“水桶模型”,為了完成一步訓(xùn)練,Worker端需要和所有PS完成交互,因此PS之間的平衡就顯得非常重要。但是在實(shí)踐中,我們發(fā)現(xiàn)多個(gè)PS的耗時(shí)并不均衡,其中的原因,既包括TensorFlow PS架構(gòu)簡(jiǎn)單的切圖邏輯(Round-Robin)帶來(lái)的負(fù)載不均衡,也有異構(gòu)機(jī)器導(dǎo)致的不均衡。對(duì)于推薦模型來(lái)說(shuō),我們的主要優(yōu)化策略是,把所有稀疏參數(shù)和大的稠密參數(shù)自動(dòng)、均勻的切分到每個(gè)PS上,可以解決大多數(shù)這類問(wèn)題。而在實(shí)踐過(guò)程中,我們也發(fā)現(xiàn)一個(gè)比較難排查的問(wèn)題:原生Adam優(yōu)化器,實(shí)現(xiàn)導(dǎo)致PS負(fù)載不均衡。下面會(huì)詳細(xì)介紹一下。在Adam優(yōu)化器中,它的參數(shù)優(yōu)化過(guò)程需要兩個(gè)β參與計(jì)算,在原生TensorFlow的實(shí)現(xiàn)中,這兩個(gè)β是所有需要此優(yōu)化器進(jìn)行優(yōu)化的Variabl(或HashTable)所共享的,并且會(huì)與第一個(gè)Variable(名字字典序)落在同一個(gè)PS上面,這會(huì)帶來(lái)一個(gè)問(wèn)題:每個(gè)優(yōu)化器只擁有一個(gè)和一個(gè),且僅位于某個(gè)PS上。因此,在參數(shù)優(yōu)化的過(guò)程中,該P(yáng)S會(huì)承受遠(yuǎn)高于其他PS的請(qǐng)求,從而導(dǎo)致該P(yáng)S成為性能瓶頸。 圖7 Adam優(yōu)化算法但是通過(guò)觀察Adam的優(yōu)化算法,我們可以看到和都是常量,且藍(lán)色高亮的部分都是相對(duì)獨(dú)立的計(jì)算過(guò)程,各個(gè)PS之間可以獨(dú)立完成?;谶@樣的發(fā)現(xiàn),優(yōu)化的方法也就非常直觀了,我們?yōu)槊恳粋€(gè)PS上的Adam優(yōu)化器冗余創(chuàng)建了β參數(shù),并在本地計(jì)算t和alpha值,去除了因此負(fù)載不均導(dǎo)致的PS熱點(diǎn)問(wèn)題。該優(yōu)化所帶來(lái)的提升具備普適性且效果明顯,在美團(tuán)內(nèi)部某業(yè)務(wù)模型上,通過(guò)β熱點(diǎn)去除可以帶來(lái)9%左右的性能提升。此外,由于擺脫了對(duì)β的全局依賴,該優(yōu)化還能提高PS架構(gòu)的可擴(kuò)展性,在擴(kuò)增Worker數(shù)量的時(shí)候相比之前會(huì)帶來(lái)更好的加速比。3.3 通信優(yōu)化通過(guò)2.2章節(jié)的分析可知,系統(tǒng)的通信壓力也非常大,我們主要基于RDMA做了通信優(yōu)化的工作。首先簡(jiǎn)單介紹一下RDMA,相比較于傳統(tǒng)基于套接字TCP/IP協(xié)議棧的通信過(guò)程,RDMA具有零拷貝、內(nèi)核旁路的優(yōu)勢(shì),不僅降低了網(wǎng)絡(luò)的延遲,同時(shí)也降低了CPU的占用率,RDMA更適合深度學(xué)習(xí)模型的相關(guān)通信過(guò)程。RDMA主要包括三種協(xié)議Infiniband、RoCE(V1, V2)、iWARP。在美團(tuán)內(nèi)部的深度學(xué)習(xí)場(chǎng)景中,RDMA通信協(xié)議使用的是RoCE V2協(xié)議。目前在深度學(xué)習(xí)訓(xùn)練領(lǐng)域,尤其是在稠密模型訓(xùn)練場(chǎng)景(NLP、CV等),RDMA已經(jīng)是大規(guī)模分布式訓(xùn)練的標(biāo)配。然而,在大規(guī)模稀疏模型的訓(xùn)練中,開源系統(tǒng)對(duì)于RDMA的支持非常有限,TensorFlow Verbs[4]通信模塊已經(jīng)很長(zhǎng)時(shí)間沒(méi)有更新了,通信效果也并不理想,我們基于此之上進(jìn)行了很多的改進(jìn)工作。經(jīng)過(guò)優(yōu)化后的版本,在1TB Click Logs[5]公開數(shù)據(jù)集、DLRM[6]模型、100個(gè)Worker以上的訓(xùn)練,性能提升了20%~40%。在美團(tuán)的多個(gè)業(yè)務(wù)模型上,對(duì)比TensorFlow Seastar[7]改造的通信層實(shí)現(xiàn)也有10%~60%的速度提升。同時(shí)也把我們的工作回饋給了社區(qū)。3.3.1 Memory Registration優(yōu)化RDMA有三種數(shù)據(jù)傳輸?shù)姆绞絊END/RECV、WRITE、READ,其中WRITE、READ類似于數(shù)據(jù)發(fā)送方直接在遠(yuǎn)程Memory進(jìn)行讀寫,Receiver無(wú)法感知,WRITE和READ適用于批量數(shù)據(jù)傳輸。在TensorFlow內(nèi)部,基于RDMA的數(shù)據(jù)傳輸方式使用的是WRITE單邊通信模式。 圖8 RDMA傳輸方式在RDMA傳輸數(shù)據(jù)時(shí),需要提前開辟內(nèi)存空間并將其注冊(cè)到網(wǎng)卡設(shè)備上(Memory Registration過(guò)程,下稱MR),使得這片空間可以被網(wǎng)卡直接操作。開辟新的內(nèi)存并注冊(cè)到設(shè)備上,整個(gè)過(guò)程是比較耗時(shí)的。下圖9展示了不同大小的內(nèi)存綁定到網(wǎng)卡設(shè)備上的耗時(shí),可以看到隨著注冊(cè)內(nèi)存的增大,綁定MR的耗時(shí)迅速增加。 圖9 MR過(guò)程開銷社區(qū)版Tensorflow RDMA實(shí)現(xiàn),Tensor創(chuàng)建依舊沿用了統(tǒng)一的BFC Allocator,并將所有創(chuàng)建的Tensor都注冊(cè)到MR上。正如上面所提到的,MR的注冊(cè)綁定具有性能開銷,高頻、大空間的MR注冊(cè)會(huì)帶來(lái)顯著的性能下降。而訓(xùn)練過(guò)程中的Tensor,只有那些涉及到跨節(jié)點(diǎn)通信的Tensor有必要進(jìn)行MR,其余Tensor并不需要注冊(cè)到MR。因此,優(yōu)化的方法也就比較直接了,我們識(shí)別并管理那些通信Tensor,僅對(duì)這些跨節(jié)點(diǎn)通信的Tensor進(jìn)行MR注冊(cè)就好了。3.3.2 RDMA靜態(tài)分配器RDMA靜態(tài)分配器是上一個(gè)MR注冊(cè)優(yōu)化的延伸。通過(guò)Memory Registration優(yōu)化,去除非傳輸Tensor的MR注冊(cè),我們降低了MR注冊(cè)數(shù)量。但是在稀疏場(chǎng)景大規(guī)模的訓(xùn)練下,并行訓(xùn)練的Worker常有幾百上千個(gè),這會(huì)帶來(lái)新的問(wèn)題:- PS架構(gòu)中的PS和Worker互為Client-Server,這里以PS端為例,當(dāng)Worker數(shù)目增加到上千個(gè)時(shí),Worker數(shù)目的增多,造成PS端MR注冊(cè)頻次還是非常高,增加了內(nèi)存分配注冊(cè)的耗時(shí)。
- 由于稀疏場(chǎng)景不同Step之間同一個(gè)算子輸出Tensor的形狀可能發(fā)生變化,導(dǎo)致了創(chuàng)建的MR可復(fù)用性較差,帶來(lái)了較高的內(nèi)存碎片和重復(fù)注冊(cè)MR開銷。
針對(duì)上面的問(wèn)題,我們引入了MR靜態(tài)分配器的策略。 圖10 MR靜態(tài)分配器
- 雖然稀疏場(chǎng)景同一個(gè)算子輸出Tensor的Shape存在變化的可能,但是整體變化幅度可控,通過(guò)監(jiān)控與分析,是可以找到一個(gè)較為穩(wěn)定的內(nèi)存大小,滿足多Step間Tensor的存儲(chǔ)需求。
- 基于上面的信息,我們修改了原有逐Tensor(Request)的MR申請(qǐng)策略,通過(guò)一次性預(yù)申請(qǐng)一塊較大的空間并注冊(cè)到網(wǎng)卡端,后續(xù)通過(guò)自己維護(hù)的分配策略進(jìn)行空間的分配,大大降低了MR申請(qǐng)的頻率,絕大多數(shù)情況下,訓(xùn)練全過(guò)程中只需要一次MR注冊(cè)申請(qǐng)即可。
- 我們引入了一種簡(jiǎn)單的交換協(xié)議,將傳輸Tensor的Shape,Data打包到一起,寫到Client端。Client端根據(jù)協(xié)議,解析出Tensor大小,并最終讀取Data,避免了原生實(shí)現(xiàn)中因Tensor的Shape變化而產(chǎn)生的多次協(xié)商過(guò)程。
 圖11 MR靜態(tài)分配器構(gòu)造流程具體到實(shí)現(xiàn)中,我們引入了Allocation Analysis模塊,在訓(xùn)練開始的一段時(shí)間,我們會(huì)對(duì)分配的歷史數(shù)據(jù)進(jìn)行分析,以得到一個(gè)實(shí)際預(yù)開辟M(fèi)R大小以及各個(gè)Tensor的預(yù)留空間大小。然后我們會(huì)暫停訓(xùn)練的進(jìn)程,啟動(dòng)Allocator的構(gòu)造過(guò)程,包括MR的創(chuàng)建以及通信雙端的信息同步。利用相關(guān)信息構(gòu)造MR Info Map,這個(gè)Map的Key是傳輸Tensor的唯一標(biāo)記(ParsedKey,計(jì)算圖切圖時(shí)確定),Info結(jié)構(gòu)體中包含了本地地址指針、offset大小、ibv_send_wr相關(guān)信息等。然后恢復(fù)訓(xùn)練,后續(xù)Tensor的傳輸就可以使用靜態(tài)開辟好的MR進(jìn)行收發(fā),也免去了因Shape變化而產(chǎn)生的多次協(xié)商過(guò)程。3.3.3 Multi RequestBuffer與CQ負(fù)載均衡TensorFlow社區(qū)版的RDMA通信過(guò)程,不僅僅包含上面Tensor數(shù)據(jù)的發(fā)送和接收過(guò)程,還包括傳輸相關(guān)的控制消息的發(fā)送和接收過(guò)程,控制消息的發(fā)送和接收過(guò)程同樣是使用了ibv_post_send和ibv_post_recv原語(yǔ)。原生的控制流實(shí)現(xiàn)存在一些瓶頸,在大規(guī)模訓(xùn)練時(shí)會(huì)限制控制流的吞吐,進(jìn)而影響數(shù)據(jù)收發(fā)的效率。具體體現(xiàn)在:- 請(qǐng)求的發(fā)送通過(guò)同一片RequestBuffer內(nèi)存進(jìn)行寫出,多個(gè)Client的請(qǐng)求均依賴這一片Buffer,也就導(dǎo)致到控制流信息實(shí)際是串行發(fā)送的,只有等到對(duì)端的Ack信息后,才可以下一個(gè)Request的寫出,限制了請(qǐng)求的發(fā)送吞吐。
- 在Client端需要輪詢RDMA Completion Queue來(lái)獲得請(qǐng)求的到達(dá),以及相關(guān)狀態(tài)的變更。原生實(shí)現(xiàn)僅有一個(gè)Completion Queue,單線程進(jìn)行輪詢處理,在大規(guī)模分布式訓(xùn)練中,限制了應(yīng)答的效率。
針對(duì)上面的問(wèn)題,我們采用了Multi RequestBuffer與CQ負(fù)載均衡優(yōu)化,破除了在請(qǐng)求發(fā)送和請(qǐng)求應(yīng)答環(huán)節(jié)可能存在的吞吐瓶頸。3.3.4 Send-Driven & Rendezvous-Bypass對(duì)于Tensorflow PS架構(gòu)熟悉的同學(xué)會(huì)了解,一整張計(jì)算圖被切割為Worker端和PS端后,為了使兩張計(jì)算圖能夠彼此交換數(shù)據(jù),建立了基于Rendezvous(匯合點(diǎn))機(jī)制的異步數(shù)據(jù)交換模式。如下圖12所示: 圖12 TensoFlow切圖之Send-Recv對(duì)添加基于上圖的切圖邏輯,Recv算子代表著這一側(cè)計(jì)算圖有Tensor的需求,而Tensor的生產(chǎn)者則位于與之配對(duì)的另一設(shè)備上的Send算子背后。在具體實(shí)現(xiàn)上,Tensorflow實(shí)現(xiàn)了Recv-Driven的數(shù)據(jù)交換模式,如上圖所示,位于DeviceA和DeviceB的兩張計(jì)算圖會(huì)異步并發(fā)的執(zhí)行,位于DeviceB的Recv執(zhí)行時(shí)會(huì)發(fā)起一條RPC請(qǐng)求發(fā)往DeviceA,DeviceA收到請(qǐng)求后,會(huì)將請(qǐng)求路由到Rendezvous中,如果在當(dāng)中發(fā)現(xiàn)所需要的數(shù)據(jù)已經(jīng)生產(chǎn)好,并被Send算子注冊(cè)了進(jìn)來(lái),那么就地獲取數(shù)據(jù),返回給DeviceB;如果此時(shí)數(shù)據(jù)還沒(méi)有生產(chǎn)好,則將來(lái)自于DeviceB的Recv請(qǐng)求注冊(cè)在Rendezvous中,等待后續(xù)DeviceA生產(chǎn)好后,由Send算子發(fā)送過(guò)來(lái),找到注冊(cè)的Recv,觸發(fā)回調(diào),返回?cái)?shù)據(jù)給DeviceB。我們看到,匯合點(diǎn)機(jī)制優(yōu)雅地解決了生產(chǎn)者消費(fèi)者節(jié)奏不同情況下數(shù)據(jù)交換的問(wèn)題。不過(guò)Recv-Driven的模式也引入了兩個(gè)潛在的問(wèn)題:- 據(jù)我們的觀察,在實(shí)際業(yè)務(wù)模型中,在Rendezvous中Recv算子等待Send算子的比例和Send算子等待Recv算子的比例相當(dāng),也就是說(shuō)對(duì)于Send等到Recv的數(shù)據(jù),在Send準(zhǔn)備好的那一剎那就可以發(fā)給對(duì)端,但是由于機(jī)制實(shí)現(xiàn)問(wèn)題,還是等待Recv算子過(guò)來(lái),才將數(shù)據(jù)拉取回去,通信過(guò)程耗時(shí)較長(zhǎng)。
- Rendezvous作為一個(gè)數(shù)據(jù)交換的熱點(diǎn),它內(nèi)部的邏輯開銷并不低。
針對(duì)上面提到的問(wèn)題,我們?cè)赗DMA上實(shí)現(xiàn)了另外一種數(shù)據(jù)交換的模式,叫做Send-Driven模式。與Recv-Driven模式相對(duì),顧名思義就是有Send算子直接將數(shù)據(jù)寫到Recv端,Recv端接收數(shù)據(jù)并注冊(cè)到本地Rendezvous中,Recv算子直接從本地的Rendezvous中獲取數(shù)據(jù)。具體流程如下圖13所示: 圖13 原生的Recv-Driven與補(bǔ)充的Send-Driven機(jī)制從圖中可以看到,相較于Recv-Driven模式,Send-Driven模式的通信流程得到了比較大的簡(jiǎn)化,另外在數(shù)據(jù)ready后立即發(fā)送的特性,跳過(guò)了一側(cè)的Rendezvous,并且對(duì)于生產(chǎn)者先于消費(fèi)者的情況,可以加快消費(fèi)端數(shù)據(jù)獲取的速度。3.4 延遲優(yōu)化這部分優(yōu)化,也是分布式計(jì)算的經(jīng)典優(yōu)化方向。整個(gè)流程鏈路上那些可以精簡(jiǎn)、合并、重疊需要不斷去挖掘。對(duì)于機(jī)器學(xué)習(xí)系統(tǒng)來(lái)說(shuō),相比其它的系統(tǒng),還可以用一些近似的算法來(lái)做這部分工作,從而獲得較大的性能提升。下面介紹我們?cè)趦蓚€(gè)這方面做的一些優(yōu)化實(shí)踐。3.4.1 稀疏域參數(shù)聚合在啟用HashTable存儲(chǔ)稀疏參數(shù)后,對(duì)應(yīng)的,一些配套參數(shù)也需要替換為HashTable實(shí)現(xiàn),這樣整個(gè)計(jì)算圖中會(huì)出現(xiàn)多張HashTable以及大量的相關(guān)算子。在實(shí)踐中,我們發(fā)現(xiàn)需要盡量降低Lookup/Insert等算子的個(gè)數(shù),一方面降低PS的負(fù)載,一方面降低RPC QPS。因此,針對(duì)稀疏模型的常見(jiàn)用法,我們進(jìn)行了相關(guān)的聚合工作。以Adam優(yōu)化器為例,需要?jiǎng)?chuàng)建兩個(gè)slot,以保存優(yōu)化中的動(dòng)量信息,它的Shape與Embedding相同。在原生優(yōu)化器中,這兩個(gè)Variable是單獨(dú)創(chuàng)建的,并在反向梯度更新的時(shí)候會(huì)去讀寫。同理,使用HashTable方案時(shí),我們需要同時(shí)創(chuàng)建兩張單獨(dú)的HashTable用來(lái)訓(xùn)練m、v參數(shù)。那么在前向,反向中需要分別對(duì)Embedding、 m、v進(jìn)行一次Lookup和一次Insert,總共需要三次Lookup和三次Insert。這里一個(gè)優(yōu)化點(diǎn)就是將Embedding、 m、v,以及低頻過(guò)濾的計(jì)數(shù)器(見(jiàn)下圖14的Counting HashTable)聚合到一起,作為HashTable的Value,這樣對(duì)稀疏參數(shù)的相關(guān)操作就可以聚合執(zhí)行,大大減少了稀疏參數(shù)操作頻次,降低了PS的壓力。 圖14 基于HashTable的參數(shù)融合策略該特性屬于一個(gè)普適型優(yōu)化,開啟聚合功能后,訓(xùn)練速度有了顯著的提高,性能提升幅度隨著模型和Worker規(guī)模的變化,效果總是正向的。在美團(tuán)內(nèi)部真實(shí)業(yè)務(wù)模型上,聚合之后性能相比非聚合方式能提升了45%左右。3.4.2 Embedding流水線優(yōu)化流水線,在工業(yè)生產(chǎn)中,指每一個(gè)生產(chǎn)單位只專注處理某個(gè)片段的工作,以提高工作效率及產(chǎn)量的一種生產(chǎn)方式。在計(jì)算機(jī)領(lǐng)域內(nèi),更為大家熟知的是,流水線代表一種多任務(wù)之間Overlap執(zhí)行的并行化技術(shù)。例如在典型的RISC處理器中,用戶的程序由大量指令構(gòu)成,而一條指令的執(zhí)行又可以大致分為:取指、譯碼、執(zhí)行、訪存、寫回等環(huán)節(jié)。這些環(huán)節(jié)會(huì)利用到指令Cache、數(shù)據(jù)Cache、寄存器、ALU等多種不同的硬件單元,在每一個(gè)指令周期內(nèi),這5個(gè)環(huán)節(jié)的硬件單元會(huì)并行執(zhí)行,得以更加充分的利用硬件能力,以此提高整個(gè)處理器的指令吞吐性能。處理器的指令流水線是一套復(fù)雜而系統(tǒng)的底層技術(shù),但其中的思想在分布式深度學(xué)習(xí)框架中也被大量的使用,例如:- 如果將分布式訓(xùn)練簡(jiǎn)單的抽象為計(jì)算和通信兩個(gè)過(guò)程,絕大多數(shù)主流的深度學(xué)習(xí)框架都支持在執(zhí)行計(jì)算圖DAG時(shí),通信和計(jì)算的Overlap。
- 如果將深度模型訓(xùn)練簡(jiǎn)單的分為前向和反向,在單步內(nèi),由于兩者的強(qiáng)依賴性,無(wú)法做到有效并行,字節(jié)BytePS[8]中引入的通信調(diào)度打破了step iteration間的屏障,上一輪的部分參數(shù)更新完畢后,即可提前開始下輪的前向計(jì)算,增強(qiáng)了整體視角下前反向的Overlap。
- 百度AIBox[9]為了解決CTR場(chǎng)景GPU訓(xùn)練時(shí),參數(shù)位于主存,但計(jì)算位于GPU的問(wèn)題,巧妙調(diào)度不同硬件設(shè)備,搭建起了主要利用CPU/主存/網(wǎng)卡的參數(shù)預(yù)準(zhǔn)備階段和主要利用GPU/NVLink的網(wǎng)絡(luò)計(jì)算階段,通過(guò)兩個(gè)階段的Overlap達(dá)到更高的訓(xùn)練吞吐。
我們看到,在深度學(xué)習(xí)框架設(shè)計(jì)上,通過(guò)分析場(chǎng)景,可以從不同的視角發(fā)掘可并行的階段,來(lái)提高整體的訓(xùn)練吞吐。對(duì)于大規(guī)模稀疏模型訓(xùn)練時(shí),核心模型流程是:先執(zhí)行稀疏參數(shù)的Embedding,然后執(zhí)行稠密部分子網(wǎng)絡(luò)。其中稀疏參數(shù)Embedding在遠(yuǎn)端PS上執(zhí)行,主要耗費(fèi)網(wǎng)絡(luò)資源,而稠密部分子網(wǎng)絡(luò)在本地Worker執(zhí)行,主要耗費(fèi)計(jì)算資源。這兩部分占了整個(gè)流程的大部分時(shí)間,在美團(tuán)某實(shí)際業(yè)務(wù)模型上分別耗時(shí)占比:40%+、50%+。那我們是否可以提前執(zhí)行稀疏參數(shù)的Embedding,來(lái)做到通信和計(jì)算的Overlap,隱藏掉這部分時(shí)間呢?從系統(tǒng)實(shí)現(xiàn)上肯定是可行的,但從算法上講,這樣做會(huì)引入?yún)?shù)Staleness的問(wèn)題,可能會(huì)導(dǎo)致模型精度受到影響。但在實(shí)際的生產(chǎn)場(chǎng)景中,大規(guī)模異步訓(xùn)練時(shí)本身就會(huì)帶來(lái)幾十到幾百個(gè)步的滯后性問(wèn)題。經(jīng)過(guò)我們測(cè)試,提前獲取一兩步的稀疏參數(shù),模型精度并未受到影響。在具體實(shí)現(xiàn)上,我們把整個(gè)計(jì)算圖拆分為Embedding Graph(EG)和Main Graph(MG)兩張子圖,兩者異步獨(dú)立執(zhí)行,做到拆分流程的Overlap(整個(gè)拆分過(guò)程,可以做到對(duì)用戶透明)。EG主要覆蓋從樣本中抽取Embedding Key,查詢組裝Embedding向量,Embedding向量更新等環(huán)節(jié)。MG主要包含稠密部分子網(wǎng)絡(luò)計(jì)算、梯度計(jì)算、稠密參數(shù)部分更新等環(huán)節(jié)。 圖15 Embedding流水線模塊交互關(guān)系兩張子圖的交互關(guān)系為:EG向MG傳遞Embeding向量(從MG的視角看,是從一個(gè)稠密Variable讀取數(shù)值);MG向EG傳遞Embedding參數(shù)對(duì)應(yīng)的梯度。上述兩個(gè)過(guò)程的表達(dá)都是TensorFlow的計(jì)算圖,我們利用兩個(gè)線程,兩個(gè)Session并發(fā)的執(zhí)行兩張計(jì)算圖,使得兩個(gè)階段Overlap起來(lái),以此到達(dá)了更大的訓(xùn)練吞吐。 圖16 Embedding流水線架構(gòu)流程圖上圖是Embedding流水線的架構(gòu)流程圖。直觀來(lái)看分為左側(cè)的樣本分發(fā)模塊,頂部的跨Session數(shù)據(jù)交換模塊,以及自動(dòng)圖切分得到的Embedding Graph和Main Graph,藍(lán)色的圓圈代表新增算子,橙色箭頭代表EG重點(diǎn)流程,藍(lán)色箭頭代表MG重點(diǎn)流程,紅色箭頭代表樣本數(shù)據(jù)重點(diǎn)流程。- 以對(duì)用戶透明的形式引入了一層名為Pipeline Dataset的抽象層,這一層的產(chǎn)生是為了滿足EG/MG兩張計(jì)算圖以不同節(jié)奏運(yùn)行的需求,支持自定義配置。另外,為了使得整個(gè)流水線中的數(shù)據(jù)做到彼此的配套,這里還會(huì)負(fù)責(zé)進(jìn)行一個(gè)全局Batch ID的生成及注冊(cè)工作。Pipeline Dataset對(duì)外暴露兩種Iterator,一個(gè)供EG使用,一個(gè)供MG使用。Pipeline Dataset底部共享TensorFlow原生的各層Dataset。
- 頂部的ExchangeManager是一個(gè)靜態(tài)的,跨Session的數(shù)據(jù)交換媒介,對(duì)外暴露數(shù)據(jù)注冊(cè)和數(shù)據(jù)拉取的能力。抽象這個(gè)模塊的原因是,EG和MG原本歸屬于一張計(jì)算圖,因?yàn)榱魉€的原因拆解為拆為兩張圖,這樣我們需要建立一種跨Session的數(shù)據(jù)交換機(jī)制,并準(zhǔn)確進(jìn)行配套。它內(nèi)部以全局Batch ID做Key,后面管理了樣本數(shù)據(jù)、Embeding向量、Embedding梯度、Unique后的Index等數(shù)據(jù),并負(fù)責(zé)這些數(shù)據(jù)的生命周期管理。
- 中間的Embedding Graph由獨(dú)立的TF Session運(yùn)行于一個(gè)獨(dú)立的線程中,通過(guò)a算子獲得樣本數(shù)據(jù)后,進(jìn)行特征ID的抽取等動(dòng)作,并進(jìn)行基于HashTable方法的稀疏參數(shù)查詢,查詢結(jié)果通過(guò)c算子放置到ExchangeManager中。EG中還包含用于反向更新的f算子,它會(huì)從ExchangeManager中獲取Embedding梯度和與其配套的前向參數(shù),然后執(zhí)行梯度更新參數(shù)邏輯。
- 下面的Main Graph負(fù)責(zé)實(shí)際稠密子網(wǎng)絡(luò)的計(jì)算,我們繼承并實(shí)現(xiàn)一種可訓(xùn)練的EmbeddingVariable,它的構(gòu)建過(guò)程(d算子)會(huì)從ExchangeManager查找與自己配套的Embedding向量封裝成EmbeddingVariable,給稠密子網(wǎng)絡(luò)。此外,在EmbeddingVariable注冊(cè)的反向方法中,我們添加了e算子使得Embedding梯度得以添加到ExchangeManager中,供EG中的f算子消費(fèi)。
通過(guò)上面的設(shè)計(jì),我們就搭建起了一套可控的EG/MG并發(fā)流水線訓(xùn)練模式??傮w來(lái)看,Embedding流水線訓(xùn)練模式的收益來(lái)源有:- 經(jīng)過(guò)我們對(duì)多個(gè)業(yè)務(wù)模型的Profiling分析發(fā)現(xiàn),EG和MG在時(shí)間的比例上在3:7或4:6的左右,通過(guò)將這兩個(gè)階段并行起來(lái),可以有效的隱藏Embedding階段,使得MG網(wǎng)絡(luò)計(jì)算部分幾乎總是可以立即開始,大大加速了整體模型的訓(xùn)練吞吐。
- TensorFlow引擎中當(dāng)使用多個(gè)優(yōu)化器(稀疏與非稀疏)的時(shí)候,會(huì)出現(xiàn)重復(fù)構(gòu)建反向計(jì)算圖的問(wèn)題,一定程度增加了額外計(jì)算,通過(guò)兩張子圖的拆分,恰好避免了這個(gè)問(wèn)題。
- 在實(shí)施過(guò)程中的ExchangeManager不僅負(fù)責(zé)了Embedding參數(shù)和梯度的交換,還承擔(dān)了元數(shù)據(jù)復(fù)用管理的職責(zé)。例如Unique等算子的結(jié)果保存,進(jìn)一步降低了重復(fù)計(jì)算。
另外,在API設(shè)計(jì)上,我們做到了對(duì)用戶透明,僅需一行代碼即可開啟Embedding流水線功能,對(duì)用戶隱藏了EG/MG的切割過(guò)程。目前,在美團(tuán)某業(yè)務(wù)訓(xùn)練中,Embedding流水線功能在CPU PS架構(gòu)下可以帶來(lái)20%~60%的性能提升(而且Worker并發(fā)規(guī)模越大,性能越好)。3.5 單實(shí)例PS并發(fā)優(yōu)化經(jīng)過(guò)2.2章節(jié)的分析可知,我們不能通過(guò)持續(xù)擴(kuò)PS來(lái)提升分布式任務(wù)的吞吐,單實(shí)例PS的并發(fā)優(yōu)化,也是非常重要的優(yōu)化方向。我們主要的優(yōu)化工作如下。3.5.1 高性能的HashTablePS架構(gòu)下,大規(guī)模稀疏模型訓(xùn)練對(duì)于HashTable的并發(fā)讀寫要求很高,因?yàn)槊總€(gè)PS都要承擔(dān)成百乃至上千個(gè)Worker的Embedding壓力,這里我們綜合速度和穩(wěn)定性考慮,選用了tbb::concurrent_hash_map[10]作為底層HashTable表實(shí)現(xiàn),并將其包裝成一個(gè)新的TBBConcurrentHashTable算子。經(jīng)過(guò)測(cè)試,在千億規(guī)模下TBBConcurrentHashTable比原生MutableDenseHashTable訓(xùn)練速度上快了3倍。3.5.2 HashTable BucketPool對(duì)于大規(guī)模稀疏模型訓(xùn)練來(lái)說(shuō),Embedding HashTable會(huì)面對(duì)大量的并發(fā)操作,通過(guò)Profiling我們發(fā)現(xiàn),頻繁動(dòng)態(tài)的內(nèi)存申請(qǐng)會(huì)帶來(lái)了較大性能開銷(即使TensorFlow的Tensor有專門的內(nèi)存分配器)。我們基于內(nèi)存池化的思路優(yōu)化了HashTable的內(nèi)存管理。我們?cè)贖ashTable初始化時(shí),會(huì)先為Key和Value分別創(chuàng)造兩個(gè)BucketPool,每個(gè)池子都會(huì)先Malloc較大一塊內(nèi)存?zhèn)溆茫紤]到可能會(huì)有對(duì)HashTable進(jìn)行中的Key和Value進(jìn)行Remove的場(chǎng)景(如Online Learning訓(xùn)練時(shí)),需要對(duì)從HashTable中刪除的Key和Value所使用的內(nèi)存進(jìn)行回收,因此每個(gè)BucketPool還有一個(gè)ReuseQueue來(lái)負(fù)責(zé)維護(hù)回收的內(nèi)存。每次向內(nèi)部的哈希表數(shù)據(jù)結(jié)構(gòu)中Insert Key和Value的時(shí)候,Key和Value內(nèi)存和釋放分配都進(jìn)行池化管理。用這種方式降低了大規(guī)模稀疏訓(xùn)練中遇到稀疏內(nèi)存分配開銷,整體端到端訓(xùn)練性能提升了5%左右。 圖17 HashTable內(nèi)存優(yōu)化3.6 單位算力吞吐優(yōu)化經(jīng)過(guò)2.2章節(jié)的分析,Worker的計(jì)算壓力也非常大,如果不優(yōu)化Worker,同時(shí)要保持吞吐,需要橫向擴(kuò)展更多的Worker,給PS帶來(lái)更大的壓力。而對(duì)于用戶來(lái)說(shuō),如果能在有限的計(jì)算資源下帶來(lái)性能提升,對(duì)業(yè)務(wù)價(jià)值更高。我們通過(guò)CAT統(tǒng)計(jì)出了一些高頻算子,并進(jìn)行了專項(xiàng)優(yōu)化。這里選取Unique&DynamicPartition算子融合案例進(jìn)行分享。在TensorFlow PS架構(gòu)中,包括Embedding向量在內(nèi)的共享參數(shù)都存儲(chǔ)在PS上,并通過(guò)網(wǎng)絡(luò)與Worker交互,在進(jìn)行Embedding查詢過(guò)程中,往往會(huì)涉及如下兩個(gè)環(huán)節(jié):- 由于稀疏參數(shù)的性質(zhì),從樣本中抽取得到的待查詢Embedding ID,它的重復(fù)率往往高達(dá)70%~90%,如果不進(jìn)行去重查詢,不論是對(duì)HashTable的查詢還是網(wǎng)絡(luò)的傳輸,都會(huì)帶來(lái)不小的壓力。因此,通常會(huì)在查詢前進(jìn)行Unique操作。
- 在大規(guī)模稀疏場(chǎng)景中,為了存儲(chǔ)千億規(guī)模的參數(shù),會(huì)有多個(gè)PS機(jī)器共同承載。而Worker端會(huì)負(fù)責(zé)對(duì)查詢請(qǐng)求按照設(shè)定的路由規(guī)則進(jìn)行切分,這里通常會(huì)在查詢前進(jìn)行DynamicPartition動(dòng)作。
通常這兩個(gè)過(guò)程會(huì)利用TensorFlow既有的算子進(jìn)行搭建,但在實(shí)際使用中,我們發(fā)現(xiàn)它并不是很高效,主要問(wèn)題在于:- Unique算子原生實(shí)現(xiàn),它內(nèi)部使用的內(nèi)存分配策略較為低效。使用了兩倍輸入?yún)?shù)(Embedding ID)的大小進(jìn)行內(nèi)存分配,但由于輸入?yún)?shù)較大,而且重復(fù)率高,導(dǎo)致HashTable創(chuàng)建過(guò)大且非常稀疏。幾乎每次插入都會(huì)產(chǎn)生一次minor_page_fault,導(dǎo)致HashTable性能下降。我們使用Intel Vtune驗(yàn)證了這一點(diǎn)(參見(jiàn)圖18)。
- Unique和Dynamic Partition算子存在冗余數(shù)據(jù)遍歷,這些操作其實(shí)可以在一次數(shù)據(jù)遍歷中全部做完,節(jié)省掉算子切換、冗余數(shù)據(jù)遍歷的耗時(shí)。
 圖18 Unique算子內(nèi)部出現(xiàn)DRAM Bound問(wèn)題總結(jié)來(lái)說(shuō),HashTable開辟過(guò)大會(huì)導(dǎo)致大量的minor_page_fault,導(dǎo)致訪存的時(shí)間增加,HashTable過(guò)小又可能會(huì)導(dǎo)致擴(kuò)容。我們采用了基于啟發(fā)式算法的內(nèi)存自適應(yīng)Unique算子實(shí)現(xiàn),通過(guò)對(duì)訓(xùn)練歷史重復(fù)率的統(tǒng)計(jì),我們可以得到一個(gè)相對(duì)合理的HashTable大小,來(lái)提高訪存的性能;另外Unique算子內(nèi)HashTable的具體選擇上,經(jīng)過(guò)我們的多種測(cè)試,選擇了Robin HashTable替換了原生TF中的實(shí)現(xiàn)。進(jìn)一步,我們對(duì)圍繞Embedding ID的Unique和Partition環(huán)節(jié)進(jìn)行了算子合并,簡(jiǎn)化了邏輯實(shí)現(xiàn)。經(jīng)過(guò)上述的優(yōu)化,Unique單算子可以取得51%的加速,在真實(shí)模型端到端上可以獲得10%左右的性能提升,算子總數(shù)量降低了4%。在整個(gè)關(guān)鍵算子優(yōu)化的過(guò)程中,Intel公司的林立凡、張向澤、高明進(jìn)行大量的技術(shù)支持,我們也復(fù)用了他們的部分優(yōu)化工作,在此深表感謝!4 大規(guī)模稀疏算法建模大規(guī)模稀疏能力在業(yè)務(wù)落地的過(guò)程中,算法層面還需要從特征和模型結(jié)構(gòu)上進(jìn)行對(duì)應(yīng)升級(jí),才能拿到非常好的效果。其中外賣廣告從業(yè)務(wù)特點(diǎn)出發(fā),引入大規(guī)模稀疏特征完成外賣場(chǎng)景下特征體系的升級(jí),提供了更高維的特征空間和參數(shù)空間,增強(qiáng)了模型的擬合能力。重新設(shè)計(jì)了面向高維稀疏場(chǎng)景的特征編碼方案,解決了特征編碼過(guò)程中的特征沖突問(wèn)題,同時(shí)編碼過(guò)程去掉了部分冗余的特征哈希操作,一定程度上簡(jiǎn)化了特征處理邏輯,并降低了特征計(jì)算的耗時(shí)。在系統(tǒng)層面,面對(duì)百億參數(shù)、百億樣本以上量級(jí)的大規(guī)模稀疏模型的訓(xùn)練,會(huì)帶來(lái)訓(xùn)練迭代效率的大大降低,單次實(shí)驗(yàn)從一天以內(nèi),增長(zhǎng)到一周左右。美團(tuán)機(jī)器學(xué)習(xí)平臺(tái)訓(xùn)練引擎團(tuán)隊(duì),除了上述TensorFlow框架層面的優(yōu)化、還針對(duì)業(yè)務(wù)模型進(jìn)行了專項(xiàng)優(yōu)化,整體吞吐優(yōu)化了8到10倍(如果投入更多計(jì)算資源,可以進(jìn)一步加速),大大提升業(yè)務(wù)的迭代效率,助力外賣廣告業(yè)務(wù)取得了較為明顯的提升。5 總結(jié)與展望TensorFlow在大規(guī)模推薦系統(tǒng)中被廣泛使用,但由于缺乏大規(guī)模稀疏的大規(guī)模分布式訓(xùn)練能力,阻礙了業(yè)務(wù)的發(fā)展。美團(tuán)基于TensorFlow原生架構(gòu),支持了大規(guī)模稀疏能力,并從多個(gè)角度進(jìn)行了深度優(yōu)化,做到千億參數(shù)、千億樣本高效的分布式訓(xùn)練,并在美團(tuán)內(nèi)部進(jìn)行了大規(guī)模的使用。對(duì)于這類關(guān)鍵能力的缺失,TensorFlow社區(qū)也引起了共鳴,社區(qū)官方在2020年創(chuàng)建了SIG Recommenders[11],通過(guò)社區(qū)共建的方式來(lái)解決此類問(wèn)題,美團(tuán)后續(xù)也會(huì)積極的參與到社區(qū)的貢獻(xiàn)當(dāng)中去。美團(tuán)推薦系統(tǒng)場(chǎng)景的模型訓(xùn)練,目前主要運(yùn)行在CPU上,但隨著業(yè)務(wù)的發(fā)展,有些模型變得越來(lái)越復(fù)雜,CPU上已經(jīng)很難有優(yōu)化空間(優(yōu)化后的Worker CPU使用率在90%以上)。而近幾年,GPU的計(jì)算能力突飛猛進(jìn),新一代的NVIDIA A100 GPU,算力達(dá)到了156TFLOPS(TF32 Tensor Cores)、80G顯存、卡間帶寬600GB/s。對(duì)于這類復(fù)雜模型的Workload,我們基于A100 GPU架構(gòu),設(shè)計(jì)了下一代的分布式訓(xùn)練架構(gòu),經(jīng)過(guò)初步優(yōu)化,在美團(tuán)某大流量業(yè)務(wù)推薦模型上也拿到了較好的效果,目前還在進(jìn)一步優(yōu)化當(dāng)中,后續(xù)我們會(huì)進(jìn)行分享,敬請(qǐng)期待。6 作者簡(jiǎn)介逸帆、家恒、崢少、鵬鵬、永宇、正陽(yáng)、黃軍等,來(lái)自美團(tuán)基礎(chǔ)研發(fā)平臺(tái),機(jī)器學(xué)習(xí)平臺(tái)訓(xùn)練引擎組,主要負(fù)責(zé)美團(tuán)分布式機(jī)器學(xué)習(xí)訓(xùn)練系統(tǒng)的性能優(yōu)化與能力建設(shè)。 海濤,來(lái)自美團(tuán)外賣廣告策略團(tuán)隊(duì),主要負(fù)責(zé)美團(tuán)外賣廣告業(yè)務(wù)的算法探索和策略落地工作。7 參考文獻(xiàn)
|