Apache SparkとAmazon DSSTNEを使った、Amazon規模のレコメンデーション生成

Amazonのパーソナライゼーションでは、お客様毎の製品レコメンデーションを生成するためにニューラルネットワークを使っています。Amazonの製品カタログは、あるお客様が購入した製品の数に比較して非常に巨大なので、データセットは極端に疎になってしまいます。そして、お客様の数と製品の数は何億にものぼるため、我々のニューラルネットワークのモデルは複数のGPUで分散しなければ、空間や時間の制約を満たすことができません。

そのため、GPU上で動作するDSSTNE (the Deep Scalable Sparse Tensor Neural Engine)を開発しオープンソースにしました。我々はDSSTNEを使ってニューラルネットワークを学習しレコメンデーションを生成していて、ECのウェブサイトやAmazonのデバイスにおいて個人毎に異なる多様な経験を生み出しています。

一方で、学習や予測のタスクに使われるデータはCPUのクラスタ上でApache Sparkを使って処理され生成されています。そのため、データ処理はCPU上で起こるのに学習と予測はGPU上で起こるという、基本的な課題が発生します。

データ生成とその分析はモデリングをするなかでしばしば見落とされます。特にデータの準備にはモデルのデザインと少なくとも同じくらいの時間がかかるので、流動的で素早いモデルのプロトタイプのために、我々は1つのツールの中でデータを分析し学習しモデルを評価したいと考えました。また、DSSTNEは疎なネットワークとスケーラビリティに最適化されているだけでなく、現在のDSSTNEに欠けているリカレントニューラルネットワーク構造といった機能を持った他のライブラリ、例えばKerasを利用することもできます。

CPUとGPUのインスタンスの両方が含まれるハイブリッドクラスタを管理することは、YARNやMesosといったクラスタ管理がネイティブにはGPUをサポートしていないため、非常に難しいです。また、もしネイティブなGPUサポートがあったとしても、オープンソースのディープラーニングライブラリをクラスタ管理のAPIで動作するように書きなおさなければなりません。

この記事では、別の解法として、CPUとGPUを別のクラスタで実行し、Apache Sparkを使ってエンドツーエンドでモデリング処理を実行する手法をご紹介します。

構成概要

我々は、1つのツールからCPU上とGPU上のどちらにもタスクを実行でき、様々な言語やAPIを持ったディープラーニングのライブラリをアルゴリズムの書き換え無しに実行できるような構成が必要でした。Sparkをメインの開始地点として残しておくために、ニューラルネットワークによる学習と予測は、特別なGPUクラスタに委譲可能な荒い粒度のタスクであると考えました。これは、行列の掛け算といった低レベルなタスクをプリミティブなタスクとして切り出す様な、より古典的なアプローチとは大きく異なります。

Spark (CPU)のクラスタはAmazon EMR上で実行され、GPUのインスタンスはAmazon ECSで管理されます。言い換えると、我々はECSをGPUのマスターと見なしています。ECSは、タスクをDockerコンテナとして実行し、そのイメージはAmazon ECRに保存されているので、ディープラーニングライブラリはDockerイメージをECRに入れるだけで簡単に取り入れることができます。

以下の図が、ハイレベルな構成図になります。

この構成では、データ分析と処理(つまりCPUジョブ)は素のSparkで実行されますが、その際ジョブはタスクに分解されSpark executorとして実行されます。前述のGPUジョブは学習もしくは予測用のニューラルネットワークを参照しています。ジョブのためのデータのパーティショニングはSparkによって実施されますが、これらジョブの実行はECSに委譲されていて、GPUスレーブ群の上のDockerコンテナ内で実行されます。2つのクラスタ間のデータ転送は、Amazon S3を介しています。

GPUジョブが実行されると、1つ以上のGPUタスクに分解されます(詳しくは後半のセクションをご覧下さい)。Spark同様に、GPUタスクもデータのRDDの各パーティション毎に割り当てられます。Spark executorは関連するパーティションをS3に保存し、ECS上のあるタスク定義を、インプットのパーティションがあるS3の場所の指定と、実行するDockerイメージを指定する様にコンテナの上書きをしつつ実行するように呼び出します。その後、各GPUタスクのステータスを監視するためにECSをロングポーリングします。

GPUノードでは、各タスクは以下のように実行されます:

1. S3からデータパーティションをダウンロードします
2. 指定されたコマンドを実行します
3. コマンドのアウトプットをS3にアップロードします

学習と予測はDockerコンテナ内で実行されるので、新しいライブラリをサポートするときにやることは、Dockerイメージを作成し、ECRにアップロードし、そして正しいイメージにマップしたタスク定義を作成するだけです。

次のセクションでは、DSSTNEを使ってどういった種類のGPUタスクを実行しているかの詳細をご説明します。

DSSTNEでディープラーニング

我々のニューラルネットワークのモデルは、しばしば数10万のノードがインプットとアウトプットのレイヤに存在します(つまり、幅広のネットワークです)。この規模では、たとえ薄くても、全接続のネットワークでは重みの数は簡単に数兆に到達します。そのため、我々のモデルは1つのGPUのメモリには収まらないことが多いです。

前述の様に、こうした大規模な疎なレイヤをサポートするために我々はDSSTNEを開発しました。それを行うための1つの方法として、DSSTNEではモデル並列学習をサポートしています。モデル並列学習では、モデルはN個のGPUに分散されますが、データセット(RDD)は全てのGPUノードに複製されます。これに対して各GPUがデータの部分集合だけで学習するデータ並列学習では、パラメータサーバ等の同期手法を使って重みを各ノードが共有しています。

モデルが学習された後は、各お客様用に予測(すなわちレコメンデーション)を生成します。各お客様のレコメンデーションは独立して生成可能なので、これはEmbarrasingly parallelなタスクになります。なので、我々はデータ並列で予測を実施していて、各GPUがお客様のバッチ単位で予測を行っています。これによって単にGPUを追加するだけで線形にスケールすることができます。そのため、パーティションの数(GPUの数)を2倍にすれば予測生成の時間は1/2になります。Auto Scalingの機能を使って、ワークロードとSLAの制約を基準にGPUクラスタをスケールさせることができています。

以下の図で、モデル並列学習とデータ並列予測を表現しています。

以上でGPUのワークロードの種類を御覧頂きました。ここからはSparkからどの様にエンドツーエンドでモデリングの繰り返しを構成しているのかの例をご覧ください。

Apache Sparkでのオーケストレーション

前にご説明した様に、ディープラーニングアルゴリズムはECS管理の分離されたGPUクラスタ上のDockerコンテナ内で実行されるので、プログラム的にリモートクラスタ上での実行エンジンにアクセスが可能です。ユーザの入り口はSparkになります。ZeppelinやJupyterといったノートブックの力で、ユーザはノートブックから離れることなしに、対話的にデータを取り出して分析し、ニューラルネットワークを学習し、予測を生成することができます。

ここからの2つのパートでは、DSSTNEのGitHubページで紹介されているサンプルのニューラルネットワークを、Zeppelinのノートブックから先に紹介した2つのクラスタを使って、どの様に学習・予測するかをご紹介します。

モデル並列学習

このセクションでは、MovieLensのデータセットを4つのGPUにまたがったモデル並列学習を使いながら、どの様に3層のオートエンコーダを学習するかをご説明します。データ、モデルの設定、そしてモデルのトポロジーはZeppelinのノートブック内で定義され、S3をファイルシステムとして使いGPUスレーブ群に伝達されます。ECS上で実行されるコマンドの一覧としては、データと設定をダウンロード、学習の開始、モデルの成果物をS3にアップロードする、となります。

GPUタスクをECSに委譲するために、ECSクラスタとタスク定義をセットアップする必要があります。以下のステップがECSのセットアップの要約になります。

1. Amazon ECSのためのAWS Java SDKをSparkのclasspathに追加します
2. EMRクラスタのEC2インスタンスがECSを呼び出す権限を持ったインスタンスプロファイルを持っていることを確認します
3. ECSクラスタをGPUインスタンスで作成し、DSSTNEと互換のあるNVIDIAドライバーをインストールします
4. 特権フラグをセットしDSSTNEのDockerイメージを指定したECSのタスク定義を、AWS CLI等で作成します

EMRクラスタにはSparkとZeppelin (Sandbox)がインストールされていて、Zeppelinの%shインタプリタを使って必要なファイルをダウンロードできます。EMRクラスタのサイズはデータセットのサイズに依存しますが、今回のMovieLensのデータセットでは2つのc3.2xlargeインスタンスの利用をお勧めします。

学習を開始するためにsparkインタプリタに切り替えます。runOnECS()メソッドは、与えられたコマンドリストを指定されたタスク定義のDockerコンテナ上で実行してくれるものです。このメソッドはZeppelinのノートブック上でも定義することができ、AmazonECSClientをInstanceProfileCredentialsProviderを使って初期化し、コマンドのリストをcontainer overridesとして使ってRunTaskRequestを送信することで実現できます。

以下の図が、Zeppelinノートブックからモデル並列学習を実行するデモになります。

データ並列予測

上述の学習されたモデルを使って、データ並列予測を実行して、各お客様向けに上位100個のレコメンデーションを生成します。predict_input RDDの各パーティションは別のGPUで並列実行されます。予測の生成後には、テストデータセットと比較することでその予測を評価することがえきます。もしモデルのパフォーマンスが十分でなければ、同じノートブック内でネットワークの学習に戻って、データのスライスやモデルのパラメータ、もしくはネットワークトポロジを変更して、高速に繰り返し実行することができます。

以下の図が、データ並列予測とその予測結果をプロットしているデモになります。

まとめと今後

この記事では、パーソナライゼーションチームがSparkとDSSTNEを使ってどの様に巨大なディープラーニングの課題を高速に繰り返し実行しているか、そして何百万ものお客様へのレコメンデーションを毎日どの様に生成しているかをご説明してきました。

SparkとDSSTNEは両方ともそれらが持つ並列化の手法によって、巨大なデータを扱える様にスケールすることができます。ECSを使ってGPUのハードウェア上でミニバッチジョブを実行することで、CPUとGPUの両方のノードが混ざったクラスタを管理するという複雑さを回避しつつ、複数のディープラーニングライブラリのインストールが共存できるようになりました。さらに、ディープラーニングライブラリを簡単に差し替えできるので、科学者にも技術者にも自分の問題に最も適したライブラリを選択できる自由を与えてくれました。

EMRを使ってSparkクラスタを作成することで、ワークロードに応じてCPUクラスタのサイズとクラスタ数を簡単にスケールさせることができます。しかし、ECS側では現在1つのGPUクラスタを共有しているのですが、ECSにはタスクのキューをする機能がありません。その結果、ユーザはGPU不足に苦しんでいます。これを緩和するために、ECS Mesos Scheduler Driverのプロトタイプをベースにして、ECS上でMesosのスケジューラを実装しています。

全ての機械学習のライフサイクルをSparkから行える様にしていることで、クリエイティブなアプリケーションを生み出すためにSparkやディープラーニングのライブラリを組み合わせていくことが可能になります。例えば、Spark Streamingの力によって、ストリーミングデータからGPUで予測タスクを実行し、お客様のレコメンデーションを準リアルタイムに更新することができるでしょう。MLlibがサポートするモデルはニューラルネットワークの起点にもできるし、調和して一緒に動作させることもできるでしょう。Sparkの並列実行エンジンを使って、ハイパーパラメータの最適化を並列実行もできるでしょう。この様に、無限に考えられます。

AWSを使っておけば、このスタックの管理とセットアップは、シンプルで、繰り返し可能で、スケール可能なものとなります。これらのツールを使って作られたあなたのアプリケーションを見るのが楽しみです。

あなたのアプリケーションと解決可能な課題について、コメント欄を使ってぜひ教えて下さい。

大規模なディープラーニングとレコメンデーションに興味がありますか？であれば、パーソナライゼーションのキャリアページもぜひご覧下さい。

原文: Generating Recommendations at Amazon Scale with Apache Spark and Amazon DSSTNE (翻訳: SA岩永)