Scratch to SOTA: Build Famous Classification Nets 3 (Train/[Distributed] Data Parallelism)
Tags: Pytorch, Computer Vision, Step By Step Guide, Image Classification, Machine Learning

Some Perspective on Training with Multiple GPUs

Training a neural network with PyTorch is really easy. In its core, there are just three steps — forward pass to compute loss, use loss.backward() to compute gradients of the parameters and then optimizer.step() to update the parameters. Of course, there are some details and tips to this process that we will go through later.

The slightly trickier problem we are trying to address in this article is on using multiple GPUs for training.

As the model we use gets larger, it incurs more memory footprint, limiting the number of images that we can feed to the model in a batch. The mini-batch’s statistics deviate away more from that of the whole dataset and stochastic gradient descent may be less effective. Smaller batch also affects the running mean and standard deviation for batch normalization. We need to use data parallelism by scattering data to different GPUs in order to feed larger batch (this does not solve the batch normalization issue though, stay tuned for why later).

If the model gets so large that we cannot even fit a single image on a GPU, we need to use model parallelism by putting different parts of the model on different GPUs. We will not go through model parallelism here, but you can check out this simple tutorial.

Suppose that we want to use GPUs from multiple nodes (a fancy name for servers or machines in this context) to feed an even larger batch, we need to use distributed data parallelism in PyTorch. It allows different nodes to communicate with each other and collectively train a model. In fact, it is the recommended method for data parallelism even for a single server, i.e. there is little reason to use the data parallelism mentioned two paragraphs ago…, except that data parallelism is just easier to implement…

We will go through using distributed data parallelism on a single node in this article. Using it across multiple nodes is a simple extension without needing to modify the code.

Training with Data Parallelism

For data parallelism, we will use PyTorch’s nn.DataParallel class. How it works is as follow.

The model is stored on the default GPU. On forward pass, the model is replicated and sent to other GPUs. The tensor inputs to the model will be split along the specified dimension (0th dimension by default, which is the default batch dimension). Inputs of other type will be shallow copied. We wiil only send labels to the default GPU. The split data and copied data will be sent to the respective GPU devices and fed to the model replicas. As the input data passes through the model/replicas, any in-place updates to the parameters or buffers are only recorded on the default GPU (for this reason, batchnorm’s running statistics are only derived from a portion of the input).

All the outputs are sent back to the default GPU and concatenated along the specified dimension (batch dimension by default). Loss is then computed. The diagram below illustrates the forward passing process.

For backward pass (when we run loss.backward() ), the loss is scattered back to the the respective GPUs and gradients are computed for each replica according to the respective input data they received. Then, all the gradients are summed on the default GPU and parameter update is conducted there. The diagram below demonstrates this process.

This whole flow may seem complex, but its implementation is quite trivial with PyTorch’s nn.DataParallel class. All we need to do is just wrap our model with this wrapper.

However, before we jump into coding the training script, there is an issue with this data parallelization. You may have noticed that the default GPU has taken up more work than other GPUs, especially when hosting the gathered output and compute the final loss. This extra memory and computational load is insignificant for image classification, but can post a memory issue for tasks like object detection or text generation whose outputs are much larger. You can check out the discussion about it in PyTorch forum.

To solve the memory imbalance problem, we can either parallelize the loss computation (like this GitHub Gist), or embed the loss computation in the model’s forward pass, so that the model replicas are only sending scalar loss back to the default GPU, instead of huge tensors. We can achieve this with a model wrapper like this one.

After we wrap our model with this class, our wrapped model will take both the input data and labels as inputs, and output predictions and loss. In this way, the loss is directly computed on the respective GPUs.

Note: this wrapper will still gather all the outputs ( preds ) back onto the default GPU, but they do not take up much space for our task. We can remove preds from the return statement, if memory is an issue.

Code

We will use a YAML file to record all the training arguments. It makes it easier to manage a project and track different training configurations. Below is the YAML configuration file we will use for the training script. Most of them are common hyper-parameters, we will apply them along the way. distributed controls whether we will use distributed data parallelism or just data parallelism. gpus_per_node and total_nodes are only used for the distributed training.

Accompanying this config file, we will create a utility class for easy access of the arguments.

After wrapping our YAML file with this class, we can access the arguments by using dot notation like params.gpus , params.lr etc. If the argument does not exist in our YAML file, the dot notation will return None .

To start our training script proper, let’s conduct the import routine. We have coded all the utility classes in previous articles, except add_weight_decay , we will explain a bit later.

We define a train function, whose argument, model , is an instance of ModelWithLoss class and params is an instance of YAML file wrapper we coded just now.

Then, there are some helper functions. print_log() will append the strings we print to train_log.txt for training progress recording. save_checkpoint() saves a generic checkpoint.

We also creates the work directory, store the configuration YAML file we use for this training for record keeping.

Then we can initiate our training dataset and validation dataset. We will use center crop to evaluate our model’s accuracy on the validation set.

Next, we will decide whether to move the model to GPU. If our params.gpus has length one, we will move our model to GPU with that id. If it has lenth more than one, meaning that we are using multiple GPU, we first move the model to params.gpus[0] (as the model’s parameters and buffers must be placed on the default GPU) and wrap the model with nn.DataParallel(model, params.gpus) .

We will then define an optimizer along with weight decay. Weight decay is effective in preventing overfiting as it reduces the norm of the weights. Though not entirely accurate, it intuitively means that a lot of the weights are near zero, it is sort of like using less parameters in a model. Though not entirely accurate, it can also mean that the optimizer will prefer the weights being more spread out instead of concentrating on a few terms (as the weight decay is usually implemented using l2 regularization, one large parameter will contribute to loss a lot more than several not so large parameters). This forces the model to make decision based on the aggregation of many features instead of memorizing just a few features.

However, if our model has batch normalization layers, we should exclude the parameters from these layers when applying weight decay. BN shifts the output’s statistics to be centered around beta with a standard deviation of gamma (where beta and gamma are learnt along with other parameters). It does not make sense to always reduce them towards zero. The same in fact should be applied to all the bias terms.

Therefore, we implement the add_weight_decay() function to separate these parameters from the rest.

We can then create the optimizer and learning rate scheduler. Also, we should write code catering to resuming training from a specific checkpoint.

Finally, the standard training loop is as following. I wrote quite detailed comments to discuss the rationale behind most lines. It helps clear up relevant concepts regarding both the theories and coding practice.

Sanity Check

It has been quite a long script so far. Luckily, because of the extra lines dedicated to ensure flexibility, our sanity check is pretty simple. Almost all the training regimes can be set up by altering the YAML configuration file.

At the bottom of the training scripts, append this portion.

As can be seen, we will try to train an AlexNet.

Running the training script only requires this line

python train.py ./train_configs/train_args.yaml

As mentioned above, I was able to train an AlexNet to a 56.67% top 1 accuracy with single center crop evaluation after 92 epochs of training. This is 0.1% higher than the official torchvision pretrained model.

As I try to explain the training script clearly, the train() function spans multiple code blocks. If you find it hard to merge them back into a complete function, feel free to look up the scripts in my GitHub Repo.

Training with Distributed Data Parallelism

According to PyTorch documentation, there is little reason not to use distributed data parallelism instead of data parallelism (except, of course, distributed training is a bit more troublesome to set up). We will in this section go through PyTorch’s distributed data parallelism mechanism.

For distributed data parallelism, there are three distinct stages to take care of — construction, forward pass and backward pass.

To initialize distributed training, we need to initiate a model on each GPU on every machine. Each GPU in this context has a rank from 0 to worldsize-1 . World size here is the total number of GPU we are using for distributed training. If each node (server) equally has num_gpus GPUs, then the world size is simply num_nodes * num_gpus .

We should take note of Node 0’s IP address and the port we are using for distributed training. We need to provide this information to other Nodes so that they know where to access the “master” Node via network connection.

As parameters are never sent between GPUs, we can use seed to ensure the models are initialized with the same weights ( torch.manual_seed(0) ).

For forward pass, the batch normalization parameters are first broadcast to all the GPUs. Each model then get a slice of input to compute their own loss. In contrast to the forward pass of data parallelism, we do not send the model and parameters from GPU with rank 0 to other GPUs, but only broadcast the batch normalization parameters. Also, on each server, there should be a copy of the data set as we do not send the input data across the network (instead, we justuse torch.utils.data.distributed.DistributedSampler to ensure each process works on different input data).

For backward pass, gradients are computed separately on each GPU. They are then broadcast across all Nodes and averaged across all GPUs for parameter update. As the training settings for all ranks should be identical, the model parameters for all ranks are also updated the same way, all processes should thus arrive at the same weights each step.

Despite putting gradient computation and broadcasting as two separate steps, they in fact happen concurrently. As gradients are being computed for layer K, the gradients for layer K + 1 are broadcast for parameter update. This reduces communication bottleneck.

Note: For distributed training, we should imagine that we are opening a console to run the training script for every GPU across a few servers (which is in fact the case as each GPU has its own process for loading data and running the model). However, instead of getting different gradients and updates, at each step, the gradients are communicated across all GPUs and averaged, so that all model get the same gradients and update the weights in the same fashion. After training, all models on all the GPUs should have the same weights.

code

Multi-processing and communication between nodes in the context of training a neural network may seem daunting in the beginning. However, PyTorch has well-developed modules to make it a lot simpler.

Firstly, we will do a bunch of imports in additional to the data parallelism’s imports.

We will define a distributed_train() function that takes in an initiated model, path to the YAML configuration file and nr that is the Node’s id (rank).

We should also record Node 0’s IP address ( os.environ[“MASTER_ADDR”] )and a free port that we are using on Node 0 ( os.environ[“MASTER_PORT”] ).

We will then use mp.spawn(process_fn, nprocs = how_many_process_to_run, args = (inputs for process_fn)) to start several processes to run the process_fn that we defined.

We will will initiate num_gpus_per_node processes each process running train_process(i, model, train_params_file, nr) . i here will be the process index. mp.spawn will prepend the args tuple with the process index and send it the function.

Here, I send an instance of the model to the train_process() will is sort of equivalent to broadcasting the model parameters across the GPUs, as the args will be deep copied. Most articles I browsed do not do this. Instead, the model initiation is done in the train_process() . If that is the case, then we need to use torch.manual_seed() to ensure that the models in different processes are initialized with the same weights.

As the train_process() function is similar to the train() we wrote in the last section, I will only elaborate on the parts that are relevant to distributed training in the code snippets below. The full script can be accessed here.

In the train_process() , we will first compute the world size and the GPU(process)’s rank. Then we will initialize the process group. Finally, we send the model to the GPU we have been referring to and wrap the model with DistributedDataParallel() wrapper.

We will create the work directory and construct a SummaryWriter only if we are in the first GPU of the node. This ensures that each server only has one work directory and event folder.

For dataset, validation dataset can remain the same as before. However, for training dataset, we need to apply DistributedSampler() to ensure that each process is getting different data for training. We should also turn off shuffle in the data loader.

Resuming from checkpoint is similar as before. However, we need to ensure that the checkpoint is mapped to the right GPU before loading.

The training loop is quite standard, but for logging and saving checkpoint, we need to be conscious of the process’ GPU id and rank to avoid repeated work.

Finally, we should clean up.

Sanity Check

The sanity check is the same as before, all we need to do is set distributed to True in our YAML configuration file. The configuration file was only for using distributed training with one Node. However, the training scripts are written with training across multiple Nodes in mind. By changing the total_nodes , we should get it work across multiple servers.

Conclusion

This article is quite long. However, training is a key process in deep learning, so it does deserve more pages.

In the next article, we will go through the implementation of GoogLeNet and ResNet. Unlike AlexNet and VGG family we implemented in the previous article, GoogLeNet and ResNet has fancier structures. It is going to be fun to write code for them.

The GitHub repo for this series is here.

References

To clear up my concepts about data parallelism and distributed training. I read through a few articles. The links below along with PyTorch’s official documentation should give a comprehensive idea of these two concepts.