TensorFlow
GPUs
Deep Learning
GPU Utilization
Machine Learning

How to force tensorflow to use all available GPUs?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

TensorFlow can see multiple GPUs without automatically training across all of them in the way people often expect. To make a model use every visible GPU on one machine, the practical answer is to ensure the GPUs are visible and then train under tf.distribute.MirroredStrategy(), which uses all available GPUs by default when you do not pass an explicit device list.

Step 1: Confirm TensorFlow Can See the GPUs

Start by checking what TensorFlow detects:

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4print("GPUs:", gpus)

If this prints an empty list, the problem is not TensorFlow strategy configuration. It is usually the environment: drivers, CUDA setup, container runtime, or a visibility restriction such as CUDA_VISIBLE_DEVICES.

Step 2: Do Not Hide GPUs by Accident

TensorFlow can only use GPUs that are visible to the process. Two common ways to restrict visibility are:

  • the CUDA_VISIBLE_DEVICES environment variable
  • 'tf.config.set_visible_devices(...)'

If you want all GPUs, do not narrow the visible device set before TensorFlow initializes them.

You may also want to enable memory growth so TensorFlow does not grab nearly all memory up front on every visible GPU:

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4for gpu in gpus:
5    tf.config.experimental.set_memory_growth(gpu, True)

This does not make training distributed by itself, but it makes multi-GPU setups easier to coexist with other processes.

Step 3: Train Under MirroredStrategy

For single-machine multi-GPU training, this is the standard TensorFlow pattern:

python
1import tensorflow as tf
2
3strategy = tf.distribute.MirroredStrategy()
4print("Replicas:", strategy.num_replicas_in_sync)
5
6with strategy.scope():
7    model = tf.keras.Sequential([
8        tf.keras.layers.Dense(128, activation="relu"),
9        tf.keras.layers.Dense(10, activation="softmax")
10    ])
11    model.compile(
12        optimizer="adam",
13        loss="sparse_categorical_crossentropy",
14        metrics=["accuracy"]
15    )
16
17dataset = tf.data.Dataset.from_tensor_slices((
18    tf.random.normal((1024, 32)),
19    tf.random.uniform((1024,), maxval=10, dtype=tf.int32)
20)).batch(32)
21
22model.fit(dataset, epochs=3)

According to the TensorFlow API docs, MirroredStrategy() uses all available GPUs when no device list is provided. That is the clearest answer to "use all available GPUs" for synchronous training on one machine.

Why Plain TensorFlow Code Is Not the Same Thing

People often see GPU memory allocated on several devices and assume the model is already training across all of them. That is not necessarily true.

Without a distribution strategy, TensorFlow may place some operations automatically, but standard Keras training on one model instance does not magically become multi-GPU data parallel training. You need the distribution strategy to replicate the model and coordinate gradient updates across devices.

When You Want a Specific Subset Instead

If you want only selected GPUs, pass them explicitly:

python
strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])

This is useful when a machine has four GPUs but a specific experiment should only use two.

Input Pipelines Still Matter

Even with all GPUs visible and mirrored training enabled, you may not see good utilization if the input pipeline is slow. To keep multiple GPUs busy, the dataset usually needs batching, prefetching, and enough CPU-side throughput.

python
dataset = dataset.shuffle(1000).batch(256).prefetch(tf.data.AUTOTUNE)

A multi-GPU strategy cannot fix a starved pipeline.

Common Pitfalls

The biggest pitfall is thinking there is a single switch called "use all GPUs." In practice, there are two separate requirements: the GPUs must be visible, and the training code must use a distribution strategy that actually replicates work across them.

Another pitfall is setting visible devices after TensorFlow has already initialized the GPUs. TensorFlow requires visibility configuration before device initialization.

A third pitfall is equating memory allocation with useful utilization. A GPU showing allocated memory does not guarantee it is actively participating in training.

Finally, be realistic about scaling. Some models, batch sizes, and input pipelines will not benefit much from more GPUs unless the workload is large enough.

Summary

  • First make sure TensorFlow can see all physical GPUs
  • Do not hide devices accidentally through visibility settings
  • Use tf.distribute.MirroredStrategy() for single-machine training across all visible GPUs
  • Enable memory growth if you want more cooperative GPU memory behavior
  • Multi-GPU training also depends on a fast enough input pipeline and a workload that can scale

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.