Multi-worker Training with Keras

View on Run in Google Colab View source on GitHub Download notebook


This tutorial demonstrates multi-worker distributed training with Keras model using tf.distribute.Strategy API. With the help of the strategies specifically designed for multi-worker training, a Keras model that was designed to run on single-worker can seamlessly work on multiple workers with minimal code change.

Distributed Training in TensorFlow guide is available for an overview of the distribution strategies TensorFlow supports for those interested in a deeper understanding of tf.distribute.Strategy APIs.


First, setup TensorFlow and the necessary imports.

from __future__ import absolute_import, division, print_function, unicode_literals
!pip install -q tensorflow==2.0.0-beta1
import tensorflow_datasets as tfds
import tensorflow as tf

Preparing dataset

Now, let's prepare the MNIST dataset from TensorFlow Datasets. The MNIST dataset comprises 60,000 training examples and 10,000 test examples of the handwritten digits 0–9, formatted as 28x28-pixel monochrome images.


# Scaling MNIST data from (0, 255] to (0., 1.]
def scale(image, label):
  image = tf.cast(image, tf.float32)
  image /= 255
  return image, label

datasets, info = tfds.load(name='mnist',
train_datasets_unbatched = datasets['train'].map(scale).shuffle(BUFFER_SIZE)
train_datasets = train_datasets_unbatched.batch(BATCH_SIZE)

Build the Keras model

Here we use tf.keras.Sequential API to build and compile a simple convolutional neural networks Keras model to train with our MNIST dataset.

def build_and_compile_cnn_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
      tf.keras.layers.Dense(64, activation='relu'),
      tf.keras.layers.Dense(10, activation='softmax')
  return model

Let's first try training the model for a small number of epochs and observe the results in single worker to make sure everything works correctly. You should expect to see the loss dropping and accuracy approaching 1.0 as epoch advances.

single_worker_model = build_and_compile_cnn_model(), epochs=3)
Epoch 1/3

WARNING: Logging before flag parsing goes to stderr.
W0614 15:36:27.263422 140668346640128] From /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/ add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.where in 2.0, which has the same broadcast rule as np.where

938/938 [==============================] - 20s 22ms/step - loss: 2.0755 - accuracy: 0.4123
Epoch 2/3
938/938 [==============================] - 18s 19ms/step - loss: 1.1850 - accuracy: 0.7703
Epoch 3/3
938/938 [==============================] - 17s 19ms/step - loss: 0.6343 - accuracy: 0.8458

<tensorflow.python.keras.callbacks.History at 0x7fef8c442470>

Multi-worker Configuration

Now let's enter the world of multi-worker training. In TensorFlow, TF_CONFIG environment variable is required for training on multiple machines, each of which possibly has a different role. TF_CONFIG is used to specify the cluster configuration on each worker that is part of the cluster.

There are two components of TF_CONFIG: cluster and task. cluster provides information about the training cluster, which is a dict consisting of different types of jobs such as worker. In multi-worker training, there is usually one worker that takes on a little more responsibility like saving checkpoint and writing summary file for TensorBoard in addition to what a regular worker does. Such worker is referred to as the 'chief' worker, and it is customary that the worker with index 0 is appointed as the chief worker (in fact this is how tf.distribute.Strategy is implemented). task on the other hand provides information of the current task.

In this example, we set the task type to "worker" and the task index to 0. This means the machine that has such setting is the first worker, which will be appointed as the chief worker and do more work than other workers. Note that other machines will need to have TF_CONFIG environment variable set as well, and it should have the same cluster dict, but different task type or task index depending on what the roles of those machines are.

For illustration purposes, this tutorial shows how one may set a TF_CONFIG with a single worker on localhost. In practice, users would create multiple workers on external IP addresses/ports, and set TF_CONFIG on each worker appropriately.

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:12345", "localhost:23456"]
    'task': {'type': 'worker', 'index': 0}

Note that while the learning rate is fixed in this example, in general it may be necessary to adjust the learning rate based on the global batch size.

Choose the right strategy

In TensorFlow, distributed training consists of synchronous training, where the steps of training are synced across the workers and replicas, and asynchronous training, where the training steps are not strictly synced.

MultiWorkerMirroredStrategy, which is the recommended strategy for synchronous multi-worker training, will be demonstrated in this guide. To train the model, use an instance of tf.distribute.experimental.MultiWorkerMirroredStrategy. MultiWorkerMirroredStrategy creates copies of all variables in the model's layers on each device across all workers. It uses CollectiveOps, a TensorFlow op for collective communication, to aggregate gradients and keep the variables in sync. The tf.distribute.Strategy guide has more details about this strategy.

strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
W0614 15:37:22.454128 140668346640128] Collective ops is not configured at program startup. Some performance features may not be enabled.
W0614 15:37:22.456581 140668346640128] There is non-GPU devices in <a href="../../../versions/r2.0/api_docs/python/tf/distribute/Strategy"><code>tf.distribute.Strategy</code></a>, not using nccl allreduce.

MultiWorkerMirroredStrategy provides multiple implementations via the CollectiveCommunication parameter. RING implements ring-based collectives using gRPC as the cross-host communication layer. NCCL uses Nvidia's NCCL to implement collectives. AUTO defers the choice to the runtime. The best choice of collective implementation depends upon the number and kind of GPUs, and the network interconnect in the cluster.

Train the model with MultiWorkerMirroredStrategy

With the integration of tf.distribute.Strategy API into tf.keras, the only change you will make to distribute the training to multi-worker is enclosing the model building and model.compile() call inside strategy.scope(). The distribution strategy's scope dictates how and where the variables are created, and in the case of MultiWorkerMirroredStrategy, the variables created are MirroredVariables, and they are replicated on each of the workers.

# Here the batch size scales up by number of workers since 
# `` expects the global batch size. Previously we used 64, 
# and now this becomes 128.
train_datasets = train_datasets_unbatched.batch(GLOBAL_BATCH_SIZE)
with strategy.scope():
  multi_worker_model = build_and_compile_cnn_model(), epochs=3)
Train on None steps
Epoch 1/3
469/469 [==============================] - 15s 32ms/step - loss: 2.2209 - accuracy: 0.2955
Epoch 2/3
469/469 [==============================] - 11s 24ms/step - loss: 1.9541 - accuracy: 0.5819
Epoch 3/3
469/469 [==============================] - 11s 24ms/step - loss: 1.4593 - accuracy: 0.7229

<tensorflow.python.keras.callbacks.History at 0x7fef9f6ad0b8>

In multi-worker training, sharding data into multiple parts is needed to ensure convergence and performance. However, note that in above code snippet, the datasets are directly sent to without needing to shard; this is because tf.distribute.Strategy API takes care of the dataset sharding automatically in multi-worker trainings.

Another thing to notice is the batch size for the datasets. Here we change the batch size to be twice as large as the case it was for single worker, because the effective per worker batch size is the global batch size (the parameter passed in divided by the number of workers, and with this change we are keeping the per worker batch size same as before.


You now have a Keras model that is all set up to run in multiple workers with MultiWorkerMirroredStrategy. You can try the following techniques to tweak performance of multi-worker training.

  • MultiWorkerMirroredStrategy provides multiple collective communication implementations. RING implements ring-based collectives using gRPC as the cross-host communication layer. NCCL uses Nvidia's NCCL to implement collectives. AUTO defers the choice to the runtime. The best choice of collective implementation depends upon the number and kind of GPUs, and the network interconnect in the cluster. To override the automatic choice, specify a valid value to the communication parameter of MultiWorkerMirroredStrategy's constructor, e.g. communication=tf.distribute.experimental.CollectiveCommunication.NCCL.
  • Cast the variables to tf.float if possible. The official ResNet model includes an example of how this can be done.

See also

  1. Distributed Training in TensorFlow guide provides an overview of the available distribution strategies.
  2. Official ResNet50 model, which can be trained using either MirroredStrategy or MultiWorkerMirroredStrategy.