Horovod
Horovod was originally developed by Uber to make distributed deep learning fast and easy to use, bringing model training time down from days and weeks to hours and minutes. With Horovod, an existing training script can be scaled up to run on hundreds of GPUs in just a few lines of Python code. Once Horovod has been configured, the same infrastructure can be used to train models with any framework, making it easy to switch between TensorFlow, PyTorch, MXNet, and future frameworks as machine learning tech stacks continue to evolve.
import tensorflow as tf # Choose the right backend import horovod.tensorflow as hvd # Initialize Horovod hvd.init() # Pin GPU to be used to process local rank (one GPU per process) config = tf.ConfigProto() config.gpu_options.visible_device_list = str(hvd.local_rank()) # Build model... loss = ... opt = ... # Add Horovod Distributed Optimizer opt = hvd.DistributedOptimizer(opt) # Add hook to broadcast variables from rank 0 to all other processes during initialization. hooks = [hvd.BroadcastGlobalVariablesHook(0)] # Make training operation train_op = opt.minimize(loss)
The runscript is shown as following
#!/bin/bash #SBATCH --job-name="horovod_tensorflow_mnist" #SBATCH --output="horovod_tensorflow_mnist_%j.out" #SBATCH --partition=gpu #SBATCH --nodes=2 #SBATCH --ntasks-per-node=4 #SBATCH --cpus-per-task=40 #SBATCH --cores-per-socket=20 #SBATCH --threads-per-core=4 #SBATCH --sockets-per-node=2 #SBATCH --mem-per-cpu=1200 #SBATCH --export=ALL #SBATCH --gres=gpu:v100:4 #SBATCH --time=4:00:00 module load opence mpirun -n 8 python 05_horovod_tensorflow_mnist.py