Gradient synchronization coordinates gradient aggregation across data-parallel workers so every worker applies consistent model updates. Training adds synchronization after local gradient computation; all-reduce is one implementation approach. section 4.2.1
Choose the update behavior according to consistency and throughput requirements:
- Synchronous updates pause every worker at gradient synchronization until aggregated gradients arrive, keeping models aligned but allowing slow machines or network problems to idle the group. section 4.2.1
- Asynchronous updates let each worker update immediately without waiting, increasing training steps per minute but generally taking longer and producing less-accurate models. section 4.2.1
PyTorch distributed training
Register workers with a shared master address, port, rank, and world size; initialize the process group; then wrap the model with DistributedDataParallel. After backward returns, each parameter gradient contains the synchronized tensor. section 4.2.3
import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP # create model and move it to GPU model = DpModel().to(device) # wrap the model with DDP ddp_model = DDP(model, device_ids=[rank]) outputs = ddp_model(data) #1 # compute the loss and sync gradient with other workers. # when 'backward' function returns, the param.grad already # contains synchronized gradient tensor loss_fn(outputs, labels).backward()
Inspect the compact PyTorch listing showing model wrapping, forward execution, loss computation, and synchronized gradients after backward. 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 the Horovod listing showing DistributedGradientTape, averaged gradients, optimizer application, and initial state broadcasts. section 4.2.3
