TensorFlow
Multi-node Computing
Parallel Processing
Distributed Systems
CPU Optimization

How to run TensorFlow on multiple nodes with several CPUs each

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

To run TensorFlow across multiple machines with several CPUs on each machine, the usual high-level tool is tf.distribute.MultiWorkerMirroredStrategy. Each worker runs the same training program, TensorFlow coordinates the cluster through TF_CONFIG, and the strategy keeps model variables synchronized during training.

The Mental Model: Same Script on Every Worker

Multi-worker TensorFlow is not a master script pushing work to passive helpers. Each worker is a full training process running the same code. The difference between workers comes from environment configuration, not from separate programs.

The key pieces are:

  • identical training script on every node
  • 'TF_CONFIG describing the full cluster and the current worker'
  • a distribution strategy created before the model

That is why distributed TensorFlow often feels simpler than people expect. You do not write separate networking code for each node. You describe the cluster and let the strategy handle synchronization.

Use MultiWorkerMirroredStrategy

For synchronous data-parallel training across multiple CPU workers, this is the normal starting point:

python
1import tensorflow as tf
2
3strategy = tf.distribute.MultiWorkerMirroredStrategy()
4
5with strategy.scope():
6    model = tf.keras.Sequential([
7        tf.keras.layers.Dense(64, activation="relu"),
8        tf.keras.layers.Dense(10),
9    ])
10    model.compile(
11        optimizer="adam",
12        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
13        metrics=["accuracy"],
14    )

Creating the model inside strategy.scope() ensures variables are created in a distributed-aware way.

Configure the Cluster with TF_CONFIG

Before the script starts, each worker needs a TF_CONFIG environment variable. It lists all workers and identifies which one the current process represents.

Worker 0:

bash
export TF_CONFIG='{"cluster":{"worker":["host1:12345","host2:12345"]},"task":{"type":"worker","index":0}}'
python train.py

Worker 1:

bash
export TF_CONFIG='{"cluster":{"worker":["host1:12345","host2:12345"]},"task":{"type":"worker","index":1}}'
python train.py

Both workers run the same train.py. The only difference is the task.index field.

Without a correct TF_CONFIG, TensorFlow does not know how to join the processes into one cluster.

A Minimal End-to-End Example

Here is a small runnable example that uses synthetic data and normal Keras fit.

python
1import numpy as np
2import tensorflow as tf
3
4tf.config.threading.set_intra_op_parallelism_threads(8)
5tf.config.threading.set_inter_op_parallelism_threads(2)
6
7strategy = tf.distribute.MultiWorkerMirroredStrategy()
8
9x = np.random.rand(1000, 20).astype("float32")
10y = np.random.randint(0, 2, size=(1000,))
11
12dataset = tf.data.Dataset.from_tensor_slices((x, y))
13dataset = dataset.shuffle(1000).batch(32)
14
15with strategy.scope():
16    model = tf.keras.Sequential([
17        tf.keras.layers.Dense(32, activation="relu"),
18        tf.keras.layers.Dense(2),
19    ])
20    model.compile(
21        optimizer="adam",
22        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
23        metrics=["accuracy"],
24    )
25
26model.fit(dataset, epochs=3)

This is still one program. Distributed behavior comes from the strategy and the cluster configuration, not from changing the training loop into a different style.

CPU Tuning Still Matters

Multiple machines do not automatically mean efficient CPU use. Each TensorFlow worker also uses local threads for kernels and input work, so per-node tuning matters.

Typical levers include:

  • batch size per worker
  • total global batch size
  • input pipeline throughput
  • intra-op thread count
  • inter-op thread count
  • network quality between workers

If the workload is small or communication-heavy, multi-node training may not help much. Distributed CPU training pays off most when the model and dataset are large enough to justify the coordination overhead.

Input Pipelines Must Keep Up

It is easy to focus on the model and forget the data path. In distributed training, every worker needs data fast enough to keep its CPUs busy.

That means you should think about:

  • whether every worker can reach the dataset efficiently
  • whether dataset sharding is happening the way you expect
  • whether preprocessing is the actual bottleneck

If the input pipeline stalls, adding more nodes only spreads the bottleneck around. Good distribution strategy does not rescue a slow data pipeline automatically.

Checkpointing and Coordination

Real training jobs also need disciplined checkpoint handling. Even though every worker runs the same script, some outputs should usually be written in a coordinated way to avoid collisions or duplicated work.

Keras and TensorFlow handle much of this well, but the operational rule still matters: distributed training is not only about model code. File paths, shared storage, and restart behavior also need to be designed for multiple workers.

That is especially important once training moves from a toy cluster to a production environment where failures and retries are normal.

Common Pitfalls

One common mistake is trying to write separate master and worker training programs instead of running the same script on every node with different TF_CONFIG values.

Another pitfall is forgetting to create the model inside strategy.scope(). That can lead to confusing distributed-variable behavior.

A third issue is ignoring input pipeline throughput and focusing only on the model. On CPU clusters, data feeding can become the real bottleneck very quickly.

Finally, do not assume more workers always means faster training. Network overhead, synchronization cost, and thread settings still determine whether the cluster scales well.

Summary

  • Use tf.distribute.MultiWorkerMirroredStrategy for synchronous multi-node TensorFlow training across CPUs.
  • Run the same script on every worker and configure the cluster with TF_CONFIG.
  • Build and compile the model inside strategy.scope().
  • Tune thread settings, batch sizes, and input pipelines for CPU-heavy workloads.
  • Distributed training helps most when the workload is large enough to justify coordination overhead.

Course illustration
Course illustration

All Rights Reserved.