Multi-GPU distributed training with PyTorch

Author: fchollet Date created: 2023/06/29 Last modified: 2023/06/29 Description: Guide to multi-GPU training for Keras models with PyTorch.
View in Colab • GitHub source
There are generally two ways to distribute computation across multiple devices:
Data parallelism, where a single model gets replicated on multiple devices or multiple machines. Each of them processes different batches of data, then they merge their results. There exist many variants of this setup, that differ in how the different model replicas merge results, in whether they stay in sync at every batch or whether they are more loosely coupled, etc.
Model parallelism, where different parts of a single model run on different devices, processing a single batch of data together. This works best with models that have a naturally-parallel architecture, such as models that feature multiple branches.
This guide focuses on data parallelism, in particular synchronous data parallelism, where the different replicas of the model stay in sync after each batch they process. Synchronicity keeps the model convergence behavior identical to what you would see for single-device training.
Specifically, this guide teaches you how to use PyTorch's DistributedDataParallel module wrapper to train Keras, with minimal changes to your code, on multiple GPUs (typically 2 to 16) installed on a single machine (single host, multi-device training). This is the most common setup for researchers and small-scale industry workflows.
Let's start by defining the function that creates the model that we will train, and the function that creates the dataset we will train on (MNIST in this case).
Next, let's define a simple PyTorch training loop that targets a GPU (note the calls to .cuda()).
In this setup, you have one machine with several GPUs on it (typically 2 to 16). Each device will run a copy of your model (called a replica). For simplicity, in what follows, we'll assume we're dealing with 8 GPUs, at no loss of generality.
At each step of training:
In practice, the process of synchronously updating the weights of the model replicas is handled at the level of each individual weight variable. This is done through a mirrored variable object.
To do single-host, multi-device synchronous training with a Keras model, you would use the torch.nn.parallel.DistributedDataParallel module wrapper. Here's how it works:
Here's the flow, where each step is split into its own utility function:
Time to start multiple processes: