TensorFlow
MirroredStrategy
Session Implementation
Non-Keras Usage
Machine Learning Algorithms

tf.distribute.MirroredStrategy implementation with sessions not with Keras?

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

tf.distribute.MirroredStrategy is not limited to Keras fit, but modern TensorFlow expects you to use it through strategy.scope, distributed datasets, and strategy.run rather than through an old hand-managed Session training loop. If you are asking whether multi-GPU training is possible without Keras, the answer is yes; if you are asking whether the best pattern is still raw sess.run, the answer is usually no.

What MirroredStrategy Actually Does

MirroredStrategy creates replicated variables across the available GPUs on one machine and coordinates synchronous updates. Each replica runs the same step on a different slice of the batch, gradients are reduced, and the optimizer applies the merged update.

That distribution model works with high-level Keras APIs, but it also works with custom training loops. The important part is that the computation must run through the strategy APIs so TensorFlow knows which operations belong on each replica.

Custom Training Loop Without Keras fit

A low-level but modern approach uses a custom step function with tf.GradientTape and strategy.run.

python
1import tensorflow as tf
2
3strategy = tf.distribute.MirroredStrategy()
4GLOBAL_BATCH_SIZE = 64
5
6features = tf.random.normal((256, 10))
7labels = tf.random.uniform((256,), maxval=2, dtype=tf.int32)
8dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(GLOBAL_BATCH_SIZE)
9dist_dataset = strategy.experimental_distribute_dataset(dataset)
10
11with strategy.scope():
12    model = tf.keras.Sequential([
13        tf.keras.layers.Dense(16, activation="relu"),
14        tf.keras.layers.Dense(2)
15    ])
16    optimizer = tf.keras.optimizers.Adam()
17    loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
18
19    def compute_loss(labels, predictions):
20        per_example_loss = loss_obj(labels, predictions)
21        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)
22
23    @tf.function
24    def train_step(dist_inputs):
25        def step_fn(inputs):
26            x, y = inputs
27            with tf.GradientTape() as tape:
28                predictions = model(x, training=True)
29                loss = compute_loss(y, predictions)
30            grads = tape.gradient(loss, model.trainable_variables)
31            optimizer.apply_gradients(zip(grads, model.trainable_variables))
32            return loss
33
34        per_replica_losses = strategy.run(step_fn, args=(dist_inputs,))
35        return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
36
37    for batch in dist_dataset:
38        loss = train_step(batch)
39        print(float(loss))

This is not Keras fit, but it is still the supported distributed pattern in modern TensorFlow.

Where Sessions Fit In

If you are still on graph-mode TensorFlow 1 style code, you may see tf.compat.v1.distribute.MirroredStrategy. That compatibility API exists, but it is legacy and substantially more awkward than the TensorFlow 2 workflow.

The key point is that MirroredStrategy is not “feed placeholders into one session and hope the GPUs synchronize.” The strategy needs to participate in variable creation and per-replica execution. That is why the APIs revolve around strategy.scope, distributed datasets, and replica-aware step functions.

If you are maintaining an older codebase, you can sometimes bridge through tf.compat.v1, but for new code the cleaner answer is to stop designing around sessions altogether.

When Non-Keras Still Makes Sense

Avoiding Keras fit is reasonable when you need:

  • unusual loss aggregation
  • manual gradient accumulation
  • custom metric timing
  • multi-stage training logic
  • tighter control over checkpoint or step boundaries

Those are valid reasons to use a custom loop. They are not reasons to go back to a pure session-driven training style unless you are constrained by legacy infrastructure.

Common Pitfalls

  • Assuming MirroredStrategy requires Keras fit when it also supports custom training loops.
  • Assuming a traditional Session loop is the normal modern interface for distributed training.
  • Creating variables outside strategy.scope and then wondering why replication behaves incorrectly.
  • Forgetting to distribute the dataset before calling the per-replica training step.
  • Mixing TensorFlow 1 session-era patterns with TensorFlow 2 distribution APIs in a way that obscures control flow.

Summary

  • 'MirroredStrategy works without Keras fit, but not as a simple old-style sess.run pattern.'
  • The modern approach is strategy.scope plus distributed datasets plus strategy.run.
  • Custom training loops are the right low-level option when you need more control.
  • 'tf.compat.v1 exists for legacy graph-mode code, but it is not the best default for new work.'
  • If you want manual control, use the distribution APIs directly rather than bypassing them.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the 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.