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.

1. 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

2. 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

3. 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

4. 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.

5. 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

6. Remove the temporary smoke-test file after the distributed run has been proven.

   $ rm distributed-gpu-training.py