How to run multi-worker training in TensorFlow

A synchronous TensorFlow cluster advances only when every worker reaches the same collective operations, so a local two-process launch can expose mismatched topology, missing peers, and broken gradient reduction before a model moves to separate hosts.

Each process receives its own TF_CONFIG value. The cluster mapping stays identical across the pair, but task.index selects the worker identity that MultiWorkerMirroredStrategy uses when it starts the collective runtime. Separate loopback ports keep the two local endpoints unambiguous.

The training program shards one deterministic dataset, mirrors model and optimizer state, reduces every replica loss, and refuses to print a result unless both workers complete the same synchronized updates.

1. Encode worker 0 and the shared two-port topology in worker-0.sh.

   worker-0.sh

   TF_CONFIG='{"cluster":{"worker":["127.0.0.1:12345","127.0.0.1:23456"]},"task":{"type":"worker","index":0}}' exec python3 multiworker_train.py

   The ports must be unused and reachable by both workers. A multi-host cluster uses routable hostnames or IP addresses instead of loopback addresses.

2. Encode worker 1 against the same cluster mapping in worker-1.sh.

   worker-1.sh

   TF_CONFIG='{"cluster":{"worker":["127.0.0.1:12345","127.0.0.1:23456"]},"task":{"type":"worker","index":1}}' exec python3 multiworker_train.py

   Only task.index changes between the worker launchers; changing the cluster list would make the processes form different collectives.

3. Initialize multiworker_train.py with CPU isolation and the worker-local cluster contract.

   multiworker_train.py

   import json
   import os
    
   os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import numpy as np
   import tensorflow as tf
    
   tf_config = json.loads(os.environ["TF_CONFIG"])
   task_index = tf_config["task"]["index"]
   worker_count = len(tf_config["cluster"]["worker"])
   per_worker_batch_size = 16
   global_batch_size = per_worker_batch_size * worker_count
   steps_per_epoch = 4

   TF_CONFIG must exist before MultiWorkerMirroredStrategy is created because strategy construction starts the worker's collective runtime.

4. Write dataset_fn() below the cluster contract to shard repeatable input across worker pipelines.

   def dataset_fn(input_context):
       rng = np.random.default_rng(7)
       features = rng.normal(size=(256, 8)).astype("float32")
       scores = (
           features[:, 0] * 0.8
           + features[:, 1] * 0.6
           - features[:, 2] * 0.4
           + features[:, 3] * 0.2
       )
       labels = (scores > 0).astype("float32")[:, None]
       batch_size = input_context.get_per_replica_batch_size(global_batch_size)
       return (
           tf.data.Dataset.from_tensor_slices((features, labels))
           .shard(
               input_context.num_input_pipelines,
               input_context.input_pipeline_id,
           )
           .shuffle(256, seed=7, reshuffle_each_iteration=False)
           .repeat()
           .batch(batch_size)
           .prefetch(tf.data.AUTOTUNE)
       )

   InputContext gives each worker a distinct input-pipeline index and calculates the per-replica batch from the global batch size.

5. Construct the mirrored model state below dataset_fn() inside the strategy scope.

   strategy = tf.distribute.MultiWorkerMirroredStrategy()
   tf.keras.utils.set_random_seed(7)
    
   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()
    
   distributed_dataset = strategy.distribute_datasets_from_function(dataset_fn)
   distributed_iterator = iter(distributed_dataset)

   The strategy scope mirrors model variables, optimizer state, and metrics across the cluster.

6. Define train_step() below the distributed iterator to reduce each synchronized replica update.

   @tf.function
   def train_step(iterator):
       def replica_step(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_size,
               )
           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(replica_step, args=(next(iterator),))
       return strategy.reduce(
           tf.distribute.ReduceOp.SUM,
           per_replica_losses,
           axis=None,
       )

   strategy.run() executes one replica_step() per replica, while strategy.reduce() returns one global loss for the update.

7. Define the three-epoch training driver below train_step() with fail-capable collective result checks.

   initial_loss = None
   final_loss = None
    
   for _ in range(3):
       train_accuracy.reset_state()
       total_loss = 0.0
       for _ in range(steps_per_epoch):
           total_loss += train_step(distributed_iterator)
       final_loss = total_loss / steps_per_epoch
       if initial_loss is None:
           initial_loss = tf.identity(final_loss)
    
   tf.debugging.assert_equal(strategy.num_replicas_in_sync, worker_count)
   tf.debugging.assert_less(final_loss, initial_loss)
   tf.debugging.assert_greater(train_accuracy.result(), 0.75)
    
   print(
       f"worker={task_index} replicas={strategy.num_replicas_in_sync} "
       f"initial_loss={float(initial_loss):.4f} "
       f"final_loss={float(final_loss):.4f} "
       f"final_accuracy={float(train_accuracy.result()):.4f}",
       flush=True,
   )

   The process exits with an assertion error if a worker is missing, the loss does not fall, or final accuracy stays at or below 0.75.

8. Launch worker 0 in the background through its worker-local topology file.

   $ bash worker-0.sh > worker-0.log 2>&1 &

   Worker 0 waits at collective operations until worker 1 joins the matching cluster.

9. Start worker 1 in the foreground through its matching topology file.

   $ bash worker-1.sh > worker-1.log 2>&1

10. Wait for the background worker to leave every collective successfully.

    $ wait

11. Confirm worker 0 joined two replicas and passed the training assertions.

    $ cat worker-0.log
    worker=0 replicas=2 initial_loss=0.7459 final_loss=0.4168 final_accuracy=0.8281

12. Confirm worker 1 independently recorded the same synchronized result.

    $ cat worker-1.log
    worker=1 replicas=2 initial_loss=0.7459 final_loss=0.4168 final_accuracy=0.8281