喜馬拉雅基於 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 地址:

http://github.com/alibaba/HybridBackend