Introduction

Running the TensorFlow MNIST example in distributed mode usually means using tf.distribute so the same model trains across multiple devices or workers. The simplest entry point is MirroredStrategy on one machine with multiple GPUs, then MultiWorkerMirroredStrategy if you want the same pattern across several machines.

Start with the Single-Machine Distributed Pattern

For local distributed training, TensorFlow’s normal pattern is:

1. create a distribution strategy
2. build the model inside the strategy scope
3. feed a batched dataset into model.fit

MNIST is a good fit because the dataset is small and already packaged with Keras:

1import tensorflow as tf
2
3(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
4
5x_train = x_train.astype("float32") / 255.0
6x_test = x_test.astype("float32") / 255.0
7
8x_train = x_train[..., tf.newaxis]
9x_test = x_test[..., tf.newaxis]
10
11strategy = tf.distribute.MirroredStrategy()
12
13global_batch_size = 256
14
15train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
16train_ds = train_ds.shuffle(10000).batch(global_batch_size).prefetch(tf.data.AUTOTUNE)
17
18test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
19test_ds = test_ds.batch(global_batch_size).prefetch(tf.data.AUTOTUNE)
20
21with strategy.scope():
22    model = tf.keras.Sequential(
23        [
24            tf.keras.layers.Input(shape=(28, 28, 1)),
25            tf.keras.layers.Conv2D(32, 3, activation="relu"),
26            tf.keras.layers.MaxPooling2D(),
27            tf.keras.layers.Flatten(),
28            tf.keras.layers.Dense(128, activation="relu"),
29            tf.keras.layers.Dense(10, activation="softmax"),
30        ]
31    )
32    model.compile(
33        optimizer="adam",
34        loss="sparse_categorical_crossentropy",
35        metrics=["accuracy"],
36    )
37
38model.fit(train_ds, epochs=3, validation_data=test_ds)

If multiple GPUs are visible on the same machine, MirroredStrategy replicates the model on each one and reduces gradients automatically.

Build the Model Inside strategy.scope()

This is the most important structural rule. Variables, optimizer state, and model weights should be created inside the distribution scope so TensorFlow knows how to mirror or shard them correctly.

If you create the model outside the scope and then try to train under a strategy, you can end up with placement errors or unexpectedly non-distributed variables.

Understand the Batch Size Rule

With distribution strategies, the batch size you pass to the dataset is usually the global batch size, not the per-device batch size. TensorFlow splits that batch across the replicas.

For example, with two GPUs and global_batch_size = 256, each replica processes 128 examples per step. If you accidentally treat 256 as per replica and scale again, the effective batch size becomes much larger than intended and training dynamics change.

That matters even for MNIST, because optimizers and learning rates are sensitive to batch size.

Multi-Worker Training Uses the Same Model Code

If you move from one machine to multiple workers, the model code can stay almost the same. The main extra step is setting TF_CONFIG so each process knows the cluster layout:

1export TF_CONFIG='{
2  "cluster": {
3    "worker": ["host1:12345", "host2:12345"]
4  },
5  "task": {
6    "type": "worker",
7    "index": 0
8  }
9}'

Then in Python:

strategy = tf.distribute.MultiWorkerMirroredStrategy()

Everything else still follows the same broad idea: build inside the scope, feed datasets, call fit.

Make the Input Pipeline Efficient

The MNIST example is small enough that data loading is rarely the bottleneck, but the habit still matters. A good distributed input pipeline should:

• batch after shuffling
• prefetch data
• avoid unnecessary Python work during training

That is why prefetch(tf.data.AUTOTUNE) is present in the example. The data pipeline stays out of the way while the devices train.

Common Pitfalls

The most common mistake is building the model or optimizer outside strategy.scope(). That breaks the variable-placement assumptions of the distribution API.

Another pitfall is confusing global batch size with per-replica batch size. The number you batch in the dataset is usually the total batch seen across all replicas together.

It is also easy to expect MirroredStrategy to distribute work without actual multiple visible devices. If TensorFlow sees only a CPU or one GPU, the code still runs, but the “distributed” effect is limited.

Finally, multi-worker training fails quickly if TF_CONFIG is inconsistent across workers. When moving beyond one machine, cluster configuration becomes just as important as model code.

Summary

• Use MirroredStrategy for one machine with multiple GPUs and MultiWorkerMirroredStrategy for multiple workers.
• Build the model and optimizer inside strategy.scope().
• Treat the dataset batch size as the global batch size unless you are doing something custom.
• Use a batched, prefetched tf.data pipeline even for simple MNIST examples.
• For multi-worker runs, correct TF_CONFIG is part of the setup, not an optional extra.