[论文解读] Snap ML: A Hierarchical Framework for Machine Learning

Celestine Dünner, T. A. Parnell+6 more|arXiv (Cornell University)|Mar 16, 2018

Stochastic Gradient Optimization Techniques被引用 6

一句话总结

Snap ML 是一个分层机器学习框架，通过利用 GPU 加速、流水线化的磁盘外计算以及多节点集群中的优化通信模式，加速广义线性模型（GLMs）的训练。在 Criteo Terabyte Click Logs 数据集上，其训练速度比 TensorFlow 快 500 倍，仅用 1.5 分钟即可训练完一个逻辑回归模型——比以往任何报告的结果快一个数量级以上。

ABSTRACT

We describe a new software framework for fast training of generalized linear models. The framework, named Snap Machine Learning (Snap ML), combines recent advances in machine learning systems and algorithms in a nested manner to reflect the hierarchical architecture of modern computing systems. We prove theoretically that such a hierarchical system can accelerate training in distributed environments where intra-node communication is cheaper than inter-node communication. Additionally, we provide a review of the implementation of Snap ML in terms of GPU acceleration, pipelining, communication patterns and software architecture, highlighting aspects that were critical for achieving high performance. We evaluate the performance of Snap ML in both single-node and multi-node environments, quantifying the benefit of the hierarchical scheme and the data streaming functionality, and comparing with other widely-used machine learning software frameworks. Finally, we present a logistic regression benchmark on the Criteo Terabyte Click Logs dataset and show that Snap ML achieves the same test loss an order of magnitude faster than any of the previously reported results, including those obtained using TensorFlow and scikit-learn.

研究动机与目标

• 为云环境和企业环境中大规模机器学习模型的快速、可扩展且成本效益高的训练提供支持。
• 克服现有框架（如 scikit-learn、Spark MLlib 和 TensorFlow）在内存限制和 GPU 利用效率方面的性能瓶颈。
• 设计一种分层系统，模拟现代计算集群的硬件堆栈，优化通信和资源使用。
• 通过高效的磁盘外流水线，实现在 GPU 内存无法容纳的大型数据集上对大规模 GLMs 进行高性能训练。

提出的方法

• 提出 CoCoA 框架的分层扩展，支持节点、GPU 和 CPU-GPU 数据流中的嵌套优化。
• 采用流水线化的磁盘外随机坐标下降算法，最大化利用 CPU、GPU 和互连资源。
• 使用基于 GPU 的本地求解器，结合流式 CUDA 操作，分块处理数据，减轻内存压力。
• 实现分层通信模式，优先使用成本更低的节点内通信而非节点间通信，提升可扩展性。
• 将优化的软件架构与底层 GPU 内核及高效数据流结合，最小化空闲时间。
• 支持单节点和多节点训练，具备动态负载均衡和容错能力。

实验结果

研究问题

• RQ1 能否通过模拟硬件堆栈的分层机器学习框架，在分布式 GLM 训练中实现更快的收敛速度和更高的资源利用率？
• RQ2 在多节点集群中，优先使用节点内通信而非节点间通信的分层通信模式如何影响训练性能？
• RQ3 GPU 加速和流水线化的磁盘外处理在多大程度上能减少大规模 GLMs 的训练时间？
• RQ4 在真实世界基准测试中，Snap ML 与 TensorFlow、Vowpal Wabbit 和 Spark MLlib 等先进框架相比，在性能和准确性方面表现如何？

主要发现

• Snap ML 在 Criteo Terabyte Click Logs 数据集上训练逻辑回归模型仅用 1.5 分钟，测试损失为 0.128。
• 与此前最快结果相比，提速达 46 倍，此前结果由 TensorFlow 在 60 个工作节点和 29 个参数服务器上实现。
• 在相同硬件配置下，Snap ML 的训练速度比 TensorFlow 快超过 500 倍，且使用相同的优化器和输入函数。
• 分层通信模式显著降低了节点间通信开销，提升了多节点环境中的训练速度。
• 流水线化的磁盘外处理管道最大化了 GPU 利用率，实现了对超过 GPU 内存容量的数据集的高效处理。
• Snap ML 在单节点和多节点部署中均表现出卓越性能，训练速度超过 scikit-learn、Vowpal Wabbit、Spark MLlib 和 TensorFlow。