Fix Keras model.fit Failing Under MirroredStrategy (Multi-GPU)

tf.distribute.MirroredStrategy runs data-parallel training across the GPUs on one machine. With Keras, model.fit handles the distribution for you — if you set it up correctly.

The correct pattern

Create and compile the model inside strategy.scope() so its variables are mirrored:

import tensorflow as tf
 
strategy = tf.distribute.MirroredStrategy()
print("Replicas:", strategy.num_replicas_in_sync)
 
with strategy.scope():
    model = create_model()          # build here
    model.compile(                  # compile here
        optimizer="adam",
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"],
    )
 
# fit OUTSIDE the scope is fine — the model is already distributed
model.fit(train_ds, epochs=10)

The three usual failures

1. Model/optimizer built outside the scope. Distributed (mirrored) variables must be created inside strategy.scope(). If you build the model first and only wrap fit, you'll get errors about variables not belonging to the strategy. Fix: move construction and compile inside the scope.

2. Global batch size not scaled. fit splits the global batch across replicas. With 2 GPUs, a global batch of 64 runs 32 per GPU. Scale the global batch up with the replica count so each GPU stays well-fed:

per_replica = 32
global_batch = per_replica * strategy.num_replicas_in_sync
train_ds = train_ds.batch(global_batch)

(And typically scale the learning rate with the global batch.)

3. A non-replica-aware custom step. If you wrote a custom train_step or metric, it must reduce correctly across replicas. For standard Keras training you don't need this — plain model.fit is already replica-aware.

Don't "test" with strategy.run(model, ...)

A common wrong fix. strategy.run executes a per-replica step function, not a model:

# ❌ wrong — strategy.run takes a function, not a model
strategy.run(model, args=())
 
# ✅ if you really need a custom loop:
@tf.function
def train_step(inputs):
    x, y = inputs
    with tf.GradientTape() as tape:
        loss = compute_loss(y, model(x, training=True))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss
 
strategy.run(train_step, args=(next(iter(dist_ds)),))

Common traps

• OOM under MirroredStrategy — the per-replica batch (global ÷ GPUs) is still too large.
• Metrics look wrong — a custom metric that isn't reduced across replicas.
• Building the model outside the scope — the single most frequent cause.

Related guides

• AI/ML interview & engineering hub

Sources

• TensorFlow — tf.distribute.MirroredStrategy