[论文解读] OneFlow: Redesign the Distributed Deep Learning Framework from Scratch

Jinhui Yuan, Xinqi Li+10 more|arXiv (Cornell University)|Oct 28, 2021

Advanced Neural Network Applications参考文献 40被引用 34

一句话总结

OneFlow 引入 SBP（split, broadcast, partial-value）及演员模型运行时，以在分布式深度学习中实现灵活的数据、模型与流水线并行性，目标相较于传统框架提供更简单、运行效率更高的实现。

ABSTRACT

Deep learning frameworks such as TensorFlow and PyTorch provide a productive interface for expressing and training a deep neural network (DNN) model on a single device or using data parallelism. Still, they may not be flexible or efficient enough in training emerging large models on distributed devices, which require more sophisticated parallelism beyond data parallelism. Plugins or wrappers have been developed to strengthen these frameworks for model or pipeline parallelism, but they complicate the usage and implementation of distributed deep learning. Aiming at a simple, neat redesign of distributed deep learning frameworks for various parallelism paradigms, we present OneFlow, a novel distributed training framework based on an SBP (split, broadcast and partial-value) abstraction and the actor model. SBP enables much easier programming of data parallelism and model parallelism than existing frameworks, and the actor model provides a succinct runtime mechanism to manage the complex dependencies imposed by resource constraints, data movement and computation in distributed deep learning. We demonstrate the general applicability and efficiency of OneFlow for training various large DNN models with case studies and extensive experiments. The results show that OneFlow outperforms many well-known customized libraries built on top of the state-of-the-art frameworks. The code of OneFlow is available at: https://github.com/Oneflow-Inc/oneflow.

研究动机与目标

• 为需要能够自动支持多样并行策略的通用、简单的分布式 DL 框架提供动力.
• 提出 SBP（split, broadcast, partial-value）作为统一的张量/运算并行抽象。
• 引入具备显式资源依赖处理的演员模型运行时，以实现稳定的分布式执行。
• 提供编译器，将逻辑图转换为以 SBP 驱动并行的物理图。
• 通过与最先进系统的广泛实验，证明通用性和效率。

提出的方法

• 将 SBP 定义为跨设备和跨节点的全局张量到局部张量的多维映射（S、B、P 签名）。
• 对运算符输入/输出推断 SBP 以指定并行性（例如数据并行与模型并行）。
• 引入 boxing 运算以转换 SBP，使数据在不同 SBP 签名之间路由。
• 采用演员模型作为运行时，其中每个运算符都是一个带有显式输入/输出寄存器和基于消息的依赖机制的演员。
• 实现显式的资源依赖计数器（输入/输出/引用计数）以实现编译阶段规划和运行时背压。
• 使用统一的基于演员的消息总线来路由节点间和节点内通信（基于拉取的节点间数据传输）。
• 提供一个编程接口，使单设备与分布式 API 保持一致性，依赖放置/ SBP 注释而非底层通信原语。

实验结果

研究问题

• RQ1 SBP 能否提供统一、灵活的抽象来表达不同异构硬件上的数据、模型和流水线并行吗？
• RQ2 基于演员的运行时如何比传统调度器更健壮地处理复杂依赖和资源约束？
• RQ3 从零开始设计（编译器 + 运行时）是否能实现与在现有框架之上定制库的同等性能？

主要发现

• OneFlow 在代表性的大模型场景下，其性能与基于最先进框架的定制库相当或略有超越。
• SBP 使混合并行（数据+模型）编程更易于实现，且有自动编译器能够生成物理图。
• 具有显式资源计数器和背压的演员运行时支持稳定执行和跨设备、跨节点的自然流水线。
• 在合成数据和真实数据场景下，数据加载/流水线可以在无需额外插件（如 DALI）的情况下达到近乎理想的吞吐量。
• 实验表明 OneFlow 在 FP32/FP16 条件下的数据并行 ResNet 与 BERT 场景中优于官方 TensorFlow、PyTorch 和 MXNet。