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.

Steps to run TensorFlow multi-worker training with MultiWorkerMirroredStrategy:

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.

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.

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.
Wait for the background worker to finish before reading the saved logs.
```
$ wait
```

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.

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
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.

Remove the temporary smoke-test files.

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