【導讀】:近期,快手宣布將在2020年春節前實現3億DAU,快手商業化營收步伐也隨之加速。快手從2018年“商業化元年”開始推行個性化的廣告推薦。截止5月底,快手DAU已經突破2億。隨著用戶和使用時長的迅速增長,為了更好地挖掘海量用戶和實時數據的核心價值,推薦模型需要快速迭代,從而對用戶興趣遷移的做出迅捷的反應。因此,模型訓練效率成為連接商業效率和用戶興趣的關鍵一環。
作者:快手FeDA智能決策實驗室
來源:AI前線(微信ID:ai-front)
基于歷史原因,行業內推薦模型的訓練大都通過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種這樣的id: {idi}i?=?1m1,它們對應的Embedding層 {Ei}i?=?1m1將會輸出m1個向量:{Ei(idi)}i?=?1 m1。圖像信息、LDA等實數向量特征。這部分將會與id對應的Embedding vector 組合在一起,輸入到DNN中預測點擊率等。假設我們有m2種這樣的向量:{densei}i?=?1m2。DNN。這部分是一個傳統神經網絡,接受Embedding vector和實數向量特征,輸出點擊率等希望預測的量:prediction?=?DNN([E1(idi),?E2(id2),?…,?Em1(idm1),?dense1,?dense2,?…,?densem2])。
Persia使用多種技術訓練廣告模型,我們將在接下來幾節依次介紹。
大模型Embedding分片訓練
廣告模型的Embedding部分占模型體積和計算量的大部分。很有可能無法放入單個GPU的顯存中。為了使用GPU運算以解決CPU運算速度過慢的問題,但又不受制于單GPU顯存對模型大小的限制,Persia系統使用多GPU分散存儲模型,每個GPU只存儲模型一部分,并進行多卡協作查找Embedding向量訓練模型的模式。
Persia將第i個Embedding層Ei 放入第 (i%總顯卡數) 個顯卡中,從而使每個顯卡只存放部分Embedding。與此同時,實數向量特征和DNN部分則置于第0個顯卡中。在使用Persia時,它將自動在各個顯卡中計算出 {Ei}i?=?1m1的值(如果對于一個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 上,使得每組∑i?∈?Sjli,?∀j 大致相等。這等價于如下優化問題:
minS1,?…,?SkVariancej[∑i?∈?Sjli],
s.t. ∑i?∈?SjVi?≤?C,
其中Vi 是第i個模型的大小,C是單個GPU的顯存大小。Persia使用貪心算法得到該問題的一個近似解,并依此將不同Embedding均勻分散在不同GPU上,以達到充分利用GPU的目的。當需要精確求解最優的Embedding放置位置時,Persia還可以通過integer optimization給出精確解。
簡化小模型多 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的更新為:
xt?+?1?←?xt?−?learning rate?×?g(xt?−?τt),
其中g(xt?−?τt)是訓練樣本的損失函數在τt 個更新之前的模型上的梯度。而 τt 的大小一般與計算機數量成正比,當計算機數量增多,xt?−?τt 與 xt 相差就越大,不可避免地導致模型質量的降低。Persia的訓練模式在Embedding分片存儲時沒有這種延遲問題,而在AllReduce模式下也僅在Embedding層有常數量級的延遲,因此模型質量也有所提升。
優化算法:與此同時,Persia還可以使用Adam等momentum optimizer,并為其實現了sparse版本的更新方式,比PyTorch/TensorFlow內置的dense版本更新在廣告任務上快3x-5x。這些算法在很多時候可以在同樣時間內得到比使用 SGD或Adagrad更好的模型。
訓練數據分布式實時處理
快手Persia的高速GPU訓練,需要大量數據實時輸入到訓練機中,由于不同模型對樣本的需求不同,對于每個新實驗需要的數據格式可能也不同。因此 Persia需要:
簡單靈活便于修改的數據處理流程,可以輕易并行的程序架構,節約帶寬的數據傳輸方式。
為此,Persia系統實現了基于Hadoop集群的實時數據處理系統,可以應不同實驗需求從HDFS中使用任意多計算機分布式讀取數據進行多級個性化處理傳送到訓練機。傳輸使用高效消息隊列,并設置多級緩存。傳輸過程實時進行壓縮以節約帶寬資源。
并行數據處理
數據處理pipeline:為了使Persia獲取數據的方式更靈活,Persia使用dataflow構建數據處理pipeline。在Persia中可以定義每一步處理,相當于一個函數,輸入為上一個處理步驟的輸出,輸出提供給下一個處理步驟。我們定義這些函數為 {fi}i?=?1p。在Persia中,這些函數可以單獨定義修改。在每個函數的入口和出口,Persia有數據隊列緩存,以減少每個函數獲取下一個輸入的時間。這些函數的運行可以完全并行起來,這也是pipeline的主要目的。以在食堂就餐為例,pipeline的運行就像這樣:
數據壓縮和傳輸:全部處理之后,數據處理任務會將數據組成mini-batch并使用zstandard高速壓縮每個batch,通過ZeroMQ將壓縮數據傳輸給訓練機進行訓練。定義batching操作為函數 B,壓縮操作為函數C,則每個數據處理任務相當于一個函數C(B(fp(fp?−?1(?f1(raw data from HDFS))))) 。
Queue server:在Hadoop集群中Persia將啟動多個數據處理任務,每個數據處理任務之間完全獨立。數據處理任務本身并不知道處理哪些數據,而是通過請求訓練機得知訓練數據的位置。這樣的好處是,在Persia中訓練機可以應自己需求動態控制使用什么樣的訓練數據,而數據處理任務相當于一個無狀態的服務,即使訓練機更換了新的訓練任務也不需要重啟數據處理任務。具體來說,在Persia中訓練機會啟動一個queue server進程,該queue server將會應數據處理任務的請求返回下一個需要讀取的數據文件。Persia的每個數據處理任務會同時從queue server請求多個文件,并行從HDFS讀取這些文件。
整個系統的構造如下圖:
實時訓練
由于Persia的數據處理任務在獲取數據時完全依賴于訓練機的指示,Persia支持對剛剛生成的數據進行在線訓練的場景,只需要使queue server返回最近生成的數據文件即可。因此,Persia在訓練時的數據讀取模式上非常靈活,對queue server非常簡單的修改即可支持任意數據讀取的順序,甚至可以一邊訓練一邊決定下一步使用什么數據。
更快的數據讀取速度:訓練機共享內存讀取數據
由于訓練機要同時接收從不同數據處理任務發送來的大量數據,并進行解壓縮和傳輸給訓練進程進行實際訓練的操作,接收端必須能夠進行并行解壓和高速數據傳輸。為此,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系統在多機情景下預計還將在單機基礎上做到數倍到數十倍效率提升。