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.

主要参考资料

相关推荐
奔跑中的小象1 小时前
统信UOS + 天数AI卡部署SGLang服务手册
人工智能·uos·sglang·天数智芯
DevSecOps选型指南1 小时前
中国版Mythos,为何是悬镜安全灵脉CodeAI?
人工智能·安全
Shockang1 小时前
LangGraph 状态机实战
人工智能
DLYSB_2 小时前
存储运维实战:基于 Ceph Event 监听与 Python 适配器的分布式存储健康度物理声光响应架构
运维·ceph·python·报警灯
刹那芳华19922 小时前
循环神经网络的从零开始实现(RNN)
人工智能·rnn·深度学习
阿童木写作2 小时前
跨境图片翻译工具多合一,批量图片视频字幕翻译加智能抠图
人工智能·python·音视频·语音识别
科技之门2 小时前
AI3D从建模到贴图、绑骨和动画的完整流程怎么做?V2Fun完整工作流指南
人工智能·3d·贴图
宇宙第一小趴菜2 小时前
二、机器学习的应用领域和发展史
人工智能·机器学习
前端开发江鸟2 小时前
我写过 MCP Server,却一直以为 MCP 只有 Tool
人工智能
天国梦2 小时前
自习室智能化升级避坑指南:天学网AI智习室方案实测与选型建议
大数据·人工智能