腾讯太极机器学习平台|大规模训练加速框架Light 在广告粗排场景的落地

背景介绍

太极机器学习平台由腾讯云机智平台和tesla平台协同共建而成，太极联合团队在深度学习训练加速上有深厚的技术累积，曾两次刷新了 ImageNet 训练速度的世界记录，并发表相应论文。

为使团队沉淀的训练加速技术赋能鹅厂更多业务场景并创造更大价值，Light 训练产品应运而生。Light 是云帆Oteam基于当前社区主流深度学习框架开发的一套多机多卡深度学习训练加速框架，用户只需要做几行代码即可接入并获得高性能加速能力。

从去年开始，太极团队针对广告训练场景进行了专项的性能优化，并针对业务模型迭代中遇到的痛点问题，基于 Light 通用框架融合数十种广告场景的技术和功能，打造了专门服务于广告粗排、预排序、召回场景的 Light 广告训练框架。

该框架主要具备以下特点：

• 易用性：兼容原 nabu 粗排系统的使用习惯，性能优化对用户层透明，并支持了一系列易用性功能以提高算法用户的研发效率。
• 可运营：集成了定时调度、指标监控、训练告警、数据统计等功能，框架版本持续快速迭代。
• 保持接口和模型文件兼容社区 Tensorflow：用户可使用社区主流接口进行模型定义，无须任何转换即可上线在线服务。
• 高性能：在原训练框架的基础上将训练吞吐提高了数倍，并行扩展率接近线性，不断打磨突破性能瓶颈。

Light 广告训练框架致力于为广告业务同学提供可迭代、好运营、易用、高效的训练能力，当前已助力广告粗排、预排序和召回业务在几十个模型上进行算法和模型实验并上线、取得了较为显著的 GMV 提升。

架构层次

在太极机器学习平台上，可通过拖拉拽的方式使用 Light 广告训练组件创建一个模型训练 pipeline 的工作流任务，用户只需要将相关的配置、代码文件上传，即可启动对应模型的训练：

（某广告 LiteCTR 模型的工作流界面）

训练启动后，可在太极页面上快捷地查看训练指标的变化趋势：

从实现架构上看，Light 广告训练框架是构建于 Light 通用框架 API 之上的业务专用训练框架：    

• 在接入层中，为适配原 nabu 系统的多种训练模式，团队从所有广告粗排/预排序/召回的训练流程中提取出了通用的过程，并使用基类将其封装为抽象接口，现有的训练模式只需要继承该基类并实现各自的训练逻辑即可：

• 在引擎层的加速模块中，lightIO、lightCalc、lightCC 以及 lightET 通过封装、劫持底层计算框架等方式，为上层训练提供了易用甚至透明的训练优化技术。
• 在引擎层的计算框架模块中，包含了当前 Light 支持的计算框架 Tensorflow/TTensorflow 和 Pytorch/TPytorch，广告推荐当前使用的是 Tensorflow 框架进行训练。在性能优化与落地的过程中，我们也沉淀了许多针对框架本身进行优化的优化代码，并将部分通用的优化向社区进行 PR 贡献，合到了社区 master 分支中。

训练加速技术

通过 out-of-the-box 的训练性能提升为广告推荐模型创建更大的业务价值，是 Light 广告训练框架的核心出发点。以下列出了在组件迭代过程中，Light 广告加速组件包含的几项主要训练加速技术。

3.1. I/O 优化

老训练系统使用独立的下载进程将训练数据分批下载到本地。在这种模式下，由于模型结构变化、数据文件大小的变化、HDFS 集群的负载状态变化等原因，I/O 耗时无法被很好的 overlap，在训练过程中经常出现由于数据下载导致训练卡住的情况。

为了提升 I/O 性能，Light 组件首先改变了数据下载的模式：采用数据流式下载不落盘的方式，即数据直接下载到内存，从而减少了不必要的内存和磁盘操作。其次，针对训练数据异地导致下载缓慢的问题，结合 tensorflow 框架进行下载 buffer 优化、多级缓存预取优化，最终在训练中基本隐藏了 I/O 的耗时影响。

3.2. 通信优化

在原 nabu 训练系统中，只支持 PS 模式的多卡训练，而这种通信模式由于存在单点问题，同步训练的瓶颈比较明显，扩展性仅有 25% 左右。而在 Light 广告组件中，底层通信依赖于云帆 Oteam 通信组件 LightCC，使用 Ring AllReduce 同步规约通信：

在这种模式下，每个 worker 上有全量的参数，各个 worker 完成一个 batch 数据的前向计算并得到所有 variables 的相应梯度后，基于 NCCL 进行梯度规约通信，并将获得的梯度更新到本地的参数上。在当前模型规模不大（百兆级别）的粗排场景，Ring AllReduce 去中心化的调度模式有效地规避了老 nabu 框架中传统 PS 模式的单点瓶颈问题，大大提升了训练的并行扩展性。

除此之外，借助稠密通信、梯度融合、多流通信等优化方式，Light广告组件在典型模型上将多机扩展性提升到了近 90%。

3.3. 计算优化

典型的训练流程包括了前向计算、反向计算和梯度规约，这里计算优化包含的是前向、反向计算中所涉及的多种优化思路。

• 算子优化：

由于广告推荐是一个 CPU bound 的训练场景，依赖于大量的 CPU 资源来做数据解析，若训练相关计算占用了大量 CPU 资源，则会让解析数据用 CPU 资源变少，从而拖慢训练速度，因此算子优化的基本思路，就是将计算任务交由 GPU 来完成，充分利用 GPU 的并行计算能力，同时解放 CPU 来完成数据处理的工作。

在这种基本思路的指导下，我们通过等价算子替换、实现新的稀疏算子、优化部分已有算子等方式，在广告粗模型上取得了符合预期的明显效果；同时也向社区贡献了不少优化 PR。

• 整型特征拷贝优化：

由于整型到字符串类型的转化运算需要在 CPU 上完成，且在一个 string Tensor 中，各个元素的地址相差很大，散落在非连续空间中，进而导致存在大量的 glibc memcpy 调用，占用大量 CPU 资源。因此，为避免在 Host 内存中碎片化拷贝导致的性能问题，我们选择直接将整数特征拷贝到 GPU 显存，然后用 GPU 替代 CPU 完成 AsString 和 Hash 的计算，从而避免碎片化拷贝问题。

• 字符串特征 Hash 优化：

在训练时，我们需要将数据文件解压读取，并依次地将整型和字符串特征反序列化出来再进行 Hash。对于整型特征 Hash，可以用 GPU 来实现 atoi 的过程，但由于 string 特征本来的地址已经是离散的，无法使用和整数特征相同的方式来做优化；因此，我们通过 custom OP 的方式，使用 Variant 构造了自定义 Tensor 类型，直接实现将字符串拷贝到 GPU 显存做进一步运算。

• Embedding Fusion：

在推荐模型中，通常会对每个特征类单独进行 embedding，得到结果向量后进行拼接再送入 DNN；其中，每个 embedding 操作都会产生相应的 gather/lookup ops。

一方面，多个 embedding table 的 lookup 算子会带来多次 Op/Kernel 启用的 overhead；另一方面，经过实际测试，随着输入 size 的增加，embedding lookup 的执行耗时呈现对数增长趋势。因此，我们只需要维护好每个特征对应的索引位置，就可以将多个 embedding table 进行 fusion，只进行一次 lookup 操作，从而减少大量的算子调度开销和 host-device 拷贝开销。

• T4 机器上线训练：

针对广告推荐场景中明显 CPU bound 的情况，团队推动了 T4 推理机器的上线和置换，当使用拥有更多 CPU 核心数的机型进行训练时，线程间 CPU 抢占的情况能有所缓解，CPU 瓶颈造成的性能瓶颈也能得到突破。

经过测试，在几个典型的召回、预排序、粗排模型上，使用 188 核心的 T4 推理机型进行训练时，4 卡 T4 的训练性能与 8 卡 P40 训练性能基本持平。

4. 易用性提升

Light 广告训练组件在粗排模型上落地实践的过程中，在业务需求的驱动下，经过持续近一年的迭代，团队也不断丰富了许多扩展功能，力图提升组件的易用性。

4.1. 使用模式

由于原有广告训练系统中不同模型使用的构图配置模式、数据配置的解析规则以及具体训练流程都各不相同，为降低广告算法同学的使用成本和框架的维护成本，我们在 light 广告训练组件中将原有的多种训练模式中能相互兼容的模式进行融合统一；除此之外，还提供了更加灵活的自定义构图模式。

• 使用 study_option.pbtxt 配置构图的训练模式

使用这种配置模式构图时，用户通过 protobuf 配置定义所需的输入特征信息（包括名称、类型、定长或变长等）、各种预定义 OP 的输入输出关系，最终在框架中经过拓扑排序进而组成完整的 graph。

e.g.

op {  name: "ad_industry_id_concat"  concat {    operand: "ad_industry_id_embedding_0"    operand: "ad_industry_id_embedding_1"    axis: 1  }}

在这种模式下，不同模型使用相同的算子库，管理维护都比较方便；但当模型变得越来越复杂时，成千上万行的配置文件使可读性急剧下降；用户无法使用条件判断、循环控制等基本编程语义，难以实现更灵活的模型逻辑；此外，若模型所需逻辑在预定义算子库中缺失，只能在框架层面添加新的算子，生产效率较低。

• 使用 graph_meta.python 代码构图的训练模式

使用这种配置模式构图时，用户可以直接使用 python 编程语义，继承框架定义的 ModelBuilder 基类，使用原生的 Tensorflow 接口，来完成模型的构图。

e.g.

model_builder = DeepShallowModelBuilder(nodes, env)# user code# ....model_builder.deep_shallow_model(...)graph_builder = GraphBuilder(...)graph_builder.add_calculate_nodes(...)graph_builder.add_serving_nodes(...)graph_builder.build(...)

在这种模式下，算法用户可以使用原生 tensorflow 接口来完成模型的定义，使用更为灵活。但由于构图还是要依赖于框架提供的预定义函数来实现，并按照约定创建 ModelBuilder 和 GraphBuilder 对象，因此用户还是需要先理解框架层面的代码结构和约束条件，才能较好的构建自定义模型；且当框架代码迭代升级时，用户的模型代码也需要做相应的改动，明显存在不便。

• 自定义代码构图

以上两种使用模式，主要是迁移自老系统的基本配置方案，虽在 Light 组件中做了易用性优化，但由于提及的种种问题，也逐渐无法满足部分算法工程师日益复杂的模型迭代需求。因此我们提供了更加灵活的自定义构图模式，进一步解耦算法与训练。

e.g.

import lightimport tensorflow as tfimport user_code_1import user_code_2def build_graph(*args, **kwargs):# ...return [train_op_1, train_op_2], [a_loss, g_loss], [metric_op]light.run(build_graph)

使用这种模式时，用户可以完全自由地编写模型代码，几乎不存在任何跟框架相关的依赖或约束，当完成模型构图后，只需 import 框架的指定接口，并传递训练时要执行的必要 op（如：init op、train op 等）即可调用 light 训练框架在平台上启动分布式模型训练。框架与平台提供的业务接入、数据拉取、性能优化、资源调度等工程能力，会通过透明劫持或 API 封装的方式，提供给模型层自主调用。

4.2. 实时训练

Light 广告训练组件最初主要落地于离线训练的场景，但在算法模型迭代的过程中，业务提出了分钟级更新的训练延迟要求；若通过平台的定时任务进行调度，考虑到平台资源分发、框架初始化的耗时 overhead，难以满足分钟级的更新延迟。因此，我们在原有实现架构的基础上进行扩展，支持了使用实时数据进行在线训练的 streaming 训练模式，并通过一系列针对性延时优化，将模型的更新延时有效控制在了分钟级别。

4.2.1. 方案设计

Light 框架的训练模式是基于 MPI 的 Ring-AllReduce 分布式同步训练，不同于原 nabu 系统采用单进程 PS 架构，因此在实现实时训练时面临着新的困难挑战。

1.为了支持流式训练，需要在每轮训练数据都消耗完后，立即使用新生产的数据开启新一轮训练。最容易想到的实现方案是在一轮训练结束后，直接让所有进程退出，再从 MPI 外层将所有 ranks 重新拉起，使用新生产的数据开启新一轮训练：

这种方式实现逻辑简单，对原有框架改动也非常小；且能保证所有 ranks 在训练时能保证所用数据的时间窗口完全一致。但其最大的问题在于反复拉起 MPI 训练程序、训练初始化阶段所产生的 overhead 耗时，特别在数据分钟级更新，每轮训练的实际训练时长都非常短时，这些 overhead 占总耗时的占比将变得非常大，无法满足模型分钟级更新的要求。

2.因此我们考虑了另一种思路，即在一轮训练结束后，当某个 rank “率先”消耗完所有训练样本时，不直接退出进程，而是立即拉取新生产的数据无缝的“接上”继续训练，而其它 ranks 同理：

这种方式规避了反复启动 MPI 训练程序带来的 overhead 耗时，在数据生产效率有保证的前提下，可以实现模型更新的分钟级延迟；且训练效率高，训练样本的利用率也得到了充分提升，i.e. 原本在某个 rank 训练结束后，由于样本分配不均，其它 ranks 只能”抛弃“的少量样本在这种方案下都能被有效的利用起来。

因此，该方案最终被敲定为 Light 流式训练的实现模式。

4.2.2. 延时优化

• 数据异步拉取

通过分析发现，在每轮训练中，最耗时的 overhead 操作就是将路径 patterns 转换为 files 路径的过程，该操作阻塞了单个 rank 将近一分钟，导致整体的同步训练都在等待。

实际上对于每个 rank，都只是需要得到当前最新的数据情况，因此我们在 local chief 中创建一个全局的数据分配器，使用一个独立的后台 daemon 线程实时拉取最新的数据情况缓存到本地，从而规避多个 ranks 各自去做匹配对整体训练造成的阻塞耗时。

• 异步并行推送

每轮训练结束之后，chief rank 需要进行模型的前向图导出、打包、上传到 HDFS 这一系列操作，当模型文件较大或网络延迟明显时，这些操作会造成较大的阻塞耗时。

由于新一轮的训练实际上并不依赖于模型上传这个操作，因此我们将原有同步阻塞的模型上传操作，修改为异步完成，触发上传逻辑后该 rank 立即返回，随即进入下一轮训练；同时将所有文件并行上传，充分利用网络带宽。

• 定时推送

应用以上优化措施后，在样本数据流稳定的情况下，模型更新延迟已能被稳定控制在分钟级别。但由于使用的是实时数据流，流量高峰时相同时间窗口大小内的数据量可达非高峰时段的几十倍，导致这种情况下模型更新的时长被大大延长。

因此，我们通过模型定时推送来规避在不同时间段相同时间窗口大小的数据量差异对模型更新延迟造成的影响，进一步保障了在线模型的稳定更新。

通过以上延时优化，在一个典型的业务模型上，顺利将模型更新延迟控制在 1-2 分钟，实现了分钟级别的实时训练。

4.3. Tensorboard 支持

TensorBoard 是 Tensorflow 生态中一个强大的 debug/profling 工具，因此我们在 Light 广告训练组件中也支持了 tensorboard 的使用。

通过简单的配置启用 tensorboard 后，Light 会默认为模型的添加 batch_loss、mean_loss 图表，用户也可以使用原生的 tensorboard python api 来定义任意自定义的 summary，在模型训练启动后即可在太极平台上打开 tensorboard 视图，实时查看模型的训练状态：

4.4. 训练时 Evaluation

原系统下离线训练 (train) 与预估 (evaluate/calculate) 是分开进行的模式，但在落地过程中算法同学提出了在训练进行期间使用测试集进行 evaluate 的迫切需求，因此我们也进行了专项支持。

我们在 tensorflow 的 graph 中同时构造了两个 dataset pipeline，并在训练期间通过数据集和切换来分别执行 train和 evaluate。因此考虑的主要问题在于如何控制 dataset 的切换，考虑以下方式：

• 使用条件判断 OP 结合 placeholder 控制进行切换，但由于引入了分支判断，会在性能上带来一定影响。
• 使用 tf.data.Iterator.from_string_handle 进行切换，结合 placeholder 创建 feedable iterator，最后在实际训练时指定使用哪一个 handle 即可实现数据集的切换。

经过实测，这种切换方式对 I/O 性能几乎没有影响，因此最终选择了这种方式进行实现。

完成该功能的支持后，用户仅需配置好 evaluation 的数据集，即可使用该功能。若在启用 evaluate 功能的同时启用 tensorboard，组件还会将相同 tensor 分别在 train 数据集和 test 数据集下的变化曲线放在同一张图表中，以方便用户进行对比：

5. 落地情况

目前在广告召回、预排序、粗排推荐链路的模型中，已有 50% 的模型使用 Light 广告加速组件结合业务特性（扩充特征、增加数据，优化模型等) 全量上线，40+ 模型实验还在灰度当中，助力线上业务效果取得明显的提升。

后续腾讯太极团队会进一步打造组件的训练性能和应用性，将分布式 GPU 训练技术应用到更多推荐场景中，与鹅厂更多业务合作整体提升业务更大的价值。

近期热文推荐

  你“在看”我吗？