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

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

基于历史原因，行业内推荐模型的训练大都通过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分布式训练、高速数据读取等多个方面。

一、 GPU 分布式运算加速模型训练效率

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

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

l 用户id、广告id 等构成的Embedding层。每个id对应一个预设大小的向量，由于id数量往往十分巨大，这些向量常常会占据整个模型体积的99%以上。假设我们有m1种这样的id: {idi}i=1m1，它们对应的Embedding层 {Ei}i=1m1将会输出m1个向量：{Ei(idi)}i=1 m1。

l 图像信息、LDA等实数向量特征。这部分将会与id对应的Embedding vector 组合在一起，输入到DNN中预测点击率等。假设我们有m2种这样的向量：{densei}i=1m2。

l DNN。这部分是一个传统神经网络，接受Embedding vector和实数向量特征，输出点击率等希望预测的量：prediction=DNN([E1(idi),E2(id2),…,Em1(idm1),dense1,dense2,…,densem2])。

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

1. 大模型Embedding分片训练

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

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

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

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

给定k个GPU，当模型的m1 个Embedding层对应GPU负载分别为 l1,l2,…,lm1，Persia将会尝试将Embedding分为k 组S1,S2,…,Sk，并分别存放在对应GPU 上，使得每组∑i∈Sjli,∀j 大致相等。这等价于如下优化问题：

minS1,…,SkVariancej[∑i∈Sjli],

s.t. ∑i∈SjVi≤C,

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

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

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

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

此外，在TensorFlow上，Persia还支持TensorFlow的"Replicated