[论文解读] Enabling Fast and Flexible Distributed Deep Learning with Programmable Switches

Heng Pan, Penglai Cui+12 more|arXiv (Cornell University)|May 11, 2022

Brain Tumor Detection and Classification被引用 7

一句话总结

本文提出 Libra，一种基于可编程交换机的系统，通过将频繁更新的（“热”）参数的梯度聚合卸载到 Intel Tofino 交换机，加速分布式稀疏深度学习。通过采样-based 热参数识别、热感知参数放置以及基于查找表的实时浮点加法机制，Libra 在稀疏模型上的训练速度相比传统参数服务器训练提升了 1.5× 至 4×。

ABSTRACT

Deep learning has been used in a wide range of areas and made a huge breakthrough. With the ever-increasing model size and train-ing data volume, distributed deep learning emerges which utilizes a cluster to train a model in parallel. Unfortunately, the performance is often far from linear speedup due to the communication overhead between cluster nodes. To address this challenge, this paper designs and implements Libra, a network aggregator, that utilizes in-network computation to optimize the communication for distributed DL training in two aspects: 1) reduce active connections and 2) aggregate exchanged network packets. We implemented our Libra on Intel Tofino switches, customized a lightweight host stack and integrated it into an open-source training framework PS-lite. The experimental result shows that our Libra can achieve 1.5~4 times speedup.

研究动机与目标

• 为解决分布式稀疏深度学习训练中的通信瓶颈问题，其中密集且突发的通信限制了可扩展性。
• 利用工业级稀疏深度学习模型中参数更新频率的高度偏斜特性，即少数参数被频繁更新。
• 设计一种将这些“热”参数的梯度聚合卸载到可编程交换机上的系统，减轻参数服务器的负载，提升训练吞吐量。
• 通过新颖的系统级优化，克服可编程交换机的硬件与协议限制，如寄存器写冲突和缺乏原生浮点支持。
• 通过数据包丢失恢复与交换机故障切换机制，确保在生产规模部署中的可靠性。

提出的方法

• 提出一种基于采样的机制，通过在完整数据集的 4–8% 上训练来识别热参数，实现在无需全规模性能分析的情况下高效发现热参数。
• 在交换机上设计基于热度的参数放置策略，以最小化梯度聚合过程中的寄存器冲突，提升流水线效率。
• 在工作节点设计参数布局感知的梯度打包技术，降低单个流水线阶段中对同一寄存器的多次写入概率。
• 开发基于查找表的机制，在交换机上实现对 32 位浮点数的实时求和，尽管原生浮点支持有限，仍能实现网络内梯度聚合。
• 实现双路径通信模型：热参数在交换机中聚合，冷参数仍由参数服务器处理，确保正确性与可扩展性。
• 通过数据包丢失恢复与交换机故障切换机制增强系统可靠性，确保在交换机故障时训练连续进行。

实验结果

研究问题

• RQ1 与传统参数服务器架构相比，稀疏深度学习模型的网络内梯度聚合能否实现显著的加速？
• RQ2 在不引入完整训练开销的前提下，如何高效识别稀疏深度学习工作负载中的热参数？
• RQ3 在可编程交换机的硬件限制下，需要哪些系统级优化才能实现高效且正确的梯度聚合？
• RQ4 结合交换机对热参数的聚合与服务器对冷参数的聚合，能否在提升性能的同时保持训练的正确性？
• RQ5 Libra 在真实网络环境（如数据包丢失或交换机故障）下如何确保可靠性？

主要发现

• 与基线参数服务器实现相比，Libra 在分布式稀疏深度学习训练中实现了 1.5× 至 4× 的加速。
• 基于采样的热参数识别机制能准确检测出更新频率最高的前 30,000 个参数，其更新量约占全部更新的 50%。
• 基于热度的参数放置与布局感知打包技术使寄存器写冲突减少超过 70%，显著提升了交换机流水线效率。
• 基于查找表的机制在交换机上实现了低性能开销的实时 32 位浮点数求和，支持高精度梯度聚合。
• 系统通过有效的数据包丢失恢复与交换机故障切换机制，保持了高可靠性，确保训练不间断。
• Libra 已与 PS-lite 集成，并部署于 Intel Tofino 交换机，验证了端到端兼容性与真实世界可部署性。