作者 | 焦學武、李淼
導讀
本文介紹百度基于GPU的超大規模離散DNN訓練架構PaddleBox。PaddleBox 打造了業界首個階層化GPU稀疏參數伺服器，結合高效的流水線排程架構，以及多機多卡的分布式架構，支援單機10TB級/多機數十TB模型訓練，具備低成本、高性能、高穩定、靈活易用的多重優勢。PaddleBox于2019在百度鳳巢首次上線，目前已在百度廣告系統全面落地，覆寫搜尋廣告、資訊流廣告、聯盟廣告業務，相比傳統CPU解決方案，資源成本效益提升5~40倍。
注：本文基于2020年百度内網文章修改而成。
全文5246字，預計閱讀時間24分鐘。
01 超大離散DNN模型訓練的背景與挑戰
精準的廣告是很多網際網路産品重要的盈利方式之一，基于海量資料的廣告系統能夠很好的比對使用者的意圖，為使用者帶來良好産品體驗，同時提升商業變現效率。CTR預估模型是廣告系統中非常重要的一環，被譽為“鑲嵌在網際網路技術上的明珠”，而高效穩定的訓練架構則是CTR模型日常疊代更新的基石。
近年來深度學習迅猛發展，百度早在2014年就已經将DNN模型應用到廣告系統，是國内首個将大規模DNN用在廣告領域并取得了良好效果的公司。Google 此前釋出了萬億級模型 Switch Transformer，參數量達到了 1.6 萬億，其實鳳巢廣告模型的參數量比Google的Switch Transformer還要大，廣告模型采用了十億級别的使用者ID、Query、廣告ID以及多種組合特征構模組化型，特征數量可以達到千億級别，模型訓練過程中需要将千億級的高維稀疏特征轉化為低維的稠密特征向量（8、16以及64維），模型參數規模高達萬億維。
△圖1. 大規模離散DNN
大規模離散DNN CTR模型以Embedding Layer和 FC Layer為主，雖然與CV和NLP領域的經典模型相比網絡結構并不複雜，但是龐大的參數規模、樣本規模和大規模高維稀疏的特性為模型的高效穩定訓練提出了多重挑戰。(下文把高維稀疏特征對應的Embedding參數稱為Sparse參數，DNN參數為Dense參數)
- 存儲挑戰。Sparse參數量可達萬億級别，存儲空間占用達到10TB級别，遠超單機記憶體，一般需要分布式全記憶體的存儲系統。
- IO挑戰。IO挑戰來自訓練樣本和模型參數兩個次元，很多場景每天産出億級使用者的數十億條訓練樣本；每個mini-batch模型需要從千億特征的Embedding Table中查找和更新對應的參數，需要應對超高頻的參數查詢和更新。
- 計算挑戰。CTR模型與NLP和CV模型不同，大部分計算不是矩陣運算，樣本解析、稀疏參數的查詢和更新等非矩陣運算占比超過70%。
02 傳統分布式CPU解決方案
基于多台cpu伺服器的參數伺服器架構是大規模模型訓練的成熟方案，該方案采用了資料并行和模型并行的混合方式，資料并行是指将訓練樣本分散到多個計算節點上，模型并行是指将稀疏參數分散到多個存儲節點上。參數伺服器架構有兩個角色Worker 和 Server。其中Worker 用于執行模型的前向和反向計算；Server 負責從各個Worekr 收集彙總梯度并更新參數。
△圖2. 參數伺服器的分布式訓練架構
參數伺服器以Scale-out的思路成功解決了參數規模和樣本規模帶來的系統可伸縮性問題，可以通過調整CPU伺服器的數量來應對參數規模和樣本規模的伸縮性需求，但也有自身的限制：
- 成本問題。訓練任務占用上百台CPU伺服器，此前百度大商業用于CTR模型訓練的伺服器接近2萬台，帶來大量的硬體采購和維護成本。
- 通信長尾和穩定性問題。訓練過程中上百台CPU伺服器需要進行mini-batch級的超高頻網絡通信（參數拉取和梯度更新），通信長尾導緻嚴重的性能惡化，也導緻了梯度過期問題，影響政策效果；此外，由于故障機率的指數級上升，基于上百台CPU伺服器的訓練任務存在嚴重的穩定性問題。
- 算力問題。随着Gate Network、Attention等網絡結構被引入到CTR模型，模型網絡越來越複雜，對算力的要求越來越高，CPU機器很難滿足模型複雜度增長的算力需求。
NVIDIA近年來一直在推動GPU通用計算能力的發展，從2006年推出CUDA程式設計平台以來，面向通用計算的GPU架構先後經曆了Fermi、Kepler、Maxwell、Pascal、Volta和Ampere等6代産品，最新一代的Tesla A100單精度浮點數運算性能達156 TFLOPS, 比Intel Xeon E7系列高出2個數量級，算力的提升使得語音、圖像、NLP等領域取得了革命性的進展，但在超大規模離散DNN場景卻一直沒有落地。
直接采用GPU做參數伺服器架構的Worker來解決算力問題并不可行，因為模型訓練過程中高頻的參數拉取與梯度更新需要頻繁地進行CPU/GPU通信，難以發揮GPU算力優勢，且存儲超大模型需要較多GPU伺服器，成本過高。
03 PaddleBox GPU 解決方案
為了解決傳統參數伺服器架構面臨的問題，本架構設計了基于GPU的PaddleBox解決方案，整體架構如下圖所示：
△圖3. PaddleBox整體架構
PaddleBox推出了異構階層化參數伺服器，通過SSD、MEM、HBM三級參數伺服器高效協同的方案來提升模型參數容量和訪存速度，采用高效的流水線排程流程使得異構硬體最大化并行，充分發揮異構硬體優勢，實作單台GPU伺服器支援10TB級模型高效訓練。為了支援更大規模的模型和樣本，PaddleBox基于多機進行擴充，設計了性能近線性加速比的分布式架構。而且PaddleBox與Paddle共建開源生态，可支援複雜模型的快速接入與調研，具備低成本、高性能、高穩定、靈活易用的多重優勢。
PaddleBox打造了業界首個分布式GPU 稀疏參數伺服器，下面将詳細介紹PaddleBox是如何解決模型訓練的多重挑戰的。
3.1 異構階層化參數伺服器
△圖4. 異構參數伺服器
3.1.1 支撐超大模型參數存儲的SSD參數伺服器
GPU伺服器擁有超強的算力，從算力角度，單台或幾台GPU參數伺服器的算力就可以支援超大規模模型的訓練。但是從存儲角度，幾台機器的記憶體和顯存無法存儲10TB級的模型參數，是以也就無法進行模型訓練。為了應對超大規模模型的存儲挑戰，在HDD、NVMe SSD、Optane 和 AEP等多種硬體儲存設備間，綜合考慮訓練所需的存儲容量、性能和成本等多方面因素，最終選擇SSD作為全量參數儲存設備。
測試顯示在Raid0 下的SSD的讀寫性能依舊落後于記憶體1-2個數量級，業務場景對讀寫性能要求非常高，傳統的SSD索引一次檢索需要多次IO，性能太差，為此系統進行多重性能優化。
- IO優化：通過建構多級全記憶體hash索引，實作了對SSD資料一次性的準确讀寫，每次查詢最多一次IO。
- 查詢剪枝：結合資料通路的冷熱比例，MEM Cache 存儲熱資料；使用BloomFilter判斷目前key 是否在SSD上進一步降低無效的訪盤次數。這兩種剪枝政策使得SSD查詢次數降低一個數量級。
- 底層優化：通過異步IO和資料對齊等優化技術，SSD讀寫性能提升5倍，直達SSD理論極限5\*3GB/s。
△圖5. SSD參數伺服器結構
這是業界首個基于SSD的超大規模異構存儲稀疏參數伺服器，成功支援單機10TB、萬億維參數存儲，使得采用單GPU伺服器進行10 TB級模型訓練成為現實。
3.1.2 支撐參數高效訪存的HBM參數伺服器
采用SSD參數伺服器後，單台GPU伺服器即可進行超大模型訓練，但是其訓練速度卻并沒有得到很大提升。其原因在于訓練過程中Sparse參數的超高頻CPU-GPU通信，無法充分發揮GPU超強算力。
為了解決這個問題，架構設計實作了業界首個多機多卡分布式GPU稀疏參數伺服器。HBM、MEM和SSD三層參數伺服器自動化協同，既具備SSD的大容量，又具備HBM的高性能。
機器内GPU卡間需要應對超高頻的萬億次稀疏參數通路，傳統通信技術無法滿足性能需求，我們通過軟硬體雙重創新來優化。
- 軟體創新：NVLink的帶寬高達300GB/s, 性能超過PCIe 一個數量級。傳統卡間通信技術采用低速的PCIe, 難以滿足海量的參數通信需求。對此我們根據GPU的拓撲結構實作了Multi-Hop 通信政策，實作GPU跨卡參數通路性能提升7倍。(百度自研XMAN2.0機型)
△圖6. 非NVLink全互連的GPU拓撲結構（來自NVIDIA官網）
- 硬體創新：首次引入NVSwitch全互聯總線架構，任意GPU間可以通過NVLink進行P2P高速通信，GPU跨卡參數通路性能提升10倍。(百度自研XMAN3.0機型)
△圖7. NVSwitch 全互聯總線架構
3.2 高效的訓練流水線架構
PaddleBox 使用的機型裝置具備多種異構硬體，怎樣可以充分發揮硬體能力呢？
PaddleBox 模型的訓練流程包括樣本讀取、樣本解析、參數拉取和模型訓練四個階段，每個階段對硬體資源類型的需求是不同的。
- 樣本讀取：從分布式檔案系統讀取資料，是網絡密集型任務。
- 樣本解析：把文本資料解析為特定的資料結構，是CPU密集型任務。
- 參數拉取：從SSD中拉取Sparse參數，是SSD IO密集型任務。
- 模型訓練：在GPU硬體上訓練模型，是GPU密集型任務。
△圖8.高效的流水線結構
3.3 近線性加速比的分布式版本
PaddleBox單機版本能夠支撐10TB模型的高效訓練。為了支援更大規模的模型和樣本，以及更快地訓練，PaddleBox 進一步實作了分布式架構。
△圖9. 分布式PaddleBox架構
- 分布式SSD存儲引擎：通過參數分區存儲的方式将稀疏參數sharding到多台機器上，建構分布式SSD參數伺服器來實作更大規模的稀疏參數存儲。
- 超高的多機通信效率：超強的GPU算力，要求更加高效的多機通信效率，否則會嚴重制約分布式的加速比。對此，PaddleBox架構在網卡拓撲、通信協定等方面進行了軟硬一體的整體設計。
- 算法創新：梯度聚合和量化通信結合，通信量降為原來的1/4， 極大的提升了模型訓練效率。我們通過多個次元的優化榨幹異構硬體的性能，最終使得PaddleBox的分布式架構具備近似線性的加速比。
△圖10. 分布式PaddleBox性能測試結果
3.4 Paddle 開源生态
PaddlePaddle 源于産業實踐，是業界公認的國内唯一功能完備的開源深度學習平台，具備靈活的組網能力和豐富的算法庫，提供高效的并行訓練能力和良好的可擴充性，目前已經支援公司内外多個業務上線。
PaddleBox 提供基于GPU的超大規模異構參數伺服器，結合Paddle靈活的組網能力和豐富的算法庫，同時具備算力群組網的雙重優勢，可以快速引入來自學術界的CNN、RNN、Attention、Bert 等模型和機制，為業務場景引入開源複雜模型算法提供了廣闊空間。
04 落地與收益
PaddleBox目前已經覆寫了百度搜尋廣告、推薦廣告、網盟廣告和手百推薦等場景的上百個模型，支撐各業務取得收入上的巨大提升。與此同時，公司内部豐富的業務也在驅動着PaddleBox的更新與成熟。擷取的收益主要來源于三方面：
- 極高的成本效益：PaddleBox以更低的成本提供了更多的訓練資源，相對采用MPI參數伺服器，取得了5~40倍的成本效益提升。
- 算力和靈活性雙重優勢：支援複雜的模型，PaddleBox 使得CTR不再隻是全連接配接模型，融入開源生态後提供語義模型、注意力模型、多模态聯合訓練的能力。
- 支援多種業務場景：PaddleBox 不僅僅是CTR模型訓練架構，更是通用的大規模離散模型訓練架構，适合多種場景。除點選率模型外，架構在轉化率CVR、TDM模型、圖模型等場景也有廣泛的落地。
05 總結
Paddlebox設計之初就開始潛心研究GPU分布式訓練技術，以應對大規模離散模型的訓練任務，在豐富的廣告推薦業務驅動下，PaddleBox推出了業内首創的異構參數伺服器，支援多種場景而且可以極大地提升硬體資源使用率，具有極高的成本效益。
同時，PaddleBox也在進行架構的更深層次更新，包括模型混布、異構叢集、昆侖芯新硬體的探索等，進一步提升模型訓練性能和資源使用率。PaddleBox也在場景上擴充到了特征抽取，目前已經支援特征抽取、模型訓練一體化任務，極大提升了模型訓練效率。後續，也希望能跟大家分享特征抽取架構FeaBox和大規模圖引擎PGLBox。