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.

主要参考资料

相关推荐
大数据魔法师2 小时前
Streamlit(二十三)- 教程(二)- 动态导航
python·web
HackTorjan4 小时前
2026年5月29日:全球首个通用人工智能操作系统正式发布,开启人机协同新纪元
人工智能
刘大猫.4 小时前
智造短剧新引擎:火山引擎上线「火山剧创 1.0」,制作效率提升 80%
人工智能·ai·chatgpt·机器人·大模型·火山引擎·短剧新引擎
AI人工智能+电脑小能手4 小时前
【大白话说Java面试题 第87题】【Mysql篇】第17题:分布式事务的实现原理?
java·数据库·分布式·mysql·面试
红尘散仙5 小时前
我把终端小说阅读器接上了 AI Agent:TRNovel 现在能用 skill 生成书源了
人工智能·后端·rust
雅菲奥朗5 小时前
企业级 AI 自动化|OpenClaw 龙虾实战与认证
运维·人工智能·自动化·openclaw
HIT_Weston5 小时前
99、【Agent】【OpenCode】task 工具提示词(Slash command)(一)
人工智能·agent·opencode
25 Hz5 小时前
Mind 爱好者时空表征刊 第24期 | 时间结构学习、空间对时间表征的补偿、事件内部的时间扭曲……
人工智能
心中有国也有家5 小时前
GE图引擎深度解析——CANN的计算图优化与执行引擎
人工智能·pytorch·python·学习·numpy
海兰5 小时前
【文字三国志:第一篇】天命重构,大语言模型(LLM)动态生成文言风格的叙事文本的文字游戏
人工智能·游戏·语言模型