ディープ ラーニングをベースとする レコメンダー モデルのトレーニング、Google Cloud により 100 兆個のパラメータ規模で実現

※この投稿は米国時間 2022 年 5 月 24 日に、Google Cloud blog に投稿されたものの抄訳です。

100 兆個のパラメータによるレコメンダー モデル トレーニング

レコメンダー システムは、今日のインターネット サービスの重要な構成要素です。10 億ドル規模の収益ビジネスが、大手テクノロジー企業のレコメンデーション サービスによって直接推進されています。現状、本番環境のレコメンダー システムではディープ ラーニングをベースとするアプローチが主流となっています。このアプローチでは、はじめに埋め込み層を採用し、固定長の埋め込みベクトルにきわめて大規模な ID タイプの機能をマッピングします。そこから複雑なニューラル ネットワーク アーキテクチャによって埋め込みが利用され、レコメンデーションが生成されます。レコメンダー モデルの継続的な進歩は、多くの場合、モデルサイズの増加によって推進されます。10 億個のパラメータを持ついくつかのモデルが過去にリリースされており、ごく最近では 1 兆個にまで達しています。モデルの容量が急増するたびに品質が大幅に改善されてきており、100 兆個のパラメータの時代も間近に迫っています。

レコメンダー モデルにおけるトレーニング タスクの規模は、独自の課題を生み出しました。トレーニングの計算には驚異的な異質性があります。モデルのサイズ全体の 99.99% 以上がモデルの埋め込み層に含まれる可能性があり、これはメモリを大いに消費します。一方、複雑で高密度な残りのニューラル ネットワークは、トレーニングの反復ごとに 100 を超える TFLOP を使用するため、ますますコンピューティング集約型になっていきます。そのため、異種リソースを使用してクラスタを管理する洗練されたメカニズムをこのようなトレーニング タスク向けに準備することが重要になります。

最近、スイス連邦工科大学チューリッヒ校の Kwai Seattle AI 研究室と DS3 研究室が協力して、トレーニング アルゴリズムとトレーニング システムの両方を慎重を期して共同設計し、この問題に取り組む「Persia」という新しいシステムを提案しました。アルゴリズム レベルでは、Persia にハイブリッド型のトレーニング アルゴリズムを採用し、埋め込み層と高密度ニューラル ネットワーク モジュールを異なる方法で処理します。埋め込み層はトレーニング サンプルのスループットを向上させるために非同期的にトレーニングし、残りのニューラル ネットワークは統計効率を維持するために同期的にトレーニングします。システム レベルでは、ハイブリッド型のアルゴリズムの可能性を最大限に引き出すために、メモリ管理と通信量の削減のためのさまざまなシステム最適化が実装されています。

Google Cloud に大規模なトレーニングをデプロイ

Persia が必要とする規模の大きさは、コンポーネント全体で必要なネットワーク帯域幅から、埋め込みを保存するために必要な RAM メモリの量まで、さまざまな課題をもたらしました。さらに、パイプラインを最小限に抑えてコストを最適化するために、デプロイ、自動化、およびオーケストレーションを行う必要のある仮想マシンが相当数ありました。具体的には、ワークロードは次の異種リソースで実行されます。

• 3,000 個のコアを備えたコンピューティング集約型仮想マシン

• 合計 64 個の A100 Nvidia GPU を追加した 8 台の A2 仮想マシン

• それぞれ 12 TB（合計 360 TB）の RAM を搭載した 30 台のハイメモリ仮想マシン

• Kubernetes とのオーケストレーション

ネットワーク レイテンシを最小限に抑えるために、すべてのリソースを同じゾーンで同時に起動する必要がありました。Google Cloud は、ほとんど予告することなく、必要な容量を提供することができました。

トレーニングのバースト性を考慮して、Google Kubernetes Engine（GKE）を使用して、138 台の仮想マシンとソフトウェア コンテナのデプロイをオーケストレートしました。ワークロードをコンテナ化することで、トレーニングの移植と反復性も実現できます。

チームは、トレーニング中にすべての埋め込みをメモリに保持することを選択しました。これには、比較的短期間ではあるものの、高度に専門化された「Ultramem」仮想マシンを使用する必要があります。これは、処理のコストと期間を管理しながら、トレーニングを 100 兆個のパラメータまでスケールさせるために欠かせないものでした。

結果と結論

Google Cloud インフラストラクチャのサポートにより、チームは Persia を 100 兆個のパラメータまでスケールできることを実証しました。ハイブリッド分散型トレーニング アルゴリズムにより、バニラ SGD と同じ速さで収束しながら、異種クラスタを効率的に利用するための精巧なシステム リラクゼーションを導入しました。Google Cloud はオンプレミス ハードウェアの制限を克服するために不可欠であり、大規模な分散型機械学習トレーニングに最適なコンピューティング環境であることが証明されました。

Persia は、Google Cloud のセットアップ手順を含むオープンソース プロジェクトとして GitHub にリリースされました。学界と産業界の誰もがディープ ラーニング レコメンダー モデルを 100 兆個のパラメータ規模で簡単にトレーニングできます。