Skip to content

Argonne Leadership Computing Facility

Distributed Training on ThetaGPU Using Data Parallelism

There are two schemes for distributed learning:

  1. Model parallelization: in this scheme, disjoint subsets of a neural network are assigned to different devices. Therefore, all the computations associated to the subsets are distributed. Communication happens between devices whenever there is dataflow between two subsets. Model parallelization is suitable when the model is too large to be fitted into a single device (CPU/GPU) because of the memory capacity. However, partitioning the model into different subsets is not an easy task, and there might potentially introduce load imbalance issues limiting the scaling efficiency. 

  2. Data parallelization: in this scheme, all the workers own a replica of the model. The global batch of data is split into multiple minibatches, and processed by different workers. Each worker computes the corresponding loss and gradients with respect to the data it posseses. Before the updating of the parameters at each epoch, the loss and gradients are averaged among all the workers through a collective operation. This scheme is relatively simple to implement. MPI_Allreduce is the only commu

Our recent presentation about the data parallel training can be found here: https://youtu.be/930yrXjNkgM

In this documentation, we would like to show how to do data parallel training on ThetaGPU.

Software environment setup

We are still in the process of setting up the software stacks on ThetaGPU. Currently, one can get TensorFlow, PyTorch, and Horovod with the following setup script.

source /lus/theta-fs0/software/datascience/thetagpu/anaconda3/setup.sh

TensorFlow with Horovod

1. Initialize Horovod

import horovod.tensorflow as hvd hvd.init()
After this initialization, the rank ID and the number of processes can be refered as hvd.rank() and hvd.size(). Besides, one can also call hvd.local_rank() to get the local rank ID within a node. This is useful when we are trying to assign GPUs to each rank.

2. Assign GPU to each rank

gpus = tf.config.experimental.list_physical_devices('GPU') 
for gpu in gpus: 
    tf.config.experimental.set_memory_growth(gpu, True) 
if gpus: 
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
In this case, we set one GPU per process: ID=hvd.local_rank()

3. Scale the learning rate

Typically, since we use multiple workers, the global batch is usually increases n times (n is the number of workers). The learning rate should increase proportionally as follows (assuming that the learning rate initially is 0.01).

opt = tf.train.AdagradOptimizer(0.01*hvd.size())

4. Wrap the optimizer with Distributed Optimizer

opt = hvd.DistributedOptimizer(opt)

5. Broadcast the model from rank 0

This is to make sure that all the workers will have the same starting point.

hooks = [hvd.BroadcastGlobalVariablesHook(0)]

6. Loading data according to rank ID

TensorFlow has some functions for parallel distribution of data. But for specific applications, the user might have to write their own data loader.

In general, one has two ways to deal with the data loading:

  1. Each worker randomly select one batch of data from the dataset at each step. In such case, each worker can see the entire dataset. It is important to make sure that the different worker have different random seeds so that they will get different data at each step.

  2. Each worker accesses a subset of dataset. One manually partition the entire dataset into different partions, and each rank access one of the partions.

In both cases, the total number of steps per epoch is nsamples / hvd.size().

7. Checkpointing on root rank

It is important to let only one process to do the checkpointing I/O lest perhaps the file been corrupted.

if hvd.rank() == 0: 
   checkpoint.save(checkpoint_dir)

8. Average metric across all the workers

Notice that in the distributed training, any tensor are local to each worker. In order to get the global averaged value, one can use Horovod allreduce. Below is an example on how to do the average.

def tensor_average(val, name): 
     tensor = torch.tensor(val) 
     if (with_hvd): 
         avg_tensor = hvd.allreduce(tensor, name=name) 
     else: 
         avg_tensor = tensor 
   return avg_tensor.item()
We provided some examples in: https://github.com/argonne-lcf/sdl_ai_workshop/blob/master/01_distributedDeepLearning/Horovod/tensorflow2_mnist.py

PyTorch with DDP

PyTorch has its own native parallelization library called DDP. We will provide more details on how to run this on ThetaGPU. The current PyTorch on ThetaGPU does not have DDP built in. We will update to our users once we have DDP.

For now, please refer to https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

MPI Profiling for data parallel training

We support two ways for profling the performance of data parallel training.

  1. mpitrace library MPI trace allows us to get a flat profiling of all the MPI function calls involved during the training. To enable this, one can set the environment variable
    export LD_PRELOAD=/lus/theta-fs0/software/datascience/thetagpu/hpctw/lib/libmpitrace.so
    
    Then run the application as usual. MPI profiling results will be generated after the run finishes mpi_profile.XXXX.[rank_id].

Below is an example output:

Data for MPI rank 0 of 8: 
Times and statistics from MPI_Init() to MPI_Finalize(). 
----------------------------------------------------------------------- 
MPI Routine #calls avg. bytes time(sec) 
----------------------------------------------------------------------- 
MPI_Comm_rank 3 0.0 0.000 
MPI_Comm_size 3 0.0 0.000 
MPI_Bcast 520 197140.6 0.518 
MPI_Allreduce 24561 208138.3 162.080 
MPI_Gather 126 4.0 0.363 
MPI_Gatherv 126 0.0 0.434 
MPI_Allgather 2 4.0 0.000 
----------------------------------------------------------------- 
MPI task 0 of 8 had the maximum communication time. 
total communication time = 163.396 seconds. 
total elapsed time = 187.298 seconds. 
user cpu time = 4127.728 seconds. 
system time = 728.100 seconds. 
max resident set size = 8403.938 MBytes. 

Rank 0 reported the largest memory utilization : 8403.94 MBytes 
Rank 0 reported the largest elapsed time : 187.30 sec 
----------------------------------------------------------------- 
Message size distributions: 
                       MPI_Bcast      #calls   avg. bytes       time(sec) 
                                         126          4.0          0.008 
                                           1          8.0          0.000 
                                         121         25.0          0.006 
                                          30        251.5          0.002 
                                          32        512.0          0.002 
                                          64       1024.0          0.005 
                                          44       2048.0          0.003 
                                          29       4092.8          0.003 
                                          16       8192.0          0.032 

                       MPI_Allreduce   #calls  avg. bytes       time(sec) 
                                        19780         8.0         90.822 
                                         4576        24.0         18.239 
                                           43      4004.0          0.295 
                                            5   2780979.2          0.469 
                                           50   8160289.2         20.893 
                                            9  11803392.0          0.964 
                                           48  28060640.0          3.293 
                                           50  64731668.5         27.105 

                       MPI_Gather      #calls   avg. bytes      time(sec) 
                                          126         4.0          0.363
The useful information here is the message size distribution.

  1. Horovod Timeline To perform Horovod timeline analysis, one has to set the environment variable HOROVOD_TIMELINE which specifies the file for the output. export HOROVOD_TIMELINE=timeline.json This file is only recorded on rank 0, but it contains information about activity of all workers. You can then open the timeline file using the chrome://tracing facility of the Chrome browser.

More details: https://horovod.readthedocs.io/en/stable/timeline_include.html