Table des matières
Jean Zay: MultiGPU and multinode distribution for training a TensorFlow or PyTorch model
The training of a neural network model is faced with two major problems:
 The training time can be very long (several weeks).
 The model can be too large for the GPU memory (32 GB or 16 GB, depending on the partition used).
The solution to these two problems can be found in multiGPU distribution. We distinguish between two techniques:
 Data parallelism which accelerates the model training by distributing each batch of data and therefore, each iteration upgrade of the model. In this case, the model is replicated on each utilised GPU to treat a data subset. It is adapted to trainings involving voluminous data or largesized batches. This is the most documented distribution process and is the easiest one to implement.
 Model parallelism which is adapted to models which are too large to be replicated and manipulated with the chosen batch size in each GPU. However, it is more difficult to implement, particularly in multinodes. It can be optimised with pipeline techniques but it cannot accelerate the training as significantly as data parallelism.
These two techniques can be used simultaneously (hybrid parallelism) to solve both problems at the same time (training time and model size) by grouping the GPUs in function of the task to be effectuated.
Builtin functionalities of TensorFlow and PyTorch
Data parallelism
Data parallelism consists of replicating the model on all the GPUs and dividing the data batch into minibatches. Each GPU takes charge of a minibatch, resulting in parallel learning on the entirety of the minibatches.
 The training script is executed in distributed mode on several workers (i.e. several GPUs or several groups of GPUs). In this way, each GPU :
 Reads a part of the data (minibatch).
 Trains the model from its minibatch.
 Calculates the parameter updates of the local model (local gradients).
 The global gradients are calculated from averaging the local gradients returned by all the GPUs. At this step, it is necessary to synchronise the different workers.
 The model is updated on all the GPUs.
 The algorithm is repeated beginning with step 1.a. until the end of the training.
Implementation
Model parallelism
Caution: This can be difficult to put in place because it requires, at minimum, the definition of a superclass for the model or the complete redefinition of the model sequentially. It would be preferable to use other optimisations to address the problems of memory occupation (Gradient Checkpointing, ZeRO of DeepSpeed, …). Only very big models (for example, for the most efficient transformers) will use model parallelism by absolute necessity.
Model parallelism consists of splitting a model onto more than one GPU in order to:
 Limit the volume of calculations to be effectuated by a single GPU during the training.
 To reduce the memory footprint of the model on each GPU (necessary for models which are too big to fit into the memory of a single GPU). In theory, the memory occupancy per GPU is thus divided by the number of GPUs.
The model is divided into groups of successive layers and each GPU takes charge of only one of these groups; this limits the number of parameters to process locally.
Unlike data parallelism, it is not necessary to synchronise the gradients: The GPUs collaborate sequentially to effectuate the training on each data input. The following figure illustrates model parallelism on 2 GPUs, thus dividing by 2 the size of the memory consumed in one GPU.
Note that model parallelism does not accelerate the learning because a GPU must wait for preceding layers of the model to be treated by another GPU before being able to execute its own layers. The learning is even a little longer because of the time of data transfer between GPUs.
Nevertheless, it is possible to decrease the time through process optimisation using a data pipeline technique: batch splitting. This consists of dividing the batch into microbatches in order to limit the wait time and have the GPUs work quasisimultaneously. However, there will still be a small wait time (the Bubble in the illustration).
The illustration below shows the quasisimultaneity and the gain obtained by using the splitting technique where the data batch is divided into 4 microbatches, for a training distributed on 4 GPUs.
The optimal number of microbatches depends on the model under consideration.
Implementation
For further information

 Optimal acceleration of learning and memory efficiency
 ZeRO, ZeRO Offload and ZeRO Infinity
 Activation checkpointing (equivalent to gradient checkpointing)
 Optimisers

 Simplification and good practices
 Sharded training (equivalent to ZeRO)
 Model parallelism
 DeepSpeed plugin
 Profiler