DeepAdvisors | Scaling Training Beyond One GPU: tf.distribute and the Keras 3 Distribution API

Most teams reach for a second GPU too early. The first question we ask on a "our training is too slow, we need a cluster" engagement is whether the single GPU they already have is actually busy - and about half the time it is not, which makes this a tf.data problem rather than a scaling problem. But the other half of the time the model genuinely does not fit or genuinely needs more throughput, and then you need distribution done properly.

This tutorial covers the two layers that matter on the modern stack: classic tf.distribute strategies, and the newer Keras 3 distribution API that works across backends.

Step 0: earn the right to scale out

Before adding hardware, confirm three things with the TensorBoard profiler:

import tensorflow as tf

tf.profiler.experimental.start("logs/profile")
model.fit(train_ds, epochs=1, steps_per_epoch=100)
tf.profiler.experimental.stop()

GPU compute utilisation is high (say 80%+). If it is not, you are input-bound or host-bound; fix that first, because distribution multiplies input pipelines and will make the problem worse.
Mixed precision is on. keras.mixed_precision.set_global_policy("mixed_float16") (or mixed_bfloat16 on Ampere and newer, or TPU) is often a 1.5-2x win for one line of code.
Batch size is at the memory limit. Small batches make collective communication dominate.

Only after those three does a second device pay for itself.

Single host, multiple GPUs: MirroredStrategy

Data parallelism on one machine is the easy, high-value case. Each GPU holds a full copy of the model, processes a different slice of the global batch, and gradients are all-reduced across devices.

import tensorflow as tf

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

PER_REPLICA_BATCH = 64
GLOBAL_BATCH = PER_REPLICA_BATCH * strategy.num_replicas_in_sync

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

train_ds = make_dataset(batch_size=GLOBAL_BATCH)
model.fit(train_ds, epochs=10)

Things that bite people here:

Everything stateful goes inside strategy.scope() - model construction, optimizer, metrics. Variables created outside the scope are not mirrored and you will get confusing errors or silently wrong training.
Batch size is global, not per replica. Keras splits the batch across replicas for you. If you keep the old batch size, each GPU gets a quarter of it and you gain nothing.
Scale the learning rate. Going from batch 64 to 256 usually needs a higher LR plus a warmup of a few hundred steps. Linear scaling with warmup is the safe default.
Custom losses must reduce correctly. If you write your own training loop, use tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH) rather than reduce_mean, or your gradients scale with replica count.

Custom training loops

When model.fit is not enough:

@tf.function
def train_step(iterator):
    def step_fn(batch):
        x, y = batch
        with tf.GradientTape() as tape:
            preds = model(x, training=True)
            per_example = loss_fn(y, preds)
            loss = tf.nn.compute_average_loss(per_example, global_batch_size=GLOBAL_BATCH)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        return loss

    per_replica = strategy.run(step_fn, args=(next(iterator),))
    return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica, axis=None)

Wrap the whole multi-step loop in one tf.function where you can; per-step Python overhead is visible once you have eight replicas waiting on the host.

Multiple hosts: MultiWorkerMirroredStrategy

Across machines, the same data-parallel model applies, but now the cluster topology comes from the TF_CONFIG environment variable:

{
  "cluster": {"worker": ["host-0:12345", "host-1:12345", "host-2:12345"]},
  "task": {"type": "worker", "index": 0}
}

Every worker runs the same script; only task.index differs. In Kubernetes this is usually a TFJob (Kubeflow training operator) that renders TF_CONFIG for you; on Vertex AI or SageMaker the managed training service does it.

strategy = tf.distribute.MultiWorkerMirroredStrategy()

with strategy.scope():
    model = build_model()
    model.compile(...)

model.fit(
    train_ds,
    epochs=50,
    callbacks=[tf.keras.callbacks.BackupAndRestore(backup_dir="gs://bucket/backup")],
)

The non-obvious requirements:

Shard the input. Give each worker a different subset of files. tf.data auto-shards by file when the dataset starts from a file list, which is why sharded TFRecords matter; if auto-sharding cannot work it falls back to sharding by element, which reads everything on every worker. Set the policy explicitly when you care:
```
options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.FILE
train_ds = train_ds.with_options(options)
```
Only chief writes. Checkpoints, TensorBoard logs, and the final export must go to a worker-specific temp path on non-chief workers, or the writes collide on shared storage. BackupAndRestore and ModelCheckpoint handle this internally; your own file writes do not.
Plan for preemption. On spot or preemptible nodes, a worker dying kills the whole job. BackupAndRestore plus frequent checkpoints turns a fatal failure into a restart that loses a few minutes.
Interconnect decides your ceiling. All-reduce over 10 GbE Ethernet with a large model will not scale; NVLink within a host and 100 Gb+ RDMA between hosts will. Measure before buying.

The Keras 3 distribution API

Keras 3 adds a backend-agnostic distribution layer (fully supported on JAX today, with TensorFlow support progressing). If you already train on JAX and serve on TensorFlow - the pattern from our Keras 3 multi-backend tutorial - this is the API to use.

Data parallelism across all visible devices:

import keras
from keras import distribution

data_parallel = distribution.DataParallel(devices=distribution.list_devices())
distribution.set_distribution(data_parallel)

model = build_model()          # variables replicated automatically
model.compile(...)
model.fit(train_ds, epochs=10)

Model parallelism, for models whose weights do not fit on one device (large language models, very wide embeddings):

mesh = distribution.DeviceMesh(
    shape=(2, 4),
    axis_names=["data", "model"],
    devices=distribution.list_devices(),
)

layout_map = distribution.LayoutMap(mesh)
layout_map[".*token_embedding.*"] = ("model", "data")
layout_map[".*attention.*kernel.*"] = ("model", None)
layout_map[".*feedforward.*kernel.*"] = (None, "model")

model_parallel = distribution.ModelParallel(layout_map=layout_map, batch_dim_name="data")
distribution.set_distribution(model_parallel)

The mesh says "split my 8 devices into 2 data-parallel groups of 4 model-parallel shards"; the layout map says which weight axis is sharded across the model axis. Getting the layout wrong is usually visible as an out-of-memory error or as communication time dominating the step - not as a wrong result, which makes it debuggable.

A useful rule: data parallel until you run out of memory, then model parallel for the layers that caused it. Sharding a model that fits comfortably on one device just adds all-gather traffic.

Prove you actually scaled

Always report scaling efficiency, not just wall clock:

efficiency = (throughput_on_N_devices) / (N * throughput_on_1_device)

Measure throughput in examples or tokens per second, after a warmup of at least 50 steps, excluding the first epoch (compilation, caching, autotuning). Numbers we consider acceptable on client work:

1 host, 2-8 GPUs, NVLink: 0.85-0.95 is normal for a decently sized model.
Multi-host over fast RDMA: 0.7-0.85.
Below 0.6: stop and profile. The usual culprits are an input pipeline that cannot feed N replicas, a batch size that is too small per replica, or gradient all-reduce on a model with many small variables (fuse them or raise batch size).

And check convergence, not only speed. A run that is 4x faster but needs 5x more epochs to reach the same validation metric is a regression. Compare validation curves against the single-device baseline at equal epochs, then at equal cost.

Cost sanity check

Before committing to a cluster, price three options for your actual model: one large GPU instance, N smaller instances, and a TPU slice with the Keras 3 or TPUStrategy path. For many of the mid-size models we see (vision backbones, tabular deep models, sub-3B fine-tunes), a single 8-GPU host beats a multi-host cluster on both cost and engineering time - and the LoRA fine-tuning approach removes the need to scale out at all for a lot of LLM work.

The official references worth having open: the TensorFlow distributed training guide, the multi-worker Keras tutorial, and the Keras distribution API docs.

When you have finished scaling

A distributed training setup is not done until the artifact it produces is servable. Export the SavedModel from the chief only, then hand it to the deployment path described in our TensorFlow Serving tutorial.

If your training jobs are the bottleneck on a delivery date, our senior TensorFlow consultants do short, focused engagements: profile the current run, decide whether scaling out is warranted, and stand up the multi-GPU or multi-host pipeline with checkpointing and monitoring in place. Get in touch with your model size, dataset size, and current step time, and we will tell you what we think is realistic.