[論文レビュー] TF-Replicator: Distributed Machine Learning for Researchers

Peter Buchlovsky, David Budden+10 more|arXiv (Cornell University)|Feb 1, 2019

Adversarial Robustness in Machine Learning参考文献 47被引用数 21

ひとこと要約

TF-Replicator は、TensorFlow をベースにしたフレームワークであり、データ並列またはモデル並列学習を用いて、CPU、GPU、TPU の分散クラスタ across で深層学習モデルを簡単にスケーリングできる。最小限のコード変更で、分散システムの専門知識が不要であり、さまざまなアーキテクチャ（ResNet-50、SN-GAN、D4PG など）で強力なスケーラビリティを達成する。TensorFlow 2.0 の一部としてオープンソース化される予定である。

ABSTRACT

We describe TF-Replicator, a framework for distributed machine learning designed for DeepMind researchers and implemented as an abstraction over TensorFlow. TF-Replicator simplifies writing data-parallel and model-parallel research code. The same models can be effortlessly deployed to different cluster architectures (i.e. one or many machines containing CPUs, GPUs or TPU accelerators) using synchronous or asynchronous training regimes. To demonstrate the generality and scalability of TF-Replicator, we implement and benchmark three very different models: (1) A ResNet-50 for ImageNet classification, (2) a SN-GAN for class-conditional ImageNet image generation, and (3) a D4PG reinforcement learning agent for continuous control. Our results show strong scalability performance without demanding any distributed systems expertise of the user. The TF-Replicator programming model will be open-sourced as part of TensorFlow 2.0 (see https://github.com/tensorflow/community/pull/25).

研究の動機と目的

• 深層学習研究におけるスケーラブルで汎用的な分散学習の増大するニーズに対応すること。
• 低レベルの分散システムの懸念を抽象化することで、分散学習の実装の複雑さを軽減すること。
• 同じモデルを CPU、GPU、TPU といった多様なハードウェアおよび同期／非同期の学習方式にシームレスにデプロイできること。
• マルチロス最適化や強化学習エージェントのような複雑な、非標準の学習パターンをサポートすること。
• 研究の反復を加速するために、既存の単一マシンモデルを学習ツールを学び直さずに分散クラスタにスケーリングできるようにすること。

提案手法

• TF-Replicator は、TensorFlow の分散実行を抽象化する高水準 API を提供し、ユーザーがモデルを一度定義するだけで、異種のクラスタにデプロイできる。
• レプリカベースのプログラミングモデルを採用しており、各レプリカがワーカー上でモデルのコピーを実行し、パラメータサーバーまたはコールレクティブ通信（例：NCCL）を介して勾配を自動同期する。
• データ並列とモデル並列学習の両方をサポートし、マシン間のデバイス間で通信を自動挿入し、デバイスの配置を透明に管理する。
• TensorFlow 2.0 のイージー実行と Keras スタイルの API と統合され、自然な制御フローと柔軟な学習ループが可能になる。
• ユーザーがデバイス配置と通信戦略を設定可能にすることで、データ並列とモデル並列を組み合わせたハイブリッド並列処理をサポートする。
• クラスタ構成は ClusterSpec を介して自動で処理され、論理デバイスを物理ハードウェアにマップすることで、単一マシンとマルチマシンの両環境でのポータビリティが実現される。

実験結果

リサーチクエスチョン

• RQ1 高水準な抽象化により、分散システムの専門知識がなくても研究者が分散学習を簡素化できるか？
• RQ2 フレームワークは、さまざまなモデルアーキテクチャ（例：CNN、GAN、RLエージェント）およびハードウェア（GPU、TPU）でどれほどスケーリングできるか？
• RQ3 TF-Replicator は、マルチロス最適化や階層的強化学習のような複雑な学習方式をサポートできるか？
• RQ4 スケーラビリティと学習速度の観点から、手動で最適化された分散システムと比較して、性能にどれほど差があるか？
• RQ5 研究者が最小限の修正で単一マシンコードを分散学習に再利用できる範囲はどの程度か？

主な発見

• TF-Replicator は ImageNet における ResNet-50 モデルで強力なスケーリングを実現し、GPU および TPU クラスタで競争力のある学習速度を達成した。
• クラス条件付き画像生成のための SN-GAN モデルは、TF-Replicator を用いて効果的にスケーリングされ、フレームワークが複雑な生成モデルをサポートしていることを示した。
• D4PG 強化学習エージェントは、学習時間に顕著な加速を達成し、8 個の V100 GPU が 1 個の TPUv2 デバイス（8 コア）の性能を同等に達成した。
• 1 個の TPUv2 デバイスが、NVLink 接続された 8 個の V100 GPU と同等の性能を発揮した。これは、フレームワークの通信およびデバイス管理の効率性を示している。
• フレームワークは、多様なクラスタアーキテクチャにわたってモデルをスケーリング可能であり、最小限のコード変更で高いスループットと低通信オーバーヘッドを維持した。
• TF-Replicator は、マルチロスや階層的 RL を含む複雑な学習ループを持つモデルの学習に成功し、低レベルの分散システムプログラミングを必要としなかった。