前言
这段时间参加了天池上的 “DeepRec CTR模型性能优化” 比赛,通过阅读 DeepRec 官方文档,可以了解 DeepRec 做了哪些优化,哪些优化可以迁移借鉴,哪些优化是针对推荐系统的。这篇文章是对 DeepRec 文档的提炼总结,不涉及代码。
一,图优化
DeepRec 提供的图优化可以分成两类,一类是通过 “重叠通信、访存、计算” 来提高性能,一类是通过算子融合来减少调度和访存的开销。
1. Auto Micro Batch
思路:利用了流水线 pipeline 的思想。在图优化阶段,复制 N 个子图出来,重叠通信、访存、计算以提升性能,多个子图计算的 micro batch 对梯度进行累加之后再更新。 如下图,原本 4 个 micro batch 需要依次完成,利用流水线就可以并行执行。
效果:单机实验数据显示,在达到相同效果的前提下,2 pipeline 吞吐提升 30%,4 pipeline 吞吐提升 17%。
2. Embedding Fusion 算子
动机:“DeepRec 及 TensorFlow 原生的 embedding lookup 相关 API,如 safe_embedding_lookup_sparse,会创建比较多的 op,因此在 GPU 上执行时容易出现 kernel launch bound 的问题。” 另外,小 op 太多会增加调度的时间。
解决方法:提供 embedding fusion ops。在用户端的 embedding_column 提供一个 do_fusion
参数,判断是否要使用融合算子 fused_embedding_lookup_sparse
。如果使用这个算子,那么就会像下图这样产生几个小 op,分别进行前处理,收集数据,后处理,反向传播。
3. Stage 流水线 && Smart 自动流水线
动机:一个训练任务由数据读取和计算两部分组成,数据读取属于 IO bound 操作,不能很好利用资源,从而拖慢训练。
解决办法:DeepRec 提供数据 prefetch 功能,提供 staged 接口,使数据读取和计算异步执行。SmartStage 可以自动化的寻找最大可以 stage 的范围,减少用户显式操作,并且可以将寻找更小粒度的算子来异步执行。
效果:CPU 下 DLRM 每秒训练 steps 可以提升 212 步,GPU 场景需要使用专门针对 GPU 优化的选项,否则性能会降低。
4. Embedding Lookup 异步化
动机:在分布式训练的场景下,模型的特征增多,worker 在 ps 节点进行大量的 Embedding Lookup 操作,在整个端到端的耗时不断增长,成为模型训练速度的瓶颈,导致不能高效地利用计算资源。
解决办法:将 Embedding Lookup 部分子图划分出来,与计算主图异步执行,实现通信和计算重叠。使用需要开启 Stage,并且有一个 IO Stage,但是和 SmartStage 不能同时开启。
正面影响:在分布式场景下,通过通信和计算重叠,提高训练速度。
负面影响:可能影响收敛速度和模型精度。因为 Embedding Lookup 和梯度反向传播更新参数是异步执行的,所以可能不能获取到最新的参数。收敛速度如果减慢了,那么即使是训练吞吐量上去了,整体效果还是降低的。
5. Auto Graph Fusion 算子融合
问题:“大量小 Op 执行中调度开销占比高,导致整个图执行效率低”
思路:自动匹配计算子图,使用一个 Macro Op 替换计算子图,使用更大粒度的 Op 好处是减少调度和访存的开销。
文档中的例子:
二,Runtime 优化
DeepRec 提供的 Runtime 优化主要是内存/显存优化,调度优化。
1. CPU Memory Optimization
动机:“在 CPU 端,高性能 malloc 库存在大内存分配时带来的 minor page fault 严重导致的性能问题”
补充:page fault 有几个类型,major,minor,invalid。major page fault,指访问的内存不在虚拟空间里面,也不在物理空间,需要从慢速设备载入。minor page fault,指访问的内存不在虚拟空间里面,但在物理空间,只需要建立 MMU 物理内存和虚拟内存的映射,当调用 malloc 获取虚拟空间地址后,首次访问该地址会发生一次 soft page fault,或者进程访问共享内存中的数据,但是还没有建立映射。invalid page fault,访问的地址不在合理的范围内。
解决办法:收集内存分配信息,指定内存分配策略。提供两个参数,一个参数为收集内存信息的步数,一个参数为策略稳定运行达到的步数。
2. GPU Memory Optimization
动机:“在 GPU 端,TensorFlow 原生 BFCAllocator 存在显存碎片过多的问题,导致显存浪费。”
解决办法:“该功能在开始运行的前 K 个 step 收集运行信息并使用 cudaMalloc 分配显存,所以会造成性能下降,在收集信息结束后,开启优化,这时的性能会有所提升。”
影响:推荐模型 DBMTL-MMOE 的性能会微小下降 1%,显存占用大大减少 23%
3. GPU 虚拟显存
动机:显存不足的问题。
解决办法:“在显存不足时使用 cuda 的统一内存地址的 cuMemAllocManaged API 来分配显存,可以使用 CPU 内存来增加显存使用”
影响:可以用利用 CPU 内存,但是性能下降。
4. Executor 优化
思路:调度优化,优化 TensorFlow 原生的 executor 策略。DeepRec 提供了两种策略。第一,Inline Executor。支持 Session Run 在同一个线程上执行,减少线程切换,减少框架的 overhead,适用于 serving 高并发场景下以及部分 PS 场景。第二,基于 CostModel 的 Executor。统计与计算多组指标,通过 CostModel 计算出一个较优的调度策略。该功能中包含了基于关键路径的调度策略和根据 CostModel 批量执行耗时短的算子的调度策略。
影响:个人实测,在训练场景下,后两种策略都会变慢。Inline Executor 变慢是因为场景不适用。基于 CostModel 的 Executor 变慢的原因不详。
三,算子优化
DeepRec 中算子优化主要是:一,老老实实优化,这对开发人员的能力要求高,而且存在重复劳动的可能;二,对接硬件厂商的高性能计算库,比如 oneDNN,cudnn 等;三,充分利用新硬件及新特性,比如 PMEM 和 NVIDIA TF32。
1. 对接 oneDNN
方案:oneDNN 是 Intel 提供的深度学习性能加速库,在图优化阶段,可以将自己实现的算子,替换成 oneDNN 内支持的原语,从而将 oneDNN 的算子用起来。oneDNN 可以充分利用 CPU 提供的指令集加速算子计算,比如 AVX512。
2. 算子优化
例子:DeepRec 文档中给出了一个具体的算子优化实例。Select 算子在 TensorFlow 中默认的实现是使用 Eigen 库提供的原语拼出来,Broadcast + ElementWise Select。DeepRec 的优化是,直接在数据上操作,既然是 Select,那么直接 Row Select 即可,并且可以利用 AVX 512 指令来优化。
3. NVIDIA TensorFloat-32
方案:“TensorFloat-32,是 Nvidia 在 Ampere 架构的 GPU 上推出的专门运用于 TensorCore 的一种计算格式。” “请注意,TF32 仅仅是在使用 TensorCore 时的一种中间计算格式,它并不是一个完全的数据类型。”
影响:使用 TF32 可以大大提升速度,但是对精度基本无影响。
4. 持久内存
简介:持久内存在内存层级中位于内存和存储器之间,速度比内存慢,容量比内存大;速度比存储器快,容量比存储器小。内存访问延迟是几十纳秒,PMEM 是几百纳秒,SSD 是几十微秒,HDD 是几毫秒。
使用方法:
- 内存模式。将 PMEM 放在内存和存储器中间,作为存储器的缓存,对程序完全透明。
- 应用直接访问模式,存在两种用法。将 PMEM 配置为块设备,当存储器使用;将 PMEM 当内存使用,可视为一个 NUMA 节点。
DeepRec 中的用处:Embedding 表通常是非常大的,Embedding 可以简单理解为用户-特征矩阵,用户很多,算法人员不断扩展新特征,这样 Embedding 会随着时间不断变大,直到放不进内存,内存容量成为了瓶颈。所以引入 PMEM,可以提高大规模分布式训练的内存存储能力,降低总体拥有成本 TCO(Total Cost of Ownership)。
5. BFloat16
方案:Intel 处理器 Cooper Lake 架构支持 BFloat16(BF16),可以使用低精度来进行推理。使用方法是,将输入数据从 float32 转为 BFloat16,模型参数会在算子调用的时候,自动被转为 BFloat16,并且缓存在算子中。“为了减小这部分开销,DeepRec 会自动将 cast 算子和就近的算子进行融合,提高运行速度。” 比如将矩阵乘法 Matmul 和 Cast 融合。
影响:对 AUC 影响小,WDL 速度可以提升达到 1.42 倍。
(正在更新中......)
四,其他优化
1. 稀疏功能
2. 分布式训练
GRPC++
问题:TensorFlow 使用 GRPC 作为通信协议,在大规模训练场景下,通信和数据拷贝有额外的开销。
解决:DeepRec 在通信上做了若干优化:“Sharing Nothing 架构、BusyPolling 机制、用户态零拷贝、Send/Recv 融合等多种优化实现。”
StarServer
问题:“在超大规模任务场景下(几百甚至上千 worker),原生 tensorflow 框架中的一些问题被暴露出来,譬如低效的线程池调度,多处关键路径上的锁开销,低效的执行引擎,频繁的 rpc 带来的开销以及内存使用效率低等等。”
解决:第一,通信使用新的模型,将 send/recv 修改为 pull/push,支持子图划分;第二,图执行过程 lock free,提高并发执行子图的效率。
SOK
SparseOperationKit 是支持稀疏训练和推理的工具集,适用于大规模推荐系统,可以将推荐系统的 Embedding 切分到多个 GPU 上,即支持模型并行。SOK 还可以和常见的数据并行框架一同工作,比如 Horovod。因此可以通过 SOK 和 Horovod 实现混合并行。