[论文解读] BAGUA: Scaling up Distributed Learning with System Relaxations

Shaoduo Gan, Jiawei Jiang+12 more|arXiv (Cornell University)|Jan 1, 2021

Stochastic Gradient Optimization Techniques参考文献 79被引用 3

一句话总结

BAGUA 是一个模块化、类似 MPI 的通信库，可高效实现分布式深度学习中的高级系统松弛技术，如量化、去中心化和异步训练。通过灵活的优化框架支持重叠、融合和分层通信，BAGUA 在各种工作负载下，端到端训练时间相比 PyTorch-DDP、Horovod 和 BytePS 最快提升 2 倍。

ABSTRACT

Recent years have witnessed a growing list of systems for distributed data-parallel training. Existing systems largely fit into two paradigms, i.e., parameter server and MPI-style collective operations. On the algorithmic side, researchers have proposed a wide range of techniques to lower the communication via system relaxations: quantization, decentralization, and communication delay. However, most, if not all, existing systems only rely on standard synchronous and asynchronous stochastic gradient (SG) based optimization, therefore, cannot take advantage of all possible optimizations that the machine learning community has been developing recently. Given this emerging gap between the current landscapes of systems and theory, we build BAGUA, a MPI-style communication library, providing a collection of primitives, that is both flexible and modular to support state-of-the-art system relaxation techniques of distributed training. Powered by this design, BAGUA has a great ability to implement and extend various state-of-the-art distributed learning algorithms. In a production cluster with up to 16 machines (128 GPUs), BAGUA can outperform PyTorch-DDP, Horovod and BytePS in the end-to-end training time by a significant margin (up to 2 times) across a diverse range of tasks. Moreover, we conduct a rigorous tradeoff exploration showing that different algorithms and system relaxations achieve the best performance over different network conditions.

研究动机与目标

• 弥合分布式训练算法理论进展与仍依赖标准同步/异步 SGD 的现有系统之间的差距。
• 设计一个灵活、模块化的通信库，原生支持多种系统松弛技术，如量化、去中心化和通信延迟。
• 通过统一的优化框架，实现对最先进分布式学习算法的高效且可扩展的实现。
• 在不同网络条件和工作负载下，实证评估不同算法与系统松弛之间的权衡。

提出的方法

• 设计一个模块化、类似 MPI 的通信库，抽象底层集体通信操作，以支持多种通信模式。
• 集成三项核心系统优化：计算-通信重叠（O）、张量融合与展平（F）以及分层 GPU 通信（H）。
• 支持多种系统松弛技术：低精度梯度（QSGD、1-bit Adam）、去中心化训练（Decen）和异步执行（Async）。
• 提供统一框架，用户可使用直接映射到系统级优化的原语来组合和扩展算法。
• 通过优化堆栈实现自动性能调优，动态适应模型和网络特性。

实验结果

研究问题

• RQ1 能否设计一个通信库，原生支持广泛的系统松弛技术，而无需硬编码特定算法的逻辑？
• RQ2 不同系统松弛技术（量化、去中心化、异步）在不同工作负载和网络条件下对端到端训练性能有何影响？
• RQ3 关键系统优化（重叠、融合、分层通信）对整体训练效率的相对影响是什么？
• RQ4 在给定模型和网络环境条件下，哪种算法配置能实现最佳性能？

主要发现

• 在 VGG16、BERT、Transformer 和 LSTM+AlexNet 等工作负载下，BAGUA 的端到端训练速度相比 PyTorch-DDP、Horovod 和 BytePS 最快提升 2 倍。
• 在低带宽网络中，压缩算法（如 QSGD 和 1-bit Adam）显著降低通信开销并提升性能。
• 在高延迟网络中，去中心化算法（Decen-32bits/8bits）因减少同步瓶颈而优于集中式算法。
• 消融研究证实，三项系统优化（重叠、融合、分层通信）均至关重要，其影响因工作负载而异：H 优化通信密集型任务，F 优化具有大量小张量的模型，O 在计算密集型场景中效果最佳。
• 异步训练（Async）在存在慢速节点时，将每个训练周期时间减少 30–50%，验证了其在异构集群中的有效性。
• 建立了经验性指导原则：使用 QSGD 适配基于 SGD 的优化器，使用 1-bit Adam 适配 Adam 优化器，当通信与计算比值较低时，优先采用异步方法。