喜馬拉雅基於 HybridBackend 的深度學習模型訓練優化實踐

語言: CN / TW / HK

喜馬拉雅作者:李超、陶雲、許晨昱、胡文俊、張爭光、趙雲鵬、張玉靜

喜馬拉雅AI雲藉助阿里雲提供的HybridBackend開源框架,實現了其推薦模型在 GPU 上的高效訓練。

業務介紹

推薦場景是喜馬拉雅app的重要應用之一,它廣泛應用於熱點、猜你喜歡、私人FM、首頁資訊流、發現頁推薦、每日必聽等模組。這些模組都依賴於喜馬拉雅AI雲,這是一套從資料、特徵、模型到服務的全流程一站式演算法工具平臺。

image.png

推薦服務的一個核心訴求是能快速捕捉和反映使用者不斷變化的興趣和當前熱點,這就要求模型能在短時間內,以可控的成本完成對海量使用者資料的訓練。使用GPU等高效能硬體來加速模型訓練已經成為CV, NLP等領域的行業標準;在使用稀疏訓練資料的推薦場景下,國內外的各大廠商也在積極轉向使用高效能GPU來替代傳統的CPU訓練叢集,以提升訓練的效率。

圖片

喜馬拉雅AI雲藉助阿里雲機器學習平臺PAI的開源框架HybridBackend,實現了其推薦模型在 GPU 上的高效訓練。在加速訓練的同時, HybridBackend 框架高度易用,幫助其演算法團隊提升了開發效率。

問題與挑戰

隨著推薦業務的底層訓練硬體逐漸從CPU向GPU轉變,我們在生產實踐中發現傳統的訓練方式存在嚴重的計算資源利用率不足的問題。經過調查與分析,我們發現計算資源利用不足主要來自於稀疏資料訪問和分散式訓練兩方面:

  • 稀疏資料訪問:我們使用經典機器學習中常用的 libsvm 資料格式來儲存資料,將多個特徵合併成一個稀疏字串表達。在訓練時,訓練節點從遠端的儲存(如 OSS)上下載字串,並從字串中切分出多個特徵輸入,然後再喂入對應的 Embedding Table。在特徵維度爆炸性增長的情況下,拼接字串的資料量很大,導致資料讀取嚴重受制於網路頻寬;同時切分字串也造成了 CPU 資源的消耗。

  • 分散式訓練:我們嘗試過多種分散式訓練方式。起初,我們使用 keras+horovod 實現多GPU分散式訓練,但在具體使用過程中發現有不少問題,比如出現加速不穩定、模型指標惡化等現象。後來,我們自研實現了一版基於引數伺服器(PS)的分散式訓練框架,通過內部的 xcache 服務實現 embedding 儲存管理並進行線上同步,並使用自研的pspull和pspush運算元進行embedding表的更新,一定程度上解決了分散式訓練的效率問題。但在進一步增大訓練資料量時發現,引入的 ps 運算元因為頻繁的 IO 互動成為了訓練速度的瓶頸,降低了GPU裝置利用率,同時 xcache 服務儲存變長embedding 的支援成本很高,限制了演算法工程師的優化空間。

HybridBackend

我們在調研如何解決上述問題和探索未來技術發展空間時發現了阿里雲正在推廣的開源框架 HybridBackend,該框架對稀疏模型訓練過程中的資料訪問、稀疏計算以及分散式訓練都進行了深度優化(見圖1),並提供了簡單易用的介面。令人驚喜的是,這個框架相容性很強,可以支援 TensorFlow、DeepRec 等多種訓練框架,可以很好地滿足我們服務不同業務客戶的需求。此外,相關的架構和系統設計已經以論文形式在ICDE2022會議上公開,並且在Github上開源了主要功能,可以直接以pip方式安裝。

圖片

HybridBackend的基本功能模組

圖2 描述了在我司模型訓練任務中落地 HybridBackend 的全景示意圖。藍色框代表了 HybridBackend 參與或加速了的流程部分。可以看到基本涵蓋了全部模型訓練流程,  下面重點介紹在資料讀取和分散式訓練上的優化成果。

圖片

HybridBackend在喜馬拉雅業務流程中的落地

稀疏資料訪問優化

HybridBackend 提供了 hb.data.Dataset 介面 ,通過支援如 Parquet 這樣的列存資料格式,可以極大加速稀疏資料的訪問。如表1所示,HybridBackend 框架稀疏資料讀取效能遠高於其他實現。

檔案格式 檔案大小(MB) 框架 執行緒數 耗時(ms)
CSV 11062.61 TensorFlow 1 8858.38
Parquet (SNAPPY) 3346.10 TensorFlow I/O 1 103056.17
Parquet (SNAPPY) 3346.10 HybridBackend 1 397.88
Parquet (SNAPPY) 3346.10 HybridBackend 20 21.67

在我們的實際應用中,HybridBackend 稀疏資料訪問功能中的一些功能效果顯著:

  • 資料列選擇性解析:我們將需要原有的類libsvm格式切換成寬表格式,其中每列對應一個特徵。HybridBackend 可以支援在讀取 Parquet 檔案時只讀取選擇的欄位,並將欄位資料解析成 TensorFlow 所需要的格式,如自動將 list 型別的資料轉換為 SparseTensor,或將 list 型別的資料進行填充截斷後轉換為 Tensor,滿足了我們資料載入的多種需要。
  • 資料讀取並行度設定:HybridBackend 可以通過設定num_parallel_reads 引數來調整讀不同檔案的並行度,通過設定num_parallel_parser_calls 引數來調整讀檔案中不同列的並行度。通過並行讀取,在充分利用機器 CPU 資源的同時,加速了資料讀取的效能。

在使用 HybridBackend 後,資料訪問不再是我們的訓練瓶頸。單卡訓練的 GPU 平均利用率提升了 3x 以上,業務模型的訓練週期顯著縮短。

分散式訓練優化

HybridBackend 提供了混合並行訓練模式(如圖3),每張 GPU 都會儲存全部的稠密引數和部分的稀疏引數,並使用可以利用 NVLink 的 NCCL 通訊協議來代替傳統 PS 訓練方式所使用的 RPC 協議。

image.png

HybridBackend提供的混合並行訓練模式

根據我們對未來一段時間內模型特徵維度和大小的預估,以及我們對訓練速度的需求,我們採用了 HybridBackend 混合並行方式進行訓練,有效地提高了訓練速度和 GPU 利用率。

我們還與 HybridBackend 社群的開發者協同工作,促進了 HybridBackend 對 Keras Model API 的支援,使我們能夠在 Keras Model API 下利用 HybridBackend 進行混合並行,並實現模型熱啟等重要功能。這些功能極大地降低了使用成本。

總體收益

整體流程改造完畢之後,我們在推薦場景中,單機多卡訓練 GPU 平均利用率提升了1.4x 以上(視具體模型不同),訓練環節整體耗時減少50%以上。目前我們已經在使用了 Tensorflow 和 DeepRec 的模型中全量推廣基於 HybridBackend 的訓練方案。

image.png

未來規劃

喜馬拉雅 AI 雲平臺目前覆蓋了喜馬拉雅多個app的推薦、廣告、搜尋推薦等核心業務場景,以及畫像產出、資料分析、BI資料生成等定製化開發場景。我們也在探索後續與 HybridBackend 社群的一些合作,以便更好地滿足業務需求:

  • 運算元優化:HybridBackend支援了embedding lookup 過程中的各種運算元的融合優化。我們會嘗試通過這種方式提升模型線上推理的效能。
  • PyTorch支援:NLP 搜推場景中有用 Pytorch 進行訓練和部署的需求。我們需要HybridBackend能夠支援該場景的實現。
  • 超大型分散式訓練:我們的模型訓練級別達到了百億樣本十億特徵維度。隨著演算法複雜度的提升,我們需要支援更大的資料量和更高的維度的訓練。

鳴謝

在合作共建階段,我們得到了 HybridBackend 社群 陳浪石、袁滿等的技術支援,他們技術高超、服務周到、響應及時。幫助我們快速完成了深度學習模型的訓練流程優化,為我們的業務指標和演算法優化空間帶來了明顯的提升。在此表示衷心的感謝!

HybridBackend 社群

歡迎在 GitHub 上 star 和提 issue,也可以直接在釘釘群中聯絡 HybridBackend 社群。

GitHub 地址:

https://github.com/alibaba/HybridBackend