Distributed Training

import tensorflow as tf
import os
import json

# ============================================================
# 1) Single host, multiple GPUs — MirroredStrategy
# ============================================================
# Each GPU has a replica of the model. Gradients are AllReduced
# across replicas after each batch. Effectively scales batch size
# by num_gpus; tune learning rate accordingly.

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

GLOBAL_BATCH = 256 * strategy.num_replicas_in_sync
ds = (tf.data.Dataset.from_tensor_slices((tf.random.normal((50_000, 32)),
                                            tf.random.uniform((50_000,), 0, 10, dtype=tf.int32)))
        .shuffle(10_000)
        .batch(GLOBAL_BATCH)
        .prefetch(tf.data.AUTOTUNE))

with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(32,)),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10),
    ])
    model.compile(optimizer='adam',
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])

model.fit(ds, epochs=3)

# ============================================================
# 2) Multi-host — MultiWorkerMirroredStrategy
# ============================================================
# Each worker runs the SAME script and reads its rank from TF_CONFIG.
# Launch via your scheduler (k8s Job, SLURM, Ray, Vertex AI Training).

# Example TF_CONFIG (on worker 0):
# {
#   "cluster":  { "worker": ["w0:12345", "w1:12345"] },
#   "task":     { "type": "worker", "index": 0 }
# }
os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {'worker': ['w0:12345', 'w1:12345']},
    'task':    {'type': 'worker', 'index': 0},
})

mw_strategy = tf.distribute.MultiWorkerMirroredStrategy()

# Sharded dataset — each worker reads its own slice
options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
ds = ds.with_options(options)

with mw_strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(32,)),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10),
    ])
    model.compile(optimizer='adam',
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])

model.fit(ds, epochs=3,
          callbacks=[
              tf.keras.callbacks.BackupAndRestore('./bk/'),    # survive worker restarts
              tf.keras.callbacks.TensorBoard('./logs/mw'),
          ])

# Save the model from ONE worker only (worker 0); writing the same file
# from every worker will corrupt the checkpoint.
if json.loads(os.environ['TF_CONFIG'])['task']['index'] == 0:
    model.save('/checkpoints/distributed.keras')

# ============================================================
# 3) TPU — TPUStrategy
# ============================================================
# resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='local')
# tf.config.experimental_connect_to_cluster(resolver)
# tf.tpu.experimental.initialize_tpu_system(resolver)
# strategy = tf.distribute.TPUStrategy(resolver)

# ============================================================
# 4) Common pitfalls
# ============================================================
# - global_batch_size vs per_replica_batch_size: confusing both is the #1 bug.
# - learning rate that worked for 1 GPU underperforms at 8; warm up + linear scale.
# - tf.data input pipeline cannot keep the GPUs fed -> use AUTOTUNE, prefetch, cache.
# - Mixed precision: with strategy.scope(): tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')
#   gives 2x speed on TPUs and recent GPUs without code changes.
# - Logging metrics: aggregate via strategy.experimental_local_results then average.

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = build_model()
    model.compile(...)
model.fit(...)  # automatically sharded across GPUs