单机训练速度提升640倍！独家解读快手商业广告模型GPU训练平台Persia

【导读】：近期，快手宣布将在2020年春节前实现3亿DAU，快手商业化营收步伐也随之加速。快手从2018年“商业化元年”开始推行个性化的广告推荐。截止5月底，快手DAU已经突破2亿。随着用户和使用时长的迅速增长，为了更好地挖掘海量用户和实时数据的核心价值，推荐模型需要快速迭代，从而对用户兴趣迁移的做出迅捷的反应。因此，模型训练效率成为连接商业效率和用户兴趣的关键一环。

作者：快手FeDA智能决策实验室

来源：AI前线（微信ID：ai-front）

基于历史原因，行业内推荐模型的训练大都通过CPU来实现。然而随着模型从Logistic Regression到深度神经网络的演化以及硬件的发展，基于CPU的训练系统或许已经不再是最合适的解决方案了。本着不盲从、不抄袭、坚持原创技术路线的原则，快手西雅图FeDA智能决策实验室推出了名为"Persia"的基于GPU的广告推荐训练系统。以往需要50台CPU机器训练20小时的系统，如今只需要一台普通的GPU机器在一到两小时完成，单机效率提升高达640倍。这意味着：

• 以往使用五十台计算机，一天只能尝试一个新想法，新系统只需一台计算机，一两个小时就能尝试一个新想法。
• 以往同时只能有一两个同学尝试新模型，新系统可以让很多同学同时尝试各自的新想法。

这套系统已经在快手商业化内部迅速推广使用，让大家可以快速试错和测试新模型以及特征。项目发起者是一位来自罗切斯特大学的实习生。他提出的GPU解决方案得到他在罗切斯特大学的导师、FeDA智能决策实验室负责人刘霁和公司内很多算法策略专家的肯定。

FeDA实验室随即成立了项目组，并决定以项目发起人最喜爱的漫画角色Persia（“佩尔西亚”）命名，展开了紧锣密鼓的开发。团队首先以PyTorch为基础平台着手解决各种技术难题，然后实现并优化TensorFlow版本。经过4个月的开发和通力合作，Persia GPU广告训练系统初步成型。系统同时支持PyTorch和TensorFlow两套方案，以方便模型开发同学的不同偏好。目前，Persia已支持多个业务项目，每位研发人员只需要一台机器便可以迅速地迭代试错。

快手AI概览

Persia背后的技术

Persia实现高效训练背后的技术包含GPU分布式训练、高速数据读取等多个方面。

1. GPU 分布式运算加速模型训练效率

近年来，GPU训练已在图像识别、文字处理等应用上取得巨大成功。GPU训练以其在卷积等数学运算上的独特效率优势，极大地提升了训练机器学习模型，尤其是深度神经网络的速度。然而，在广告模型中，由于大量的稀疏样本存在（比如用户id），每个id在模型中都会有对应的Embedding向量，因此广告模型常常体积十分巨大，以至于单GPU无法存下模型。目前往往将模型存在内存中，由CPU进行这部分巨大的Embedding层的运算操作。这既限制了训练的速度，又导致实际生产中无法使用比较复杂的模型——因为使用复杂模型会导致CPU对给定输入计算时间过长，无法及时响应请求。

广告模型的构成：在广告模型中，模型往往由下图中的三部分构成：

• 用户id、广告id 等构成的Embedding层。每个id对应一个预设大小的向量，由于id数量往往十分巨大，这些向量常常会占据整个模型体积的99%以上。假设我们有m₁种这样的id: {id_i}_i = 1^m1，它们对应的Embedding层 {E_i}_i = 1^m1将会输出m1个向量：{E_i(id_i)}_i = 1 ^m1。
• 图像信息、LDA等实数向量特征。这部分将会与id对应的Embedding vector 组合在一起，输入到DNN中预测点击率等。假设我们有m₂种这样的向量：{dense_i}_i = 1^m2。
• DNN。这部分是一个传统神经网络，接受Embedding vector和实数向量特征，输出点击率等希望预测的量：prediction = DNN([E₁(id_i), E₂(id₂), …, E_m1(id_m1), dense₁, dense₂, …, dense_m2])。

Persia使用多种技术训练广告模型，我们将在接下来几节依次介绍。

1. 大模型Embedding分片训练

广告模型的Embedding部分占模型体积和计算量的大部分。很有可能无法放入单个GPU的显存中。为了使用GPU运算以解决CPU运算速度过慢的问题，但又不受制于单GPU显存对模型大小的限制，Persia系统使用多GPU分散存储模型，每个GPU只存储模型一部分，并进行多卡协作查找Embedding向量训练模型的模式。

Persia将第i个Embedding层E_i 放入第 (i%总显卡数) 个显卡中，从而使每个显卡只存放部分Embedding。与此同时，实数向量特征和DNN部分则置于第0个显卡中。在使用Persia时，它将自动在各个显卡中计算出 {E_i}_i = 1^m1的值（如果对于一个Embedding输入了多个id，则计算其中每个值对应的Embedding vector的平均），并传送给第0个显卡。第0个显卡会合并这些Embedding vector和实数向量特征，输入DNN中进行预测。

当求解梯度时，第0个显卡会将各个Embedding层输出处的导数传回各个显卡，各个显卡各自负责各自Embedding的反向传播算法求梯度。大致结构如下图所示：

GPU分配的负载均衡：由于将 Embedding 依次分配在每个GPU上，可能导致部分GPU负载显著高于其他GPU，为了让每个GPU都能充分发挥性能，Persia训练系统还支持对Embedding运算在GPU上进行负载均衡。

给定k个GPU，当模型的m₁ 个Embedding层对应GPU负载分别为 l₁, l₂, …, l_m1，Persia将会尝试将Embedding分为k 组S₁, S₂, …, S_k，并分别存放在对应GPU 上，使得每组∑_i ∈ Sjl_i, ∀j 大致相等。这等价于如下优化问题：

min_S1, …, _SkVariance_j[∑_i ∈ Sjl_i],

s.t. ∑_i ∈ SjV_i ≤ C,

其中V_i 是第i个模型的大小，C是单个GPU的显存大小。Persia使用贪心算法得到该问题的一个近似解，并依此将不同Embedding均匀分散在不同GPU上，以达到充分利用GPU的目的。当需要精确求解最优的Embedding放置位置时，Persia还可以通过integer optimization给出精确解。

1. 简化小模型多 GPU 分布训练

当模型大小可以放入单个GPU时，Persia也支持切换为目前在图像识别等任务中流行的AllReduce分布训练模式。这样不仅可以使训练算法更加简单，在某些情景下还可以加快训练速度。

使用这种训练模式时，每个GPU都会拥有一个同样的模型，各自获取样本进行梯度计算。在梯度计算后，每个GPU只更新自己显存中的模型。需要注意的是即使模型可以置于一个GPU的显存中，往往Embedding部分也比较大，如果每次更新都同步所有GPU上的模型，会大大拖慢运算速度。因此Persia在AllReduce模式下，每次更新模型后，所有GPU使用AllReduce同步DNN部分，而Embedding部分每隔几个更新才同步一次。这样，即不会损失太多信息，又保持了训练速度。

此外，在TensorFlow上，Persia还支持TensorFlow的"Replicated", "PS", "PS" + "Asynchronous" 模式多卡训练，它们的主要区别如下图：

1. 模型准确度提升

同步更新：由于普遍使用的传统异步 SGD 有梯度的延迟问题，若有n台计算机参与计算，每台计算机的梯度的计算实际上基于n个梯度更新之前的模型。在数学上，对于第t步的模型x_t，传统异步SGD的更新为：

x_t + 1 ← x_t − learning rate × g(x_{t − τt}),

其中g(x_{t − τt})是训练样本的损失函数在τ_t 个更新之前的模型上的梯度。而 τ_t 的大小一般与计算机数量成正比，当计算机数量增多，x_{t − τt} 与 x_t 相差就越大，不可避免地导致模型质量的降低。Persia的训练模式在Embedding分片存储时没有这种延迟问题，而在AllReduce模式下也仅在Embedding层有常数量级的延迟，因此模型质量也有所提升。

优化算法：与此同时，Persia还可以使用Adam等momentum optimizer，并为其实现了sparse版本的更新方式，比PyTorch/TensorFlow内置的dense版本更新在广告任务上快3x-5x。这些算法在很多时候可以在同样时间内得到比使用 SGD或Adagrad更好的模型。

1. 训练数据分布式实时处理

快手Persia的高速GPU训练，需要大量数据实时输入到训练机中，由于不同模型对样本的需求不同，对于每个新实验需要的数据格式可能也不同。因此 Persia需要：

• 简单灵活便于修改的数据处理流程，
• 可以轻易并行的程序架构，
• 节约带宽的数据传输方式。

为此，Persia系统实现了基于Hadoop集群的实时数据处理系统，可以应不同实验需求从HDFS中使用任意多计算机分布式读取数据进行多级个性化处理传送到训练机。传输使用高效消息队列，并设置多级缓存。传输过程实时进行压缩以节约带宽资源。

1. 并行数据处理

数据处理pipeline：为了使Persia获取数据的方式更灵活，Persia使用dataflow构建数据处理pipeline。在Persia中可以定义每一步处理，相当于一个函数，输入为上一个处理步骤的输出，输出提供给下一个处理步骤。我们定义这些函数为 {f_i}_i = 1^p。在Persia中，这些函数可以单独定义修改。在每个函数的入口和出口，Persia有数据队列缓存，以减少每个函数获取下一个输入的时间。这些函数的运行可以完全并行起来，这也是pipeline的主要目的。以在食堂就餐为例，pipeline的运行就像这样：

数据压缩和传输：全部处理之后，数据处理任务会将数据组成mini-batch并使用zstandard高速压缩每个batch，通过ZeroMQ将压缩数据传输给训练机进行训练。定义batching操作为函数 B，压缩操作为函数C，则每个数据处理任务相当于一个函数C(B(f_p(f_p − 1(⋯f₁(raw data from HDFS))))) 。

Queue server：在Hadoop集群中Persia将启动多个数据处理任务，每个数据处理任务之间完全独立。数据处理任务本身并不知道处理哪些数据，而是通过请求训练机得知训练数据的位置。这样的好处是，在Persia中训练机可以应自己需求动态控制使用什么样的训练数据，而数据处理任务相当于一个无状态的服务，即使训练机更换了新的训练任务也不需要重启数据处理任务。具体来说，在Persia中训练机会启动一个queue server进程，该queue server将会应数据处理任务的请求返回下一个需要读取的数据文件。Persia的每个数据处理任务会同时从queue server请求多个文件，并行从HDFS读取这些文件。

整个系统的构造如下图：

1. 实时训练

由于Persia的数据处理任务在获取数据时完全依赖于训练机的指示，Persia支持对刚刚生成的数据进行在线训练的场景，只需要使queue server返回最近生成的数据文件即可。因此，Persia在训练时的数据读取模式上非常灵活，对queue server非常简单的修改即可支持任意数据读取的顺序，甚至可以一边训练一边决定下一步使用什么数据。

1. 更快的数据读取速度：训练机共享内存读取数据

由于训练机要同时接收从不同数据处理任务发送来的大量数据，并进行解压缩和传输给训练进程进行实际训练的操作，接收端必须能够进行并行解压和高速数据传输。为此，Persia使用ZeroMQ device接收多个任务传输而来的压缩数据，并使用多个解压进程读取该device。每个解压进程独立进行解压，并与训练进程共享内存。当结束解压后，解压进程会将可以直接使用的batch样本放入共享内存中，训练任务即可直接使用该batch进行训练，而无需进一步的序列化反序列化操作。

训练效果

Persia系统在单机上目前实现了如下训练效果：

• 数据大小：百T数据。
• 样本数量：25亿训练样本。
• 8卡V100计算机，25Gb带宽：总共1小时训练时间，每秒64万样本。
• 8卡1080Ti计算机，10Gb带宽：总共不到2小时训练时间，每秒40万样本。
• 4卡1080Ti达30万样本/秒，2卡1080Ti达20万样本/秒。
• Persia同样数据上Test AUC高于原ASGD CPU平台。
• Persia支持很大batch size，例如25k。

综上，Persia不仅训练速度上远远超过CPU平台，并且大量节省了计算资源，使得同时尝试多种实验变得非常方便。

展望：分布式多机训练

未来，Persia系统将展开分布式多GPU计算机训练。有别于成熟的计算机视觉等任务，由于在广告任务中模型大小大为增加，传统分布式训练方式面临计算机之间的同步瓶颈会使训练效率大为降低。Persia系统将支持通讯代价更小、系统容灾能力更强的去中心化梯度压缩训练算法。据快手FeDA智能决策实验室负责人刘霁介绍，该算法结合新兴的异步去中心化训练 (Asynchronous decentralized parallel stochastic gradient descent, ICML 2018) 和梯度压缩补偿算法 (Doublesqueeze: parallel stochastic gradient descent with double-pass error-compensated compression, ICML 2019），并有严格理论保证，快手Persia系统在多机情景下预计还将在单机基础上做到数倍到数十倍效率提升。