BAGUA : システムリラクゼーションによる分散学習のスケールアップ

今回は、" BAGUA: Scaling up Distributed Learning with System Relaxations "という論文について翻訳をして紹介していきます！

BAGUA: Scaling up Distributed Learning with System Relaxations

Recent years have witnessed a growing list of systems for distributed data-parallel training. Existing systems largely fit into two paradigms, i.e., parameter s...

論文情報
著者 : Shaoduo Gan, Xiangru Lian, Rui Wang, Jianbin Chang, Chengjun Liu, Hongmei Shi, Shengzhuo Zhang, Xianghong Li, Tengxu Sun, Jiawei Jiang, Binhang Yuan, Sen Yang, Ji Liu, Ce Zhang
(Submitted on on 3 Jul 2021 (v1))

ひとこと紹介
イントロダクション
この論文の構成

ひとこと紹介

　分散型機械学習における理論の発展とシステムの現状のギャップを埋める通信フレームワーク "BAGUA"を提案！

イントロダクション

分散型機械学習　

　近年の深層学習・機械学習でのモデルの性能の向上により、モデルのサイズや学習できる量が増えつつあり、計算量も急激に増加している。この向上していることの背景には、GPUやTPUなどのハードウェアの高速化と、効率的でスケーラブルな分散型学習アルゴリズムの開発という2つの大きな要因があると考えるが、今回は後者の分散型学習に焦点をあてて取り上げる。

データ並列トレーニングシステムの現在の状況

　この論文では、データ並列学習に焦点を当てている。
データ並列学習とは、最も一般的な分散学習パラダイムの1つとされていて、データセットを複数の作業者に分割し、モデルを1つのデバイスに収めるというものである。

　近年、分散型データ並列学習のためのシステムが増えてきているがここでは、既存のシステムを2つのパラダイムに分類して説明している。

Liらによるパラメータサーバの研究
Adam（Microsoft），初期のTensorFlow（Google），Poseidon（Petuum），Angel（Tencent），BytePS（ByteDance）など
SergeevらによるAllreduceなどのMPIコレクティブ操作の研究
PyTorch-DDP（Facebook），Mariana（Tencent），MALT（NEC Labs），NCCL（NVIDIA），Horovod（Uber）など

データ並列学習アルゴリズムの現状

　理論やアルゴリズムの面では、標準的な同期および非同期の確率的勾配（SG）ベースのアルゴリズムの性能向上にも積極的に取り組んでいる。しかし、「通信」部分がボトルネックになっていることもあり、主に通信のある側面を「緩和」することで、通信のオーバーヘッドを低減するため、下記のような手法を提案している。
（1）通信の圧縮（量子化、スパース化、エラー補償など）
（2）通信の分散化
（3）通信の遅延（LocalSGDなど）や非同期化などが挙げられます。

　これらの技術は、異なる作業負荷や異なるネットワーク条件に合わせて最適化されている。またこれらの技術をあわせることで、通信のオーバーヘッドを帯域幅とレイテンシーの両方の観点から大幅に削減したり、「外れ値」の存在に対する耐性を高めたりすることができると期待されている。

理論の間の新たなギャップと技術的課題

　この論文では、分散学習の理論やアルゴリズムの発展と実際に使われているシステムのギャップに着目し、最適化が十分に活かしきれていないことに注目している。

このギャップを埋めるためには、既存のシステムが持つパラメータサーバやAllreduceの抽象化を利用して、これらのアルゴリズムを実装するだけでは不十分であり、下記の2つの課題が考えられる。

① パラメータサーバやAllreduceのパラダイムにおいて、これらのシステムリラクゼーションを直接かつ自然にサポートすることは困難
　例えば、誤差補正を用いたほとんどの通信圧縮アルゴリズムで必要とされる、サーバ側にメモリと状態を必要とするアルゴリズムをサポートするために、パラメータサーバが提供するput/getの抽象化を使用することは困難である。同様に、どちらのパラダイムも分散型の通信をサポートすることは困難である。その結果、今日の多くの緩和アルゴリズムをサポートするためには、システムの抽象化の設計を根本的に見直さなければならなくなってしまう点。

② システムリラクゼーションの多様性に対応するために、モジュール式のシステム抽象化と最適化をサポートする必要がある
　Horovod やBytePS などの既存のシステムが性能を最適化する場合、教科書的なSGベースのアルゴリズムの通信パターンに焦点を当てていることが多い。大規模な学習アルゴリズムの集合をサポートしたいと考えた場合、それぞれを個別に最適化することはできない。代わりに、この多様なアルゴリズムの集合を共通のフレームワークで自動的に最適化する方法を理解する必要がある。

BAGUAシステムと我々の貢献

　この2つの課題を解決するために、この研究では通信フレームワークであるBAGUAを構築している。
BAGUAは、最先端のシステムリラクゼーション技術である分散学習をサポートすることを目的としている。

1. BAGUAのシステムデザインであり、通信のためのモジュラーデザインを提供

この研究での抽象化は、パラメータサーバやAllreduceパラダイムを超えて、MPIスタイルの集合操作のコレクションを提供し、異なる精度や集中化戦略での通信を容易にしる。この抽象化機能は、多くのアルゴリズムをサポートするのに十分な柔軟性とモジュール性を備えている。さらに、BAGUAフレームワークで実装されたアルゴリズムを高速化する、シンプルな自動最適化フレームワークも開発している。このフレームワークの鍵となるのは、通信の自動バッチ処理とスケジューリング。HorovodやBytePSのような先行研究とは異なり、我々の最適化フレームワークは、標準的なSGベースのアルゴリズムを超えて、より広く適用することができる。

2. 2つの仮説を中心とした広範な実証的研究

(1）BAGUAは、様々なシステム緩和手法をサポートすることで、実世界のアプリケーションや workloads、実世界のインフラにおいて、既存のシステムよりも大幅に改善することができる
(2）BAGUAは、多様なシステム緩和手法をサポートすることで、ユーザーが様々なアルゴリズムを選択できるように、多様なネットワーク条件でスケーラブルなMLトレーニングを提供することができる。この目的のために、我々はKwai Inc.において、ベンチマークタスクと実世界のアプリケーションを用いた大規模な実証実験を行った。

最大16台のマシン（合計128GPU、Tensor Coreで2ペタフロップス）を持つクラスタ上で、V100 GPUマシン（p3.8xlarge、p3.16xlarge、p3dn.24xlarge）がAWS上でどのように接続されているか、様々なネットワーク条件を検討した結果、BAGUAは、様々なタスクにおいてBytePS, Horovod, PyTorch-DDPを大幅に上回った（10Gbpsでは最大1.9倍、100Gbpsでは最大1.34倍）。さらに、厳密なトレードオフ分析を行った結果、異なるアルゴリズムやシステムリラクセーションが，異なるネットワーク条件で最高の性能を達成することが分かった。これは、エンドユーザにこのような多様なアルゴリズムを提供することの重要性を示している。

限界と今後の課題

現在のBAGUAシステムは、ユーザーが最適なシステムリラクゼーションを自動的に選択するための原理的な方法を提供していません。BAGUAがこれらのアルゴリズムをサポートした後は、原理的な自動チューニングシステムを構築する方法を理解することが一つの興味深い方向性である。
また、現在のBAGUAはデータ並列にのみ焦点を当てており、モデル並列（やパイプライン並列などの他の技術を統合し、システムの抽象化を理解することは、今後の興味深い課題である。