[论文解读] Scaling Distributed Machine Learning with In-Network Aggregation

Amedeo Sapio, Marco Canini+8 more|arXiv (Cornell University)|Feb 22, 2019

Stochastic Gradient Optimization Techniques参考文献 73被引用 119

一句话总结

SwitchML 通过可编程交换机在网络内实现聚合，对工作节点的更新进行求和，降低数据流量，并在分布式训练中实现高达 5.5x 的端到端加速。

ABSTRACT

Training machine learning models in parallel is an increasingly important workload. We accelerate distributed parallel training by designing a communication primitive that uses a programmable switch dataplane to execute a key step of the training process. Our approach, SwitchML, reduces the volume of exchanged data by aggregating the model updates from multiple workers in the network. We co-design the switch processing with the end-host protocols and ML frameworks to provide an efficient solution that speeds up training by up to 5.5$ imes$ for a number of real-world benchmark models.

研究动机与目标

• 推动在大规模分布式 ML 中将训练瓶颈从计算转向通信。
• 提出在可编程交换机中实现的实用网络内聚合原语。
• 对端主机协议和 ML 框架进行协同设计，以加速基于 SGD 的训练。
• 展示端到端的训练加速效果，并与传统的 all-reduce 和参数服务器方法进行比较。

提出的方法

• 将模型更新分解为适合交换机流水线限制以进行流式聚合的块。
• 实现基于 pool 的、在交换机内的整型聚合以求和更新，浮点值通过自适应缩放映射为定点表示。
• 协调整端聚合协议以驱动对交换机聚合器的同步、基于时隙的利用。
• 提供面向主机的轻量级重传和交换机端影子副本的容错机制，以应对数据包丢失。
• 将 SwitchML 与现有的 ML 框架（例如 PyTorch、TensorFlow）以及 Horovod 集成，以实现实际部署。

实验结果

研究问题

• RQ1 在可编程交换机内进行网络内聚合能否降低同步分布式 SGD 的通信开销？
• RQ2 在交换机数据平面中实现梯度聚合需要哪些设计权衡（精度、内存、可靠性）？
• RQ3 在不同网络速度和模型大小下，SwitchML 与传统的 all-reduce 和参数服务器方法相比如何？
• RQ4 将 SwitchML 集成到真实 ML 流水线中对端到端训练性能的影响是什么？

主要发现

• SwitchML 在真实世界的 DNN 基准测试中实现高达 5.5x 的端到端训练加速。
• 网络内聚合在测试的工作负载中始终优于环形 all-reduce，对于网络瓶颈模型，在合成基准中在 10 Gbps 时快出最多 80%，在 100 Gbps 时快出最多 67%。
• 与使用 RDMA 的 NCCL 相比，基于 SwitchML 的聚合速度最高可快 2.9x；在所报道的微基准中，相对于使用 TCP 的 NCCL，快了最多 9.1x。
• 梯度压缩（Top-k、QSGD）并非普遍有益；在许多情况下，压缩开销抵消数据缩减，当硬件和工作负载条件限制压缩增益时，SwitchML 常占上风。
• 该设计利用基于 pool 的流式聚合和两池影子副本机制，以在尽量降低交换机复杂度的同时容忍数据包丢失。