單機訓練速度提升640倍!獨家解讀快手商業廣告模型GPU訓練平臺Persia

2018年是快手的“商業化元年”。爲了聯結用戶體驗和商業價值,快手開始推行個性化的廣告推薦。截止5月底,快手的DAU已經突破2億,並且宣佈將在明年1月達到3億DAU的目標。隨着快手用戶和使用時長的迅速增長,爲了更好地挖掘海量用戶和實時數據的核心價值,推薦模型需要快速迭代,從而對用戶興趣遷移的做出迅捷的反應。因此,模型訓練效率成爲至關重要的一環。

基於歷史原因,行業內推薦模型的訓練大都通過CPU來實現。然而隨着模型從Logistic Regression到深度神經網絡的演化以及硬件的發展,基於CPU的訓練系統或許已經不再是最合適的解決方案了。本着不盲從、不抄襲、堅持原創技術路線的原則,快手西雅圖FeDA智能決策實驗室推出了名爲"Persia"的基於GPU的廣告推薦訓練系統。以往需要50臺CPU機器訓練20小時的系統,如今只需要一臺普通的GPU機器在一到兩小時完成,單機效率提升高達640倍。這意味着:

  • 以往使用五十臺計算機,一天只能嘗試一個新想法,新系統只需一臺計算機,一兩個小時就能嘗試一個新想法。
  • 以往同時只能有一兩個同學嘗試新模型,新系統可以讓很多同學同時嘗試各自的新想法。

這套系統已經在快手商業化內部迅速推廣使用,讓大家可以快速試錯和測試新模型以及特徵。項目發起者是一位來自羅切斯特大學的實習生。他提出的GPU解決方案得到他在羅切斯特大學的導師、FeDA智能決策實驗室負責人劉霽和公司內很多算法策略專家的肯定。

FeDA實驗室隨即成立了項目組,並決定以項目發起人最喜愛的漫畫角色Persia(“佩爾西亞”)命名,展開了緊鑼密鼓的開發。團隊首先以PyTorch爲基礎平臺着手解決各種技術難題,然後實現並優化TensorFlow版本。經過4個月的開發和通力合作,Persia GPU廣告訓練系統初步成型。系統同時支持PyTorch和TensorFlow兩套方案,以方便模型開發同學的不同偏好。目前,Persia已支持多個業務項目,每位研發人員只需要一臺機器便可以迅速地迭代試錯。

快手AI概覽

Persia背後的技術

Persia實現高效訓練背後的技術包含GPU分佈式訓練、高速數據讀取等多個方面。

GPU 分佈式運算加速模型訓練效率

近年來,GPU訓練已在圖像識別、文字處理等應用上取得巨大成功。GPU訓練以其在卷積等數學運算上的獨特效率優勢,極大地提升了訓練機器學習模型,尤其是深度神經網絡的速度。然而,在廣告模型中,由於大量的稀疏樣本存在(比如用戶id),每個id在模型中都會有對應的Embedding向量,因此廣告模型常常體積十分巨大,以至於單GPU無法存下模型。目前往往將模型存在內存中,由CPU進行這部分巨大的Embedding層的運算操作。這既限制了訓練的速度,又導致實際生產中無法使用比較複雜的模型——因爲使用複雜模型會導致CPU對給定輸入計算時間過長,無法及時響應請求。

廣告模型的構成:在廣告模型中,模型往往由下圖中的三部分構成:

  • 用戶id、廣告id 等構成的Embedding層。每個id對應一個預設大小的向量,由於id數量往往十分巨大,這些向量常常會佔據整個模型體積的99%以上。假設我們有m1**(注:因排版原因,此處1爲下標,下同)種這樣的id:,它們對應的Embedding層將會輸出m1個向量。

  • 圖像信息、LDA等實數向量特徵。這部分將會與id對應的Embedding vector 組合在一起,輸入到DNN中預測點擊率等。假設我們有m2種這樣的向量:

  • DNN。這部分是一個傳統神經網絡,接受Embedding vector和實數向量特徵,輸出點擊率等希望預測的量:prediction = 

Persia使用多種技術訓練廣告模型,我們將在接下來幾節依次介紹。

1. 大模型Embedding分片訓練

廣告模型的Embedding部分佔模型體積和計算量的大部分。很有可能無法放入單個GPU的顯存中。爲了使用GPU運算以解決CPU運算速度過慢的問題,但又不受制於單GPU顯存對模型大小的限制,Persia系統使用多GPU分散存儲模型,每個GPU只存儲模型一部分,並進行多卡協作查找Embedding向量訓練模型的模式。

Persia將第i個Embedding層Ei放入第 (i%總顯卡數) 個顯卡中,從而使每個顯卡只存放部分Embedding。與此同時,實數向量特徵和DNN部分則置於第0個顯卡中。在使用Persia時,它將自動在各個顯卡中計算出 的值(如果對於一個Embedding輸入了多個id,則計算其中每個值對應的Embedding vector的平均),並傳送給第0個顯卡。第0個顯卡會合並這些Embedding vector和實數向量特徵,輸入DNN中進行預測。

當求解梯度時,第0個顯卡會將各個Embedding層輸出處的導數傳回各個顯卡,各個顯卡各自負責各自Embedding的反向傳播算法求梯度。大致結構如下圖所示:

GPU分配的負載均衡:由於將 Embedding 依次分配在每個GPU上,可能導致部分GPU負載顯著高於其他GPU,爲了讓每個GPU都能充分發揮性能,Persia訓練系統還支持對Embedding運算在GPU上進行負載均衡。

給定k個GPU,當模型的m1個Embedding層對應GPU負載分別爲l1, l2, …, lm1,Persia將會嘗試將Embedding分爲k 組 S1, S2, …, Sk,並分別存放在對應GPU 上,使得每組大致相等。這等價於如下優化問題:

其中Vi是第i個模型的大小,C是單個GPU的顯存大小。Persia使用貪心算法得到該問題的一個近似解,並依此將不同Embedding均勻分散在不同GPU上,以達到充分利用GPU的目的。當需要精確求解最優的Embedding放置位置時,Persia還可以通過integer optimization給出精確解。

2. 簡化小模型多 GPU 分佈訓練

當模型大小可以放入單個GPU時,Persia也支持切換爲目前在圖像識別等任務中流行的AllReduce分佈訓練模式。這樣不僅可以使訓練算法更加簡單,在某些情景下還可以加快訓練速度。

使用這種訓練模式時,每個GPU都會擁有一個同樣的模型,各自獲取樣本進行梯度計算。在梯度計算後,每個GPU只更新自己顯存中的模型。需要注意的是即使模型可以置於一個GPU的顯存中,往往Embedding部分也比較大,如果每次更新都同步所有GPU上的模型,會大大拖慢運算速度。因此Persia在AllReduce模式下,每次更新模型後,所有GPU使用AllReduce同步DNN部分,而Embedding部分每隔幾個更新才同步一次。這樣,即不會損失太多信息,又保持了訓練速度。

此外,在TensorFlow上,Persia還支持TensorFlow的"Replicated", “PS”, “PS” + “Asynchronous” 模式多卡訓練,它們的主要區別如下圖:

模型準確度提升

同步更新:由於普遍使用的傳統異步 SGD 有梯度的延遲問題,若有n臺計算機參與計算,每臺計算機的梯度的計算實際上基於n個梯度更新之前的模型。在數學上,對於第t步的模型xt,傳統異步SGD的更新爲:

其中是訓練樣本的損失函數在τt個更新之前的模型上的梯度。而τt的大小一般與計算機數量成正比,當計算機數量增多,xt相差就越大,不可避免地導致模型質量的降低。Persia的訓練模式在Embedding分片存儲時沒有這種延遲問題,而在AllReduce模式下也僅在Embedding層有常數量級的延遲,因此模型質量也有所提升。

優化算法:與此同時,Persia還可以使用Adam等momentum optimizer,併爲其實現了sparse版本的更新方式,比PyTorch/TensorFlow內置的dense版本更新在廣告任務上快3x-5x。這些算法在很多時候可以在同樣時間內得到比使用 SGD或Adagrad更好的模型。

訓練數據分佈式實時處理

快手Persia的高速GPU訓練,需要大量數據實時輸入到訓練機中,由於不同模型對樣本的需求不同,對於每個新實驗需要的數據格式可能也不同。因此 Persia需要:

  • 簡單靈活便於修改的數據處理流程,
  • 可以輕易並行的程序架構,
  • 節約帶寬的數據傳輸方式。

爲此,Persia系統實現了基於Hadoop集羣的實時數據處理系統,可以應不同實驗需求從HDFS中使用任意多計算機分佈式讀取數據進行多級個性化處理傳送到訓練機。傳輸使用高效消息隊列,並設置多級緩存。傳輸過程實時進行壓縮以節約帶寬資源。

1. 並行數據處理

數據處理pipeline:爲了使Persia獲取數據的方式更靈活,Persia使用dataflow構建數據處理pipeline。在Persia中可以定義每一步處理,相當於一個函數,輸入爲上一個處理步驟的輸出,輸出提供給下一個處理步驟。我們定義這些函數爲 。在Persia中,這些函數可以單獨定義修改。在每個函數的入口和出口,Persia有數據隊列緩存,以減少每個函數獲取下一個輸入的時間。這些函數的運行可以完全並行起來,這也是pipeline的主要目的。以在食堂就餐爲例,pipeline的運行就像這樣:

數據壓縮和傳輸:全部處理之後,數據處理任務會將數據組成mini-batch並使用zstandard高速壓縮每個batch,通過ZeroMQ將壓縮數據傳輸給訓練機進行訓練。定義batching操作爲函數 B,壓縮操作爲函數C,則每個數據處理任務相當於一個函數

Queue server:在Hadoop集羣中Persia將啓動多個數據處理任務,每個數據處理任務之間完全獨立。數據處理任務本身並不知道處理哪些數據,而是通過請求訓練機得知訓練數據的位置。這樣的好處是,在Persia中訓練機可以應自己需求動態控制使用什麼樣的訓練數據,而數據處理任務相當於一個無狀態的服務,即使訓練機更換了新的訓練任務也不需要重啓數據處理任務。具體來說,在Persia中訓練機會啓動一個queue server進程,該queue server將會應數據處理任務的請求返回下一個需要讀取的數據文件。Persia的每個數據處理任務會同時從queue server請求多個文件,並行從HDFS讀取這些文件。

整個系統的構造如下圖:

2. 實時訓練

由於Persia的數據處理任務在獲取數據時完全依賴於訓練機的指示,Persia支持對剛剛生成的數據進行在線訓練的場景,只需要使queue server返回最近生成的數據文件即可。因此,Persia在訓練時的數據讀取模式上非常靈活,對queue server非常簡單的修改即可支持任意數據讀取的順序,甚至可以一邊訓練一邊決定下一步使用什麼數據。

3. 更快的數據讀取速度:訓練機共享內存讀取數據

由於訓練機要同時接收從不同數據處理任務發送來的大量數據,並進行解壓縮和傳輸給訓練進程進行實際訓練的操作,接收端必須能夠進行並行解壓和高速數據傳輸。爲此,Persia使用ZeroMQ device接收多個任務傳輸而來的壓縮數據,並使用多個解壓進程讀取該device。每個解壓進程獨立進行解壓,並與訓練進程共享內存。當結束解壓後,解壓進程會將可以直接使用的batch樣本放入共享內存中,訓練任務即可直接使用該batch進行訓練,而無需進一步的序列化反序列化操作。

訓練效果

Persia系統在單機上目前實現瞭如下訓練效果:

  • 數據大小:百T數據。
  • 樣本數量:25億訓練樣本。
  • 8卡V100計算機,25Gb帶寬:總共1小時訓練時間,每秒64萬樣本。
  • 8卡1080Ti計算機,10Gb帶寬:總共不到2小時訓練時間,每秒40萬樣本。
  • 4卡1080Ti達30萬樣本/秒,2卡1080Ti達20萬樣本/秒。
  • Persia同樣數據上Test AUC高於原ASGD CPU平臺。
  • Persia支持很大batch size,例如25k。

綜上,Persia不僅訓練速度上遠遠超過CPU平臺,並且大量節省了計算資源,使得同時嘗試多種實驗變得非常方便。

展望:分佈式多機訓練

未來,Persia系統將展開分佈式多GPU計算機訓練。有別於成熟的計算機視覺等任務,由於在廣告任務中模型大小大爲增加,傳統分佈式訓練方式面臨計算機之間的同步瓶頸會使訓練效率大爲降低。Persia系統將支持通訊代價更小、系統容災能力更強的去中心化梯度壓縮訓練算法。據快手FeDA智能決策實驗室負責人劉霽介紹,該算法結合新興的異步去中心化訓練 (Asynchronous decentralized parallel stochastic gradient descent, ICML 2018) 和梯度壓縮補償算法 (Doublesqueeze: parallel stochastic gradient descent with double-pass error-compensated compression, ICML 2019),並有嚴格理論保證,快手Persia系統在多機情景下預計還將在單機基礎上做到數倍到數十倍效率提升。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章