How to run multi-worker training in TensorFlow

Multi-worker TensorFlow training turns one training program into several cooperating worker processes that execute the same training step under a shared cluster description. It is useful once a model already trains on one host and the next scaling step is to add another process or machine without rewriting the model code around sockets or RPC calls.

TensorFlow reads the cluster from TF_CONFIG before tf.distribute.MultiWorkerMirroredStrategy starts. Every worker receives the same cluster list, while the task section identifies the current worker's type and index.

A local smoke test can use two worker processes on separate localhost ports. For real multi-host training, replace those addresses with hostnames or routable IP addresses that every worker can reach, and keep checkpoints, logs, and backup directories on shared storage.

1. Save the multi-worker training script as multiworker_train.py.

   multiworker_train.py

   import json
   import os
   from pathlib import Path
    
   import numpy as np
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
   os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1")
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("ERROR")
    
   tf_config = json.loads(os.environ["TF_CONFIG"])
   task = tf_config["task"]
   workers = tf_config["cluster"]["worker"]
    
   per_worker_batch = 16
   global_batch = per_worker_batch * len(workers)
   steps_per_epoch = 4
   epochs = 3
    
   rng = np.random.default_rng(7)
   features = rng.normal(size=(256, 8)).astype("float32")
   score = (
       features[:, 0] * 0.8
       + features[:, 1] * 0.6
       - features[:, 2] * 0.4
       + features[:, 3] * 0.2
   )
   labels = (score > 0).astype("float32")[:, None]
    
   strategy = tf.distribute.MultiWorkerMirroredStrategy()
    
   print(
       f"worker={task['index']} "
       f"cluster_workers={len(workers)} "
       f"replicas={strategy.num_replicas_in_sync} "
       f"global_batch={global_batch}",
       flush=True,
   )
    
    
   def dataset_fn(input_context):
       batch_size = input_context.get_per_replica_batch_size(global_batch)
       dataset = tf.data.Dataset.from_tensor_slices((features, labels))
       dataset = dataset.shard(
           input_context.num_input_pipelines,
           input_context.input_pipeline_id,
       )
       dataset = dataset.shuffle(128, seed=7, reshuffle_each_iteration=False)
       return dataset.repeat().batch(batch_size).prefetch(tf.data.AUTOTUNE)
    
    
   with strategy.scope():
       model = tf.keras.Sequential(
           [
               tf.keras.layers.Input(shape=(8,)),
               tf.keras.layers.Dense(16, activation="relu"),
               tf.keras.layers.Dense(1, activation="sigmoid"),
           ]
       )
       optimizer = tf.keras.optimizers.Adam(learning_rate=0.03)
       loss_object = tf.keras.losses.BinaryCrossentropy(
           reduction=tf.keras.losses.Reduction.NONE
       )
       train_accuracy = tf.keras.metrics.BinaryAccuracy(name="binary_accuracy")
    
   dist_dataset = strategy.distribute_datasets_from_function(dataset_fn)
   dist_iterator = iter(dist_dataset)
    
    
   @tf.function
   def train_step(iterator):
       def step_fn(inputs):
           batch_features, batch_labels = inputs
           with tf.GradientTape() as tape:
               predictions = model(batch_features, training=True)
               per_example_loss = loss_object(batch_labels, predictions)
               loss = tf.nn.compute_average_loss(
                   per_example_loss,
                   global_batch_size=global_batch,
               )
           gradients = tape.gradient(loss, model.trainable_variables)
           optimizer.apply_gradients(zip(gradients, model.trainable_variables))
           train_accuracy.update_state(batch_labels, predictions)
           return loss
    
       per_replica_losses = strategy.run(step_fn, args=(next(iterator),))
       return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
    
    
   for epoch in range(1, epochs + 1):
       train_accuracy.reset_state()
       total_loss = 0.0
       for _ in range(steps_per_epoch):
           total_loss += train_step(dist_iterator)
       print(
           f"worker={task['index']} "
           f"epoch={epoch} "
           f"loss={float(total_loss / steps_per_epoch):.4f} "
           f"binary_accuracy={float(train_accuracy.result()):.4f}",
           flush=True,
       )
    
   if task["type"] == "worker" and task["index"] == 0:
       Path("multiworker-complete.txt").write_text("chief worker completed training\n")
    
   print(f"worker={task['index']} training_complete", flush=True)

   Create MultiWorkerMirroredStrategy before running TensorFlow tensor operations. Collective training is configured at process startup, so TF_CONFIG must be present before the Python process creates the strategy.

2. Start worker 0 and write its log to worker-0.log.

   $ TF_CONFIG='{"cluster":{"worker":["localhost:12345","localhost:23456"]},"task":{"type":"worker","index":0}}' \
     python3 multiworker_train.py > worker-0.log 2>&1 &

   The first worker waits for the other task listed in TF_CONFIG. Keep it running while worker 1 starts.

3. Start worker 1 with the same cluster list and task index 1.

   $ TF_CONFIG='{"cluster":{"worker":["localhost:12345","localhost:23456"]},"task":{"type":"worker","index":1}}' \
     python3 multiworker_train.py > worker-1.log 2>&1

   For real hosts, replace localhost:12345 and localhost:23456 with addresses reachable from every worker. Keep the cluster list identical on all workers and change only the task.index value.

4. Wait for the background worker to finish before reading the saved logs.

   $ wait

5. Read the worker 1 log.

   $ cat worker-1.log
   worker=1 cluster_workers=2 replicas=2 global_batch=32
   worker=1 epoch=1 loss=0.6817 binary_accuracy=0.6016
   worker=1 epoch=2 loss=0.4741 binary_accuracy=0.8047
   worker=1 epoch=3 loss=0.3265 binary_accuracy=0.8750
   worker=1 training_complete

   replicas=2 shows that both workers joined the cluster, and the falling loss shows that synchronized training steps completed.

6. Read the worker 0 log.

   $ cat worker-0.log
   worker=0 cluster_workers=2 replicas=2 global_batch=32
   worker=0 epoch=1 loss=0.6817 binary_accuracy=0.6016
   worker=0 epoch=2 loss=0.4741 binary_accuracy=0.8047
   worker=0 epoch=3 loss=0.3265 binary_accuracy=0.8750
   worker=0 training_complete

   Both workers report the same epoch metrics because each worker participates in the same synchronized update. In production, write checkpoints and TensorBoard logs to storage visible to every worker.
   Related: How to save and restore a TensorFlow checkpoint
   Related: How to profile TensorFlow training in TensorBoard

7. Check the chief completion marker.

   $ cat multiworker-complete.txt
   chief worker completed training

   The marker is written only by worker 0 after the training loop exits.

8. Remove the temporary smoke-test files.

   $ rm -f multiworker_train.py worker-0.log worker-1.log multiworker-complete.txt