LayerGenius

Scaling Keras Training to Multiple GPUs: What Actually Gets Faster

Adding GPUs is the first thing teams try when training is slow, and often the last thing that helps. Before you write a line of distribution code, measure where the time goes: a step that is input-bound at one GPU stays input-bound at four, and you will have paid for three idle accelerators. We covered that measurement in the tf.data performance checklist, and in review work it settles the question more often than not. Mixed precision is usually the cheaper win too — one policy line, frequently 1.5–2x on Ampere-class and newer hardware, no topology changes.

When you have exhausted those, data parallelism is the next step. This post covers how it works in Keras 3, the settings that are easy to get wrong, and how to tell whether it worked.

What data parallelism does and does not fix

In data-parallel training, every device holds a full copy of the model and processes a different slice of each batch. Gradients are all-reduced across devices, so every replica applies the same update. This helps when the bottleneck is compute per step and the model fits comfortably in one device's memory.

It does not fix:

  • Input starvation. If your tf.data pipeline delivers 400 images/sec and one GPU consumes 380, a second GPU gets you to roughly 400, not 760.
  • Models that do not fit. That is a model-parallel or sharding problem, not a data-parallel one.
  • Tiny models. With a small step time, communication and Python overhead dominate; two GPUs can be slower than one.

Decide which case you are in first. The honest answer is often "none of them, the pipeline is the problem."

Single host, TensorFlow backend

On the TF backend the mechanism is tf.distribute. Build and compile inside the strategy scope so variables are created as mirrored variables:

strategy = tf.distribute.MirroredStrategy()
print("replicas:", strategy.num_replicas_in_sync)

with strategy.scope():
    model = build_model()
    model.compile(
        optimizer=keras.optimizers.Adam(3e-4),
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"],
    )

model.fit(train_ds, validation_data=val_ds, epochs=30)

The scope covers model construction, optimizer creation, and compile. Data loading and fit stay outside it. Loading a checkpoint into a model built outside the scope, then wrapping it later, is a reliable way to get variables on the wrong devices.

Single host, Keras 3 distribution API

Keras 3 ships a backend-agnostic distribution API, currently most complete on JAX. It is declarative: describe the device mesh, set it globally, then write ordinary Keras code.

import keras

devices = keras.distribution.list_devices("gpu")
data_parallel = keras.distribution.DataParallel(devices=devices)
keras.distribution.set_distribution(data_parallel)

model = build_model()
model.compile(optimizer=keras.optimizers.Adam(3e-4), loss="sparse_categorical_crossentropy")
model.fit(train_ds, epochs=30)

No scope, no wrapper. ModelParallel with a LayoutMap exists for sharding large weights across devices, but reach for it only when a model genuinely does not fit — it adds a layout contract you now have to maintain through every architecture change.

If your codebase is still on the TF backend, keep MirroredStrategy. If you are already running JAX, prefer the Keras distribution API; it is the path that stays portable.

Batch size and learning rate: the part teams get wrong

In Keras, the batch_size you pass — or the batch size baked into your tf.data dataset — is the global batch size. It is split across replicas. Two consequences:

  1. Global batch size must be divisible by the replica count, or you get ragged final shards and, on some setups, an error mid-epoch.
  2. To keep per-device batch size and memory use constant when you add devices, you must scale the global batch size up. Four GPUs at a per-device batch of 64 means batch_size=256.

And once the global batch grows, the effective learning rate is wrong. A larger batch gives a lower-variance gradient, so the same step size under-trains. The standard starting point is linear scaling: multiply the base learning rate by the replica count, with a warmup of a few epochs so early steps do not blow up.

replicas = strategy.num_replicas_in_sync
global_batch = 64 * replicas
base_lr = 3e-4 * replicas

schedule = keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=0.0,
    warmup_target=base_lr,
    warmup_steps=steps_per_epoch * 3,
    decay_steps=steps_per_epoch * 30,
)

Linear scaling is a heuristic, not a law; it holds well up to moderate batch sizes and degrades past them. Treat it as the starting point and confirm with a run against your single-GPU baseline curve. If accuracy drops, the batch size grew past what the schedule supports — shorten warmup, lower the multiplier, or stop adding devices.

One more trap: BatchNorm statistics are computed per replica, over the per-device batch, not the global one. If scaling drives per-device batch down to 8 or 16, BN gets noisy and results shift. Fine-tuning with frozen BN sidesteps this; training from scratch at small per-device batches does not, and group normalization or a larger per-device batch is the usual answer.

Feed it, or none of this matters

With N replicas you need roughly N times the input throughput. Before scaling out:

  • Shard files across workers, not records: many TFRecord shards, num_parallel_reads set, interleave with AUTOTUNE.
  • prefetch(tf.data.AUTOTUNE) at the end of the pipeline, always.
  • Move decode and resize into the pipeline with num_parallel_calls=AUTOTUNE; move augmentation into the model where it runs on device.
  • cache() after decode if the decoded dataset fits in RAM. This is often the single biggest multi-GPU win, and it costs one line.

A quick check: run fit on a synthetic dataset of pre-batched random tensors of the right shape. That step time is your compute ceiling. If real training is much slower, you have an input problem and scaling out will not help yet.

Checkpoints, callbacks, and logs

ModelCheckpoint and TensorBoard behave under distribution, with two caveats. Write checkpoints to a path all workers can see, and make sure only the chief writes final artifacts — in multi-worker setups other replicas should write to temporary directories that are cleaned up. Metrics reported by fit are already aggregated across replicas; custom metrics that accumulate Python state, rather than keras.metrics objects, are not, and will report one replica's view. Use keras.metrics subclasses so aggregation happens correctly.

Also keep the evaluation path single-device unless the evaluation set is large. Distributed evaluation with a set that does not divide evenly across replicas is a common source of metrics that drift slightly between runs, and it is not worth the debugging time.

Measure scaling efficiency, then decide

Report two numbers after every scaling change:

  • Throughput: samples/sec at steady state, measured after the first epoch so tracing and warmup are excluded.
  • Scaling efficiency: throughput on N devices divided by (N x throughput on 1 device).

On a single host with a well-fed pipeline, 85–95% at 2–4 GPUs is a reasonable expectation for mid-size vision models. Below 70%, stop and find the bottleneck — it is almost always input, an under-sized per-device batch, or a model too small to amortize the all-reduce. And measure time-to-target-accuracy, not just samples/sec: a configuration that doubles throughput while needing three times the epochs to converge is a loss, and you only see it if you track both.

The order that works: profile, fix the input pipeline, enable mixed precision, then add devices — and re-measure at each step. Teams that start at the last step tend to buy hardware to work around a missing prefetch call.