OneFlow: Redesign the Distributed Deep Learning Framework from Scratch

By Jinhui Yuan, Xinqi Li, Cheng Cheng, Juncheng Liu, Ran Guo, Shenghang Cai, C. Yao, Fei Yang, Xiaodong Yi, Chuan Wu and 2 more. 2021

54 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Electrical and Electronic Engineering, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Jie Tang
2. Yuxiao Dong
3. Zhilin Yang
4. Myle Ott
5. Sam Shleifer
6. Luke S. Zettlemoyer
7. Chien-chin Huang
8. Min Xu

In Computer Vision and Pattern Recognition

1. Jie Tang
2. Sam Shleifer
3. Chien-chin Huang
4. Min Xu
5. Yuxiao Dong
6. Zhilin Yang
7. Yufei Xue
8. Xiao Xia

In Electrical and Electronic Engineering

1. Zhilin Yang

In Optimization & Training

1. Chien-chin Huang
2. Min Xu
3. Yuxiao Dong
4. Zhilin Yang
5. Jie Tang
6. Yufei Xue
7. Xiao Xia
8. Andi Wang

Citations in the corpus

• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, background): “Tensor parallelism [8, 20, 30, 31] shards model parameters, conducts partial computation on individual devices and communicates activations at required layer boundaries.”
• CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Evaluations on HumanEval-X (2023, methodology): “We also compare the inference speed with implementations in LLM.int() (Dettmers et al., 2022) and Oneflow (Yuan et al., 2021).”
• A Survey of Large Language Models (2023, unknown): “OneFlow [133]) have also provided the support for parallelism algorithms, which are commonly used for training large-scale models.”