[源码解析] 深度学习分布式训练框架 horovod (12) --- 弹性训练总体架构
0x00 摘要
Horovod 是Uber于2017年发布的一个易于使用的高性能的分布式训练框架,在业界得到了广泛应用。
本系列将通过源码分析来带领大家了解 Horovod。本文是系列第十二篇,看看horovod 如何实施弹性训练。
弹性训练使得Horovod具备运行时worker数量动态伸缩,而不需要重启 或者 只是从存储中的checkpoint恢复训练。
本系列其他文章链接如下:
[源码解析] 深度学习分布式训练框架 Horovod (1) --- 基础知识
[源码解析] 深度学习分布式训练框架 horovod (2) --- 从使用者角度切入
[源码解析] 深度学习分布式训练框架 horovod (3) --- Horovodrun背后做了什么
[源码解析] 深度学习分布式训练框架 horovod (4) --- 网络基础 & Driver
[源码解析] 深度学习分布式训练框架 horovod (5) --- 融合框架
[源码解析] 深度学习分布式训练框架 horovod (6) --- 后台线程架构
[源码解析] 深度学习分布式训练框架 horovod (7) --- DistributedOptimizer
[源码解析] 深度学习分布式训练框架 horovod (8) --- on spark
[源码解析] 深度学习分布式训练框架 horovod (9) --- 启动 on spark
[源码解析] 深度学习分布式训练框架 horovod (10) --- run on spark
[源码解析] 深度学习分布式训练框架 horovod (11) --- on spark --- GLOO 方案
0x01 总述
1.1 问题点
我们思考下,Horovod 目前遇到了什么问题?
-
无法自动调节容量(Auto Scale)。
- 因为计算资源也许会有弹性调度,所以应该考虑到如果集群缩容了怎么办?如果扩容了怎么?理想状态应该是:在训练过程中可以自动增加或者减少worker数量。而且在worker数量变化时,不会中断训练任务,做到平滑过渡。
- 目前Horovod无法在资源有限的情况下执行。假如一共需要100个GPU,暂时只有40个GPU到位,在这种情况下,Horovod就只能等待,不能用现有的40个GPU先在少量进程上开始训练,从而无法快速开始模型迭代。
- 资源充裕时,Horovod 无法自动增加进程加速训练。就上例而言,在理想状态下,Horovoid应该先用这40个GPU构建一个环来启动训练,如果发现60个新GPU到位了就自动动态扩容,从而在下一个 epoch 开始就用100个GPU构建新的环开始训练;
- 没有容错机制(Fault Tolerance)。目前如果某一个节点失败,整个训练会失败,用户只能从头开始训练。如果可以支持 auto scale,加上一些之前陆续保存的 checkpoint,则Horovod可以重新选取一个好节点启动这个worker,或者用剩下的节点构建一个环继续训练。
-
调度机制不灵活。
- 机器学习训练任务一般时间较长,占用算力大,而Horovod任务缺少弹性能力,不支持动态配置 worker,不支持高优先级抢占实例。因此当资源不足时,无法按需为其他高优先级业务腾出资源, 只能等待任务自己主动终止或者出错终止。
为了解决以上几个问题,我们会思考很多的其他具体技术问题和细节,让我们先罗列出来:
- 何时构建 checkpoint?哪一个阶段是合适的?每一个 epoch 之后自动保存?还是由用户自行控制(这样可以做到更好的)?
- 如何从 checkpoint恢复?
- checkpoint需要存储哪些东西,即,对于horovod来说,哪些状态是必须的?
- 如何监听 worker 的工作情况?怎么判断机器出了问题?假如只是网络阻塞偶尔导致的怎么办?
- 需要构建一个通知机制;
- 如何知道集群的富余资源?如何发现可用节点?
- 如何构建新的通信环 ring?
- 如果构建新ring,是由一个 master 完成?还是使用类似 gossip 这样的协议?
- 是否有优先级调度,这样可以充分利用共享集群资源空闲的资源。
- 新 worker 怎么被 sync?
- 原有的active worker 节点怎么处理?
- 出问题的 worker 节点怎么处理?
- rank 0 怎么广播?
我们在本文以及后续各篇的分析中试着解答这些问题。
注:Horovod目前的调度机制依然不灵活,不支持抢占。
1.1 角色
Horovod 在单机的多个 GPU 之上采用 NCCL 来通信,在多机(CPU或者GPU)之间通过 Ring AllReduce 算法进行通信。Horovod 的弹性训练是指多机的弹性训练。
Horovod 弹性训练有两个角色:driver和 worker。driver 进程运行在 CPU 节点上,worker 进程可以运行在 CPU 节点或者 GPU 节点之上。
Driver 进程的作用是:
- 调用 Gloo 帮助 workers 构造一个 AllReduce 通信环,或者说是通信域。Driver 不参与具体构建通信环,而是提供辅助信息,从而worker可以建立环。
- Driver 进程需要给 Gloo 创建一个带有 KVStore 的 RendezvousServer,其中 KVStore 用于存储通信域内每个节点的 host 和 其在逻辑通信环分配的序号 rank 等信息。
- 这个 RendezvousServer 运行在 Horovod 的 driver 进程里。driver 进程拿到所有 worker 进程节点的地址和 GPU 卡数信息后,会将其写入RendezvousServer 的 KVStore 中,然后 worker 就可以调用 gloo 来访问 RendezvousServer 构造通信环。
- Driver 会在 worker 节点上启动/重启 worker 进程。
- Driver 会监控系统整体状态。
worker 负责训练和模型迭代。
- 每个 worker 节点会向 RendezvousServer 发起请求来得到自己的邻居节点信息,从而构造通信环。
- 在这个通信环之中,每个 worker 节点有一个左邻居和一个右邻居,在通信过程中,每个 worker 只会向它的右邻居发送数据,只会从左邻居接受数据。
具体组网机制如下:
+-------------------------------+
| Driver |
| |
| +------------------------+ |
| | RendezvousServer | |
| | | |
| | | |
| | host1, host2, host3 | |
| +------------------------+ |
+-------------------------------+
^ ^ ^
| | |
| | |
+-------------+ | +--------------+
| | |
| | |
| | |
v v v
+--------+----+ +-------+------+ +----+--------+
| Worker | | Worker | | Worker |
+------> | +------> | +---------> | | +------+
| | host1 | | host2 | | host3 | |
| +-------------+ +--------------+ +-------------+ |
| |
| |
| v
<--------------------------------------------------------------------------------+
我们下面详细分析下各个部分。
1.2 容错机制
Horovod 的容错机制是基于 gloo 来实现的,对于错误来说,这可以被认为是一个被动操作。
Gloo 本身是不支持容错的。当众多worker之间对张量进行聚合操作时候,如果某一个worker失败,则gloo不会处理异常,而是抛出异常并且退出,这样所有worker都会报异常退出。
为了不让某一个 worker 的失败导致整体训练退出,Horovod 需要做两方面工作:
-
不让异常影响现有作业。Horovod 必须捕获 gloo 抛出的异常,于是就构建了一个python处理异常机制。
- Worker 在捕获异常之后会将异常传递给对应的 Python API 处理,API 通过判断异常类型决定是否继续训练。
- 如果异常信息中包括 “HorovodAllreduce”、“HorovodAllgather” 或者 “HorovodBroadcast” 等关键字,说明这可能是某个worker死掉导致的通信失败,这种异常被Horovod认为是可以恢复的。
-
放弃失败的worker,使用剩余可用worker继续训练。
- 其他存活的 worker 停止当前的训练,记录当前模型迭代的步数。
- 此时gloo的runtime已经出现问题,通信环已经破裂,无法在剩余的 worker 之间继续进行 AllReduce 操作。
- 为了可以继续训练,Horovod Driver 会重新初始化 gloo,启动一个新的 rendezvous server,然后获取存活的 worker 的信息,利用这些worker组成新的通信环。
- 当新的通信环构造成功后,rank 0 worker 会把自己的模型广播发给其他所有worker,这样大家就可以在一个基础上,接着上次停止的迭代开始训练。
1.4 监控机制
容错机制是被动操作,监控机制就是主动操作。
弹性就意味着分布式集群的状态会随时发生变化,而 Horovod 本身和分布式集群并没有关联,所以需要有一个外部途径来让 Horovod 随时掌握集群状态。
这个外部途径就是用户需要在 Horovod 启动命令中提供一个发现脚本 discovery_host。discovery_host 由用户编写,负责发现可用的 worker 节点拓扑信息。
Driver在运行之后会定期调用这个 bash 脚本来对集群监控,当worker发生变化时,discover_host 脚本会返回最新的worker状态,Driver 根据 discover_host 的返回值得到 worker 节点信息:
- 如果Driver发现有worker失败,就捕获异常,根据存活的worker信息来更新 RendezvousServer KVStore 的节点信息,号召大家重新建立通信环进行训练。
- 如果Driver发现有新worker节点加入集群,根据目前所有worker信息来更新 RendezvousServer KVStore 的节点信息,号召大家重新建立通信环进行训练。现有worker 节点收到通知后,会暂停当前训练,记录目前迭代步数,调用
shutdown
和init
重新构造通信环。Driver也会在新节点上启动worker,扩充进程数目。 - 当新的通信环构造成功之后,rank 0 worker 会把自己的模型广播发给其他所有worker,这样大家就可以在一个基础上,接着上次停止的迭代开始训练。
这样在训练过程中,当 worker 数量有变化时,训练依然继续进行。
1.5 官方架构图
官方的一个架构图如下,我们会在后续文章中逐步讲解图中部分:
0x02 示例代码
2.1 python代码
我们从官方文档中找出 TF v2 的示例代码看看,其关键之处是使用 @hvd.elastic.run 对 train 做了一个封装,并且传入了一个 TensorFlowKerasState。
import tensorflow as tf
import horovod.tensorflow as hvd
hvd.init()
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
dataset = ...
model = ...
optimizer = tf.optimizers.Adam(lr * hvd.size())
@tf.function
def train_one_batch(data, target, allreduce=True):
with tf.GradientTape() as tape:
probs = model(data, training=True)
loss = tf.losses.categorical_crossentropy(target, probs)
if allreduce:
tape = hvd.DistributedGradientTape(tape)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
# Initialize model and optimizer state so we can synchronize across workers
data, target = get_random_batch()
train_one_batch(data, target, allreduce=False)
# 使用 @hvd.elastic.run 对 train 做了一个封装
@hvd.elastic.run
def train(state):
for state.epoch in range(state.epoch, epochs):
for state.batch in range(state.batch, batches_per_epoch):
data, target = get_random_batch()
train_one_batch(data, target)
if state.batch % batches_per_commit == 0:
state.commit()
state.batch = 0
def on_state_reset():
optimizer.lr.assign(lr * hvd.size())
# 这里是新修改处,传入了一个 TensorFlowKerasState
state = hvd.elastic.TensorFlowKerasState(model, optimizer, batch=0, epoch=0)
state.register_reset_callbacks([on_state_reset])
train(state)
2.2 脚本执行
弹性训练依然使用 horovodrun 这个命令行工具跑,和普通分布式训练不同的是,弹性训练不会在启动命令中明确指定节点列表,而是是使用一个 发现机制 来在运行时发现节点。通用的做法是在启动 Job 时候提供一个发现脚本:
horovodrun -np 18 --host-discovery-script discover_hosts.sh python train.py
此脚本用以实时反馈当前可用的 hosts 以及每个 hosts 上的 slots(下文使用 discover_hosts.sh
指代该脚本,但其无需命名为 discover_hosts.sh
)。
discover_hosts.sh 脚本必须有可执行权限,在被执行时返回可用节点列表,一行一个节点信息,结构为: ,例如:
$ sh ./discover_hosts.sh # 运行脚本,输出节点信息
host-1:4
host-2:4
host-3:4
如果这个发现脚本运行失败(没有可执行权限)或者运行时返回非0错误码,则训练进程会立刻失败,否则会一直重试直到超时(返回的slot列表不满足最小可运行数)。
弹性训练会一直等到所需最小slots数(-np)准备好之后,才会开始运行训练进程,用户可以通过 --min-np 和 --max-np 指定最小和最大的slots数,如:
horovodrun -np 8 --min-np 4 --max-np 12 --host-discovery-script discover_hosts.sh python train.py
如果可用slots数小于 --min-np 指定的数量时(比如某些节点故障,任务被抢占等),任务会被暂停等待,直到更多的节点变为活跃,或者超时时间 HOROVOD_ELASTIC_TIMEOUT(默认设置为600秒)达到。另外,如果不指定 --min-np ,则最小slots数会被默认为 -np 所配置的数目。
需要 --max-np 的原因是为了限制进程数目(防止过度使用可用资源),另外在学习率和数据分区方面也可以作为参考点(在这些情况下需要有一个固定的参考配置)。同样,如果不指定此参数,也会默认为 --np 。
0x03 逻辑流程
3.1 逻辑流程
我们先解析下弹性训练的逻辑流程(为了实现弹性训练的能力,Horovod Elastic 对 Horovod 的架构和实现进行了一定的修改),最大的差别就是:弹性训练需要在增删worker时候可以跟踪和同步worker的状态,具体修改如下。
-
聚合操作需要被定义在
hvd.elastic.run
函数之下。- 将你的主训练进程代码(初始化之后的所有代码)用一个函数(我们暂时命名为 train_func)封装起来,然后使用装饰器 hvd.elastic.run 装饰这个函数。
- 对于这个装饰器修饰的 train_func 函数,它第一个参数,必须是 hvd.elastic.State 的实例。因为某些新加入的worker可能会处于某些不确定的状态之中,所以在运行这个被装饰函数 train_func 之前,这个状态对象需要在所有worker中进行同步,以此确保所有的worker都达到一致状态。
- 因为同步函数会用到集合通信操作,并且添加worker后,活跃worker不会在此函数之前重置,所以不要在同步函数之前使用Horovod的集合操作(比如broadcast, allreduce, allgather)。
-
每个 worker 都有自己的状态(state)。
- 把所有需要在workers之间同步的变量都放进 hvd.elastic.State (比如model parameters,optimizer state,当前epoch和batch进度等等)对象之中。
- 对于TensorFlow,Keras和PyTorch,已经提供默认的标准状态实现。然而,如果用户需要在某些场景广播特殊类型,可以重载定制 hvd.elastic.State 这个对象。
- 在运行
hvd.elastic.run
函数前,此状态对象将在所有workers中同步一次,用于保持一致性。
-
周期性调用 state.commit() 来把状态(state)备份到内存。
- 定期备份非常有用。在某些worker发生意外错误时,定期备份可以避免因为状态被损坏而在重新训练时候无法恢复现场。比如,如果一个worker刚好在更新参数过程中突然出错,此时部分梯度更新完毕,部分梯度可能只更新到一半,这个状态是不可逆转而又无法继续。因此,当此状态发生时,会抛出一个 HorovodInternalError 异常,当 hvd.elastic.run 捕获到这个异常后,会利用最新一次commit中恢复所有状态。
- 因为commit状态代价高昂(比如如参数量太大会导致耗时过长),所以需要在"每个batch的处理时间"与"如果出错,训练需要从多久前的状态恢复"之间选取一个平衡点。比如,如果你每训练10个batches就commit一次,你就把复制时间降低了10倍。但是当发生错误时,你需要回滚到10个batches前的状态。
- Elastic Horowod可以通过执行我们称之为“优雅地移除worker”操作来避免这些回滚。如果driver进程发现主机已可用或标记为删除,它将向所有workers推送一个通知。于是在下次调用
state.commit()
或更轻量级的state.check_host_updates()
时,一个HostsUpdatedInterrupt
异常将被抛出。此异常的处理方式与“HorovodInternalError”类似,只是参数状态不会还原到上次commit,而是从当前实时参数中恢复。 - 一般来说,如果你的硬件设施是可靠与稳定的,并且你的编排系统会在任务节点移除时提供足够的告警,你就可低频次调用 state.commit() 函数,同时只在每个batch结束时调用相对不耗时的 state.check_host_updates() 来检查节点变更情况。
-
在 hvd.elastic.State 对象中注册一些回调函数,以便当worker成员发生变化时给予响应
- 比如回调函数可以处理如下情况:
- 当worker数量发生改变时,学习率需要根据新的world size进行相应改变。
- 对数据集进行重新分区。
- 这些回调函数会在"Horovod被重启之后"和"状态在节点间同步之前"这两个阶段中间被调用。
- 比如回调函数可以处理如下情况:
-
worker 的增减会触发其他 worker 上的重置(reset)事件,重置事件会激活以下几个操作(具体执行依据情况决定,不一定全部执行):
- 判断该 worker 是否可以继续运行。
- 将失效的 worker host 加入到黑名单,下一次组网不会使用blacklist中的host。
- 在新的 hosts 上启动 worker 进程。
- 更新每个 worker 的 rank 信息。
- 在重置之后,每个 worker 的状态会被同步
3.2 入口点
从如下代码可知 hvd.elastic.run 就是 horovod/tensorflow/elastic.py 之中的 run 函数。
import horovod.tensorflow as hvd
@hvd.elastic.run
所以我们去这个文件中探寻。
def run(func):
from tensorflow.python.framework.errors_impl import UnknownError
def wrapper(state, *args, **kwargs):
try:
return func(state, *args, **kwargs)
except UnknownError as e:
if 'HorovodAllreduce' in e.message or \
'HorovodAllgather' in e.message or \
'HorovodBroadcast' in e.message:
raise HorovodInternalError(e)
return run_fn(wrapper, _reset)
3.3 主逻辑
run_fn 函数是关于用户代码的主要逻辑所在,位于 horovod/common/elastic.py。
其主要逻辑是:
- 初始化 notification_manager;
- 在 notification_manager 注册 state;
- 运行 func 函数,就是用户的训练代码 train;
- 在worker进程出现 HorvodInternalError 错误或者 HostsUpdateInterrupt 节点增删时,会捕获这两个错误,调用 reset 来进行容错处理;
def run_fn(func, reset):
@functools.wraps(func)
def wrapper(state, *args, **kwargs):
notification_manager.init()
notification_manager.register_listener(state)
skip_sync = False
try:
while True:
if not skip_sync:
state.sync()
try:
return func(state, *args, **kwargs)
except HorovodInternalError:
state.restore()
skip_sync = False
except HostsUpdatedInterrupt as e:
skip_sync = e.skip_sync
reset()
state.on_reset()
finally:
notification_manager.remove_listener(state)
return wrapper
3.4 出错处理
在出错状态下,当worker进程出现 HorvodInternalError (代表出现错误)或者 HostsUpdateInterrupt (代表有节点增删)时,Horovod 会执行如下流程:
- 在 hvd.elastic.run 装饰器中捕获上述两个错误;
- 如果抛出的是 HorvodInternalError 错误,则会从最后的一次 commit 状态中恢复;
- 重新初始化 Horovod context,然后启动新的一轮的rendezvous,在rendezvous过程中,旧的worker会被优先被选举为新的rank-0,因为旧的worker具有上次训练中的最近状态;
- 新的 rank-0 worker 会把状态同步到其它workers;
- 继续训练;
至此,我们已经分析了horovod 弹性训练基本架构,下一篇我们分析最主要的部件:Driver。
0xEE 个人信息
★★★★★★关于生活和技术的思考★★★★★★
微信公众账号:罗西的思考
如果您想及时得到个人撰写文章的消息推送,或者想看看个人推荐的技术资料,敬请关注。
0xFF 参考
ElasticDL调用 Horovod 在Kubernetes上实现弹性 AllReduce(一)
kubernetes 培训_在Kubernetes上使用horovod进行分布式深度学习培训
在 Kubernetes 上弹性深度学习训练利器 -- Elastic Training Operator