分散機械学習プラットフォームの比較: Spark、PMLS、TensorFlow、MXNet

[[200819]]

本稿では、分散システムの観点から現在の機械学習プラットフォームのいくつかを研究し、これらのプラットフォームで使用されるアーキテクチャ設計を要約し、通信と制御におけるこれらのプラットフォームのボトルネック、フォールトトレランス、開発の難しさなどを分析および比較し、分散機械学習プラットフォームに関する将来の研究に向けていくつかの提案を行います。この論文の研究は、ムラト・デミルバス教授と大学院生のクオ・チャン氏、サレム・アルカハタニ氏によって行われました。

機械学習、特にディープラーニングは、音声認識、画像認識、自然言語処理、そして最近ではレコメンデーションや検索エンジンなどの分野で革命的な成功を収めています。これらのテクノロジーは、自動運転車、デジタル医療システム、CRM、広告、モノのインターネットなどの分野で大きな応用の見込みがあります。もちろん、テクノロジーの急速な進歩を牽引したのは資金であり、それが近年のいくつかの機械学習プラットフォームの立ち上げにつながっています。

トレーニングに関係するデータ セットとモデルの規模が膨大であることを考慮すると、機械学習プラットフォームは通常、モデルをトレーニングするために並列実行される数十または数百のコンピューティング ノードを展開する分散プラットフォームです。近い将来、データセンターのほとんどのタスクは機械学習タスクになると予想されています。

私は分散システムの研究分野出身なので、分散システムの観点からこれらの機械学習プラットフォームの研究を行い、通信と制御におけるこれらのプラットフォームのボトルネックを分析することを検討しました。また、これらのプラットフォームのフォールト トレランスとプログラミングの容易さも考慮しました。

機械学習プラットフォームは、設計アプローチに基づいて、基本データフロー、パラメーター サーバー モデル、高度なデータフローの 3 つの基本カテゴリに分類されます。

以下では、基本的なデータフローの紹介としてApache Spark、パラメータサーバーモデルの紹介としてPMLS（Petuum）、高度なデータフローの紹介としてTensorFlowとMXNetを例に、それぞれの方式について簡単に紹介します。上記のプラットフォームのパフォーマンスを比較し、一連の評価結果を提示しました。詳細な評価結果については、弊社の論文を参照してください。残念ながら、私たちは小規模な研究チームであるため、大規模な評価を行うことができません。

このブログ投稿の最後では、いくつかの結論を述べ、分散機械学習プラットフォームに関する将来の研究作業についていくつかの提案をします。これらの分散機械学習プラットフォームについてすでにある程度理解している読者は、この記事の最後まで直接ジャンプできます。

スパーク

Spark では、計算は有向非巡回グラフ (DAG) としてモデル化され、各頂点は RDD を表し、各エッジは RDD 上の操作を表します。 RDD は一連のパーティション化されたオブジェクトで構成され、メモリに保存および計算され、シャッフル プロセス中にディスクにオーバーフローします。

DAG では、頂点 A から B への有向エッジ E は、RDD B が RDD A に対して操作 E を実行した結果であることを示します。操作は「変換」と「アクション」の 2 つのカテゴリに分かれています。変換操作 (マップ、フィルター、結合など) が RDD に適用され、変換操作の出力は新しい RDD になります。

Spark ユーザーは、RDD で実行される変換とアクションを表す DAG として計算をモデル化します。その後、DAG は複数のステージにコンパイルされます。各ステージは、並行して実行される一連のタスクとして実行され、各パーティションはタスクに対応します。ここで、依存関係が狭いと計算の効率的な実行が容易になりますが、依存関係が広いと、通信集約型のシャッフル操作が導入され、操作のフローが中断されるため、ボトルネックが発生します。

Spark の分散実行は、DAG ステージを異なるコンピューティング ノードに分割することによって実現されます。上の図は、この「マスター ワーカー」アーキテクチャを明確に示しています。ドライバーには、DAG スケジューラとタスク スケジューラという 2 つのスケジューラ コンポーネントが含まれています。スケジューラは、作業者にタスクを割り当て、作業者を調整します。

Spark は一般的なデータ処理用に設計されており、機械学習タスクに特化していません。 Spark で機械学習タスクを実行するには、Spark 用の MLlib を使用できます。 Spark を基本設定で使用する場合、モデル パラメーターはドライバー ノードに保存され、ワー​​カーとドライバー間の通信を通じて各反復後に更新されます。機械学習タスクを大規模に展開する場合、ドライバーはすべてのモデルパラメータを保存できない可能性があります。この場合、すべてのパラメータに対応するために RDD を使用する必要があります。更新されたモデル パラメータに対応するために、各反復で新しい RDD を作成する必要があるため、多くの追加オーバーヘッドが発生します。モデルを更新するには、マシンとディスク間でデータをシャッフルする必要があるため、Spark のスケーラビリティが制限されます。これはまさに基本的なデータフロー モデル (つまり DAG) の欠点です。 Spark は機械学習における反復操作をあまりサポートしていません。

PMLS

PMLS は機械学習タスク用に特別に設計されています。反復的な集中トレーニング プロセスをサポートするために、「パラメータ サーバー」(PS) と呼ばれる抽象化を導入します。

PS (図では緑色のボックスで表示) は、複製可能でシャーディングされた分散キー値データ テーブルとしてメモリ内に存在します。各ノードはモデル内のシャードのプライマリ ノード (パラメーター空間) であり、他のシャードのセカンダリ ノードまたはレプリカ ノードとして機能します。このように、PS はノード数に関して優れたスケーラビリティを備えています。

PS ノードはモデル パラメータを保存および更新し、ワーカーからの要求に応答します。ワーカーは、PS のローカル コピーから最新のモデル パラメータを要求し、割り当てられたデータセット パーティションで計算を実行します。

PMLS は SSP (Stale Synchronous Parallelism) モデルも採用しています。 BSP (Bulk Synchronous Parallelism) モデルと比較すると、SSP では、各反復の終了時に各マシンを同期する必要があるという要件が緩和されます。同期を実現するために、SSP はワーカー間である程度の非同期を許可し、最速のワーカーが最も遅いワーカーより s 回の反復を超えて先行しないことを保証します。処理プロセスは許容される誤差範囲内であるため、この非厳密な一貫性モデルは機械学習に適用できます。このメカニズムについてはブログ記事を書きました。

テンソルフロー

Google は、「DistBelief」と呼ばれるパラメータ サーバー モデル ベースの分散機械学習プラットフォームを開発しました (DistBelief の論文に関する私のレビューはこちら)。私の知る限り、DistBelief の主な不満点は、機械学習アプリケーションを作成するときに低レベルのコードを混在させる必要があることです。 Google は、分散実行の専門家でなくても、従業員全員が機械学習コードを書けるようにしたいと考えています。 Google がビッグデータ処理用の MapReduce フレームワークを作成したのも同じ理由です。

TensorFlow はまさにそれを実現するために設計されたプラットフォームです。より高度なデータフロー処理パラダイムを採用しており、計算を表すグラフは DAG である必要がなくなり、グラフにサイクルを含めることができ、可変状態をサポートします。 TensorFlow の設計は、ある程度 Naiad の設計哲学に影響を受けていたと思います。

TensorFlow は、計算をノードとエッジで構成される有向グラフとして表現します。ノードは計算操作または変更可能な状態 (変数など) を表し、エッジはノード間の通信のための多次元配列を表します。この多次元データは「テンソル」と呼ばれます。 TensorFlow では、ユーザーが論理計算グラフを静的に宣言する必要があり、グラフをマシンに書き換えて分割することで分散コンピューティングを実装します。 MXNet、特に DyNet は動的に定義されたグラフを使用することに注意してください。これによりプログラミングが簡素化され、プログラミングの柔軟性が向上します。

上の図に示すように、TensorFlow では、分散機械学習のトレーニングにパラメータ サーバー アプローチが使用されます。 TensorFlow で PS 抽象化を使用する場合、サーバーおよびデータ並列性のパラメータが使用されます。 TensorFlow はより複雑なタスクを実行できると主張していますが、そのためにはユーザーが未踏の領域に到達するためにコードを記述する必要があります。

MXネット

MXNet は、2015 年に登場した CXXNet、Minverva、Purines などのディープラーニング プロジェクトから生まれた共同オープン ソース プロジェクトです。 TensorFlow と同様に、MXNet も、変更可能な状態を持つ循環計算グラフをサポートし、パラメータ サーバー モデルを使用したトレーニング計算をサポートするデータ フロー システムです。同様に、MXNet は複数の CPU/GPU 上のデータ並列処理を適切にサポートし、モデルの並列処理を実現できます。 MXNet は同期と非同期の両方のトレーニング計算をサポートします。次の図は、MXNet の主なコンポーネントを示しています。このうち、ランタイム依存関係エンジンは、計算処理における依存関係を分析し、相互依存関係のない計算については並列処理を実行します。 MXNet は、計算グラフとメモリの最適化のために、ランタイム依存エンジンの上に中間レイヤーを提供します。

MXNet は、チェックポイント メカニズムを使用して基本的なフォールト トレランスをサポートし、モデルの保存および読み込み操作を提供します。保存操作ではモデル パラメータがチェックポイント ファイルに書き込まれ、ロード操作ではモデル パラメータがチェックポイント ファイルから読み取られます。

MXNet は、記述型プログラミング言語を使用して計算を有向グラフとして表現し、テンソル計算と制御フローを定義するための命令型プログラミングの使用もサポートします。 MXNet は、C++、Python、R、Scala などの複数の言語で記述されたクライアント API を提供します。 TensorFlow と同様に、MXNet のバックエンド コア エンジン ライブラリも C++ で記述されています。

評価結果の一部

実験では、Amazon EC2 m4.xlarge インスタンスが使用されました。各インスタンスには、Intel Xeon E5-2676 v3 で実装された 4 つの vCPU、16 GB のメモリがあり、EBS 帯域幅は 750 Mbps です。私たちは、2 つの一般的な機械学習タスク、つまりバイナリ ロジスティック回帰と多層ニューラル ネットワークを使用した画像分類に関する実験を実施しました。この記事で示されているのは結果グラフの一部のみです。 より詳細な実験結果については、論文を参照してください。私たちの実験にはいくつかの制限があることに留意してください。まず、使用したマシンの数が少なかったため、スケーラビリティをテストできませんでした。次に、CPU 計算のみを使用し、GPU の使用はテストしませんでした。

このグラフは、各プラットフォームがロジスティック回帰を計算する速度を示しています。 PMLS と MXNet は最も高速な 2 つのシステムであり、TensorFlow は最も低速で、Spark はその中間です。 私たちは、これには複数の原因があると分析しました。まず、PMLS は Spark や TensorFlow と比較して軽量なシステムです。Spark は高性能な C++ プログラミング言語で実装されていますが、Spark は JVM 上で実行される Scala で記述されています。第二に、PMLS には抽象化が少なく、TensorFlow には抽象化が多すぎます。抽象化によりシステムの複雑さが増し、実行時のオーバーヘッドが発生します。

この図は、さまざまなプラットフォームでの DNN の速度を示しています。単層ニューラル ネットワークと比較すると、2 層ニューラル ネットワークに発展すると、より多くの反復計算が必要になるため、Spark のパフォーマンスは低下します。ドライバーはこれらのパラメータに対応できるため、モデル パラメータを Spark ドライバーに保存します。しかし、RDD を使用してパラメータを保存し、各反復後にパラメータを更新すると、結果が悪くなる可能性があります。

このグラフは、プラットフォームの CPU 使用率を示しています。 Spark アプリケーションは CPU 使用率が高いように見えますが、これは主にシリアル化のオーバーヘッドが原因です。以前の研究でも同じ問題が指摘されていました。

結論と今後の研究の方向性

機械学習やディープラーニングのアプリケーションは高度に並列化可能ですが、並行アルゴリズムの観点から見ると、そうするのはそれほど簡単ではありません。おそらく、パラメータ サーバー アプローチは、分散機械学習プラットフォームでのトレーニング タスクに最適です。

ボトルネックを考慮すると、分散型機械学習アプリケーションでは依然としてネットワークがボトルネックとなっています。より高度な一般的なデータフロー プラットフォームと比較すると、データとモデルのより優れた組織化と参加のメカニズムを提案することがより重要です。データとモデルは第一級のものとして扱われるべきです。

さらに、プラットフォームに微妙な影響を及ぼす他の要因もあります。 Spark の場合、CPU オーバーヘッドはネットワークの制限よりも深刻なボトルネックになります。 Spark で使用されるプログラミング言語は Scala などの JVM 言語であり、Spark のパフォーマンスに大きな影響を与えます。したがって、分散機械学習プラットフォームのより優れた監視ツールとパフォーマンス予測が必要です。最近、Ernest や CherryPick など、Spark データ処理アプリケーションのいくつかの問題を解決するツールが登場しました。

機械学習をサポートするランタイム分散システムには、リソースのスケジューリングやランタイムパフォーマンスの向上など、まだ未解決の問題がいくつかあります。アプリケーションの実行時監視とパフォーマンス分析を考慮して、次世代の分散機械学習プラットフォームは、プラットフォーム上で実行されるタスクのコンピューティング、メモリ、ネットワーク リソースに対して、実行時に柔軟な構成とスケジューリング機能を提供する必要があります。

プログラミングとソフトウェア エンジニアリングのサポートにも未解決の問題があります。機械学習アプリケーションに適した分散プログラミングの抽象化とは何でしょうか? これらの質問には、問題固有の入力による DNN のテストなど、分散機械学習アプリケーションのさらなる調査、テスト、検証が必要です。