阿里巴巴稀疏模型訓練引擎-DeepRec

DataFunTalk發表於2022-04-17


導讀:DeepRec從2016年起深耕至今,支援了淘寶搜尋、推薦、廣告等核心業務,沉澱了大量優化的運算元、圖優化、Runtime優化、編譯優化以及高效能分散式訓練框架,在稀疏模型的訓練方面有著優異效能的表現。並且沉澱了稀疏場景下的動態彈性特徵、動態維度彈性特徵、多Hash彈性特徵等功能,能夠不同程度的提高稀疏模型的效果。作為阿里巴巴集團內稀疏場景的統一訓練引擎,是AOP團隊、XDL團隊、PAI團隊、AIS團隊合作共建的專案。除此之外,DeepRec得到了Intel、NV相關團隊的支援,針對稀疏場景下的運算元、子圖、以及針對硬體特點進行了深度定製優化。

本文將圍繞下面三點展開:

  • DeepRec背景

  • DeepRec功能介紹

  • DeepRec開源

--

01

DeepRec背景

1. 為什麼需要稀疏模型訓練引擎

首先是DeepRec的背景,做搜尋推廣的同學肯定感同身受,流行的Tensorflow、PyTorch等,其實有很多痛點,這些痛點也就是我們為什麼需要做稀疏模型的訓練引擎,我們這邊總結了主要分為三個方面:

第一個方面,已有的這些開源的深度學習引擎缺少了對這種稀疏模型訓練功能的支援,大家也都知道像基礎的這種動態彈性特徵,也是在稀疏場景裡常用的一個功能,這種功能在模型效果上其實會有很好的補充和提升。

第二就是訓練效能,已有的一些開源深度學習框架並沒有特別針對稀疏模型的訓練有很好的優化。無論是在PS/worker模式還是AllReduce模式,以及不同的device CPU、GPU上的訓練,大家都做了非常多的工作,而這些工作也源於我們在使用開源深度學習框架訓練稀疏模型時遇到的一些問題。

第三點就是稀疏模型很具有特點、部署和Serving,稀疏模型的部署跟其他的像CV類、NLP類的一些模型的部署區別比較大,像ODL這種場景下,模型的更新很可能是需要分鐘級別的、甚至秒級別的,而這些模型通常可能是幾百GB、幾個TB、10TB的這種超大模型,這種模型在分鐘級別模型更新或者秒級別模型更新時使用開源的框架其實是沒有辦法做到的。

2. 什麼是DeepRec

針對上面提到的問題,我們是在基於Tensorflow上面構建了這種稀疏模型的訓練引擎,並且在這三個方面進行了很多優化和功能的沉澱,其實這些也是基於阿里巴巴內部的各大核心業務方的一些功能的沉澱。下面我分別來先簡單說一下這三個方面分別有什麼樣的功能。

① 稀疏功能

在稀疏功能方面,像動態彈性特徵,這其實是大家基礎常用的一個功能,存在稀疏引數的維度特別大的情況。動態彈性特徵在特徵准入、特徵淘汰上都能夠在一定程度上解決過擬合、訓練不充分的問題。特徵淘汰也是稀疏場景特別有特點的,比如說某些商品下架,對應的特徵就需要被淘汰掉。在DeepRec裡面針對動態彈性特徵,支援的是非常完備的,因為我們各大業務方,他們對於准入淘汰有各種不同的需求,比如我們是有基於布隆過濾器的准入,基於精準Counter的准入,種類是非常多的,也非常豐富。

此外,像基於特徵頻率的動態彈性維度,針對每個特徵的冷熱會自動地伸縮它的維度。訓練不充分的時候,過擬合會相對比較嚴重,動態彈性的維度是根據引數出現的頻度,自動的每個引數,每個特徵都會有自己的維度,這樣對於低頻的特徵,它可以用更低頻的維度來表達,對高頻特徵可以用更高頻的維度來表達。自適應的動態彈性特徵,結合了有衝突和無衝突的引數,也是在一定程度上提高了模型的效果。此外,DeepRec還支援多個雜湊的組合。

② 訓練效能

在訓練效能上,我們做了大量的工作,包括分散式訓練框架,在非同步訓練上,我們在PS/worker上面支援超大規模的非同步訓練框架。同步訓練,我們是基於GPU的一個硬體,實現的同步的訓練框架HybridBackend。在Runtime上,我們對Tensorflow針對的稀疏模型的場景進行了一個深度的重寫,包括記憶體、視訊記憶體、執行、執行緒池。圖優化上面,我們這邊包括有自動的多階段的pipeline,然後自動得像剛剛老師提到的這種maffer batch,包括結構化特徵,然後子圖的Fusion這些圖優化的工作。運算元優化的話包括了很多大量的feature OP的重寫,以及大量稀疏運算元的一些重寫。

③ 部署及Serving

在模型部署和Serving上面,包括增量模型的匯出載入,超大模型的Serving,以及多層的混合儲存,還有多Backend的支援,以及ODL的支援。

3. DeepRec業務場景

DeepRec在阿里內部使用的核心的業務場景主要就是猜你喜歡、主搜尋,還有廣告的直通車和定向。

4. DeepRec研發團隊

DeepRec專案其實是各個團隊合作共建的。在開發團隊上,主要有像主搜工程的公司,搜尋工程團隊的AOP團隊,還有RTP、XDL、PAI、AIS,以及英特爾和英偉達都在參與DeepRec開發。

--

02

DeepRec功能介紹

1. 動態彈性特徵

動態彈性特徵是稀疏模型訓練裡面非常基礎的一個功能,解決的就是像原生的Tensorflow裡面,在訓練的時候,固定的引數bucket size和固定dimension帶來的問題。稀疏模型訓練,通常這種增量訓練可能一個模型要持續訓練很久,模型引數固定之後,自然而然就有上圖中列出的幾個問題,像特徵衝突、記憶體浪費,以及可能因為靜態shape設定得比較大造成的無效的IO等問題。

動態彈性特徵,在底層使用了hashtable的結構,可以基於訓練的情況,自動插入刪除這種特徵,這樣的話底層的hashtable,自然而然就會解決掉像前面提到的這種問題。當然在動態彈性特徵之上有一些基礎的功能,比如特徵的准入、特徵的淘汰。

2. 基於特徵頻率的動態彈性維度

這個功能是在動態彈性特徵基礎之上,使用者在使用的時候也不用配dimension,可以配一個最大可能出現的dimension上限,這樣的話動態彈性維度功能會自動根據特徵本身出現的頻率,自動進行伸縮。

低頻的特徵像圖中表示的就是這種用16位來表達,越高頻的特徵用更高的維度來表達,這樣去解決低頻特徵,因為它更新的訓練的機率、機會比較少,可能出現的這種過擬合問題,當然也是一定程度上提高高頻特徵本身的表達能力,這樣的話,大家也都有一個common sense,就是在我們的稀疏模型裡面,高頻特徵佔比通常也是比較少的,但其實通過這個功能也一定程度上能夠降低模型的size,在我們的一些場景裡邊模型size能夠降低幾倍。

3. 非同步訓練框架StarServer

下面講兩個關於分散式訓練框架的功能,一個是非同步訓練框架StarServer。

StarServer主要有四個特點,包括通訊方面的優化,通訊協議的優化,圖層面的優化,以及Runtime執行引擎的優化。

第一個在通訊協議上來講,我們實現了Share Nothing的這種通訊協議,並且實現了使用者態的純零拷貝的資料傳輸。此外我們會基於圖的拓撲關係對圖進行一些Fusion操作,然後將圖做重新的partition。在Runtime,我們會在整個執行引擎的執行方式上進行變化,執行引擎在執行的過程中實現了Run To Completion的執行策略,以及整個圖執行過程中的完全的無鎖化。

非同步訓練框架StarServer,它整個支撐的規模是比較大。在超大規模worker的上面存在著過期梯度的問題,DeepRec裡面也有一些Optimizer的實現來解決過期梯度的問題。

4. 同步訓練框架HybridBackend

另一個分散式訓練框架就是同步訓練框架HybridBackend。它是基於GPU實現的一套同步的訓練框架。

其有三個特點:

  • 資料並行於模型並行混合分散式訓練

  • 面向混合硬體的訓練編排

  • 針對高維稀疏特徵大批量訪問優化

展開講一下混合並行的這種分散式訓練。不同的稀疏特徵,比如說ID型別的特徵,相對來講會非常大,對於有些特徵,它的特徵的量非常少,比如說省份、性別這種型別的特徵。不同的特徵,它的同步方式是有區別的,就像大的這種稀疏的特徵,使用AllToAllv的方式進行同步;對於小size的特徵列和稠密引數直接使用AllReduce進行同步梯度。此外會基於GPU的硬體特點,進行訓練的編排,而這裡其實主要還是基於硬體,GPU底下NVLink以及RDMA本身的硬體的排布,去做同步的方案。

此外,像高維特徵的訪問的優化,為什麼大家逐漸利用GPU來做稀疏模型的訓練,其實利用了幾點:一個就是GPU的這種高併發,再就是HBM以及NVLink的這種高頻寬,其中高維稀疏的批量的訪問優化其實就是需要這種高並行、精確的高併發以及HBM的這種高的訪存頻寬,對於高維稀疏特徵的訪問優化是非常明顯的。

5. Runtime優化-PRMalloc

上圖是挑選了Runtime優化中的PRMalloc工作,我們針對這種模型訓練的一個特點,將模型的訓練前面,比如訓練過程前面的K輪進行一些統計,通過Malloc的這樣一個模組,對記憶體的分配,包括分配的時間點、分配的大小,統計好分配的時間點和大小後,在K輪結束之後會使用啟發式的一些演算法規劃出一個較優的快取,而這個快取會用於K輪之後的記憶體分配。

整個架構來講,PRMalloc解決了在稀疏場景裡面在訪存上存在的一些問題,主要有幾個問題,一個問題就是在稀疏場景裡邊,有大量的小記憶體的分配,這些小記憶體分配很可能就是特徵的處理,包括一些特徵的拼接,或者在做一些交叉特徵,這裡會存在大量的小記憶體的分配。同樣在模型訓練也存在很多大的記憶體,特別是像後面的一些結構,比如說從embedding層cat起來之後的一些結構裡邊,可能attention的、RNN的、或者說全連線的,它會有一些大記憶體的分配。

在這些小記憶體和大記憶體的分配上,傳統的malloc庫,比如jemalloc、tcmalloc,其實並不適用於深度學習的訓練。PRMalloc一方面就是解決在小記憶體分配以及大記憶體分配上存在的這些問題,其實大家如果是稀疏模型訓練,在CPU上有些很明顯的特點就是可以使用我們PRMalloc,比如minor perfort比較高,說明有非常多大記憶體分配行為,這些分配行為源自於tcmalloc或者是jemalloc在大記憶體分配管理上存在的一些缺陷,導致每次分配後使用觸發大量minor pagefault。

6. 圖優化-結構化特徵

圖優化方面,我挑了其中一個相對比較有特點的結構化特徵,針對我們樣本的一些特點,它包含兩個層面的優化,一個就是樣本的儲存,另一個就是模型的訓練效能。

樣本儲存,大家可以看到我們通常一次曝光,user測的這些特徵是一樣的,只有item和label會不一樣,實際上在儲存上面,當我們把user側和item側重新做一個儲存格式的劃分,我們可以存一條user對應多條item以及label,這樣在儲存上可以節省大量的user側特徵的儲存。因為user側特徵儲存,特別是使用者的行為,它的量會比較大,尤其是隨著使用者行為序列越來越長,使用者特徵裡面使用者行為佔比是非常多的。

在圖優化層面,當我們把輸入變成了一個user對多個item以及多個label的形式之後,user側的網路其實也會簡化,因為user側網路的輸入,它的user側的特徵可能變得很少,這個時候它的計算以及它的embedding lookup都會變得更少,當需要和item側進行結合的時候,要做tensor的展開。

7. 模型增量匯出及載入

在模型的Serving和匯出方面,我們支援模型的增量匯出和載入,而這裡是針對稀疏模型典型的一個需求,即每隔一段時間都需要去快速產生一個上線的模型,也就是像ODL。在ODL場景下,分鐘級別秒級別的模型更新的時候,需要快速產生一個模型,並且這個模型首先要達到上線的需求,其次它的大小也應該是足夠小的。增量模型就是記錄了在一個時間段內的稀疏引數的變化情況,然後將變化的這些引數匯出到模型中,這樣增量的模型通常非常小。

8. Embedding多級混合儲存

此外,在這種超大的embedding場景下,對於多級混合儲存的需求是非常急迫的。因為本身embedding的這些特徵,存在很明顯的冷熱區分。這種明顯的冷熱區分,可以通過多級混合儲存既滿足了效能的要求,同時滿足了成本的需求。此外這種多級混合儲存更好的一點就是它更能夠大大地降低對分散式的依賴。

舉個最簡單的場景,比如說分散式Serving的一個場景,一個大模型的分散式Serving可能會增加很多的網路開銷,但是通過多級混合儲存去實現一個單節點的Serving,一個大模型通過單節點就可以Serving的話,其實在效率上要提高很多。此外在training階段也是這樣,當使用GPU進行稀疏模型訓練的時候,也是需要多級混合儲存去支援解決這種在GPU下,算力足夠強、稀疏引數又太大的矛盾的問題。在DeepRec內部包括GPU上的多級混合儲存,以及CPU上的多級混合儲存,我們分四級,包括HBM、DRAM、PMEM以及SSD,這裡其中HBM的管理以及PMEM的管理是和英特爾以及英偉達的相關的同學一起合作開發的。

--

03

DeepRec開源

下面我介紹一下DeepRec開源的狀態和一些思考。

首先DeepRec在阿里內部已經在最核心的一些業務場景以及大量的稀疏場景下得到了很好的錘鍊,並且沉澱了很多對於效能,包括對於模型效果都有提升的一些功能。

我們也希望後面在開源的過程中,能跟更多的公司去合作,去接觸到更多的業務場景,把DeepRec用得更廣。因為不同的業務場景對於模型的效果、模型的訓練,以及他們的樣本、模型差異都是非常大的。


現在我們跟英特爾、英偉達已經在DeepRec內部沉澱了非常多的功能,包括像稀疏分子的加速,HugeCTR也已經在DeepRec接入了。包括像GPU的一些OP,GPU的hashtable,以及Runtime優化、圖優化運算元優化的一些實現,以及軟硬體一體化的一些設計,另外在稀疏場景模型訓練中,也正在做一些硬體的工作。

總結:

大規模稀疏模型的應用是搜、推、廣等多業務領域所面臨的重要課題,阿里巴巴DeepRec模型在業界已有訓練引擎、框架的基礎上探索出了一套行之有效的實踐方案,在特徵使用、模型訓練、線上推理、儲存等方面都進行了一定的探索和優化。DeepRec希望在自身開源的過程中與業界夥伴一起更好地解決大規模稀疏模型應用的問題。

--

04

問答環節

Q:目前基於哪個TF的版本?

A:基於TF的1.15,並且加入了NV TF和Intel TF的功能。

Q:維度的動態變化會不會與TF網路層引數不匹配,或無法計算?

A:維度的變化,當Embedding被查出來要送入後面的網路層時,會padding到最高維度,在初始設定時一般會設定一個最高維度,動態維度在後續計算中會統一到一個最高維上。


今天的分享就到這裡,謝謝大家。

閱讀更多技術乾貨文章,請關注微信公眾號“DataFunTalk”

關於我們:
DataFun:專注於大資料、人工智慧技術應用的分享與交流。發起於2017年,在北京、上海、深圳、杭州等城市舉辦超過100+線下和100+線上沙龍、論壇及峰會,已邀請近1000位專家和學者參與分享。其公眾號 DataFunTalk 累計生產原創文章500+,百萬+閱讀,13萬+精準粉絲。


注:歡迎轉載,轉載請留言或私信。

相關文章