[论文解读] Parallax: Sparsity-aware Data Parallel Training of Deep Neural Networks

Soojeong Kim, Gyeong-In Yu+7 more|arXiv (Cornell University)|Aug 8, 2018

Advanced Neural Network Applications参考文献 37被引用 6

一句话总结

Parallax 是一种感知稀疏性的数据并行训练框架，根据模型参数的稀疏度动态结合参数服务器（Parameter Server）与AllReduce架构，以最小化通信开销。在48张GPU上训练NLP模型时，其速度比Horovod快6.02倍，比TensorFlow快2.8倍，同时在稠密模型上保持或超越现有框架的性能。

ABSTRACT

The employment of high-performance servers and GPU accelerators for training deep neural network models have greatly accelerated recent advances in deep learning (DL). DL frameworks, such as TensorFlow, MXNet, and Caffe2, have emerged to assist DL researchers to train their models in a distributed manner. Although current DL frameworks scale well for image classification models, there remain opportunities for scalable distributed training on natural language processing (NLP) models. We found that current frameworks show relatively low scalability on training NLP models due to the lack of consideration to the difference in sparsity of model parameters. In this paper, we propose Parallax, a framework that optimizes data parallel training by utilizing the sparsity of model parameters. Parallax introduces a hybrid approach that combines Parameter Server and AllReduce architectures to optimize the amount of data transfer according to the sparsity. Experiments show that Parallax built atop TensorFlow achieves scalable training throughput on both dense and sparse models while requiring little effort from its users. Parallax achieves up to 2.8x, 6.02x speedup for NLP models than TensorFlow and Horovod with 48 GPUs, respectively. The training speed for the image classification models is equal to Horovod and 1.53x faster than TensorFlow.

研究动机与目标

• 为解决现有深度学习框架在训练自然语言处理（NLP）模型时因稀疏模型参数处理不当而导致可扩展性差的问题。
• 通过利用模型变量的稀疏性，降低分布式数据并行训练中的通信开销，尤其针对具有大嵌入层的NLP模型。
• 设计一种混合通信架构，根据每个变量的稀疏度自动选择最优通信模式（参数服务器或AllReduce），且无需用户干预。
• 通过自动将单GPU计算图转换为可扩展的分布式图，实现无缝、高性能的分布式训练，同时最大限度减少用户工作量。
• 在统一框架中提升稀疏模型（如NLP）与稠密模型（如图像分类）的训练吞吐量与扩展效率。

提出的方法

• Parallax 引入一种混合通信架构，基于每个变量的稀疏度估计，对稀疏变量使用参数服务器（PS），对稠密变量使用AllReduce。
• 在训练时对计算图执行自动图变换，根据变量稀疏度在图中插入适当的通信操作（如allreduce、push/pull）。
• 对于稀疏变量，Parallax 将其划分为更小的块，并使用PS风格的push/pull操作仅传输被访问的元素，从而最小化网络流量。
• 该框架在图分析阶段估算每个变量的稀疏度，并基于成本模型选择最优通信模式——稀疏变量使用PS，稠密变量使用AllReduce。
• 它与TensorFlow集成，支持同步SGD，在保持与现有深度学习工作流兼容的同时实现可扩展训练。
• Parallax 采用一种新颖的稀疏变量分区策略，以最大化并行性并最小化通信与计算开销。

实验结果

研究问题

• RQ1 如何使数据并行训练在稀疏模型（尤其是NLP模型）中更具可扩展性，其中大型嵌入层主导了通信成本？
• RQ2 在不同类型的模型参数（稀疏与稠密）中，参数服务器或AllReduce通信架构哪一个能提供最佳性能？
• RQ3 能否通过为每个变量选择最优架构的混合通信模型，实现比单体式方法更好的整体可扩展性？
• RQ4 自动图变换在多大程度上可使用户在不修改代码的情况下实现高性能分布式训练？
• RQ5 Parallax 在稀疏与稠密深度学习模型上与Horovod和TensorFlow等先进框架相比，性能如何？

主要发现

• 在48张GPU上训练NLP模型（如语言模型和神经机器翻译模型）时，Parallax 相较于Horovod实现最高6.02倍的加速，相较于TensorFlow实现2.8倍的加速。
• 在图像分类模型上，Parallax 的性能与Horovod相当，且在48张GPU上比TensorFlow快1.53倍，表明其在各类模型中均具备出色的可扩展性。
• 在48张GPU上，Parallax 将NLP模型的扩展效率提升至19.0%，而Horovod仅为7.0%，TensorFlow为19.0%，显示出分布式训练效率的显著提升。
• Parallax 通过选择性地对稀疏变量使用PS、对稠密变量使用AllReduce，有效降低通信开销，同时不牺牲训练精度。
• 自动图变换机制使用户仅需极少代码修改即可在分布式环境中训练其原有单GPU模型，实现接近最优的性能。
• Parallax 已开源并公开发布于 https://github.com/snuspl/parallax，有助于推动稀疏感知分布式训练的采用与进一步研究。