Topolearntf.distribute scales Keras and custom training with few code changes, including on TPUs.
A `tf.distribute.Strategy` decides how variables and computation are spread across devices. You create the model, optimizer and metrics inside with strategy.scope():, and the strategy takes care of replication and gradient aggregation.
The main strategies:
MirroredStrategy: synchronous data parallelism across GPUs on one machine. Each replica gets a slice of every batch; gradients are all-reduced.MultiWorkerMirroredStrategy: the same across machines, configured through theTF_CONFIGenvironment variable.TPUStrategy: data parallelism on TPU cores.ParameterServerStrategy: asynchronous training with parameter servers, for very large clusters.
With Keras, model.fit works unchanged once the model was built in scope. The batch size you feed is the global batch, split across strategy.num_replicas_in_sync replicas, so scale it up (and usually the learning rate).
In custom loops, distribute the dataset with strategy.experimental_distribute_dataset(ds), run the step with strategy.run(step_fn, args=(batch,)), and combine results with strategy.reduce. Compute per-example losses and divide by the global batch size (for example with tf.nn.compute_average_loss) so gradients are scaled correctly when summed across replicas.
Example
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy() # all local GPUs (or the CPU)
print("replicas:", strategy.num_replicas_in_sync)
global_batch = 64 * strategy.num_replicas_in_sync
with strategy.scope(): # variables are mirrored
model = tf.keras.Sequential([
tf.keras.Input((20,)),
tf.keras.layers.Dense(64, activation="relu"),
tf.keras.layers.Dense(1),
])
model.compile(optimizer="adam", loss="mse")
x = tf.random.normal([4096, 20])
y = tf.reduce_sum(x, axis=1, keepdims=True)
ds = (tf.data.Dataset.from_tensor_slices((x, y))
.shuffle(4096).batch(global_batch).prefetch(tf.data.AUTOTUNE))
model.fit(ds, epochs=3, verbose=0) # each global batch is split across replicasRemember
- Build model, optimizer and metrics inside
strategy.scope() - Batches are global and split across replicas
- Custom loops: distribute the dataset,
strategy.run, scale losses by global batch
Watch out for
- Creating the optimizer outside the scope, which places its variables incorrectly.