DeepRec 大規模稀疏模型訓練推理引擎

導讀：

本文將以下三個方面展開介紹：

• DeepRec背景（我們為什麼要做DeepRec）
• DeepRec功能（設計動機和實現）
• DeepRec社群（最新發布的2206版本主要功能）

DeepRec背景介紹

我們為什麼需要稀疏模型引擎？TensorFlow目前的社群版本是能夠支援稀疏場景的，但是在以下三個方面存在一些功能上的短板：

• 提升模型效果的稀疏訓練功能；
• 提升模型迭代效率的訓練效能；
• 稀疏模型的部署。

因此我們提出了DeepRec，其功能定位在稀疏場景做深度的優化。

DeepRec所做的工作主要在四大方面：稀疏功能、訓練效能、Serving、以及部署& ODL。

DeepRec在阿里巴巴內部的應用主要在推薦（猜你喜歡）、搜尋（主搜）、廣告（直通車和定向）等幾個核心場景。我們也給雲上的一些客戶提供了部分稀疏場景的解決方法，為其模型效果和迭代效率的提升帶來了很大幫助。

DeepRec功能介紹

DeepRec的功能主要分為以下五大方面：稀疏功能Embedding，訓練框架（非同步、同步），Runtime（Executor、PRMalloc），圖優化（結構化模型，SmartStage），serving部署相關功能。

1. Embedding

Embedding部分將介紹以下5個子功能：

1.1 動態彈性特徵（EV）

上圖的左邊是TensorFlow支援稀疏功能的主要方式。使用者首先定義固定的shape的Tensor，稀疏的特徵通過Hash+Mod的方式map到剛剛定義的Tensor上。這個邏輯上有4個問題：

• 稀疏特徵的衝突，Hash+Mod的方式容易引入特徵衝突，這會導致有效特徵的消失，進而影響效果；
• 儲存部分會導致記憶體的浪費，有部分記憶體空間不會被使用到；
• 固定的shape，一旦Variable的shape固定了，未來無法更改；
• 低效的IO，假如使用者用這種方式定義Variable，必須通過全量的方式匯出，如果Variable的維度很大，那麼無論匯出還是載入都是十分耗時的，但我們在稀疏的場景其實變化的部分是很少的。

在這種情況下，DeepRec定義的EmbeddingVariable設計的原理是：將靜態的Variable轉化為動態的類似HashTable的儲存，每來一個key，新建立一個Embedding，這樣就天然地解決了特徵衝突的問題。經過這樣的設計，當特徵特別的多的時候，EmbeddingVariable無序的擴張，記憶體消耗也會變得很大，因此DeepRec引入了以下兩個功能：特徵准入和特徵淘汰。它們都能有效的防止特徵擴充套件到很大的維度。在搜尋和推薦這樣的稀疏場景，有些長尾特徵被模型訓練的次數十分少。因此特徵准入能通過CounterFilter或者BloomFilter的方式對特徵進入EmbeddingVariable設定一個門檻；在模型匯出Checkpoint的時候也會有特徵淘汰的功能，時間上比較老的特徵也會被淘汰。這在阿里內部某個推薦業務AUC提升5‰，在雲上某推薦業務AUC提升5‰，pvctr也有提升4%。

1.2基於特徵頻率的動態彈性維度特徵（FAE）

通常情況下同一個特徵對應的EmbeddingVariable會被設定為同一個維度，如果EmbeddingVariable被設定一個較高的維度，低頻的特徵內容容易導致過擬合，並且會消耗大量的記憶體。相反的如果維度設定的過低，高頻的特徵內容則有可能因為表達的能力不足而影響模型的效果。FAE的功能則提供了對於同一個特徵裡，根據不同特徵冷熱來配置不同的維度。這樣讓模型自動進行訓練時第一個是模型的效果能得到保證，第二個也能解決訓練對資源的使用。這是對於FAE功能的出發點的介紹。這個功能的使用目前是讓使用者傳入一個維度和統計的演算法，FAE自動根據實現的演算法來產生不同的EmbeddingVariable；後面DeepRec計劃在系統內部自適應的發現去分配特徵的維度，從而提高使用者的易用性。

1.3自適應EmbeddingVariable

這個功能和第二個功能有些類似，都是以定義高低頻的關係作為出發點。當前面提到的EV特別大時，我們會看到記憶體佔用特別高。在Adaptive Embedding Variable中我們用兩個Variable來表達，如右圖展示。我們會定義其中一個Variable為靜態的，低頻的特徵會盡可能對映到這個Variable上；另外一個則定義為動態彈性維度特徵，用於高頻部分的特徵。Variable的內部支援低頻和高頻特徵動態的轉換，這樣的優點是極大降低了系統對記憶體的使用。例如某個特徵訓練後第一維可能有接近10億，而重要的特徵只有20%-30%，通過這種自適應的方式後，可以不需要那麼大的維度，進而極大的降低了對記憶體的使用。我們在實際應用發現對模型的精度影響是很小的。

1.4 Multi-Hash Variable

這個功能是為了解決特徵衝突的問題。我們原來是通過一個Hash+Mod的方式解決特徵衝突，現在用兩個或多個Hash+Mod去得到Embedding，並且隨後對得到的Embedding做Reduction，這樣的好處是能用更少的記憶體來解決特徵衝突的問題。

1.5 Embedding多級混合儲存

這一功能的出發點同樣也是發現EV在特徵個數多的時候，記憶體開銷十分大，訓練的時候worker佔用的記憶體可能達到了幾十上百G。我們發現，特徵實際上遵循典型的冪律分佈。考慮到這個特徵點，我們將熱點特徵放到CPU這樣更寶貴的資源，而相對長尾低頻的特徵則放到相對廉價的資源中。如右圖，有DRAM、PMEM、SSD三種結構，PMEM是英特爾提供的速度介於DRAM和SSD之間，但容量很大。我們目前支援DRAM-PMEM、DRAM-SSD、PMEM-SSD的混合，也在業務上取得了效果。雲上有個業務 原來用200+多CPU分散式訓練，現在使用多級儲存後改成了單機GPU訓練。

以上是對Embedding所有功能的介紹。我們做這些功能的動機是由於TensorFlow的幾個問題（主要是特徵衝突），我們解決的方案是動態彈性特徵和Multi-Hash特徵，針對動態彈性特徵記憶體開銷較大的問題，我們又開發了特徵准入和特徵淘汰的功能；針對特徵頻次，我們開發了3組功能：動態彈性維度和自適應動態彈性特徵是從維度的方向解決的問題，多級混合儲存則是從軟硬體的方向解決的問題。

2.  訓練框架

第二個要介紹的功能是訓練框架，分為非同步和同步兩個方向來介紹。

2.1非同步訓練框架StarServer

在超大規模任務情況下，上千個worker，原生TensorFlow存在的問題是：執行緒排程十分低效，關鍵路徑開銷凸顯，另外小包通訊十分頻繁，這些都成為了分散式通訊的瓶頸。

StarServer在圖的執行緒排程、記憶體的優化方面做得很好，將框架中Send/Recv修改為了Push/Pull語義，PS在執行的時候使用了lockless的方法，極大地提高了執行的效率。我們對比原生框架有數倍的效能提升，並且在內部3Kworker左右的數量能達到線性的擴充套件。

2.2同步訓練框架HybridBackend，

這是我們為同步訓練開發的方案，它支援資料並行和模型並行混合分散式訓練。資料讀取通過資料並行來完成，模型並行能支援大引數量訓練，最後使用資料並行做稠密計算。我們針對不同EmbeddingLookup的特徵，做了多路Lookup合併的優化，分組優化，還利用了GPU Direct RDMA的優點，基於網路拓撲的感知，設計整個同步的框架。

3.  Runtime

第三個大方面的功能是Runtime，主要介紹PRMalloc和Executor優化。

3.1 PRMalloc

首先是記憶體分配，記憶體分配在TensorFlow和DeepRec中都是無處不在的，我們首先在稀疏訓練中發現，大塊記憶體分配造成了大量的minorpagefault，此外在多執行緒的分配中也存在併發分配的問題。我們在DeepRec中針對稀疏訓練前向反向的特點，設計了針對深度學習的記憶體分配方案，稱為PRMalloc。它提高了記憶體使用率和系統的效能。在圖中可以看到主要的一塊是MemoryPlanner，它的作用是在模型訓練的前k輪的minibatch先統計當前訓練的特點，每次需要分配多少Tensor，將這些行為記錄通過bin的buffer記錄下來，並且做相應的優化。在k步後，我們將其應用，從而極大減少上述的問題。我們在DeepRec的使用中發現，這能大大減少minorpagefault的出現，減少了記憶體的使用，訓練速度也得到了1.6倍的加速。

3.2 Executor優化

TensorFlow原生的Executor的實現十分簡單，首先對DAG做拓撲排序，隨後將Node插入到執行佇列中，通過Task利用Executor排程。這樣的實現沒有結合業務考慮，ThreadPool預設使用了Eigen執行緒池，若執行緒負載不均勻，會發生大量的執行緒間搶佔Steal，帶來極大開銷。我們在DeepRec中定義排程更均勻，同時定義了關鍵路徑使得在排程的時候有一定的優先順序順序，來執行Op。最終DeepRec也提供了多種包括基於Task，SimpleGraph的排程策略。

4. 圖優化相關的功能

4.1結構化特徵

這是從業務啟發的一個功能。我們發現在搜尋場景下，不管是訓練還是推理，樣本往往是1個user對應多個item，多個label的特點。原來的處理方式會視為多個樣本，這樣user的儲存是冗餘的，我們為了節省這部分開銷，自定義了儲存格式來做這部分優化。如果這些樣本在一個minibatch中是同一個user，部分user網路和item網路會分別計算，最後在做相應的邏輯計算，這樣能節省計算開銷。所以我們分別從儲存和計算端做了結構化的優化。

4.2 SmartStage

我們看到稀疏模型的訓練通常包括樣本的讀取，EmbeddingLookup，還有MLP的網路計算。樣本的讀取和Embedding查詢往往不是計算密集型的，並不能有效利用計算資源。原生框架提供的prefetch介面雖然能一定程度上完成非同步操作，但是我們在EmbeddingLookup過程中設計部分複雜的子圖，這些不能通過TensorFlow的prefetch實現流水線。TensorFlow提供的流水線功能，實際使用中需要使用者顯示的指定stage邊界，一方面會提高使用難度，另一方面由於stage的精度不夠，無法精確到op級別。對於High Level的API使用者無法手動插入，會導致很多步伐並行化。下圖是SmartStage的具體操作，它會將Op自動的歸類到不同的Stage，使得併發的流水線能得到效能的提升。我們在ModelZoo裡模型的測試效果最大加速比能達到1.1-1.3。

5. Serving

5.1模型增量匯出及載入

一開始在介紹Embedding的時候其中一個重要的點是低效的IO，如果將前面提到動態彈性功能應用後，我們天然能做增量的匯出。只要在圖中加入曾經訪問的稀疏ID，那麼在增量匯出的時候就能準確的匯出這部分我們需要的ID。我們做這個功能有兩個出發點：首先，模型訓練時我們原有的方法，在每個step匯出全量的模型匯出，在程式中斷restore時候也是restore checkpoint，最差的時候可能損失兩個checkpoint區間所有的結果，有了增量匯出，我們對於dense部分會全量匯出，sparse部分是增量匯出，這在實際場景10分鐘的增量匯出能很大程度節約restore帶來的損失；另外，增量匯出的場景是線上serving，如果每次都全量載入，那麼對於稀疏場景，模型十分大，每次載入都需要耗費很長時間，如果要做線上學習會很困難，所以增量匯出也會用到ODL場景。

5.2 ODL

最左邊是樣本處理，上下兩部分是離線和線上的訓練，右邊是serving。這裡面應用了很多PAI的元件來完成Pipeline的構造。

DeepRec 社群

社群方面，我們在6月份釋出了新版本2206，主要包括以下新功能：