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.
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:
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_DEVICESenvironment 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:
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:
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:
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.
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
- How to freeze weights in certain layer with Keras?
- How to freeze weights in certain layer with Keras?
- How to freeze/lock weights of one TensorFlow variable e.g., one CNN kernel of one layer
- How to get a tensorflow op by name?
- How to forecast using the Tensorflow model?
- How to freeze lf-net tensorflow model to use it with opencv dnn?
- How to generate a train-test-split based on a group id?
- How to generate random number in a given range as a Tensorflow variable
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.