Distributed Training with TensorFlow

A single GPU eventually runs out of room, either in compute or in memory. Distributed training spreads the work across multiple GPUs, multiple machines, or even TPU pods, letting you train larger models on larger datasets in a fraction of the time. TensorFlow makes this surprisingly approachable through the tf.distribute API.

Introduction to Distributed Training

There are two broad ways to distribute training:

• Data parallelism: the model is copied onto every device, and each copy processes a different slice of the batch. Gradients are then combined (usually averaged) before being applied. This is by far the most common approach and the one TensorFlow's built-in strategies focus on.

• Model parallelism: the model itself is split across devices, with different layers or parameter shards living on different machines. This is used for models too large to fit on a single device, such as very large language models.

TensorFlow's tf.distribute.Strategy API abstracts away most of the complexity of data parallelism. You wrap your model creation and training step inside a strategy's scope, and TensorFlow handles variable replication, gradient aggregation, and synchronisation for you.

MirroredStrategy for Multi-GPU Training

MirroredStrategy is the strategy you reach for when training on a single machine with multiple GPUs. It creates a complete copy (a mirror) of every model variable on each GPU. During training, each GPU computes gradients on its own slice of the batch, and those gradients are summed across all GPUs using an efficient all-reduce algorithm before being applied identically to every copy.

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

print(f'Number of devices: {strategy.num_replicas_in_sync}')

with strategy.scope():

model = tf.keras.Sequential([

tf.keras.layers.Dense(256, activation='relu', input_shape=(784,)),

tf.keras.layers.Dense(10, activation='softmax'),

])

model.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

# Global batch size is split evenly across all GPUs

model.fit(x_train, y_train, epochs=10, batch_size=256)

Everything inside strategy.scope() is automatically replicated. The batch size you pass to fit() is the global batch size; TensorFlow divides it across the available GPUs so each one receives an equal slice.

TPU Training on Google Colab

Tensor Processing Units (TPUs) are custom chips built by Google specifically for matrix-heavy deep learning workloads. Colab offers free TPU access, and connecting to one requires a small setup step before you can use TPUStrategy.

import tensorflow as tf

# Detect and connect to the TPU cluster

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()

tf.config.experimental_connect_to_cluster(resolver)

tf.tpu.experimental.initialize_tpu_system(resolver)

strategy = tf.distribute.TPUStrategy(resolver)

print(f'Number of TPU cores: {strategy.num_replicas_in_sync}')

with strategy.scope():

model = tf.keras.Sequential([

tf.keras.layers.Dense(256, activation='relu', input_shape=(784,)),

tf.keras.layers.Dense(10, activation='softmax'),

])

model.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

Quick Tip

TPUs are pickiest about input pipelines. Always feed them through tf.data.Dataset rather than NumPy arrays directly, and avoid Python-side control flow inside the training step, since TPUs compile the entire step into a static graph ahead of execution.

Multi-Worker Training Strategies

When a single machine is not enough, MultiWorkerMirroredStrategy extends the mirrored approach across several physical machines, each potentially with its own GPUs. Every worker needs to know about every other worker through a TF_CONFIG environment variable, which defines the cluster topology before the script even starts.

import os, json

import tensorflow as tf

# Set on every worker BEFORE the script starts (cluster topology)

os.environ['TF_CONFIG'] = json.dumps({

'cluster': {

'worker': ['worker0.example.com:12345', 'worker1.example.com:12345']

},

'task': {'type': 'worker', 'index': 0}   # changes per machine

})

strategy = tf.distribute.MultiWorkerMirroredStrategy()

with strategy.scope():

model = tf.keras.Sequential([

tf.keras.layers.Dense(256, activation='relu', input_shape=(784,)),

tf.keras.layers.Dense(10, activation='softmax'),

])

model.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

Each worker runs the exact same script. The only difference between machines is the task.index value in TF_CONFIG, which tells that particular process which worker it is. For very large clusters where parameters do not fit on every machine, ParameterServerStrategy distributes the model's variables across dedicated parameter server machines instead of mirroring them everywhere.

Performance Optimization Technique

Adding more devices only helps if the input pipeline can keep them fed. A common mistake is to distribute training across four GPUs and discover they spend half their time idle, waiting on data. The following practices keep devices saturated:

# Prefetching and parallel preprocessing

AUTOTUNE = tf.data.AUTOTUNE

dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))

.shuffle(10000)

.batch(256)

.prefetch(AUTOTUNE))

# Mixed precision

tf.keras.mixed_precision.set_global_policy('mixed_float16')

Profiling with TensorBoard's Profiler is the fastest way to find out whether your bottleneck is the GPU, the input pipeline, or communication between devices, rather than guessing.