tensorflow 异步训练及其优化

目前遇到了 tensorflow 进行分布式训练中出现 worker 训练不均的情况，这里记录一下解决问题查找的一些资料和想法推测。

关于分布式的原理以及源码说明，可以参考最后的“主要参考资料”部分。

目前框架实现部分使用了 ParameterServerStrategy 分布式策略，主要对应的文章就是 TensorFlow 分布式之 ParameterServerStrategy V1 。

该算法默认采用了 Round-Robin算法, 链接当中有说明，那我们看一下这个调度算法是否会造成 worker分布不均的情况。

Round-Robin Scheduling

参考：wiki

时间片轮转调度(Round-Robin Scheduling) 是进程和网络调度程序常用的算法之一。这一方法将相等长度的时间片按照不变的顺序依次分配给每个进程[2]，且在处理所有进程时不考虑任何优先级。这一算法简单并易于实现，并且不会产生饥饿问题。时间片轮转调度可以应用于其他调度问题，例如计算机网络中的数据包调度。它是一个操作系统概念。

该算法的名称来自于其他领域通用的循环制原则，即每个参与者轮流获得相同分量的物品。

为了公平地调度进程，循环调度程序通常采用分时机制，为每个作业分配一个时间片或时间量（CPU 时间），如果用完这一分配的时间还没有完成，则中断该进程。下次为该进程分配时间时，该进程将恢复执行。如果进程在其时间片内终止或将其状态更改为等待（或阻塞），则调度程序会选择就绪队列中的第一个进程来执行。

循环算法是一种抢占式算法，因为一旦时间片用尽，调度程序就会强制性的暂停进程的执行。

例如，如果时间片为100毫秒，而进程1完成的总时间为250毫秒，则循环调度程序将在100毫秒后暂停该进程，并让其他进程在CPU上占用时间。一旦其他线程都使用过一次相同的时间片（100毫秒），进程1将获得另一次CPU时间分配。这个过程一直将持续循环到进程结束并且不需要更多的CPU时间。

美团技术团队也有一篇文章提到了这部分的优化，TensorFlow在推荐系统中的分布式训练优化实践，这篇文章是这么理解的：

这部分优化，是分布式计算的经典优化方向。PS架构是一个典型的“水桶模型”，为了完成一步训练，Worker端需要和所有PS完成交互，因此PS之间的平衡就显得非常重要。但是在实践中，我们发现多个PS的耗时并不均衡，其中的原因，既包括TensorFlow PS架构简单的切图逻辑（Round-Robin）带来的负载不均衡，也有异构机器导致的不均衡。
对于推荐模型来说，我们的主要优化策略是，把所有稀疏参数和大的稠密参数自动、均匀的切分到每个PS上，可以解决大多数这类问题。而在实践过程中，我们也发现一个比较难排查的问题：原生Adam优化器，实现导致PS负载不均衡。下面会详细介绍一下。在Adam优化器中，它的参数优化过程需要两个β参与计算，在原生TensorFlow的实现中，这两个β是所有需要此优化器进行优化的Variabl（或HashTable）所共享的，并且会与第一个Variable（名字字典序）落在同一个PS上面，这会带来一个问题：每个优化器只拥有一个β_1和一个β_2，且仅位于某个PS上。因此，在参数优化的过程中，该PS会承受远高于其他PS的请求，从而导致该PS成为性能瓶颈。

所以，这里的 round-robin 调度算法极有可能造成 worker 之间利用率不均匀的情况。要解决这个问题，

1. 可以更换调度算法
2. 对优化器部分进行优化

Dataset 与 变量分片

这里引用一下 tensorflow 当中对变量分片的解释：

变量分片是指将一个变量拆分为多个较小的变量，这些变量称为分片。在访问这些分片时，变量分片可能有助于分配网络负载。这对在多个参数服务器之间分布计算和存储普通变量也很有用，例如，当使用可能不适合单个机器内存的非常大的嵌入时。

要启用变量分片，您可以在构造 ParameterServerStrategy 对象时传入 variable_partitioner。每次创建变量时都会调用 variable_partitioner，它预计会返回该变量每个维度上的分片数。提供了一些开箱即用的 variable_partitioner，例如 tf.distribute.experimental.partitioners.MinSizePartitioner。建议使用基于大小的分区程序（如 tf.distribute.experimental.partitioners.MinSizePartitioner）以避免对小变量进行分区，否则可能会对模型训练速度产生负面影响。

steps_per_exectuion

在模型构建过程中会使用这个函数，大意为每个 worker在执行过程中所每次执行所运算的 batch 数量，经过实践证明，这个参数调小之后会使得 worker cpu 利用率尖峰数量变多更加密集，缩短了执行时间，与参数服务器的交换速度加快。

但是并不能解决尖峰的问题。

这里是 keras 官方文档 解释：

steps_per_execution: Int. The number of batches to run during each a single compiled function call. Running multiple batches inside a single compiled function call can greatly improve performance on TPUs or small models with a large Python overhead. At most, one full epoch will be run each execution. If a number larger than the size of the epoch is passed, the execution will be truncated to the size of the epoch. Note that if steps_per_execution is set to N, Callback.on_batch_begin and Callback.on_batch_end methods will only be called every N batches (i.e. before/after each compiled function execution). Not supported with the PyTorch backend.

主要参考资料

• 随笔分类 - 011_分布式机器学习
• TensorFlow在推荐系统中的分布式训练优化实践