[Paper Review] OneFlow: Redesign the Distributed Deep Learning Framework from Scratch

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

Advanced Neural Network Applications40 references34 citations

TL;DR

OneFlow introduces SBP (split, broadcast, partial-value) and an actor-model runtime to enable flexible data, model, and pipeline parallelism in distributed deep learning, aiming for simpler, efficient execution compared to legacy frameworks.

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.

Motivation & Objective

• Motivate the need for a generic, simpler distributed DL framework capable of automatically supporting diverse parallelism strategies.
• Propose SBP (split, broadcast, partial-value) as a unifying tensor/operator parallelism abstraction.
• Introduce an actor-model runtime with explicit resource dependency handling for stable distributed execution.
• Provide a compiler that converts a logical graph into a physical graph with SBP-driven parallelism.
• Demonstrate general applicability and efficiency through extensive experiments against state-of-the-art systems.

Proposed method

• Define SBP as a multi-dimensional mapping from global tensors to local tensors across devices and nodes (S, B, P signatures).
• Infer SBP for operator inputs/outputs to specify parallelism (e.g., data vs model parallelism).
• Introduce boxing ops to transform SBP, enabling data routing between differing SBP signatures.
• Adopt the actor model for runtime, where each op is an actor with explicit in/out registers and a message-based dependency mechanism.
• Implement explicit resource dependency counters (in/out/reference counters) to enable compile-time planning and runtime back pressure.
• Use a unified actor-based message bus to route inter-node and intra-node communications (pull-based inter-node data transfer).
• Provide a programming interface that keeps single-device and distributed APIs aligned, relying on placement/SBP annotations rather than low-level communication primitives.

Experimental results

Research questions

• RQ1 Can SBP provide a unified, flexible abstraction to express data, model, and pipeline parallelism across heterogeneous hardware?
• RQ2 How can an actor-based runtime handle complex dependencies and resource constraints more robustly than traditional schedulers?
• RQ3 Does a from-scratch design (compiler + runtime) achieve competitive performance to customized libraries built on top of existing frameworks?

Key findings

• OneFlow achieves performance comparable to or slightly better than major customized libraries built on state-of-the-art frameworks in representative large-model scenarios.
• SBP enables easier programming of hybrid parallelisms (data + model) than existing frameworks, with an automatic compiler generating physical graphs.
• The actor-runtime with explicit resource counters and back-pressure supports stable execution and natural pipelining across devices and nodes.
• Data loading/pipelining can reach near-ideal throughput without extra plugins (e.g., DALI) in synthetic and real data cases.
• Experiments show OneFlow outperforms official TensorFlow, PyTorch, and MXNet in data-parallel ResNet and BERT scenarios under FP32/FP16.