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.

主要参考资料

相关推荐
咸鱼桨3 分钟前
《庐山派从入门到...》PWM板载蜂鸣器
人工智能·windows·python·k230·庐山派
强哥之神15 分钟前
Nexa AI发布OmniAudio-2.6B:一款快速的音频语言模型,专为边缘部署设计
人工智能·深度学习·机器学习·语言模型·自然语言处理·音视频·openai
yusaisai大鱼19 分钟前
tensorflow_probability与tensorflow版本依赖关系
人工智能·python·tensorflow
18号房客19 分钟前
一个简单的深度学习模型例程,使用Keras(基于TensorFlow)构建一个卷积神经网络(CNN)来分类MNIST手写数字数据集。
人工智能·深度学习·机器学习·生成对抗网络·语言模型·自然语言处理·tensorflow
Biomamba生信基地22 分钟前
R语言基础| 功效分析
开发语言·python·r语言·医药
神秘的土鸡26 分钟前
神经网络图像隐写术:用AI隐藏信息的艺术
人工智能·深度学习·神经网络
数据分析能量站27 分钟前
神经网络-LeNet
人工智能·深度学习·神经网络·机器学习
Jaly_W36 分钟前
用于航空发动机故障诊断的深度分层排序网络
人工智能·深度学习·故障诊断·航空发动机
CodeClimb37 分钟前
【华为OD-E卷-木板 100分(python、java、c++、js、c)】
java·javascript·c++·python·华为od
小嗷犬38 分钟前
【论文笔记】Cross-lingual few-shot sign language recognition
论文阅读·人工智能·多模态·少样本·手语翻译