PaddleBox：百度基于GPU的超大规模离散DNN模型训练解决方案

在这里插入图片描述

作者 | 焦学武、李淼

导读
本文介绍百度基于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。

相信PaddleBox必将成为星辰大海上的那颗北极星，引领大规模离散模型的进一步创新。

————END————