[论文解读] Stanza: Layer Separation for Distributed Training in Deep Learning

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

Advanced Neural Network Applications参考文献 47被引用 4

一句话总结

Stanza 通过将计算密集型的卷积（CONV）层与参数密集型的全连接（FC）层解耦，在分布式深度学习中实现层分离，即在大多数工作节点上训练 CONV 层，在少数节点上训练 FC 层，从而将节点间通信量减少多达 10 倍，并在 10Gb 网络下实现 1.34x–13.9x 的训练加速。该方法显著降低了参数服务器系统中的数据传输开销，同时不损害收敛性或模型精度。

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.

研究动机与目标

• 为解决基于参数服务器的分布式深度学习中因每轮训练迭代中大规模参数交换导致的高通信开销问题。
• 通过最小化工作者与参数服务器之间的数据传输量，减少训练时间，从而降低现代 GPU 集群中训练延迟的主导因素。
• 利用卷积神经网络中卷积层与全连接层在计算与通信特性上的差异，实现更高效的分布式训练。
• 设计一种系统，在大幅降低网络带宽使用的同时，保持模型收敛性与性能。

提出的方法

• 将深度学习模型分解为独立的 CONV 层与 FC 层训练流水线，将集群中的大多数节点分配用于训练 CONV 层，少数节点用于训练 FC 层。
• 采用混合通信策略：CONV-FC 激活值与梯度交换采用多对一和一对多通信模式，同时在 CONV 与 FC 工作者组内使用优化的 allreduce 实现梯度同步。
• 通过重叠 CONV 与 FC 工作者的梯度交换操作，隐藏通信延迟，提升流水线效率。
• 构建性能模型，基于节点分配预测训练吞吐量，并求解离线优化问题，以确定 CONV 与 FC 工作者的最优数量。
• 将 Stanza 集成至 PyTorch，支持在 Azure 和 AWS EC2 等云平台无缝部署，仅需极少代码修改。
• 支持与标准 SGD 的向后兼容性，并通过批量同步并行（BSP）机制保持模型一致性，同时保持对异步与过时同步训练的可扩展性。

实验结果

研究问题

• RQ1 是否可以通过分离卷积层与全连接层的训练，在不降低模型精度或收敛性的情况下，减少分布式深度学习中的通信开销？
• RQ2 如何优化 CONV 层与 FC 层之间的通信模式，以最小化延迟，同时保持训练效率？
• RQ3 在给定集群规模下，何种节点分配策略（CONV 与 FC 工作者数量）能最大化训练吞吐量？
• RQ4 在不同网络带宽条件下（包括 10Gb、40Gb 和 100Gb），分层训练方法的性能表现如何？
• RQ5 即使在通信并非主要瓶颈的高带宽数据中心环境中，该系统是否仍能实现显著的加速？

主要发现

• 在搭载 Tesla V100 GPU 与 10Gb 网络的 EC2 上，Stanza 在 AlexNet 和 VGG-16 等常见模型上相比传统参数服务器系统实现了 1.34x 至 13.9x 的加速。
• 对于 VGG-16，Stanza 相比标准参数服务器训练将数据传输量减少了高达 10 倍。
• 即使在 100Gb 网络带宽下，数值模拟结果表明 Stanza 对 AlexNet 和 VGG-16 仍分别实现 1.55x 和 1.72x 的加速，表明其优势在 10Gb 以上网络中依然持续存在。
• 系统通过使用批量同步并行（BSP）机制，保持了快速收敛与模型精度，确保一致性且无性能下降。
• 性能模型可实现最优节点分配，所得到的配置在多种模型与集群规模下均持续优于基线参数服务器系统。
• Stanza 与现有通信优化技术（如梯度压缩与量化）正交，可与之结合使用，进一步降低通信成本。