作者:劉童璇煙秋

經(jīng)歷6年時(shí)間,在各團(tuán)隊(duì)的努力下,阿里巴巴集團(tuán)大規(guī)模稀疏模型訓(xùn)練/預(yù)測(cè)引擎DeepRec正式對(duì)外開源,助力開發(fā)者提升稀疏模型訓(xùn)練性能和效果。

DeepRec是什么

DeepRec(PAI-TF)是阿里巴巴集團(tuán)統(tǒng)一的大規(guī)模稀疏模型訓(xùn)練/預(yù)測(cè)引擎,廣泛應(yīng)用于淘寶、天貓、阿里媽媽、高德、淘特、AliExpress、Lazada等,支持了淘寶搜索、推薦、廣告等核心業(yè)務(wù),支撐著千億特征、萬億樣本的超大規(guī)模稀疏訓(xùn)練。

DeepRec在分布式、圖優(yōu)化、算子、Runtime等方面對(duì)稀疏模型進(jìn)行了深度性能優(yōu)化,同時(shí)提供了稀疏場(chǎng)景下特有的Embedding相關(guān)功能。

DeepRec項(xiàng)目從2016年開發(fā)至今,由阿里巴巴集團(tuán)內(nèi)AOP團(tuán)隊(duì)、XDL團(tuán)隊(duì)、PAI團(tuán)隊(duì)、RTP團(tuán)隊(duì)以及螞蟻集團(tuán)AIInfra團(tuán)隊(duì)共建,并且得到了淘寶推薦算法等多個(gè)業(yè)務(wù)算法團(tuán)隊(duì)的支持。DeepRec的研發(fā)也得到了Intel CESG軟件團(tuán)隊(duì)、Optane團(tuán)隊(duì)和PSU團(tuán)隊(duì),NVIDIA GPU計(jì)算專家團(tuán)隊(duì)及Merlin HughCTR團(tuán)隊(duì)的支持。

DeepRec架構(gòu)設(shè)計(jì)原則

在TensorFlow引擎上支持大規(guī)模稀疏特征,業(yè)界有多種實(shí)現(xiàn)方式,其中最常見的方式是借鑒了ParameterServer的架構(gòu)實(shí)現(xiàn),在TensorFlow之外獨(dú)立實(shí)現(xiàn)了一套ParameterServer和相關(guān)的優(yōu)化器,同時(shí)在TensorFlow內(nèi)部通過bridge的方式橋接了兩個(gè)模塊。這個(gè)做法有一定的好處,比如PS的實(shí)現(xiàn)會(huì)比較靈活,但也存在一些局限性。

DeepRec采取了另一種架構(gòu)設(shè)計(jì)方式,遵循“視整個(gè)訓(xùn)練引擎為一個(gè)系統(tǒng)整體”的架構(gòu)設(shè)計(jì)原則。TensorFlow是一個(gè)基于Graph的靜態(tài)圖訓(xùn)練引擎,在其架構(gòu)上有相應(yīng)的分層,比如最上層的API層、中間的圖優(yōu)化層和最下層的算子層。TensorFlow通過這三層的設(shè)計(jì)去支撐上層不同Workload的業(yè)務(wù)需求和性能優(yōu)化需求。

DeepRec也堅(jiān)持了這一設(shè)計(jì)原則,基于存儲(chǔ)/計(jì)算解耦的設(shè)計(jì)原則在Graph層面引入EmbeddingVariable功能;基于Graph的特點(diǎn)實(shí)現(xiàn)了通信的算子融合。通過這樣的設(shè)計(jì)原則,DeepRec可以支持用戶在單機(jī)、分布式場(chǎng)景下使用同一個(gè)優(yōu)化器的實(shí)現(xiàn)和同一套EmbeddingVariable的實(shí)現(xiàn);同時(shí)在Graph層面引入多種優(yōu)化能力,從而做到獨(dú)立模塊設(shè)計(jì)所做不到的聯(lián)合優(yōu)化設(shè)計(jì)。

DeepRec的優(yōu)勢(shì)

DeepRec是基于TensorFlow1.15、Intel-TF、NV-TF構(gòu)建的稀疏模型訓(xùn)練/預(yù)測(cè)引擎,針對(duì)稀疏模型場(chǎng)景進(jìn)行了定制深度優(yōu)化,主要包含以下三類功能優(yōu)化:

模型效果

DeepRec提供了豐富的稀疏功能支持,提高模型效果的同時(shí)降低稀疏模型的大小,并且優(yōu)化超大規(guī)模下Optimizer的效果。下面簡(jiǎn)單介紹Embedding及Optimizer幾個(gè)有特色的工作:

? EmbeddingVariable(動(dòng)態(tài)彈性特征):

1)解決了靜態(tài)Shape Variable的vocabulary_size難以預(yù)估、特征沖突、內(nèi)存及IO冗余等問題,并且在DeepRec中提供了豐富的EmbeddingVariable的進(jìn)階功能,包括不同的特征準(zhǔn)入方式、支持不同的特征淘汰策略等,能夠明顯提高稀疏模型的效果。

2)在訪問效率上,為了達(dá)到更優(yōu)化的性能和更低的內(nèi)存占用,EmbeddingVariable的底層HashTable實(shí)現(xiàn)了無鎖化設(shè)計(jì),并且進(jìn)行了精細(xì)的內(nèi)存布局優(yōu)化,優(yōu)化了HashTable的訪問頻次,使得在訓(xùn)練過程中前后向只需訪問一次HashTable。

? DynamicDimensionEmbeddingVariable(動(dòng)態(tài)彈性維度):

在典型的稀疏場(chǎng)景中,同類特征的出現(xiàn)頻次往往極度不均勻。通常情況下,同一個(gè)特征列的特征都被設(shè)置成統(tǒng)一維度,如果Embedding維度過高,低頻特征容易過擬合,而且會(huì)額外耗費(fèi)大量?jī)?nèi)存;如果維度設(shè)置過低,高頻部征特征可能會(huì)由于表達(dá)不夠而影響效果。

Dynamic Dimension Embedding Variable提供了同一特征列的不同特征值,根據(jù)特征的冷熱自動(dòng)配置不同的特征維度,高頻特征可以配置更高維度增強(qiáng)表達(dá)能力,而低頻特征因?yàn)榻o定低維度embedding緩解了過擬合的問題,而且可以極大程度節(jié)省內(nèi)存(低頻長(zhǎng)尾特征的數(shù)量占據(jù)絕對(duì)優(yōu)勢(shì))。

? Adaptive Embedding(自適應(yīng)Embedding):

當(dāng)使用動(dòng)態(tài)彈性特征功能時(shí),低頻特征存在過擬合問題。EmbeddingVariable中所有的特征都從initializer設(shè)定的初始值(一般設(shè)為0)開始學(xué)起,對(duì)于一些出現(xiàn)頻次從低到高的特征,也需要逐漸學(xué)習(xí)到一個(gè)較好的狀態(tài),不能共享別的特征的學(xué)習(xí)結(jié)果。AdaptiveEmbedding功能使用靜態(tài)Shape Variable和動(dòng)態(tài)EmbeddingVariable共同存儲(chǔ)稀疏特征,對(duì)于新加入的特征存于有沖突的Variable,對(duì)于出現(xiàn)頻率較高的特征存于無沖突的EmbeddingVariable,特征遷移到EmbeddingVaraible可以復(fù)用在有沖突的靜態(tài)Shape Variable的學(xué)習(xí)結(jié)果。

? Adagrad Decay Optimizer:

為支持超大規(guī)模訓(xùn)練而提出的一種改進(jìn)版Adagrad優(yōu)化器。當(dāng)模型訓(xùn)練的樣本量大,同時(shí)持續(xù)增量訓(xùn)練較長(zhǎng)時(shí)間時(shí),Adagrad優(yōu)化器的梯度會(huì)趨近于0,導(dǎo)致新增訓(xùn)練的數(shù)據(jù)無法對(duì)模型產(chǎn)生影響。已有的累積打折的方案雖然可以解決梯度趨近0的問題,但也會(huì)帶來模型效果變差的問題(通過iteration打折策略無法反映實(shí)際的業(yè)務(wù)場(chǎng)景特點(diǎn))。Adagrad Decay Optimizer基于周期打折的策略,同一個(gè)周期內(nèi)的樣本相同的打折力度,兼顧數(shù)據(jù)的無限累積和樣本順序?qū)δＰ偷挠绊憽?/p>

此外,DeepRec還提供Multi-HashEmbedding、AdamAsyncOptimizer等功能,在內(nèi)存占用、性能、模型效果等方面為業(yè)務(wù)帶來實(shí)際的幫助。

訓(xùn)練性能

DeepRec針對(duì)稀疏模型場(chǎng)景在分布式、圖優(yōu)化、算子、Runtime等方面進(jìn)行了深度性能優(yōu)化。其中,DeepRec對(duì)不同的分布式策略進(jìn)行了深度的優(yōu)化,包括異步訓(xùn)練、同步訓(xùn)練、半同步訓(xùn)練等,其中GPU同步訓(xùn)練支持HybridBackend以及NVIDIA HugeCTR-SOK。DeepRec提供了豐富的針對(duì)稀疏模型訓(xùn)練的圖優(yōu)化功能,包括自動(dòng)流水線SmartStage、結(jié)構(gòu)化特征、自動(dòng)圖Fusion等等。DeepRec中優(yōu)化了稀疏模型中數(shù)十個(gè)常見算子,并且提供了包括Embedding、Attention等通用子圖的Fusion算子。DeepRec中CPUAllocator和GPUAllocator能夠大大降低內(nèi)存/顯存的使用量并顯著加速E2E的訓(xùn)練性能。在線程調(diào)度、執(zhí)行引擎方面針對(duì)不同的場(chǎng)景提供了不同的調(diào)度引擎策略。下面簡(jiǎn)單介紹分布式、圖優(yōu)化、Runtime優(yōu)化方面幾個(gè)有特色的工作:

? StarServer(異步訓(xùn)練框架):

在超大規(guī)模任務(wù)場(chǎng)景下(幾百、上千worker),原生開源框架中的一些問題被暴露出來,譬如低效的線程池調(diào)度、關(guān)鍵路徑上的鎖開銷、低效的執(zhí)行引擎、頻繁的小包rpc帶來的開銷導(dǎo)致ParameterServer在分布式擴(kuò)展時(shí)成為明顯的性能瓶頸。StarServer進(jìn)行了包括圖、線程調(diào)度、執(zhí)行引擎以及內(nèi)存等優(yōu)化,將原有框架中的send/recv語義修改為pull/push語義,并且在子圖劃分上支持了該語義,同時(shí)實(shí)現(xiàn)了ParameterServer端圖執(zhí)行過程中的lockfree,實(shí)現(xiàn)了無鎖化的執(zhí)行,大大提高了并發(fā)執(zhí)行子圖的效率。對(duì)比原生框架,能夠提升數(shù)倍的訓(xùn)練性能,并且支持3000worker規(guī)模的線性分布式擴(kuò)展。

? SmartStage(自動(dòng)流水線):

稀疏模型訓(xùn)練通常包含樣本數(shù)據(jù)的讀取、Embedding查找、Attention/MLP計(jì)算等,樣本讀取和Embedding查找非計(jì)算密集操作,同時(shí)并不能高效利用計(jì)算資源(CPU、GPU)。原生框架中提供的dataset.prefetch接口可以異步化樣本讀取操作,但Embedding查找過程中涉及特征補(bǔ)齊、ID化等復(fù)雜的過程,這些過程無法通過prefetch進(jìn)行流水線化。SmartStage功能能夠自動(dòng)分析圖中異步流水線化的邊界并自動(dòng)插入,可以使并發(fā)流水線發(fā)揮最大的性能提升。

? PRMalloc(內(nèi)存分配器):

如何做到既高效又有效的使用內(nèi)存,對(duì)于稀疏模型的訓(xùn)練非常關(guān)鍵,稀疏場(chǎng)景模型訓(xùn)練中大塊內(nèi)存分配使用造成大量的minor pagefault,此外,多線程分配效率存在比較嚴(yán)重的并發(fā)分配效率問題。針對(duì)稀疏模型訓(xùn)練前向、后向,Graph計(jì)算模式的相對(duì)固定、多輪反復(fù)迭代的特點(diǎn),DeepRec設(shè)計(jì)了一套針對(duì)深度學(xué)習(xí)任務(wù)的內(nèi)存管理方案,提高內(nèi)存的使用效率和系統(tǒng)性能。使用DeepRec中提供的PRMalloc能夠極大降低訓(xùn)練過程中minor pagefault,提高多線程并發(fā)內(nèi)存分配、釋放的效率。

? PMEM allocator(持久內(nèi)存分配器):

基于PMDK的底層libpmem庫(kù)實(shí)現(xiàn)的PMEM allocator將從PMEM map出的一塊空間分為若干segment,每個(gè)segment又分成若干blocks,block是allocator的最小分配單元。分配block的線程為避免線程競(jìng)爭(zhēng),緩存一些可用空間,包括一組segment和free list?？捎每臻g中為每種record size(若干個(gè)block)維護(hù)一個(gè)free list和segment。各record size對(duì)應(yīng)的segment只分配該大小的PMEM空間,各record size對(duì)應(yīng)的free list中的所有指針均指向?qū)?yīng)record size的空閑空間。此外,為了均衡各thread cache的資源,由一個(gè)后臺(tái)線程周期地將thread cache中的free list移動(dòng)到后臺(tái)的pool中,pool中的資源由所有前臺(tái)線程共享。實(shí)驗(yàn)證明,基于持久內(nèi)存實(shí)現(xiàn)的內(nèi)存分配器在大模型的訓(xùn)練性能方面與基于DRAM的訓(xùn)練性能差別很小,但是TCO會(huì)有很大的優(yōu)勢(shì)。

部署及Serving

? 增量模型導(dǎo)出及加載:

時(shí)效性要求高的業(yè)務(wù),需要頻繁的線上模型更新,頻率往往達(dá)到分鐘級(jí)別甚至秒級(jí)。對(duì)于TB-10TB級(jí)別的超大模型而言,分鐘級(jí)別的模型生成到上線很難完成。此外,超大模型的訓(xùn)練和預(yù)測(cè)存在著資源浪費(fèi)、多節(jié)點(diǎn)Serving延時(shí)加大等問題。DeepRec提供了增量模型產(chǎn)出及加載能力,極大加速了超大模型生成和加載。

? Embedding多級(jí)混合存儲(chǔ):

稀疏模型中特征存在冷熱傾斜的特性,這產(chǎn)生了某些冷門特征很少被訪問和更新導(dǎo)致的內(nèi)存/顯存浪費(fèi)問題,以及超大模型內(nèi)存/顯存放不下的問題。DeepRec提供了多級(jí)混合存儲(chǔ)(支持最多四級(jí)的混合存儲(chǔ)HBM+DRAM+PMEM+SSD)的能力,自動(dòng)將冷門特征存放到廉價(jià)的存儲(chǔ)介質(zhì)中,將熱門特征存放到訪問更快、更貴的存儲(chǔ)介質(zhì)上,通過多級(jí)混合存儲(chǔ),使得單節(jié)點(diǎn)可以進(jìn)行TB-10TB模型的Training和Serving。

通過多級(jí)混合存儲(chǔ),能夠更大發(fā)揮GPU訓(xùn)練稀疏模型的能力,同時(shí)降低由于存儲(chǔ)資源限制造成的計(jì)算資源浪費(fèi),可以使用更少的機(jī)器進(jìn)行相近規(guī)模的模型訓(xùn)練,或者使用相同數(shù)量的機(jī)器進(jìn)行更大規(guī)模的訓(xùn)練。多級(jí)混合存儲(chǔ)也能使得單機(jī)進(jìn)行超大模型預(yù)測(cè)時(shí)避免分布式Serving帶來的latency增大問題,提高大模型的預(yù)測(cè)性能的同時(shí)降低成本。多級(jí)混合存儲(chǔ)功能也擁有自動(dòng)發(fā)現(xiàn)特征的訪問特性,基于高效的熱度統(tǒng)計(jì)策略,將熱度高的特征放置到快速的存儲(chǔ)介質(zhì)中,將低頻的特征offload到低速存儲(chǔ)介質(zhì)中,再通過異步方式驅(qū)動(dòng)特征在多個(gè)介質(zhì)之間移動(dòng)。

為什么開源DeepRec

開源深度學(xué)習(xí)框架都不能很好地支持稀疏場(chǎng)景中對(duì)于稀疏Embedding功能的需求、模型訓(xùn)練性能需求、部署迭代和線上服務(wù)的需求。DeepRec經(jīng)過阿里巴巴集團(tuán)搜索、推薦、廣告等核心業(yè)務(wù)場(chǎng)景及公有云上各種業(yè)務(wù)場(chǎng)景的打磨,能夠支持不同類型的稀疏場(chǎng)景訓(xùn)練效果和性能需求。

阿里巴巴希望通過建立開源社區(qū),和外部開發(fā)者開展廣泛合作,進(jìn)一步推動(dòng)稀疏

模型訓(xùn)練/預(yù)測(cè)框架的發(fā)展,為不同業(yè)務(wù)場(chǎng)景中的搜推廣模型訓(xùn)練和預(yù)測(cè)帶來業(yè)務(wù)效果和性能提升。

今天DeepRec的開源只是我們邁出的一小步。我們非常期待得到您的反饋。最后,

如果你對(duì)DeepRec有相應(yīng)的興趣,你也可以來轉(zhuǎn)轉(zhuǎn),為我們的框架貢獻(xiàn)一點(diǎn)你的

代碼和意見,這將是我們莫大的榮幸。

開源地址:https://github.com/alibaba/DeepRec

釘釘答疑交流群:34375411

想了解更多AI開源項(xiàng)目,請(qǐng)點(diǎn)擊:

https://www.aliyun.com/activity/bigdata/opensource_bigdata__ai