阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec

阿里雲大資料AI技術發表於2022-03-22

作者:煙秋

經歷6年時間,在各團隊的努力下,阿里巴巴集團大規模稀疏模型訓練/預測引擎DeepRec正式對外開源,助力開發者提升稀疏模型訓練效能和效果。

DeepRec是什麼

DeepRec(PAI-TF)是阿里巴巴集團統一的大規模稀疏模型訓練/預測引擎,廣泛應用於淘寶、天貓、阿里媽媽、高德、淘特、AliExpress、Lazada等,支援了淘寶搜尋、推薦、廣告等核心業務,支撐著千億特徵、萬億樣本的超大規模稀疏訓練。

DeepRec在分散式、圖最佳化、運算元、Runtime等方面對稀疏模型進行了深度效能最佳化,同時提供了稀疏場景下特有的Embedding相關功能。

DeepRec專案從2016年開發至今,由阿里巴巴集團內AOP團隊、XDL團隊、PAI團隊、RTP團隊以及螞蟻集團AIInfra團隊共建,並且得到了淘寶推薦演算法等多個業務演算法團隊的支援。DeepRec的研發也得到了Intel CESG軟體團隊、Optane團隊和PSU團隊,NVIDIA GPU計算專家團隊及Merlin HughCTR團隊的支援。

DeepRec架構設計原則

在TensorFlow引擎上支援大規模稀疏特徵,業界有多種實現方式,其中最常見的方式是借鑑了ParameterServer的架構實現,在TensorFlow之外獨立實現了一套ParameterServer和相關的最佳化器,同時在TensorFlow內部透過bridge的方式橋接了兩個模組。這個做法有一定的好處,比如PS的實現會比較靈活,但也存在一些侷限性。

DeepRec採取了另一種架構設計方式,遵循“視整個訓練引擎為一個系統整體”的架構設計原則。TensorFlow是一個基於Graph的靜態圖訓練引擎,在其架構上有相應的分層,比如最上層的API層、中間的圖最佳化層和最下層的運算元層。TensorFlow透過這三層的設計去支撐上層不同Workload的業務需求和效能最佳化需求。

DeepRec也堅持了這一設計原則,基於儲存/計算解耦的設計原則在Graph層面引入EmbeddingVariable功能;基於Graph的特點實現了通訊的運算元融合。透過這樣的設計原則,DeepRec可以支援使用者在單機、分散式場景下使用同一個最佳化器的實現和同一套EmbeddingVariable的實現;同時在Graph層面引入多種最佳化能力,從而做到獨立模組設計所做不到的聯合最佳化設計。

DeepRec的優勢

DeepRec是基於TensorFlow1.15、Intel-TF、NV-TF構建的稀疏模型訓練/預測引擎,針對稀疏模型場景進行了定製深度最佳化,主要包含以下三類功能最佳化:


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec

模型效果

DeepRec提供了豐富的稀疏功能支援,提高模型效果的同時降低稀疏模型的大小,並且最佳化超大規模下Optimizer的效果。下面簡單介紹Embedding及Optimizer幾個有特色的工作:

  • EmbeddingVariable(動態彈性特徵):

1)解決了靜態Shape Variable的vocabulary_size難以預估、特徵衝突、記憶體及IO冗餘等問題,並且在DeepRec中提供了豐富的EmbeddingVariable的進階功能,包括不同的特徵准入方式、支援不同的特徵淘汰策略等,能夠明顯提高稀疏模型的效果。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec

2)在訪問效率上,為了達到更最佳化的效能和更低的記憶體佔用,EmbeddingVariable的底層HashTable實現了無鎖化設計,並且進行了精細的記憶體佈局最佳化,最佳化了HashTable的訪問頻次,使得在訓練過程中前後向只需訪問一次HashTable。

  • DynamicDimensionEmbeddingVariable(動態彈性維度):

在典型的稀疏場景中,同類特徵的出現頻次往往極度不均勻。通常情況下,同一個特徵列的特徵都被設定成統一維度,如果Embedding維度過高,低頻特徵容易過擬合,而且會額外耗費大量記憶體;如果維度設定過低,高頻部徵特徵可能會由於表達不夠而影響效果。

Dynamic Dimension Embedding Variable提供了同一特徵列的不同特徵值,根據特徵的冷熱自動配置不同的特徵維度,高頻特徵可以配置更高維度增強表達能力,而低頻特徵因為給定低維度embedding緩解了過擬合的問題,而且可以極大程度節省記憶體(低頻長尾特徵的數量佔據絕對優勢)。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec

  • Adaptive Embedding(自適應Embedding):

當使用動態彈性特徵功能時,低頻特徵存在過擬合問題。EmbeddingVariable中所有的特徵都從initializer設定的初始值(一般設為0)開始學起,對於一些出現頻次從低到高的特徵,也需要逐漸學習到一個較好的狀態,不能共享別的特徵的學習結果。AdaptiveEmbedding功能使用靜態Shape Variable和動態EmbeddingVariable共同儲存稀疏特徵,對於新加入的特徵存於有衝突的Variable,對於出現頻率較高的特徵存於無衝突的EmbeddingVariable,特徵遷移到EmbeddingVaraible可以複用在有衝突的靜態Shape Variable的學習結果。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec

  • Adagrad Decay Optimizer:

為支援超大規模訓練而提出的一種改進版Adagrad最佳化器。當模型訓練的樣本量大,同時持續增量訓練較長時間時,Adagrad最佳化器的梯度會趨近於0,導致新增訓練的資料無法對模型產生影響。已有的累積打折的方案雖然可以解決梯度趨近0的問題,但也會帶來模型效果變差的問題(透過iteration打折策略無法反映實際的業務場景特點)。Adagrad Decay Optimizer基於週期打折的策略,同一個週期內的樣本相同的打折力度,兼顧資料的無限累積和樣本順序對模型的影響。

此外,DeepRec還提供Multi-HashEmbedding、AdamAsyncOptimizer等功能,在記憶體佔用、效能、模型效果等方面為業務帶來實際的幫助。

訓練效能

DeepRec針對稀疏模型場景在分散式、圖最佳化、運算元、Runtime等方面進行了深度效能最佳化。其中,DeepRec對不同的分散式策略進行了深度的最佳化,包括非同步訓練、同步訓練、半同步訓練等,其中GPU同步訓練支援HybridBackend以及NVIDIA HugeCTR-SOK。DeepRec提供了豐富的針對稀疏模型訓練的圖最佳化功能,包括自動流水線SmartStage、結構化特徵、自動圖Fusion等等。DeepRec中最佳化了稀疏模型中數十個常見運算元,並且提供了包括Embedding、Attention等通用子圖的Fusion運算元。DeepRec中CPUAllocator和GPUAllocator能夠大大降低記憶體/視訊記憶體的使用量並顯著加速E2E的訓練效能。線上程排程、執行引擎方面針對不同的場景提供了不同的排程引擎策略。下面簡單介紹分散式、圖最佳化、Runtime最佳化方面幾個有特色的工作:

  • StarServer(非同步訓練框架):

在超大規模任務場景下(幾百、上千worker),原生開源框架中的一些問題被暴露出來,譬如低效的執行緒池排程、關鍵路徑上的鎖開銷、低效的執行引擎、頻繁的小包rpc帶來的開銷導致ParameterServer在分散式擴充套件時成為明顯的效能瓶頸。StarServer進行了包括圖、執行緒排程、執行引擎以及記憶體等最佳化,將原有框架中的send/recv語義修改為pull/push語義,並且在子圖劃分上支援了該語義,同時實現了ParameterServer端圖執行過程中的lockfree,實現了無鎖化的執行,大大提高了併發執行子圖的效率。對比原生框架,能夠提升數倍的訓練效能,並且支援3000worker規模的線性分散式擴充套件。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec


  • SmartStage(自動流水線):

稀疏模型訓練通常包含樣本資料的讀取、Embedding查詢、Attention/MLP計算等,樣本讀取和Embedding查詢非計算密集操作,同時並不能高效利用計算資源(CPU、GPU)。原生框架中提供的dataset.prefetch介面可以非同步化樣本讀取操作,但Embedding查詢過程中涉及特徵補齊、ID化等複雜的過程,這些過程無法透過prefetch進行流水線化。SmartStage功能能夠自動分析圖中非同步流水線化的邊界並自動插入,可以使併發流水線發揮最大的效能提升。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec


  • PRMalloc(記憶體分配器):

如何做到既高效又有效的使用記憶體,對於稀疏模型的訓練非常關鍵,稀疏場景模型訓練中大塊記憶體分配使用造成大量的minor pagefault,此外,多執行緒分配效率存在比較嚴重的併發分配效率問題。針對稀疏模型訓練前向、後向,Graph計算模式的相對固定、多輪反覆迭代的特點,DeepRec設計了一套針對深度學習任務的記憶體管理方案,提高記憶體的使用效率和系統效能。使用DeepRec中提供的PRMalloc能夠極大降低訓練過程中minor pagefault,提高多執行緒併發記憶體分配、釋放的效率。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec


  • PMEM allocator(持久記憶體分配器):

基於PMDK的底層libpmem庫實現的PMEM allocator將從PMEM map出的一塊空間分為若干segment,每個segment又分成若干blocks,block是allocator的最小分配單元。分配block的執行緒為避免執行緒競爭,快取一些可用空間,包括一組segment和free list。可用空間中為每種record size(若干個block)維護一個free list和segment。各record size對應的segment只分配該大小的PMEM空間,各record size對應的free list中的所有指標均指向對應record size的空閒空間。此外,為了均衡各thread cache的資源,由一個後臺執行緒週期地將thread cache中的free list移動到後臺的pool中,pool中的資源由所有前臺執行緒共享。實驗證明,基於持久記憶體實現的記憶體分配器在大模型的訓練效能方面與基於DRAM的訓練效能差別很小,但是TCO會有很大的優勢。

部署及Serving

  • 增量模型匯出及載入:

時效性要求高的業務,需要頻繁的線上模型更新,頻率往往達到分鐘級別甚至秒級。對於 TB-10TB 級別的超大模型而言,分鐘級別的模型生成到上線很難完成。此外,超大模型的訓練和預測存在著資源浪費、多節點 Serving 延時加大等問題。DeepRec 提供了增量模型產出及載入能力,極大加速了超大模型生成和載入。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec


  • Embedding多級混合儲存:

稀疏模型中特徵存在冷熱傾斜的特性,這產生了某些冷門特徵很少被訪問和更新導致的記憶體/視訊記憶體浪費問題,以及超大模型記憶體/視訊記憶體放不下的問題。DeepRec 提供了多級混合儲存(支援最多四級的混合儲存HBM+DRAM+PMEM+SSD)的能力,自動將冷門特徵存放到廉價的儲存介質中,將熱門特徵存放到訪問更快、更貴的儲存介質上,透過多級混合儲存,使得單節點可以進行 TB-10TB 模型的 Training 和 Serving。


阿里巴巴開源大規模稀疏模型訓練/預測引擎DeepRec

透過多級混合儲存,能夠更大的發揮 GPU 訓練稀疏模型的能力,同時降低由於儲存資源限制造成的計算資源浪費,可以使用更少的機器進行相近規模的模型訓練,或者使用相同數量的機器進行更大規模的訓練。多級混合儲存也能使得單機進行超大模型預測時避免分散式 Serving 帶來的 latency 增大問題,提高大模型的預測效能的同時降低成本。多級混合儲存功能也擁有自動發現特徵的訪問特性,基於高效的熱度統計策略,將熱度高的特徵放置到快速的儲存介質中,將低頻的特徵 offload 到低速儲存介質中,再透過非同步方式驅動特徵在多個介質之間移動。

為什麼開源DeepRec

開源深度學習框架都不能很好地支援稀疏場景中對於稀疏 Embedding 功能的需求、模型訓練效能需求、部署迭代和線上服務的需求。DeepRec 經過阿里巴巴集團搜尋、推薦、廣告等核心業務場景及公有云上各種業務場景的打磨,能夠支援不同型別的稀疏場景訓練效果和效能需求。

阿里巴巴希望透過建立開源社群,和外部開發者開展廣泛合作,進一步推動稀疏模型訓練/預測框架的發展,為不同業務場景中的搜推廣模型訓練和預測帶來業務效果和效能提升。

今天 DeepRec 的開源只是我們邁出的一小步。我們非常期待得到您的反饋。最後,如果你對DeepRec有相應的興趣,你也可以來轉轉,為我們的框架貢獻一點你的程式碼和意見,這將是我們莫大的榮幸。


歡迎加入釘釘答疑交流群:34375411


來自 “ ITPUB部落格 ” ,連結:http://blog.itpub.net/70004426/viewspace-2882892/,如需轉載,請註明出處,否則將追究法律責任。

相關文章