作者：烟秋

经历 6 年时间，在各团队的努力下，阿里巴巴集团大规模稀疏模型训练/预测引擎 DeepRec 正式对外开源，助力开发者提升稀疏模型训练性能和效果。

DeepRec 是什么

DeepRec(PAI-TF)是阿里巴巴集团统一的大规模稀疏模型训练/预测引擎，广泛应用于淘宝、天猫、阿里妈妈、高德、淘特、AliExpress、Lazada 等，支持了淘宝搜索、推荐、广告等核心业务，支撑着千亿特征、万亿样本的超大规模稀疏训练。

DeepRec 在分布式、图优化、算子、Runtime 等方面对稀疏模型进行了深度性能优化，同时提供了稀疏场景下特有的 Embedding 相关功能。

DeepRec 项目从 2016 年开发至今，由阿里巴巴集团内 AOP 团队、XDL 团队、PAI 团队、RTP 团队以及蚂蚁集团 AIInfra 团队共建，并且得到了淘宝推荐算法等多个业务算法团队的支持。DeepRec 的研发也得到了 Intel CESG 软件团队、Optane 团队和 PSU 团队，NVIDIA GPU 计算专家团队及 Merlin HughCTR 团队的支持。

DeepRec 架构设计原则

在 TensorFlow 引擎上支持大规模稀疏特征，业界有多种实现方式，其中最常见的方式是借鉴了 ParameterServer 的架构实现，在 TensorFlow 之外独立实现了一套 ParameterServer 和相关的优化器，同时在 TensorFlow 内部通过 bridge 的方式桥接了两个模块。这个做法有一定的好处，比如 PS 的实现会比较灵活，但也存在一些局限性。

DeepRec 采取了另一种架构设计方式，遵循“视整个训练引擎为一个系统整体”的架构设计原则。TensorFlow 是一个基于 Graph 的静态图训练引擎，在其架构上有相应的分层，比如最上层的 API 层、中间的图优化层和最下层的算子层。TensorFlow 通过这三层的设计去支撑上层不同 Workload 的业务需求和性能优化需求。

DeepRec 也坚持了这一设计原则，基于存储/计算解耦的设计原则在 Graph 层面引入 EmbeddingVariable 功能；基于 Graph 的特点实现了通信的算子融合。通过这样的设计原则，DeepRec 可以支持用户在单机、分布式场景下使用同一个优化器的实现和同一套 EmbeddingVariable 的实现；同时在 Graph 层面引入多种优化能力，从而做到独立模块设计所做不到的联合优化设计。

DeepRec 的优势

DeepRec 是基于 TensorFlow1.15、Intel-TF、NV-TF 构建的稀疏模型训练/预测引擎，针对稀疏模型场景进行了定制深度优化，主要包含以下三类功能优化：

image1.png

模型效果

DeepRec 提供了丰富的稀疏功能支持，提高模型效果的同时降低稀疏模型的大小，并且优化超大规模下 Optimizer 的效果。下面简单介绍 Embedding 及 Optimizer 几个有特色的工作：

• EmbeddingVariable（动态弹性特征）：

1）解决了静态 Shape Variable 的 vocabulary_size 难以预估、特征冲突、内存及 IO 冗余等问题，并且在 DeepRec 中提供了丰富的 EmbeddingVariable 的进阶功能，包括不同的特征准入方式、支持不同的特征淘汰策略等，能够明显提高稀疏模型的效果。

image2.png

2）在访问效率上，为了达到更优化的性能和更低的内存占用，EmbeddingVariable 的底层 HashTable 实现了无锁化设计，并且进行了精细的内存布局优化，优化了 HashTable 的访问频次，使得在训练过程中前后向只需访问一次 HashTable。

• DynamicDimensionEmbeddingVariable（动态弹性维度）：

在典型的稀疏场景中，同类特征的出现频次往往极度不均匀。通常情况下，同一个特征列的特征都被设置成统一维度，如果 Embedding 维度过高，低频特征容易过拟合，而且会额外耗费大量内存；如果维度设置过低，高频部征特征可能会由于表达不够而影响效果。

Dynamic Dimension Embedding Variable 提供了同一特征列的不同特征值，根据特征的冷热自动配置不同的特征维度，高频特征可以配置更高维度增强表达能力，而低频特征因为给定低维度 embedding 缓解了过拟合的问题，而且可以极大程度节省内存（低频长尾特征的数量占据绝对优势）。

image3.png

• Adaptive Embedding（自适应 Embedding）：

当使用动态弹性特征功能时，低频特征存在过拟合问题。EmbeddingVariable 中所有的特征都从 initializer 设定的初始值（一般设为 0）开始学起，对于一些出现频次从低到高的特征，也需要逐渐学习到一个较好的状态，不能共享别的特征的学习结果。AdaptiveEmbedding 功能使用静态 Shape Variable 和动态 EmbeddingVariable 共同存储稀疏特征，对于新加入的特征存于有冲突的 Variable，对于出现频率较高的特征存于无冲突的 EmbeddingVariable，特征迁移到 EmbeddingVaraible 可以复用在有冲突的静态 Shape Variable 的学习结果。

image4.png

• Adagrad Decay Optimizer：

为支持超大规模训练而提出的一种改进版 Adagrad 优化器。当模型训练的样本量大，同时持续增量训练较长时间时，Adagrad 优化器的梯度会趋近于 0，导致新增训练的数据无法对模型产生影响。已有的累积打折的方案虽然可以解决梯度趋近 0 的问题，但也会带来模型效果变差的问题（通过 iteration 打折策略无法反映实际的业务场景特点）。Adagrad Decay Optimizer 基于周期打折的策略，同一个周期内的样本相同的打折力度，兼顾数据的无限累积和样本顺序对模型的影响。

此外，DeepRec 还提供 Multi-HashEmbedding、AdamAsyncOptimizer 等功能，在内存占用、性能、模型效果等方面为业务带来实际的帮助。

训练性能

DeepRec 针对稀疏模型场景在分布式、图优化、算子、Runtime 等方面进行了深度性能优化。其中，DeepRec 对不同的分布式策略进行了深度的优化，包括异步训练、同步训练、半同步训练等，其中 GPU 同步训练支持 HybridBackend 以及 NVIDIA HugeCTR-SOK。DeepRec 提供了丰富的针对稀疏模型训练的图优化功能，包括自动流水线 SmartStage、结构化特征、自动图 Fusion 等等。DeepRec 中优化了稀疏模型中数十个常见算子，并且提供了包括 Embedding、Attention 等通用子图的 Fusion 算子。DeepRec 中 CPUAllocator 和 GPUAllocator 能够大大降低内存/显存的使用量并显著加速 E2E 的训练性能。在线程调度、执行引擎方面针对不同的场景提供了不同的调度引擎策略。下面简单介绍分布式、图优化、Runtime 优化方面几个有特色的工作：

• StarServer（异步训练框架）：

在超大规模任务场景下(几百、上千 worker)，原生开源框架中的一些问题被暴露出来，譬如低效的线程池调度、关键路径上的锁开销、低效的执行引擎、频繁的小包 rpc 带来的开销导致 ParameterServer 在分布式扩展时成为明显的性能瓶颈。StarServer 进行了包括图、线程调度、执行引擎以及内存等优化，将原有框架中的 send/recv 语义修改为 pull/push 语义，并且在子图划分上支持了该语义，同时实现了 ParameterServer 端图执行过程中的 lockfree，实现了无锁化的执行，大大提高了并发执行子图的效率。对比原生框架，能够提升数倍的训练性能，并且支持 3000worker 规模的线性分布式扩展。

image5.png

• SmartStage（自动流水线）：

稀疏模型训练通常包含样本数据的读取、Embedding 查找、Attention/MLP 计算等，样本读取和 Embedding 查找非计算密集操作，同时并不能高效利用计算资源（CPU、GPU）。原生框架中提供的 dataset.prefetch 接口可以异步化样本读取操作，但 Embedding 查找过程中涉及特征补齐、ID 化等复杂的过程，这些过程无法通过 prefetch 进行流水线化。SmartStage 功能能够自动分析图中异步流水线化的边界并自动插入，可以使并发流水线发挥最大的性能提升。

image6.png

• PRMalloc（内存分配器）：

如何做到既高效又有效的使用内存，对于稀疏模型的训练非常关键，稀疏场景模型训练中大块内存分配使用造成大量的 minor pagefault，此外，多线程分配效率存在比较严重的并发分配效率问题。针对稀疏模型训练前向、后向，Graph 计算模式的相对固定、多轮反复迭代的特点，DeepRec 设计了一套针对深度学习任务的内存管理方案，提高内存的使用效率和系统性能。使用 DeepRec 中提供的 PRMalloc 能够极大降低训练过程中 minor pagefault，提高多线程并发内存分配、释放的效率。

image7.png

• PMEM allocator（持久内存分配器）：

基于 PMDK 的底层 libpmem 库实现的 PMEM allocator 将从 PMEM map 出的一块空间分为若干 segment，每个 segment 又分成若干 blocks，block 是 allocator 的最小分配单元。分配 block 的线程为避免线程竞争，缓存一些可用空间，包括一组 segment 和 free list。可用空间中为每种 record size（若干个 block）维护一个 free list 和 segment。各 record size 对应的 segment 只分配该大小的 PMEM 空间，各 record size 对应的 free list 中的所有指针均指向对应 record size 的空闲空间。此外，为了均衡各 thread cache 的资源，由一个后台线程周期地将 thread cache 中的 free list 移动到后台的 pool 中，pool 中的资源由所有前台线程共享。实验证明，基于持久内存实现的内存分配器在大模型的训练性能方面与基于 DRAM 的训练性能差别很小，但是 TCO 会有很大的优势。

部署及 Serving

• 增量模型导出及加载：

时效性要求高的业务，需要频繁的线上模型更新，频率往往达到分钟级别甚至秒级。对于 TB-10TB 级别的超大模型而言，分钟级别的模型生成到上线很难完成。此外，超大模型的训练和预测存在着资源浪费、多节点 Serving 延时加大等问题。DeepRec 提供了增量模型产出及加载能力，极大加速了超大模型生成和加载。

image8.png

• Embedding 多级混合存储：

稀疏模型中特征存在冷热倾斜的特性，这产生了某些冷门特征很少被访问和更新导致的内存/显存浪费问题，以及超大模型内存/显存放不下的问题。DeepRec 提供了多级混合存储（支持最多四级的混合存储 HBM+DRAM+PMEM+SSD）的能力，自动将冷门特征存放到廉价的存储介质中，将热门特征存放到访问更快、更贵的存储介质上，通过多级混合存储，使得单节点可以进行 TB-10TB 模型的 Training 和 Serving。

image9.png

通过多级混合存储，能够更大发挥 GPU 训练稀疏模型的能力，同时降低由于存储资源限制造成的计算资源浪费，可以使用更少的机器进行相近规模的模型训练，或者使用相同数量的机器进行更大规模的训练。多级混合存储也能使得单机进行超大模型预测时避免分布式 Serving 带来的 latency 增大问题，提高大模型的预测性能的同时降低成本。多级混合存储功能也拥有自动发现特征的访问特性，基于高效的热度统计策略，将热度高的特征放置到快速的存储介质中，将低频的特征 offload 到低速存储介质中，再通过异步方式驱动特征在多个介质之间移动。

为什么开源 DeepRec

开源深度学习框架都不能很好地支持稀疏场景中对于稀疏 Embedding 功能的需求、模型训练性能需求、部署迭代和线上服务的需求。DeepRec 经过阿里巴巴集团搜索、推荐、广告等核心业务场景及公有云上各种业务场景的打磨，能够支持不同类型的稀疏场景训练效果和性能需求。

阿里巴巴希望通过建立开源社区，和外部开发者开展广泛合作，进一步推动稀疏模型训练/预测框架的发展，为不同业务场景中的搜推广模型训练和预测带来业务效果和性能提升。

今天 DeepRec 的开源只是我们迈出的一小步。我们非常期待得到您的反馈。最后，如果你对 DeepRec 有相应的兴趣，你也可以来转转，为我们的框架贡献一点你的代码和意见，这将是我们莫大的荣幸。

image10.jpeg

开源地址：https://github.com/alibaba/DeepRec

钉钉答疑交流群：34375411

想了解更多 AI 开源项目，请点击：https://www.aliyun.com/activity/bigdata/opensource_bigdata__ai