Horovod 是Uber于2017年发布的一个易于使用的高性能的分布式训练框架,在业界得到了广泛应用。
本系列将通过源码分析来带领大家了解 Horovod。本文是系列第十二篇,看看horovod 如何实施弹性训练。
弹性训练使得Horovod具备运行时worker数量动态伸缩,而不需要重启 或者 只是从存储中的checkpoint恢复训练。
本系列其他文章链接如下:
[源码解析] 深度学习分布式训练框架 Horovod (1) --- 基础知识
[源码解析] 深度学习分布式训练框架 horovod (2) --- 从使用者角度切入
[源码解析] 深度学习分布式训练框架 horovod (3) --- Horovodrun背后做了什么
[源码解析] 深度学习分布式训练框架 horovod (4) --- 网络基础 & Driver
[源码解析] 深度学习分布式训练框架 horovod (5) --- 融合框架
[源码解析] 深度学习分布式训练框架 horovod (6) --- 后台架构
[源码解析] 深度学习分布式训练框架 horovod (6) --- 线程实现
[源码解析] 深度学习分布式训练框架 horovod (7) --- DistributedOptimizer
[源码解析] 深度学习分布式训练框架 horovod (8) --- on spark
[源码解析] 深度学习分布式训练框架 horovod (9) --- 启动 on spark
[源码解析] 深度学习分布式训练框架 horovod (10) --- run on spark
[源码解析] 深度学习分布式训练框架 horovod (11) --- on spark --- GLOO 方案
我们思考下,Horovod 目前遇到了什么问题?
为了解决以上几个问题,我们会思考很多的其他具体技术问题和细节,让我们先罗列出来:
我们在本文以及后续各篇的分析中试着解答这些问题。
注:Horovod目前的调度机制依然不灵活,不支持抢占。
Horovod 在单机的多个 GPU 之上采用 NCCL 来通信,在多机(CPU或者GPU)之间通过 Ring AllReduce 算法进行通信。Horovod 的弹性训练是指多机的弹性训练。
Horovod 弹性训练有两个角色:driver和 worker。driver 进程运行在 CPU 节点上,worker 进程可以运行在 CPU 节点或者 GPU 节点之上。
Driver 进程的作用是:
worker 负责训练和模型迭代。
具体组网机制如下:
+-------------------------------+
| Driver |
| |
| +------------------------+ |
| | RendezvousServer | |
| | | |
| | | |
| | host1, host2, host3 | |
| +------------------------+ |
+-------------------------------+
^ ^ ^
| | |
| | |
+-------------+ | +--------------+
| | |
| | |
| | |
v v v
+--------+----+ +-------+------+ +----+--------+
| Worker | | Worker | | Worker |
+------> | +------> | +---------> | | +------+
| | host1 | | host2 | | host3 | |
| +-------------+ +--------------+ +-------------+ |
| |
| |
| v
<--------------------------------------------------------------------------------+
复制代码
我们下面详细分析下各个部分。
Horovod 的容错机制是基于 gloo 来实现的,对于错误来说,这可以被认为是一个被动操作。
Gloo 本身是不支持容错的。当众多worker之间对张量进行聚合操作时候,如果某一个worker失败,则gloo不会处理异常,而是抛出异常并且退出,这样所有worker都会报异常退出。
为了不让某一个 worker 的失败导致整体训练退出,Horovod 需要做两方面工作:
容错机制是被动操作,监控机制就是主动操作。
弹性就意味着分布式集群的状态会随时发生变化,而 Horovod 本身和分布式集群并没有关联,所以需要有一个外部途径来让 Horovod 随时掌握集群状态。
这个外部途径就是用户需要在 Horovod 启动命令中提供一个发现脚本 discovery_host。discovery_host 由用户编写,负责发现可用的 worker 节点拓扑信息。
Driver在运行之后会定期调用这个 bash 脚本来对集群监控,当worker发生变化时,discover_host 脚本会返回最新的worker状态,Driver 根据 discover_host 的返回值得到 worker 节点信息:
shutdown
和 init
重新构造通信环。Driver也会在新节点上启动worker,扩充进程数目。这样在训练过程中,当 worker 数量有变化时,训练依然继续进行。
官方的一个架构图如下,我们会在后续文章中逐步讲解图中部分:
我们从官方文档中找出 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)
复制代码
弹性训练依然使用 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 。
我们先解析下弹性训练的逻辑流程(为了实现弹性训练的能力,Horovod Elastic 对 Horovod 的架构和实现进行了一定的修改),最大的差别就是:弹性训练需要在增删worker时候可以跟踪和同步worker的状态,具体修改如下。
hvd.elastic.run
函数之下。
hvd.elastic.run
函数前,此状态对象将在所有workers中同步一次,用于保持一致性。state.commit()
或更轻量级的state.check_host_updates()
时,一个HostsUpdatedInterrupt
异常将被抛出。此异常的处理方式与“HorovodInternalError”类似,只是参数状态不会还原到上次commit,而是从当前实时参数中恢复。从如下代码可知 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)
复制代码
run_fn 函数是关于用户代码的主要逻辑所在,位于 horovod/common/elastic.py。
其主要逻辑是:
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
复制代码
在出错状态下,当worker进程出现 HorvodInternalError (代表出现错误)或者 HostsUpdateInterrupt (代表有节点增删)时,Horovod 会执行如下流程:
至此,我们已经分析了horovod 弹性训练基本架构,下一篇我们分析最主要的部件:Driver。
作者:罗西的思考
链接:https://juejin.cn/post/6982903090471010311
来源:掘金
著作权归作者所有。商业转载请联系作者获得授权,非商业转载请注明出处。