Distributed Training with tf.distribute

Data parallelism is a three-line change: wrap the model construction in a strategy scope and keep everything else the same. What that change does to your batch size, your number of updates and therefore your effective learning rate is where the difficulty lives, and none of it requires a GPU to demonstrate.

What a strategy does to your batch

The three-line change

strategy = tf.distribute.MirroredStrategy()
 
with strategy.scope():
    model = build_model()                    # variables created once per replica
    model.compile(optimizer, loss, metrics=["accuracy"])
 
model.fit(x, y, batch_size=64 * strategy.num_replicas_in_sync)

That last line is the part people miss. The batch_size you pass is the global batch: it is split across replicas, so each one sees batch_size / num_replicas.

Replicas	Global batch	Per replica	Gradient updates	Accuracy
1	64	64	282	0.8813
2	128	64	141	0.8840
2 (batch not scaled)	64	32	282	0.8813

Scaling the batch with the replica count keeps each device as busy as it was — but it halves the number of gradient updates for the same data. Two replicas processing 128 samples at a time take half as many steps through the epoch as one replica processing 64.

Left: paired bars of gradient updates over three epochs for one and two replicas, under a scaled global batch (282 then 141) and a fixed one (282 for both). Right: accuracy against replica count for both conventions, with all four points between 0.8813 and 0.8840. — Halving the updates is the real cost of scaling out, and it is invisible in the accuracy column here only because three epochs on MNIST is a forgiving test — 0.8840 against 0.8813. On a job where the update count is what limits convergence, this is the difference between finishing and not, and it is why the learning rate is conventionally scaled alongside the batch.

The standard remedy is the linear scaling rule: multiply the learning rate by the number of replicas, so that a step over a batch of 128 moves as far as two steps over batches of 64 would have. It is a heuristic, it needs a warm-up at large replica counts, and it stops working well past a few thousand samples per batch — but it is the right default.

The scaffolding is not free

Left: horizontal bars of milliseconds per step — 2.90 with no strategy, 4.36 with MirroredStrategy at one replica, 6.13 at two. Right: seconds for three epochs against replica count, showing a flat measured line against a dashed line for what N accelerators would ideally give. — One replica is 1.50x slower than no strategy at all, and it parallelises nothing — that is pure scaffolding: variable placement, mirrored reads, and an all-reduce that has exactly one participant. Two replicas cost 2.11x. On real accelerators these costs are repaid many times over; the point of measuring them here is that they exist and are not proportional to the work being distributed.

Configuration	ms/step	Relative
No strategy	2.90	1.00×
`MirroredStrategy` ×1	4.36	1.50×
`MirroredStrategy` ×2	6.13	2.11×

The one-replica row is the informative one. It does no distribution at all and still costs 50% more per step, because the strategy still places variables, still routes reads through mirrored copies, and still runs the all-reduce machinery over a single participant.

This is why distributing a small model is often slower than not distributing it, on any hardware: below some model size the per-step coordination outweighs what the extra devices contribute.

diagram Diagram mermaid

  flowchart TD
  B["global batch of 128"] --> S1["replica 1: 64 samples"]
  B --> S2["replica 2: 64 samples"]
  S1 --> G1["gradient 1"]
  S2 --> G2["gradient 2"]
  G1 --> AR["all-reduce: average"]
  G2 --> AR
  AR --> U["one update, applied to
both mirrored copies"]
  U -.->|"282 updates becomes 141"| L["scale the learning rate"]
  AR -.->|"equal shards: 2.2e-08 error
uneven shards: 1.8e-02"| E["shards must be equal"]

Is the averaged gradient the real gradient?

Data parallelism rests on a claim that is easy to state and easy to check: the mean of the per-replica gradients equals the gradient of the whole batch. It is true, and only when the shards are the same size.

A bar chart on a log scale of the largest disagreement with the whole-batch gradient. Two, four and eight equal shards all sit near 2e-08, while two unequal shards sit at 1.8e-02, six orders of magnitude higher. — Equal shards agree to 2.235e-08 on weights whose largest component is about 0.089 — a relative error of 2.5e-07, which is float32 rounding and nothing more. Two unequal shards of 200 and 56 samples disagree by 1.812e-02, a relative error of 20%, because averaging two per-shard means weights the 56-sample shard as heavily as the 200-sample one.

Sharding	Max difference	Relative
2 equal shards	2.235e-08	2.505e-07
4 equal shards	2.235e-08	2.505e-07
8 equal shards	1.490e-08	1.670e-07
2 unequal shards	1.812e-02	2.031e-01

This is the same bug as the per-batch metric averaging on the custom training loops page, wearing different clothes. Averaging means-of-means is only correct when the groups are equal, and in a distributed job the group that is different is usually the last batch of the epoch.

tf.distribute handles this for you — it drops or pads the remainder rather than averaging unequal shards — but a hand-written distributed loop must do so explicitly, and a custom loss that reduces with tf.reduce_mean inside a replica context has exactly this problem. The correct pattern sums the per-example losses and divides by the global batch size:

Reducing a loss correctly under a strategy

per_example = keras.losses.sparse_categorical_crossentropy(labels, logits)
loss = tf.nn.compute_average_loss(per_example, global_batch_size=GLOBAL_BATCH)

Which strategy to reach for

MirroredStrategy — several devices on one machine. Variables are mirrored, gradients all-reduced. This is the one measured above and the one most people need.
MultiWorkerMirroredStrategy — the same idea across machines, coordinated by a TF_CONFIG environment variable. The failure modes become network failure modes.
ParameterServerStrategy — asynchronous, with variables held on dedicated servers. Scales to many workers and gives up gradient synchrony to do it.
TPUStrategy — for TPUs, where the batch must be divisible by the number of cores.

All four share the interface used above, which is the genuine achievement of the API: the training code does not change, only the scope it is built inside.

sketch Replicas, batches and updates p5.js

Drag the replica count and choose whether the global batch scales with it. The readout shows what happens to the per-replica batch and the number of gradient updates.

sketch The measured table, ranked p5.js

Click a column to rank every row by it. The bars are that column's values and the highest and lowest are computed from the numbers, not written in.

Pitfalls

Forgetting the batch is global. Two replicas at batch_size=64 give each replica 32 samples, not 64.
Not scaling the learning rate. Two replicas halved the updates from 282 to 141; the same learning rate now travels half as far per epoch.
Distributing a small model. One replica cost 1.50× per step and parallelised nothing.
Building the model outside the scope. Variables must be created inside strategy.scope() or they are not mirrored.
Averaging unequal shards. 20% relative error here; use compute_average_loss with the global batch size.
Assuming more replicas is always faster. The coordination cost is per step and does not shrink as you add devices.
Measuring speed-up on virtual CPU devices. They share the same cores; only semantics and overhead are measurable that way.

Recap

The batch_size passed to fit is the global batch, split across replicas.
Scaling it with the replica count halved the gradient updates, 282 → 141, which is why the learning rate is conventionally scaled too.
MirroredStrategy cost 1.50× per step at one replica and 2.11× at two, before any parallelism.
Averaging equal shards reproduces the whole-batch gradient to 2.2e-08; unequal shards are wrong by 20%.
Reduce losses with compute_average_loss(..., global_batch_size=...) rather than reduce_mean.
The four strategies share one interface; the training code does not change, only the scope.

Training is finished. The remaining question is how anything else reaches the model: Serving Models with TensorFlow Serving.

pch.quizTag pch.quizDefaultTitle

With two replicas and the global batch scaled to 128, the number of gradient updates fell from 282 to 141. Why does that matter?

pch.quizShowAnswer

B — The same data is seen but the parameters are updated half as often, so at an unchanged learning rate the model travels half as far per epoch — which is why the linear scaling rule exists — Accuracy barely moved here (0.8840 against 0.8813) because three epochs on MNIST is forgiving, but on an update-limited job this is the difference between converging and not.
MirroredStrategy with ONE replica cost 1.50x per step against no strategy at all. What is that overhead?

pch.quizShowAnswer

B — Real scaffolding — variable placement, reads routed through mirrored copies, and the all-reduce machinery running over a single participant — It is why distributing a small model can be slower than not distributing it: below some size the per-step coordination outweighs what extra devices contribute.
Averaging gradients over equal shards matched the whole-batch gradient to 2.2e-08, but two unequal shards were wrong by 1.8e-02. What causes that?

pch.quizShowAnswer

B — Averaging per-shard means weights each shard equally regardless of how many samples it holds, so a 56-sample shard counts as heavily as a 200-sample one — It is the same bug as averaging per-batch accuracies, and in a distributed job the odd-sized group is usually the last batch of the epoch.
Why must the model be constructed inside strategy.scope()?

pch.quizShowAnswer

B — Because that is where variables get created as mirrored copies — build outside the scope and you have ordinary variables that the strategy cannot keep in sync — The training call itself stays unchanged, which is the point of the API: only the construction scope differs.
How should a custom loss be reduced inside a replica context?

pch.quizShowAnswer

B — With tf.nn.compute_average_loss and the GLOBAL batch size — reduce_mean divides by the per-replica count, which reintroduces the unequal-shard weighting — tf.distribute handles the built-in path by dropping or padding the remainder, but a hand-written reduction has to be explicit about the denominator.

🧪 Try It Yourself

pch.coffeeTagline

pch.coffeeCta

pch.feedbackHeading

pch.feedbackSubheading