作者
焦学武、李淼
导读
本文介绍百度基于GPU的超大规模离散DNN训练框架PaddleBox。PaddleBox打造了业界首个层次化GPU稀疏参数服务器，结合高效的流水线调度架构，以及多机多卡的分布式架构，支持单机10TB级/多机数十TB模型训练，具备低成本、高性能、高稳定、灵活易用的多重优势。PaddleBox于在百度凤巢首次上线，目前已在百度广告系统全面落地，覆盖搜索广告、信息流广告、联盟广告业务，相比传统CPU解决方案，资源性价比提升5~40倍。
注：本文基于年百度内网文章修改而成。
全文字，预计阅读时间24分钟。
01
超大离散DNN模型训练的背景与挑战
精准的广告是很多互联网产品重要的盈利方式之一，基于海量数据的广告系统能够很好的匹配用户的意图，为用户带来良好产品体验，同时提升商业变现效率。CTR预估模型是广告系统中非常重要的一环，被誉为“镶嵌在互联网技术上的明珠”，而高效稳定的训练框架则是CTR模型日常迭代更新的基石。
近年来深度学习迅猛发展，百度早在年就已经将DNN模型应用到广告系统，是国内首个将大规模DNN用在广告领域并取得了良好效果的公司。Google此前发布了万亿级模型SwitchTransformer，参数量达到了1.6万亿，其实凤巢广告模型的参数量比Google的SwitchTransformer还要大，广告模型采用了十亿级别的用户ID、Query、广告ID以及多种组合特征构建模型，特征数量可以达到千亿级别，模型训练过程中需要将千亿级的高维稀疏特征转化为低维的稠密特征向量（8、16以及64维），模型参数规模高达万亿维。
△图1.大规模离散DNN
大规模离散DNNCTR模型以EmbeddingLayer和FCLayer为主，虽然与CV和NLP领域的经典模型相比网络结构并不复杂，但是庞大的参数规模、样本规模和大规模高维稀疏的特性为模型的高效稳定训练提出了多重挑战。(下文把高维稀疏特征对应的Embedding参数称为Sparse参数，DNN参数为Dense参数)
存储挑战。Sparse参数量可达万亿级别，存储空间占用达到10TB级别，远超单机内存，一般需要分布式全内存的存储系统。
IO挑战。IO挑战来自训练样本和模型参数两个维度，很多场景每天产出亿级用户的数十亿条训练样本；每个mini-batch模型需要从千亿特征的EmbeddingTable中查找和更新对应的参数，需要应对超高频的参数查询和更新。
计算挑战。CTR模型与NLP和CV模型不同，大部分计算不是矩阵运算，样本解析、稀疏参数的查询和更新等非矩阵运算占比超过70%。
02
传统分布式CPU解决方案
基于多台cpu服务器的参数服务器架构是大规模模型训练的成熟方案，该方案采用了数据并行和模型并行的混合方式，数据并行是指将训练样本分散到多个计算节点上，模型并行是指将稀疏参数分散到多个存储节点上。参数服务器架构有两个角色Worker和Server。其中Worker用于执行模型的前向和反向计算；Server负责从各个Worekr收集汇总梯度并更新参数。
△图2.参数服务器的分布式训练架构
参数服务器以Scale-out的思路成功解决了参数规模和样本规模带来的系统可伸缩性问题，可以通过调整CPU服务器的数量来应对参数规模和样本规模的伸缩性需求，但也有自身的限制：
成本问题。训练任务占用上百台CPU服务器，此前百度大商业用于CTR模型训练的服务器接近2万台，带来大量的硬件采购和维护成本。
通信长尾和稳定性问题。训练过程中上百台CPU服务器需要进行mini-batch级的超高频网络通信（参数拉取和梯度更新），通信长尾导致严重的性能恶化，也导致了梯度过期问题，影响策略效果；此外，由于故障概率的指数级上升，基于上百台CPU服务器的训练任务存在严重的稳定性问题。
算力问题。随着GateNetwork、Attention等网络结构被引入到CTR模型，模型网络越来越复杂，对算力的要求越来越高，CPU机器很难满足模型复杂度增长的算力需求。
NVIDIA近年来一直在推动GPU通用计算能力的发展，从年推出CUDA编程平台以来，面向通用计算的GPU架构先后经历了Fermi、Kepler、Maxwell、Pascal、Volta和Ampere等6代产品，最新一代的TeslaA单精度浮点数运算性能达TFLOPS,比IntelXeonE7系列高出2个数量级，算力的提升使得语音、图像、NLP等领域取得了革命性的进展，但在超大规模离散DNN场景却一直没有落地。
直接采用GPU做参数服务器架构的Worker来解决算力问题并不可行，因为模型训练过程中高频的参数拉取与梯度更新需要频繁地进行CPU/GPU通信，难以发挥GPU算力优势，且存储超大模型需要较多GPU服务器，成本过高。
0
PaddleBoxGPU解决方案
为了解决传统参数服务器架构面临的问题，本框架设计了基于GPU的PaddleBox解决方案，整体架构如下图所示：
△图.PaddleBox整体架构
PaddleBox推出了异构层次化参数服务器，通过SSD、MEM、HBM三级参数服务器高效协同的方案来提升模型参数容量和访存速度，采用高效的流水线调度流程使得异构硬件最大化并行，充分发挥异构硬件优势，实现单台GPU服务器支持10TB级模型高效训练。为了支持更大规模的模型和样本，PaddleBox基于多机进行扩展，设计了性能近线性加速比的分布式架构。而且PaddleBox与Paddle共建开源生态，可支持复杂模型的快速接入与调研，具备低成本、高性能、高稳定、灵活易用的多重优势。
PaddleBox打造了业界首个分布式GPU稀疏参数服务器，下面将详细介绍PaddleBox是如何解决模型训练的多重挑战的。
.1异构层次化参数服务器
△图4.异构参数服务器
.1.1支撑超大模型参数存储的SSD参数服务器
GPU服务器拥有超强的算力，从算力角度，单台或几台GPU参数服务器的算力就可以支持超大规模模型的训练。但是从存储角度，几台机器的内存和显存无法存储10TB级的模型参数，因此也就无法进行模型训练。为了应对超大规模模型的存储挑战，在HDD、NVMeSSD、Optane和AEP等多种硬件存储设备间，综合考虑训练所需的存储容量、性能和成本等多方面因素，最终选择SSD作为全量参数存储设备。
测试显示在Raid0下的SSD的读写性能依旧落后于内存1-2个数量级，业务场景对读写性能要求非常高，传统的SSD索引一次检索需要多次IO，性能太差，为此系统进行多重性能优化。
IO优化：通过构建多级全内存hash索引，实现了对SSD数据一次性的准确读写，每次查询最多一次IO。
查询剪枝：结合数据访问的冷热比例，MEMCache存储热数据；使用BloomFilter判断当前key是否在SSD上进一步降低无效的访盘次数。这两种剪枝策略使得SSD查询次数降低一个数量级。
底层优化：通过异步IO和数据对齐等优化技术，SSD读写性能提升5倍，直达SSD理论极限5*GB/s。
△图5.SSD参数服务器结构
这是业界首个基于SSD的超大规模异构存储稀疏参数服务器，成功支持单机10TB、万亿维参数存储，使得采用单GPU服务器进行10TB级模型训练成为现实。
.1.2支撑参数高效访存的HBM参数服务器
采用SSD参数服务器后，单台GPU服务器即可进行超大模型训练，但是其训练速度却并没有得到很大提升。其原因在于训练过程中Sparse参数的超高频CPU-GPU通信，无法充分发挥GPU超强算力。
为了解决这个问题，框架设计实现了业界首个多机多卡分布式GPU稀疏参数服务器。HBM、MEM和SSD三层参数服务器自动化协同，既具备SSD的大容量，又具备HBM的高性能。
机器内GPU卡间需要应对超高频的万亿次稀疏参数访问，传统通信技术无法满足性能需求，我们通过软硬件双重创新来优化。
软件创新：NVLink的带宽高达00GB/s,性能超过PCIe一个数量级。传统卡间通信技术采用低速的PCIe,难以满足海量的参数通信需求。对此我们根据GPU的拓扑结构实现了Multi-Hop通信策略，实现GPU跨卡参数访问性能提升7倍。(百度自研XMAN2.0机型)
△图6.非NVLink全互连的GPU拓扑结构（来自NVIDIA