[Paper Review] Efficient and Robust Parallel DNN Training through Model Parallelism on Multi-GPU Platform

Chi‐Chung Chen, Chia-Lin Yang+1 more|arXiv (Cornell University)|Sep 8, 2018

Advanced Neural Network Applications49 references78 citations

TL;DR

This paper proposes SpecTrain, a weight-prediction technique for pipelined model parallelism to mitigate staleness in multi-GPU DNN training, achieving high GPU utilization with comparable accuracy to data parallelism and up to 8.91x speedups on 4 GPUs.

ABSTRACT

The training process of Deep Neural Network (DNN) is compute-intensive, often taking days to weeks to train a DNN model. Therefore, parallel execution of DNN training on GPUs is a widely adopted approach to speed up the process nowadays. Due to the implementation simplicity, data parallelism is currently the most commonly used parallelization method. Nonetheless, data parallelism suffers from excessive inter-GPU communication overhead due to frequent weight synchronization among GPUs. Another approach is pipelined model parallelism, which partitions a DNN model among GPUs, and processes multiple mini-batches concurrently. This approach can significantly reduce inter-GPU communication cost compared to data parallelism. However, pipelined model parallelism faces the weight staleness issue; that is, gradients are computed with stale weights, leading to training instability and accuracy loss. In this paper, we present a pipelined model parallel execution method that enables high GPU utilization while maintaining robust training accuracy via a novel weight prediction technique, SpecTrain. Experimental results show that our proposal achieves up to 8.91x speedup compared to data parallelism on a 4-GPU platform while maintaining comparable model accuracy.

Motivation & Objective

• Motivate and analyze the limitations of data parallelism due to inter-GPU communication overhead on multi-GPU platforms.
• Explore pipelined model parallelism to improve GPU utilization and reduce inter-GPU data transfers.
• Identify and address the weight staleness issue inherent in pipelined model parallelism.
• Introduce SpecTrain, a weight-prediction mechanism based on smoothed gradients to maintain training robustness and accuracy.
• Evaluate throughput and accuracy across diverse CNN, FCN, and RNN models on a 4-GPU platform.

Proposed method

• Adopt a pipelined model-parallel training framework derived from PipeDream and introduce SpecTrain for weight prediction.
• Predict future weights using smoothed gradients from momentum SGD to estimate weight updates during early pipeline stages.
• Compute predicted weights with equation W_hat_{t+s} = W_t - s * eta * v_{t-1}, where v_{t-1} is the smoothed gradient.
• Determine the version difference s based on GPU index and whether the mini-batch is in forward or backward pass (specific formulas provided in the paper).
• Compare data parallelism, vanilla model parallelism, PipeDream (weight stashing), and SpecTrain across six models (CNN, FCN, RNN) on CIFAR-10/IMDb datasets.
• Provide throughput, convergence behavior, and accuracy analyses to demonstrate robustness and performance gains.

Experimental results

Research questions

• RQ1 How does data parallelism compare to model parallelism in terms of inter-GPU communication and training efficiency on multi-GPU platforms?
• RQ2 Can pipelined model parallelism achieve high GPU utilization without compromising accuracy, and what mechanisms can mitigate weight staleness?
• RQ3 Does SpecTrain’s weight prediction maintain or improve model accuracy while delivering throughput gains over data parallelism and existing pipelined methods?
• RQ4 What is the performance/accuracy trade-off of SpecTrain across CNN, FCN, and RNN models on multi-GPU systems?

Key findings

• Model parallelism reduces inter-GPU communication compared to data parallelism (average 13.4x less, up to 528x, across tested models).
• Pipelined model parallelism without staleness mitigation can be unstable and degrade accuracy, especially for larger models.
• SpecTrain provides weight prediction that mitigates staleness, yielding learning curves and final accuracies comparable to data parallelism for most workloads.
• On 4-GPU systems, SpecTrain achieves up to 8.91x throughput improvement over data parallelism for FCN/RNN models and shows no accuracy drop compared to data parallelism in most cases.
• Compared with PipeDream, SpecTrain avoids additional memory overhead from weight-stashing queues and maintains stability and accuracy.
• For Transformer and other CNN/RNN models, SpecTrain maintains robustness with training loss and validation accuracy close to or matching data parallelism, while Vanilla Model P. and PipeDream can incur accuracy penalties on several models.