[Paper Review] Stanza: Layer Separation for Distributed Training in Deep Learning

Xiaorui Wu, Hong Xu+2 more|arXiv (Cornell University)|Dec 27, 2018

Advanced Neural Network Applications47 references4 citations

TL;DR

Stanza proposes layer separation in distributed deep learning by decoupling computation-heavy convolutional (CONV) layers from parameter-heavy fully connected (FC) layers, training CONV layers on most workers and FC layers on a few, reducing inter-node communication by up to 10x and accelerating training by 1.34x–13.9x on 10Gb networks. This approach significantly reduces data transfer overhead in parameter server systems without sacrificing convergence or model accuracy.

ABSTRACT

The parameter server architecture is prevalently used for distributed deep learning. Each worker machine in a parameter server system trains the complete model, which leads to a hefty amount of network data transfer between workers and servers. We empirically observe that the data transfer has a non-negligible impact on training time. To tackle the problem, we design a new distributed training system called Stanza. Stanza exploits the fact that in many models such as convolution neural networks, most data exchange is attributed to the fully connected layers, while most computation is carried out in convolutional layers. Thus, we propose layer separation in distributed training: the majority of the nodes just train the convolutional layers, and the rest train the fully connected layers only. Gradients and parameters of the fully connected layers no longer need to be exchanged across the cluster, thereby substantially reducing the data transfer volume. We implement Stanza on PyTorch and evaluate its performance on Azure and EC2. Results show that Stanza accelerates training significantly over current parameter server systems: on EC2 instances with Tesla V100 GPU and 10Gb bandwidth for example, Stanza is 1.34x--13.9x faster for common deep learning models.

Motivation & Objective

• To address the high communication overhead in parameter server-based distributed deep learning, especially due to large-scale parameter exchange in each training iteration.
• To reduce training time by minimizing data transfer between workers and parameter servers, which dominates training latency in modern GPU clusters.
• To exploit the distinct computational and communication characteristics of convolutional and fully connected layers in CNNs for more efficient distributed training.
• To design a system that maintains model convergence and performance while drastically reducing network bandwidth usage.

Proposed method

• Decompose deep learning models into separate CONV and FC layer training pipelines, assigning most cluster nodes to train CONV layers and a few to train FC layers.
• Use a hybrid communication strategy: many-to-one and one-to-many for CONV-FC activation and gradient exchange, and optimized allreduce for intra-group gradient synchronization among CONV and FC workers.
• Implement overlapping of CONV and FC worker gradient exchanges to hide communication latency and improve pipeline efficiency.
• Develop a performance model to predict training throughput based on node assignment, and solve an offline optimization problem to determine the optimal number of CONV and FC workers.
• Integrate Stanza into PyTorch, enabling seamless deployment on cloud platforms like Azure and AWS EC2 with minimal code changes.
• Support backward compatibility with standard SGD and maintain model consistency using Bulk Synchronous Parallel (BSP), while remaining extensible to asynchronous and stale-synchronous training.

Experimental results

Research questions

• RQ1 Can separating the training of convolutional and fully connected layers reduce communication overhead in distributed deep learning without degrading model accuracy or convergence?
• RQ2 How can communication patterns between CONV and FC layers be optimized to minimize latency while maintaining training efficiency?
• RQ3 What is the optimal node assignment strategy (number of CONV vs. FC workers) that maximizes training throughput for a given cluster size?
• RQ4 How does layer separation perform under varying network bandwidths, including 10Gb, 40Gb, and 100Gb?
• RQ5 Can the proposed system achieve significant speedups even when communication is not the primary bottleneck, such as in high-bandwidth datacenter environments?

Key findings

• On EC2 with Tesla V100 GPUs and 10Gb network bandwidth, Stanza achieves 1.34x to 13.9x speedup over traditional parameter server systems for common models like AlexNet and VGG-16.
• For VGG-16, Stanza reduces data transfer volume by up to 10x compared to standard parameter server training.
• Even with 100Gb network bandwidth, numerical simulations show Stanza still provides 1.55x and 1.72x speedups for AlexNet and VGG-16, respectively, indicating sustained benefits beyond 10Gb networks.
• The system maintains fast convergence and model accuracy by using Bulk Synchronous Parallel (BSP), ensuring consistency without performance degradation.
• The performance model enables optimal node assignment, and the resulting configuration consistently outperforms baseline parameter server systems across multiple models and cluster sizes.
• Stanza is orthogonal to existing communication optimization techniques such as gradient compression and quantization, and can be combined with them to further reduce communication costs.