Horovod is a distributed deep learning framework that gives PyTorch, TensorFlow, MXNet, and PySpark training code a common data-parallel execution method. It synchronizes workers while allowing a training service to treat framework-specific training code as a black box. section 4.5
A Horovod integration typically handles these responsibilities:
- Initialize Horovod so each process joins the distributed training group. section 4.2.3
- Synchronize gradients across workers with a distributed gradient tape or distributed optimizer. section 4.2.3
- Broadcast initial model parameters and optimizer state from rank 0. section 4.2.3
- Restrict progress checks and model or checkpoint persistence to the master worker, avoiding duplicated artifacts. section 4.5
TensorFlow gradient synchronization
Initialize Horovod, wrap the TensorFlow gradient tape with hvd.DistributedGradientTape, average gradients across workers, apply them locally, and broadcast initial variables from rank 0. section 4.2.3
hvd.init() #1
.. .. ..
@tf.function
def training_step(images, labels, first_batch):
with tf.GradientTape() as tape:
probs = mnist_model(images, training=True)
loss_value = loss(labels, probs)
# Wrap tape with Horovod Distributed GradientTape.
# This gradient tape averages gradients from all
# workers by using allreduce or allgather, and then
# applies those averaged gradients back to the local model.
tape = hvd.DistributedGradientTape(tape)
grads = tape.gradient(loss_value, mnist_model.trainable_variables)
opt.apply_gradients(zip(grads, mnist_model.trainable_variables))
# Broadcast initial variable states
# from rank 0 to all other processes.
if first_batch:
hvd.broadcast_variables(mnist_model.variables, root_rank=0)
hvd.broadcast_variables(opt.variables(), root_rank=0)
return loss_value
for batch, (images, labels) in \ #2
enumerate(dataset.take(10000 / hvd.size())):
loss_value = training_step(images, labels, batch == 0)
.. .. ..
# save checkpoints only on worker 0 to
# prevent other workers from corrupting it.
if hvd.rank() == 0:
checkpoint.save(checkpoint_dir)
Inspect a TensorFlow training step that combines Horovod initialization, distributed gradient synchronization, and rank-zero state broadcasting. section 4.2.3
