Module 9 — Distributed training across several accelerators

When a model fits in memory but training takes three days, distribution becomes interesting. TensorFlow makes it surprisingly simple to write, which hides two settings you absolutely must adjust by hand.

The principle: replicate the model, split the batch

The most common strategy replicates the entire model on each accelerator. Each replica receives a fraction of the batch, computes its gradients, then all gradients are averaged and applied identically everywhere. The replicas therefore stay rigorously in sync.

This is data parallelism. It assumes the model fits in a single accelerator's memory. When it no longer does — the large language models of course 16 — you need model parallelism, which splits the network itself and falls outside this module.

Three lines of code

import tensorflow as tf
from tensorflow import keras

strategy = tf.distribute.MirroredStrategy()
print(f"Replicas: {strategy.num_replicas_in_sync}")

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

model.fit(dataset, validation_data=val_dataset, epochs=20)

Everything that creates variables must sit inside the scope: model construction and compilation. fit itself stays outside. A variable created outside the scope is not replicated, and the resulting error is hard to trace back to its cause.

Strategy	Scope	Use
MirroredStrategy	several accelerators, one machine	the common case
MultiWorkerMirroredStrategy	several machines	requires network configuration
TPUStrategy	tensor processing units	after connecting to the cluster
OneDeviceStrategy	a single device	debugging distributed code

Batch size is global, which is counter-intuitive

The first setting not to miss. The batch_size you supply is the global size: it gets divided across replicas.

BATCH_PER_REPLICA = 64
GLOBAL_BATCH = BATCH_PER_REPLICA * strategy.num_replicas_in_sync

dataset = dataset.batch(GLOBAL_BATCH).prefetch(tf.data.AUTOTUNE)

Keeping batch_size=64 on four accelerators gives sixteen examples per replica. Computation becomes inefficient — the accelerators are underused — and batch normalisation statistics degrade, since they are computed per replica rather than globally. Below thirty-two examples per replica, batch normalisation becomes distinctly noisy.

The global batch must divide evenly by the replica count

A remainder produces unequally sized replicas and, depending on the version, either an error or a silent imbalance. Always derive the global batch from the per-replica batch, never the other way round.

The learning rate must follow

The second setting, and the one most often forgotten. Multiplying batch size by four divides the number of updates per epoch by four. At a constant rate the model therefore learns four times less per epoch, and training appears to regress when it is merely slowed.

Two scaling rules circulate:

• linear: multiply the rate by the replica count. Works up to batches of a few thousand examples.
• square root: multiply by the square root of the factor. More cautious on very large batches.

BASE_RATE = 1e-3
rate = BASE_RATE * strategy.num_replicas_in_sync

with strategy.scope():
    model.compile(optimizer=keras.optimizers.Adam(rate), loss="mse")

A rate multiplied by four and applied from the very first batch destabilises training. The warmup mentioned in module 6 becomes nearly indispensable here: start from the base rate and reach the scaled rate over a few hundred steps.

Mixed precision, often more profitable than distribution

Before adding accelerators, there is a cheaper lever. Mixed precision stores weights in float32 but computes in float16, which exploits dedicated hardware units and roughly halves memory use.

keras.mixed_precision.set_global_policy("mixed_float16")

One precaution comes with this setting: the output layer must stay in float32. In float16 a softmax saturates and cross-entropy loses all numerical precision.

outputs = layers.Dense(num_classes, activation="softmax", dtype="float32")(x)

Keras handles loss scaling itself, which prevents small gradients from becoming zero in the reduced float16 range. If you have overridden train_step as in module 4, that scaling becomes your responsibility: optimizer.get_scaled_loss before differentiation, get_unscaled_gradients after.

What to optimise, in what order

The first bottleneck is almost always the data pipeline from module 5, and fixing it costs nothing. Next comes mixed precision, which often delivers close to a factor of two for one line of code. Distribution comes third: it multiplies the hardware bill and introduces extra settings. Use the profiler from module 7 to confirm computation really is the bottleneck before going there.

What changes across several machines

MultiWorkerMirroredStrategy extends the principle between machines but adds constraints that did not exist before. The TF_CONFIG environment variable must describe the cluster topology on every node. Network bandwidth becomes decisive, since gradients cross the network at every step. And crucially, checkpoint writing must be coordinated: each worker writes to a distinct temporary directory, and only the chief keeps the final file. keras.callbacks.BackupAndRestore handles that coordination and allows resuming after a node failure, which becomes statistically inevitable beyond a handful of machines.

Key takeaways

• Data parallelism replicates the whole model and averages gradients; it requires the model to fit in a single accelerator.
• Everything that creates variables goes inside the strategy scope; fit stays outside.
• batch_size is global: derive it from the per-replica batch, or you underuse the hardware and degrade batch normalisation.
• The learning rate must be scaled with the replica count and paired with a warmup; before distributing, fix the data pipeline then enable mixed precision, both far cheaper.

Next module: exporting the model as a SavedModel and serving it in production.