[論文レビュー] 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（分割、ブロードキャスト、部分値）とアクターモデルのランタイムを導入し、分散深層学習におけるデータ、モデル、パイプライン並列性を柔軟に実現する。従来のフレームワークと比べて、よりシンプルで効率的な実行を目指す。

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フレームワークの必要性を動機づける。
• SB（split, broadcast, partial-value）をテンソル/演算子並列性の統一的抽象として提案する。
• 安定した分散実行のための明示的なリソース依存性の処理を備えたアクターモデルのランタイムを導入する。
• SBP駆動の並列性を用いて、論理グラフを物理グラフに変換するコンパイラを提供する。
• 最新システムに対する広範な実験を通じて一般適用性と効率を示す。

提案手法

• SBPを、デバイスとノード全体のグローバルテンソルからローカルテンソルへの多次元マッピングとして定義する（S、B、P署名）。
• オペレータの入力/出力に対してSBPを推定し、並列性（例：データ並列 vs モデル並列）を指定する。
• SBPを変換するボクシング演算を導入し、異なるSBP署名間のデータルーティングを可能にする。
• ランタイムにはアクターモデルを採用し、各演算は明示的な入出力レジスタとメッセージベースの依存機構を持つアクターである。
• コンパイル時計画とランタイムのバックプレッシャを可能にするため、明示的なリソース依存カウンター（入出力・参照カウンター）を実装する。
• ノード間/ノード内通信をルーティングする統一のアクター型メッセージバスを使用する（ノード間データ転送はプル型）。
• 単一デバイスと分散APIを整合させたプログラミングインターフェースを提供し、低レベルな通信プリミティブではなく配置/SBP注釈に依拠する。

実験結果

リサーチクエスチョン

• RQ1 SBPは異種ハードウェア全体でデータ、モデル、パイプライン並列性を表現する統一・柔軟な抽象を提供できるか。
• RQ2 伝統的なスケジューラより、アクターモデルのランタイムは複雑な依存関係とリソース制約をどのようにより堅牢に処理できるか。
• RQ3 ゼロからの設計（コンパイラ＋ランタイム）が、既存フレームワーク上に構築されたカスタマイズ済みライブラリと競争力のある性能を達成するか。

主な発見

• OneFlowは、代表的な大規模モデルの状況で、最先端フレームワーク上に構築された主要なカスタマイズライブラリと同等かわずかに優れた性能を達成する。
• SBPは既存のフレームワークよりハイブリッド並列性（データ＋モデル）のプログラミングを容易にし、物理グラフを自動的に生成するコンパイラを備える。
• 明示的なリソースカウンターとバックプレッシャを備えたアクターランタイムは、デバイスとノード間の安定した実行と自然なパイプライン処理を支援する。
• 合成データと実データのケースで、データ読み込み/パイプライン処理は、追加プラグイン（例：DALI）なしでほぼ理想的なスループットに達する。
• 実験では、FP32/FP16でのデータ並列ResNetおよびBERTのシナリオにおいて、OneFlowが公式のTensorFlow、PyTorch、MXNetを上回ることを示している。