TensorFlow在美团推荐系统中的分布式训练优化实践( 七 )


以Adam优化器为例 , 需要创建两个slot , 以保存优化中的动量信息 , 它的Shape与Embedding相同 。在原生优化器中 , 这两个Variable是单独创建的 , 并在反向梯度更新的时候会去读写 。同理 , 使用HashTable方案时 , 我们需要同时创建两张单独的HashTable用来训练m、v参数 。那么在前向 , 反向中需要分别对Embedding、 m、v进行一次Lookup和一次insert , 总共需要三次Lookup和三次insert 。
这里一个优化点就是将Embedding、 m、v , 以及低频过滤的计数器(见下图14的Counting HashTable)聚合到一起 , 作为HashTable的Value , 这样对稀疏参数的相关操作就可以聚合执行 , 大大减少了稀疏参数操作频次 , 降低了PS的压力 。

TensorFlow在美团推荐系统中的分布式训练优化实践


图14 基于HashTable的参数融合策略
该特性属于一个普适型优化 , 开启聚合功能后 , 训练速度有了显著的提高 , 性能提升幅度随着模型和Worker规模的变化 , 效果总是正向的 。在美团内部真实业务模型上 , 聚合之后性能相比非聚合方式能提升了45%左右 。
3.4.2 Embedding流水线优化
流水线 , 在工业生产中 , 指每一个生产单位只专注处理某个片段的工作 , 以提高工作效率及产量的一种生产方式 。在计算机领域内 , 更为大家熟知的是 , 流水线代表一种多任务之间Overlap执行的并行化技术 。例如在典型的RISC处理器中 , 用户的程序由大量指令构成 , 而一条指令的执行又可以大致分为:取指、译码、执行、访存、写回等环节 。这些环节会利用到指令Cache、数据Cache、寄存器、ALU等多种不同的硬件单元 , 在每一个指令周期内 , 这5个环节的硬件单元会并行执行 , 得以更加充分的利用硬件能力 , 以此提高整个处理器的指令吞吐性能 。处理器的指令流水线是一套复杂而系统的底层技术 , 但其中的思想在分布式深度学习框架中也被大量的使用 , 例如:
如果将分布式训练简单的抽象为计算和通信两个过程 , 绝大多数主流的深度学习框架都支持在执行计算图DAG时 , 通信和计算的Overlap 。如果将深度模型训练简单的分为前向和反向 , 在单步内 , 由于两者的强依赖性 , 无法做到有效并行 , 字节BytePS[8]中引入的通信调度打破了step iteration间的屏障 , 上一轮的部分参数更新完毕后 , 即可提前开始下轮的前向计算 , 增强了整体视角下前反向的Overlap 。百度AIBox[9]为了解决CTR场景GPU训练时 , 参数位于主存 , 但计算位于GPU的问题 , 巧妙调度不同硬件设备 , 搭建起了主要利用CPU/主存/网卡的参数预准备阶段和主要利用GPU/NVLink的网络计算阶段 , 通过两个阶段的Overlap达到更高的训练吞吐 。
我们看到 , 在深度学习框架设计上 , 通过分析场景 , 可以从不同的视角发掘可并行的阶段 , 来提高整体的训练吞吐 。
对于大规模稀疏模型训练时 , 核心模型流程是:先执行稀疏参数的Embedding , 然后执行稠密部分子网络 。其中稀疏参数Embedding在远端PS上执行 , 主要耗费网络资源 , 而稠密部分子网络在本地Worker执行 , 主要耗费计算资源 。这两部分占了整个流程的大部分时间 , 在美团某实际业务模型上分别耗时占比:40%+、50%+ 。

推荐阅读