MultiWorker, TF_CONFIG, and Performance Recipes

From single machine to a fleet

MultiWorkerMirroredStrategy extends MirroredStrategy across multiple machines. Each machine runs a copy of the model on all its GPUs; gradients synchronize across workers AND devices via NCCL all-reduce.

TF_CONFIG is the JSON env variable each worker needs to know its peers and its own role. The cluster field lists all worker addresses; the task field tells this worker its type and index.

Performance recipe before going distributed: always optimize single-GPU first. Profile says <5% input bound? Mixed precision on? XLA on? Batch size as large as memory allows? Each of these can give 2–5× speedup before you add distribution complexity.

Code

MultiWorker setup — TF_CONFIG·python

import os, json, tensorflow as tf

# Set on each worker before starting training
# Worker 0 (chief):
os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
    },
    "task": {"type": "worker", "index": 0},   # different on each worker
})

# Strategy setup — same code on every worker
communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL,
)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options,
)

with strategy.scope():
    model = build_and_compile_model()

NUM_WORKERS    = 3
GPUS_PER_WORKER = 2
BATCH_PER_REPLICA = 64
GLOBAL_BATCH = BATCH_PER_REPLICA * NUM_WORKERS * GPUS_PER_WORKER

model.fit(dataset, epochs=50)

Performance recipe stack·python

import tensorflow as tf
from tensorflow.keras import mixed_precision

# 1. Linear LR scaling with batch size
BASE_LR = 1e-3
GLOBAL_BATCH_SIZE = 512
WARMUP_EPOCHS = 5
lr_schedule = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=BASE_LR * (GLOBAL_BATCH_SIZE / 32),
    decay_steps=1000,
    warmup_steps=WARMUP_EPOCHS * steps_per_epoch,
)

# 2. AUTOTUNE pipeline
AUTOTUNE = tf.data.AUTOTUNE
train_ds = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(10000)
    .batch(GLOBAL_BATCH_SIZE, drop_remainder=True)   # drop_remainder for TPU
    .map(preprocess, num_parallel_calls=AUTOTUNE)
    .prefetch(AUTOTUNE)
    .cache()
)

# 3. XLA compilation
model.compile(
    optimizer=optimizer,
    loss='sparse_categorical_crossentropy',
    jit_compile=True,    # auto on TPU, opt-in on GPU
)

# 4. Mixed precision (matches hardware)
mixed_precision.set_global_policy('mixed_float16')   # GPU
# or 'mixed_bfloat16'                                  # TPU

MultiWorker, TF_CONFIG, and Performance Recipes

From single machine to a fleet

Code

Progress

Comments 0