TensorFlow
tf.distribute.Strategy
Python
Machine Learning
Parallel Computing

What has to be inside tf.distribute.Strategy.scope?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

tf.distribute.Strategy.scope() is mainly about variable creation. Anything that creates TensorFlow variables which should be mirrored, sharded, or otherwise owned by the distribution strategy should be created inside the scope. In normal Keras workflows, that means model creation, optimizer creation, and usually model.compile().

The Main Rule

TensorFlow's distributed training guides explain the core idea clearly: variables created inside the strategy scope are associated with that strategy. So the simplest rule is:

  • create distributed variables inside the scope
  • run ordinary data loading and model.fit() outside unless a specific API says otherwise

That rule covers most real projects.

What Usually Goes Inside

A typical Keras setup looks like this.

python
1import tensorflow as tf
2from tensorflow import keras
3
4strategy = tf.distribute.MirroredStrategy()
5
6with strategy.scope():
7    model = keras.Sequential([
8        keras.layers.Dense(64, activation="relu", input_shape=(32,)),
9        keras.layers.Dense(10, activation="softmax")
10    ])
11
12    optimizer = keras.optimizers.Adam()
13
14    model.compile(
15        optimizer=optimizer,
16        loss="sparse_categorical_crossentropy",
17        metrics=["accuracy"]
18    )

This is the standard placement because the model and optimizer create variables that the strategy needs to manage correctly.

What Usually Stays Outside

Your input pipeline usually does not need to be created inside the scope.

python
1x = tf.random.normal((1024, 32))
2y = tf.random.uniform((1024,), maxval=10, dtype=tf.int32)
3
4dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(32)
5
6model.fit(dataset, epochs=3)

The dataset can usually be built before or after the scope, depending on the workflow. The important point is that the dataset itself is not the thing whose variables must be strategy-owned.

Custom Training Loops

If you write a custom loop, the same principle applies. Create the model, optimizer, and any tf.Variable objects inside the scope. Then define the step function and run training logic around them.

python
1strategy = tf.distribute.MirroredStrategy()
2
3with strategy.scope():
4    model = keras.Sequential([
5        keras.layers.Dense(1, input_shape=(4,))
6    ])
7    optimizer = keras.optimizers.SGD()
8    loss_obj = keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE)

The distributed execution wrapper and per-replica step function can live outside, but the strategy-managed state should already exist.

Checkpoints And Metrics

If a checkpoint tracks model and optimizer state, create those tracked objects inside the scope as well because they depend on the distributed variables already created there.

Metrics often work most smoothly when created in the same place as the model and optimizer, especially in custom loops where they own state variables of their own.

Why model.compile() Is Commonly Inside

compile() can create optimizer slot variables later during training, but in distributed Keras examples it is commonly placed inside the scope so all strategy-related state is configured together. This also makes the setup easier to reason about.

Model Restoration Follows The Same Principle

If you restore weights into a model that must live under the strategy, create the model inside the scope first, then load the checkpoint. The restoration step depends on variables that already belong to the distributed model.

Common Pitfalls

A common mistake is creating the model outside the scope and then calling fit() inside it. By that point the important variables already exist, so the strategy cannot retroactively recreate them as distributed variables.

Another mistake is putting everything inside the scope, including unrelated data loading and preprocessing, as if the scope were a general "distributed mode" wrapper. It is mainly about strategy-managed state creation.

It is also easy to forget about custom tf.Variable objects in custom loops. If they should be mirrored or strategy-owned, create them inside the scope too.

Summary

  • Put model creation inside strategy.scope().
  • Put optimizer creation inside the scope as well.
  • In Keras workflows, model.compile() usually belongs inside too.
  • Datasets and model.fit() usually do not need to be inside.
  • The real rule is about where strategy-managed variables are created.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.