TensorFlow - Distributed Computing

TensorFlow - Introduction to Distributed Computing with tf.distribute.Strategy

Distributed training allows you to train machine learning models on multiple devices (CPUs, GPUs, TPUs) or even multiple machines. This can significantly speed up training for large models or large datasets. TensorFlow provides the tf.distribute.Strategy API as the primary mechanism for distributed training.

This chapter introduces the concept and shows a high-level overview of how one might set up a multi-worker distributed training scenario. The older tf.train.ClusterSpec and tf.train.Server APIs are legacy; tf.distribute.Strategy is the modern, recommended approach.

Overview of tf.distribute.Strategy

tf.distribute.Strategy is an API that handles the complexities of distributing model training across various hardware configurations. Key strategies include:

• MirroredStrategy: Supports synchronous distributed training on multiple GPUs on one machine. Model variables are mirrored across all GPUs, and updates are applied synchronously.
• MultiWorkerMirroredStrategy: Supports synchronous distributed training across multiple workers (machines), each potentially with multiple GPUs. Similar to MirroredStrategy but for a multi-machine setup.
• TPUStrategy: For training on Tensor Processing Units (TPUs).
• ParameterServerStrategy: Supports asynchronous parameter server training, where some workers act as parameter servers (storing variables) and others as compute workers.

Using these strategies typically involves minimal changes to your existing TensorFlow Keras training code. The main steps are:

1. Instantiate a strategy.
2. Move model creation and compilation within the strategy’s scope (with strategy.scope():).
3. Prepare your dataset using tf.data.Dataset and ensure it’s properly sharded or distributed if necessary.

Conceptual Example: MultiWorkerMirroredStrategy

Setting up MultiWorkerMirroredStrategy involves configuring the TF_CONFIG environment variable on each machine (worker) participating in the cluster. This variable tells TensorFlow about the cluster’s structure (worker addresses, current worker’s type and index).

Step 1: Define TF_CONFIG (Conceptual - on each worker machine)

Imagine you have two worker machines, worker0.example.com and worker1.example.com, listening on port 12345.

On worker0.example.com, you would set TF_CONFIG (e.g., in your shell or Python script):

import json
import os

# Example TF_CONFIG for worker 0
tf_config_worker0 = {
    'cluster': {
        'worker': ['worker0.example.com:12345', 'worker1.example.com:12345']
    },
    'task': {'type': 'worker', 'index': 0}
}

# In a real scenario, this would be set as an environment variable BEFORE the script runs.
# For demonstration, we might set it in Python if this script IS worker 0.
# os.environ['TF_CONFIG'] = json.dumps(tf_config_worker0)

On worker1.example.com, TF_CONFIG would be similar, but with 'index': 1.

# Example TF_CONFIG for worker 1
tf_config_worker1 = {
    'cluster': {
        'worker': ['worker0.example.com:12345', 'worker1.example.com:12345']
    },
    'task': {'type': 'worker', 'index': 1}
}
# os.environ['TF_CONFIG'] = json.dumps(tf_config_worker1)

Step 2: Python script for training (run on all workers)

The same Python script would be run on all worker machines after TF_CONFIG is set.

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import os
import json

# --- This part is for simulation if not running in a real multi-worker env ---
# In a real multi-worker setup, TF_CONFIG is an environment variable.
# To simulate locally for this tutorial, we can try to set it if not present.
if 'TF_CONFIG' not in os.environ:
    # Simulate a single worker setup if no TF_CONFIG is found for demonstration purposes.
    # For true multi-worker, TF_CONFIG must be set externally for each worker.
    print("TF_CONFIG not set. Simulating single-worker strategy or defaulting.")
    # strategy = tf.distribute.get_strategy() # Gets default strategy
    strategy = tf.distribute.MirroredStrategy() # For local multi-GPU, or fallback
    if not tf.config.list_physical_devices('GPU'):
        print("No GPUs found, MirroredStrategy will run on CPU. May be slow.")
else:
    # If TF_CONFIG is set, MultiWorkerMirroredStrategy will be used.
    strategy = tf.distribute.MultiWorkerMirroredStrategy()
    # You can access cluster information if needed:
    # cluster_resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()
    # print(f"Task type: {cluster_resolver.task_type}, Task ID: {cluster_resolver.task_id}")
    # print(f"Cluster spec: {cluster_resolver.cluster_spec()}")

print(f'Number of devices: {strategy.num_replicas_in_sync}')

# Hyperparameters
BUFFER_SIZE = 10000
BATCH_SIZE_PER_REPLICA = 64
GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
EPOCHS = 3

# Create a simple dataset (e.g., MNIST)
(x_train, y_train), _ = keras.datasets.mnist.load_data()
x_train = x_train.astype('float32') / 255.0
x_train = np.expand_dims(x_train, axis=-1)
y_train = keras.utils.to_categorical(y_train, num_classes=10)

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(BUFFER_SIZE).batch(GLOBAL_BATCH_SIZE)

# Model creation and compilation within strategy.scope()
with strategy.scope():
    model = keras.Sequential([
        layers.Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(28, 28, 1)),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

print("Model created and compiled within strategy scope.")

# List devices visible to TensorFlow within this strategy
# In a real multi-worker setup, this would show devices across the cluster from the perspective
# of the current worker and strategy. For a local MirroredStrategy, it lists local GPUs/CPUs.
# For TF1.x style device listing (sess.list_devices()), the closest in TF2 is checking
# physical devices and how the strategy maps to them.
print("\nPhysical Devices available to TensorFlow:")
for device_type in ['CPU', 'GPU']:
    devices = tf.config.list_physical_devices(device_type)
    if devices:
        print(f"{device_type}s: {devices}")
    else:
        print(f"No {device_type}s found.")

# Train the model
if strategy.num_replicas_in_sync > 0: # Ensure we have replicas
    print(f"\nStarting training on {strategy.num_replicas_in_sync} replicas...")
    model.fit(train_dataset, epochs=EPOCHS)
    print("Training complete.")
else:
    print("No replicas available for training. Check strategy setup.")

# For more details on distributed training:
# https://www.tensorflow.org/guide/distributed_training

Key points for MultiWorkerMirroredStrategy:

• TF_CONFIG: Essential for workers to discover each other and coordinate.
• Strategy Scope: Model definition, optimizer creation, and model.compile() must happen inside with strategy.scope():.
• Data Sharding: tf.data.Dataset handles data sharding automatically for most strategies when using distribute_datasets_from_function or if the dataset is batched correctly.
• Global Batch Size: The batch size specified in dataset.batch() should be the global batch size (batch size per replica * number of replicas).
• Saving/Loading Models: Use tf.keras.callbacks.ModelCheckpoint and tf.keras.models.load_model. For MultiWorkerMirroredStrategy, it’s recommended that only the chief worker (worker 0) saves checkpoints to avoid conflicts. This can be managed with tf.distribute.experimental.coordinator.ClusterCoordinator or custom callbacks.

This example provides a conceptual overview. Setting up a true multi-worker environment requires network configuration and ensuring all workers can communicate. The tf.distribute.Strategy API significantly simplifies the process of writing distributed training code compared to older, more manual methods.