Data Parallelism with TensorFlow
This page explains how to distribute a neural network model implemented in TensorFlow code across multiple GPUs using the data parallelism method.
A practical example is provided as a Notebook at the bottom of the page to give you access to a functional implementation of the explanations below.
Implementation of a Distribution Strategy
To distribute a model in TensorFlow, you define a distribution strategy by creating an instance of the tf.distribute.Strategy class. This strategy allows you to control how data and computations are distributed across GPUs.
Choice of Strategy
MultiWorkerMirroredStrategy
TensorFlow provides several pre-implemented strategies. In this documentation, we only present the tf.distribute.MultiWorkerMirroredStrategy strategy. This has the advantage of being generic in that it supports both multi-GPU and multi-node configurations, without performance loss compared to other tested strategies.
When using the MultiWorkerMirroredStrategy strategy, the model variables are replicated across all detected GPUs. Each GPU processes a portion of the data (mini-batch), and collective reduction operations are performed to aggregate the tensors and update the variables on each GPU at each step.
Multi-Worker Environment
The MultiWorkerMirroredStrategy is a multi-worker version of the tf.distribute.MirroredStrategy strategy. It runs on multiple tasks (or workers), each of which is associated with a hostname and a port number. The set of tasks forms a cluster on which the distribution strategy relies to synchronise the GPUs.
The concept of worker (and cluster) notably enables execution across multiple compute nodes. Each worker can be associated with one or more GPUs. On Jean Zay, we recommend defining one worker per GPU.
The multi-worker environment of an execution is automatically detected from Slurm variables defined in your submission script using the tf.distribute.cluster_resolver.SlurmClusterResolver class.
The MultiWorkerMirroredStrategy strategy can rely on two types of inter-GPU communication protocols: gRPC or NCCL. On Jean Zay, it is recommended to request the use of the NCCL communication protocol to achieve the best performance.
Declaration
The MultiWorkerMirroredStrategy strategy can be declared in a few lines:
# build multi-worker environment from Slurm variables
cluster_resolver = tf.distribute.cluster_resolver.SlurmClusterResolver(port_base=12345)
# use NCCL communication protocol
implementation = tf.distribute.experimental.CommunicationImplementation.NCCL
communication_options = tf.distribute.experimental.CommunicationOptions(implementation=implementation)
# declare distribution strategy
strategy = tf.distribute.MultiWorkerMirroredStrategy(cluster_resolver=cluster_resolver,
communication_options=communication_options)
On Jean Zay, you can use port numbers between 10000 and 20000.
There is a TensorFlow limitation regarding strategy declaration: it must be done before any other call to a TensorFlow operation.
Integration with the Training Model
To replicate a model across multiple GPUs, it must be created within the strategy.scope() context.
The .scope() method provides a context manager that captures TensorFlow variables and communicates them to each GPU according to the chosen strategy. This is where you declare the elements that create model-related variables: loading a saved model, declaring a model, the model.compile() function, the optimizer, etc.
Here is an example of declaring a model to be replicated across all GPUs:
# get total number of workers
n_workers = int(os.environ['SLURM_NTASKS'])
# define batch size
batch_size_per_gpu = 64
global_batch_size = batch_size_per_gpu * n_workers
# load dataset
dataset = tf.data.Dataset.from_tensor_slices(...)
# [...]
dataset = dataset.batch(global_batch_size)
# model building/compiling need to be within `strategy.scope()`
with strategy.scope():
multi_worker_model = tf.keras.Sequential([
tf.keras.Input(shape=(28, 28)),
tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
#...
])
multi_worker_model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.SGD(learning_rate=0.001*n_workers),
metrics=['accuracy'])
- When creating the
datasetobject, the batch size to specify is the global batch size, not the batch size per GPU. - The use of certain optimizers such as SGD requires adjusting the learning rate proportionally to the global batch size, and thus to the number of GPUs.
Distributed Training of a tf.keras.Model Model
The model.fit() function from the Keras library automatically handles the distribution of training according to the chosen strategy. Training is therefore launched in the usual way. For example:
multi_worker_model.fit(train_dataset, epochs=10, steps_per_epoch=100)
- Input data distribution across different processes (data sharding) is automatically managed within the
model.fit()function. - The evaluation step is automatically performed in distributed mode by passing the evaluation dataset to the
model.fit()function:
multi_worker_model.fit(train_dataset, epochs=3, steps_per_epoch=100, validation_data=valid_dataset)
Distributed Training with a Custom Loop
To have greater control over model training and evaluation, it is possible to bypass the .fit() method and define your own distributed training loop. To do this, you must pay attention to three things:
- correctly distribute the dataset across the workers and the associated GPUs;
- train your model only on a mini-batch for each GPU and aggregate the local gradients to form the global gradient;
- correctly distribute the calculation of metrics across each GPU and aggregate the local metrics to form the global metrics.
1. Correctly Distribute the Dataset Across the Workers and the Associated GPUs
To manually distribute the dataset, you can use the .experimental_distribute_dataset(dataset) method, which takes a tf.data.Dataset as an argument and handles splitting the dataset equally for each worker. Consequently, you will start by defining the tf.data.Dataset.
Two important operations on such a dataset are splitting into batches and shuffling.
- Splitting into batches is mandatory, as we first need to split the dataset into so-called global batches. Then, during dataset distribution, these global batches will be split into as many mini-batches as there are workers.
- Shuffling is highly recommended for neural network training, but you must be very careful because two workers might end up with different shuffles of the dataset. This is why it is imperative to set a global seed to ensure the same shuffles are performed on the workers.
- The order of operations between splitting into batches and shuffling is important! The instructions
.shuffle().batch()and.batch().shuffle()do not produce the same effect. It is recommended to apply the first instruction.shuffle().batch(). The second instruction would shuffle the order of the batches, but the content inside the batches would remain identical from one epoch to another. - To facilitate the exact calculation of gradients and metrics, it is recommended to remove the last global batch if it is not of size
global_batch_size: to do this, use the optionaldrop_remainderargument ofbatch().
In summary:
# Create a tf.data.Dataset
train_dataset = ...
# Set a global seed
tf.random.set_seed(1234)
# Shuffle tf.data.Dataset
train_dataset = train_dataset.shuffle(...)
# Batching using `global_batch_size`
train_dataset = train_dataset.batch(global_batch_size, drop_remainder=True)
If batching is performed without drop_remainder, you can manually remove the last batch:
# Remove the last global batch if it is smaller than the others
# We suppose that a global batch is a tuple of (inputs,labels)
train_dataset = train_dataset.filter(lambda x, y : len(x) == global_batch_size)
We can finally obtain our tf.distribute.DistributedDataset. When iterating over this dataset, a worker will retrieve a tf.distribute.DistributedValues containing as many mini-batches as it has GPUs, for example:
# Suppose, we have 2 workers with 3 GPUs each
# Each global batch is separated into 6 mini-batches
train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
# On one worker, distributed_batch contains 3 mini-batches from the same global batch
# On the other worker, the distributed_batch contains the 3 remaining mini-batches from the same global batch
distributed_batch = next(iter(train_dist_dataset))
In this specific context, next is a synchronous operation; if one worker does not make this call, the other workers will be blocked.
2. Train Your Model Only on a Mini-Batch for Each GPU and Aggregate the Local Gradients to Form the Global Gradient
Let us now move to the training loop. We will consider the specific and recommended case on Jean Zay of one GPU per worker. We will define the train_epoch() function, which will launch training for one epoch:
@tf.function
def train_epoch(): # ------------------------------------------ train 1 epoch
def step_fn(mini_batch): # ---------------------------------- train 1 step
x, y = mini_batch
with tf.GradientTape() as tape:
# compute predictions
predictions = model(x, training=True)
# compute the loss of each input in the mini-batch without reduction
reduction = tf.keras.losses.Reduction.NONE
losses_mini_batch_per_gpu = tf.keras.losses.SparseCategoricalCrossentropy(reduction=reduction)(
y, predictions)
# sum all the individual losses and divide by `global_batch_size`
loss_mini_batch_per_gpu = tf.nn.compute_average_loss(losses_mini_batch_per_gpu,
global_batch_size=global_batch_size)
# compute the gradient
grads = tape.gradient(loss_mini_batch_per_gpu, model.trainable_variables)
# inside MultiWorkerMirroredStrategy, `grads` will be summed across all GPUs first
# before updating the parameters
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss_mini_batch_per_gpu
# loss computed on the whole dataset
total_loss_gpu = 0.0
train_n_batches = 0
for distributed_batch in train_dist_dataset:
train_n_batches += 1
total_loss_gpu += strategy.run(step_fn, args=(distributed_batch,))
# we get the global loss across the train dataset for 1 GPU
total_loss_gpu /= float(train_n_batches)
# `strategy.reduce()` will sum the given value across GPUs and return result on current device
return strategy.reduce(tf.distribute.ReduceOp.SUM, total_loss_gpu, axis=None)
- The
@tf.functiondecorator is a recommended optimization when making calls tostrategy.run(). - In
step_fn(), one GPU is used to compute the loss for each sample in a mini-batch. It is then important to sum these losses and divide the sum by the size of the global batch rather than the mini-batch size, since TensorFlow will automatically compute the global gradient by summing the local gradients. - In the case of at least 2 GPUs per worker, the
distributed_batchcontains multiple mini-batches, andstrategy.run()distributes these mini-batches across the GPUs. total_loss_gpuonly contains the loss computed on the portion of the dataset assigned to the GPU; the.reduce()method then allows you to obtain the loss over one epoch.
3. Correctly Distribute the Calculation of Metrics Across Each GPU and Aggregate the Local Metrics to Form the Global Metrics
Regarding the calculation of metrics during model evaluation, it closely resembles training, except for the removal of the gradient computation and weight update parts.
Compute Environment Configuration
Jean Zay Modules
To achieve optimal performance in distributed execution on Jean Zay, you must load one of the following modules:
tensorflow-gpu/py3/2.4.0-noMKLtensorflow-gpu/py3/2.5.0+nccl-2.8.3- any module with version ≥ 2.6.0
In other Jean Zay environments, the TensorFlow library was compiled with certain options that can lead to significant performance losses.
Slurm Reservation Configuration
A correct Slurm reservation contains as many Slurm tasks as workers, as the Python script must be executed on each worker. Here, we associate one worker per GPU, so as many Slurm tasks as GPUs.
TensorFlow defaults to using the HTTP protocol. To prevent this, you must disable Jean Zay's HTTP proxy by removing the environment variables http_proxy, https_proxy, HTTP_PROXY, and HTTPS_PROXY.
Here is an example reservation for 2 quadri-GPU nodes:
#!/bin/bash
#SBATCH --job-name=tf_distributed
#SBATCH --nodes=2 #------------------------- number of nodes
#SBATCH --ntasks=8 #------------------------ number of tasks / workers
#SBATCH --gres=gpu:4 #---------------------- number of GPUs per node
#SBATCH --cpus-per-task=10
#SBATCH --hint=nomultithread
#SBATCH --time=01:00:00
#SBATCH --output=tf_distributed%j.out
#SBATCH --error=tf_distributed%j.err
# deactivate the HTTP proxy
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
# activate the environment
module purge
module load tensorflow-gpu/py3/2.6.0
srun python3 -u script.py
Application Example
Multi-GPU and Multi-Node Execution with MultiWorkerMirroredStrategy
An example in Notebook form can be found in $DSDIR/examples_IA/Tensorflow_parallel/Example_DataParallelism_TensorFlow-eng.ipynb on Jean-Zay.
The example is a Notebook in which the training methods presented above are implemented and executed on 1, 2, 4, and 8 GPUs. The examples are based on the ResNet50 model and the CIFAR-10 database.
You must first copy the Notebook to your personal space (for example, in your $WORK):
cp $DSDIR/examples_IA/Tensorflow_parallel/Example_DataParallelism_TensorFlow-eng.ipynb $WORK
You can then run the Notebook from a Jean Zay frontend machine by selecting a TensorFlow kernel (see our documentation on accessing JupyterHub for more information on using Notebooks on Jean Zay).