How to run distributed GPU training in TensorFlow

Large TensorFlow training jobs can outgrow one GPU when each step spends most of its time in matrix operations. On a single host with several NVIDIA GPUs, MirroredStrategy lets a Keras model train one synchronized batch across those devices instead of leaving extra accelerators idle.

MirroredStrategy creates one replica on each visible GPU, mirrors model variables, and combines gradients before the optimizer update completes. Model, optimizer, metrics, and any checkpoint restore that creates variables belong inside strategy.scope() so model.fit() runs with distributed variables.

Single-host GPU distribution is different from multi-worker training. If TensorFlow lists fewer than two GPUs, stop at the device checks and fix the GPU environment first; a one-replica run is still valid TensorFlow, but it is not distributed GPU training.

Steps to run distributed GPU training in TensorFlow with MirroredStrategy:

Confirm that the host driver sees at least two GPUs.
```
$ nvidia-smi --query-gpu=index,name,memory.total --format=csv
index, name, memory.total [MiB]
0, NVIDIA RTX 4090, 24564 MiB
1, NVIDIA RTX 4090, 24564 MiB
```
If this command fails or prints fewer than two rows, repair the NVIDIA driver, container runtime, or WSL2 GPU passthrough before changing the training script.
Related: How to enable GPU acceleration in TensorFlow
Confirm that the active TensorFlow environment sees the same GPUs.
```
$ python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'),
 PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')]
```
An empty list or a single device means the run below cannot prove distributed GPU training. Fix the GPU environment before continuing.
Related: How to fix TensorFlow not detecting a GPU

Save the training smoke test as distributed-gpu-training.py.

distributed-gpu-training.py

import os
 
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
import tensorflow as tf
 
tf.get_logger().setLevel("ERROR")
tf.keras.utils.set_random_seed(7)
 
gpus = tf.config.list_physical_devices("GPU")
if len(gpus) < 2:
    raise SystemExit(f"Need at least 2 visible GPUs; found {len(gpus)}")
 
strategy = tf.distribute.MirroredStrategy()
if strategy.num_replicas_in_sync < 2:
    raise SystemExit(
        f"Need at least 2 synchronized replicas; found {strategy.num_replicas_in_sync}"
    )
 
per_replica_batch = 16
global_batch = per_replica_batch * strategy.num_replicas_in_sync
 
features = tf.random.stateless_normal((640, 8), seed=(7, 11))
score = (
    features[:, 0] * 0.9
    + features[:, 1] * 0.5
    - features[:, 2] * 0.4
    + features[:, 3] * 0.2
)
labels = tf.cast(score > 0, tf.float32)[:, None]
 
train_ds = (
    tf.data.Dataset.from_tensor_slices((features[:512], labels[:512]))
    .shuffle(512, seed=7, reshuffle_each_iteration=True)
    .batch(global_batch)
    .prefetch(tf.data.AUTOTUNE)
)
val_ds = (
    tf.data.Dataset.from_tensor_slices((features[512:], labels[512:]))
    .batch(global_batch)
    .prefetch(tf.data.AUTOTUNE)
)
 
with strategy.scope():
    model = tf.keras.Sequential(
        [
            tf.keras.layers.Input(shape=(8,)),
            tf.keras.layers.Dense(32, activation="relu"),
            tf.keras.layers.Dense(16, activation="relu"),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ]
    )
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
        loss=tf.keras.losses.BinaryCrossentropy(),
        metrics=[tf.keras.metrics.BinaryAccuracy(name="binary_accuracy")],
    )
 
print(f"visible_gpus={len(gpus)}")
print(f"replicas={strategy.num_replicas_in_sync}")
print(f"global_batch={global_batch}")
 
history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=3,
    shuffle=False,
    verbose=2,
)
 
print(f"final_val_binary_accuracy={history.history['val_binary_accuracy'][-1]:.4f}")

The dataset is batched with the global batch size, which is the per-replica batch multiplied by the number of synchronized replicas.
Related: How to optimize TensorFlow data pipeline performance

Run the smoke test and confirm that the replica count is at least two before the epochs finish.

$ python distributed-gpu-training.py
visible_gpus=2
replicas=2
global_batch=32
Epoch 1/3
16/16 - 1s - 40ms/step - binary_accuracy: 0.7930 - loss: 0.4713 - val_binary_accuracy: 0.8750 - val_loss: 0.2806
Epoch 2/3
16/16 - 0s - 13ms/step - binary_accuracy: 0.9570 - loss: 0.1509 - val_binary_accuracy: 0.9219 - val_loss: 0.1358
Epoch 3/3
16/16 - 0s - 13ms/step - binary_accuracy: 0.9727 - loss: 0.0716 - val_binary_accuracy: 0.9609 - val_loss: 0.0717
final_val_binary_accuracy=0.9609

visible_gpus=2 confirms device discovery, replicas=2 confirms MirroredStrategy opened two synchronized replicas, and the final validation metric confirms model.fit() completed inside that distributed context.

Move the same strategy block into the project training script before model creation.

strategy = tf.distribute.MirroredStrategy()
per_replica_batch = 16
global_batch = per_replica_batch * strategy.num_replicas_in_sync
 
train_ds = make_train_dataset().batch(global_batch).prefetch(tf.data.AUTOTUNE)
val_ds = make_validation_dataset().batch(global_batch).prefetch(tf.data.AUTOTUNE)
 
with strategy.scope():
    model = build_model()
    model.compile(
        optimizer="adam",
        loss="binary_crossentropy",
        metrics=["binary_accuracy"],
    )
 
model.fit(train_ds, validation_data=val_ds, epochs=10, shuffle=False)

MirroredStrategy uses every GPU visible to TensorFlow unless a devices= list is passed. Larger global batches often need learning-rate retuning because each optimizer step sees more examples.
Related: How to compile a model in Keras
Related: How to train, evaluate, and run prediction with a Keras model

Remove the temporary smoke-test file after the distributed run has been proven.
```
$ rm distributed-gpu-training.py
```