目前遇到了 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 之间利用率不均匀的情况。要解决这个问题,
- 可以更换调度算法
- 对优化器部分进行优化
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.