Multi GPU Training
TensorFlow
Data Parallelism
feed_dict
Machine Learning

Multi GPU Training in Tensorflow Data Parallelism when Using feed_dict

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

In legacy TensorFlow 1.x code, multi-GPU data parallelism usually means building one model tower per GPU, splitting each training batch in Python, computing gradients on each device, and then averaging those gradients before applying a single update. This works with feed_dict, but it is verbose and is best thought of as maintenance knowledge rather than the preferred design for new TensorFlow code.

The Core Idea of Synchronized Data Parallelism

The high-level pattern is simple:

  • each GPU gets the same model structure,
  • each GPU receives a different slice of the batch,
  • each tower computes its own loss and gradients,
  • the host averages gradients and applies one shared optimizer step.

The last part is what makes the training synchronized. If each GPU updates its own weights independently, you no longer have the usual synchronized data-parallel setup.

Build One Tower Per GPU

A classic TensorFlow 1.x pattern uses separate placeholders per GPU shard and variable reuse after the first tower.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5num_gpus = 2
6x_parts = [
7    tf.compat.v1.placeholder(tf.float32, shape=[None, 10], name=f"x_{i}")
8    for i in range(num_gpus)
9]
10y_parts = [
11    tf.compat.v1.placeholder(tf.int32, shape=[None], name=f"y_{i}")
12    for i in range(num_gpus)
13]
14
15optimizer = tf.compat.v1.train.AdamOptimizer(1e-3)
16
17
18def tower_model(inputs, reuse):
19    with tf.compat.v1.variable_scope("net", reuse=reuse):
20        hidden = tf.compat.v1.layers.dense(inputs, 32, activation=tf.nn.relu)
21        return tf.compat.v1.layers.dense(hidden, 2)
22
23
24tower_grads = []
25for i in range(num_gpus):
26    with tf.device(f"/gpu:{i}"):
27        logits = tower_model(x_parts[i], reuse=(i > 0))
28        loss = tf.reduce_mean(
29            tf.nn.sparse_softmax_cross_entropy_with_logits(
30                labels=y_parts[i], logits=logits
31            )
32        )
33        tower_grads.append(optimizer.compute_gradients(loss))

The reuse rule matters a lot. Only the first tower should create variables. Later towers must reuse those same variables.

Average the Gradients

After each GPU computes gradients, average them and apply a single update.

python
1import tensorflow as tf
2
3
4def average_gradients(tower_grads):
5    averaged = []
6    for grad_and_vars in zip(*tower_grads):
7        grads = [g for g, _ in grad_and_vars]
8        grad = tf.reduce_mean(tf.stack(grads, axis=0), axis=0)
9        variable = grad_and_vars[0][1]
10        averaged.append((grad, variable))
11    return averaged
12
13
14train_op = optimizer.apply_gradients(average_gradients(tower_grads))

This is the synchronization step. Without it, the towers are not cooperating on one shared update.

Split the Batch in Python and Feed Each Tower

With feed_dict, the host code has to split the training batch manually.

python
1import numpy as np
2import tensorflow as tf
3
4batch_x = np.random.rand(64, 10).astype("float32")
5batch_y = np.random.randint(0, 2, size=(64,), dtype="int32")
6
7feed = {
8    x_parts[0]: batch_x[:32],
9    y_parts[0]: batch_y[:32],
10    x_parts[1]: batch_x[32:],
11    y_parts[1]: batch_y[32:],
12}
13
14config = tf.compat.v1.ConfigProto(allow_soft_placement=True)
15with tf.compat.v1.Session(config=config) as sess:
16    sess.run(tf.compat.v1.global_variables_initializer())
17    sess.run(train_op, feed_dict=feed)

This works, but it also shows why the pattern is considered legacy. Python becomes responsible for device sharding and input feeding, which can become a bottleneck.

Why Newer TensorFlow Uses tf.distribute

Modern TensorFlow usually solves this problem with tf.distribute.MirroredStrategy or related APIs. Those tools handle replication, synchronization, and device placement for you and integrate better with tf.data and Keras.

So the feed_dict tower pattern is still useful to understand when inheriting old code, but it is rarely the best design for new training pipelines.

Common Pitfalls

A common mistake is forgetting variable reuse and accidentally creating a separate model on each GPU instead of shared towers.

Another issue is averaging losses but not gradients, which is not the same optimization rule.

Teams also often underestimate the Python-side bottleneck. Even if the GPUs are configured correctly, feed_dict can become the limiting factor when input delivery is slow.

Summary

  • Legacy TensorFlow multi-GPU training with feed_dict uses one tower per GPU and averaged gradients.
  • Each tower must share the same variables rather than creating a separate model copy.
  • Python splits the batch and feeds each tower separately.
  • Gradient averaging is what makes the update synchronized across GPUs.
  • For new TensorFlow code, tf.distribute is usually the better solution.

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.