兆パラメータ規模を突破し、究極の効率性と費用対効果を追求：中国チームが初の異種並列推奨システムトレーニングフレームワークをオープンソース化 ペルシャ

パーソナライズされた推奨は、インターネット業界が DAU (Daily Active Users) と収益を増やすための中核的な技術的手段です。ディープラーニングが広く応用されるにつれて、現代の推奨システムはニューラルネットワークを使用してユーザーの行動習慣を隠して「記憶」し、ユーザーの好みを正確に予測します。モバイルインターネットの普及後、ユーザー行動データは飛躍的に増加し、単位時間あたりに生成・収集されるユーザー行動データの量は非常に大きいため、ユーザーの興味をエンコードするにはより大きなモデルが必要になります。

データ規模が大きいということは、より大きなモデル容量が必要であることを意味します。モデルパラメータの数は、5年前の10億から、しばらく前にFacebookが公開した10兆パラメータ規模まで急速に増加しています。この傾向により、より大規模なトレーニングの需要が、間違いなく克服すべき次のマイルストーンとなるでしょう。

最近、2 つの中国チームが共同でオープンソース化したトレーニング フレームワークである Persia は、効率性と精度の両方を考慮しながら、1 兆個のパラメータ (数百テラバイトのストレージが必要) のハイブリッド アーキテクチャを設計することで、Google Cloud 上でモデル規模を新たなレベルに押し上げることに成功しました。現在、このフレームワークは Pytorch エコシステム Pytorch Lightning に統合されるよう招待されています。

• GitHub: https://github.com/PersiaML/PERSIA
• 論文: https://arxiv.org/abs/2111.05897

モデルパラメータの数が指数関数的に増加するにつれて、高性能なトレーニングフレームワークの需要がますます緊急になります。従来のアーキテクチャでは、増加するパラメータ数に対応できなくなってきています。従来のアーキテクチャでは、CPU の均質並列メカニズムを採用し、対応するパラメータ分布はモデル並列性を採用しています。

比較的大型の機種にも対応できる水平拡張性を備えているのが最大の利点であり、現在でも多くの企業で広く採用されています（各社で命名方法は異なりますが、本稿では総称して mio アーキテクチャと呼んでいます）。推奨モデルが従来のロジスティック回帰からディープラーニングベースのモデルにアップグレードされ、パラメータの数が大幅に増加すると、従来のソリューションは不十分になり、非効率になり、精度のバランスを取るのが難しくなります。その後の進化版では、ディープ ネットワーク部分の計算を担うために GPU が導入されました (この記事では mio+ アーキテクチャと呼びます)。依然として同型設計コンセプト (CPU マシンを GPU 搭載 CPU マシンに置き換えただけ) を採用しているため、一定の効率向上を実現し、効率と精度の矛盾を部分的に緩和できますが、異なる規模のネットワーク構造を処理する必要がある場合、高価な GPU リソースがアイドル状態になることが多く、費用対効果に重大な損害をもたらします。

モデルサイズの継続的な拡大によって生じるこれら 2 つの問題を解決するために、Persia の中心的な設計思想は次のとおりです。

• GPU リソース使用率の問題を解決するために、異種アーキテクチャ設計が使用されます。 CPU と GPU の比率が制限されている場合、どのフレームワークでも、どのモデル構造でも同時にリソース使用率を保証することは困難です。そのため、Persia は、モデル要件に応じてリソースを割り当て、効率を確保しながらリソースを最大限に活用し、コスト効率を大幅に向上させる柔軟な異種アーキテクチャを設計しました。
• 効率性と精度のバランスをとるために、同期と非同期のハイブリッドトレーニングモードが採用されています。従来のソリューションでは、純粋に同期したトレーニングまたは純粋に非同期のトレーニングのいずれかを使用します。モデルが大きくなり、マシンの数が増えると、同期トレーニングではマシンがお互いを待つことになり、トレーニング効率が著しく低下します。非同期トレーニング方式では、マシン同士が待機する必要がなくなり、トレーニング効率が大幅に向上しますが、マシンの数が増えるとモデルの精度が大幅に低下します。超大規模モデル シナリオにおけるこのような課題に対処するために、Persia は、両方の長所を組み合わせながら短所を回避する、同期と非同期のハイブリッド トレーニング アーキテクチャを設計しました。 Persia は同期トレーニングの精度と非同期トレーニングの効率の両方を実現できることが理論的および実践的側面の両方から検証されています。

ここで簡単にいくつかのポイントを挙げます ペルシャの特徴:

• PyTorch エコシステムのネイティブ サポート: PyTorch によって研究者がモデルを定義するハードルが大幅に下がり、ディープラーニング分野全体で PyTorch がますます大きなシェアを占めるようになっていることから、Persia は、既存の推奨トレーニング フレームワーク (XDL、PaddlePaddle など) とは異なり、PyTorch エコシステムをベースとすることを決定しました。ユーザーモデルの定義などの操作は PyTorch の助けを借りて直接実装できるため、研究分野の最新かつ最先端のモデル (Transformer など) も直接呼び出すことができ、最大限の柔軟性と使いやすさを実現できます。
• 高いパフォーマンス: Criteo 標準データセットでは、同じリソース条件下で、他の一般的なオープンソースの推奨モデルトレーニングフレームワークと比較して、   Persia はパフォーマンスを 2 倍以上に向上できます。 Persia は、CPU-GPU 異種トレーニングと GPU 間の直接通信をサポートし、トレーニング コストを大幅に削減します。
• スケーラビリティ: Persia は、最大 100 兆のモデル パラメータまでの規模で高いトレーニング効率を維持します。同時に、ほとんどのシナリオでは、ほぼ線形の加速を実現できます (n 倍のリソースを投資し、トレーニング効率をほぼ n 倍向上します)。
• 産業グレードのシナリオ大規模検証：ペルシャ クラウドネイティブのデプロイメントをサポートするために、Kubernetes 用のカスタマイズされたオペレーターを実装しました。さまざまなフォールト トレラント メカニズムが実装されており、オンライン実稼働環境で 2 年以上にわたって安定して動作することが検証されています。 Persia は、数億人の DAU を持つ複数のコアビジネス シナリオでテストされ、パフォーマンスとビジネス指標の大幅な改善を達成しました。
• 安全でトラブルシューティングが簡単: Persia は、メモリの安全性、速度、並行性に重点を置いた Rust で実装されており、コンパイル時に多数のメモリの安全性の問題が排除されます。多数の監視ポイントをネイティブに提供し、Grafana と完璧に統合し、さまざまなアラーム条件をカスタマイズできます。同時に、トレースに基づいてモジュール化された階層的なログ出力が実装され、実際のシナリオでのトラブルシューティングが容易になります。
• 柔軟な機能処理: クロスフィーチャなどのさまざまな一般的な機能処理方法をサポートし、ユーザーは Python スクリプトを通じてさまざまなカスタム機能処理モードを定義できます。柔軟性と使いやすさ。
• オンラインとオフラインの一貫性: オフライン トレーニング コードとオンライン トレーニング コードが統合されているため、エンジニアが一貫性のないモデルのオンライン効果のトラブルシューティングに多くの時間を費やす必要があるという問題点が解消されます。

ペルシャデザインのアイデア

全体的なアーキテクチャ

推奨モデルでは、モデルは多くの場合、次の部分で構成されます。

• 埋め込みレイヤー:  埋め込みレイヤーは、ユーザー ID やアイテム ID などの ID 機能に対応する埋め込みで構成されます。各 ID は、あらかじめ設定されたサイズのベクトルに対応します (埋め込みと呼ばれます)。ID の数は非常に多い場合が多いため、これらのベクトルがモデル全体のボリュームの 99% 以上を占めることがよくあります。
• 非IDタイプの機能:  画像情報、LDA などの実際のベクトル特徴。この部分は、IDに対応する埋め込みベクトルと組み合わされ、DNNに入力され、クリック率などを予測します。
• 高密度ニューラルネットワーク (NN):  この部分は、埋め込みベクトルと実数ベクトルの特徴を受け入れ、クリック率などの予測したい量を出力するニューラルネットワークです。

この種の推奨モデルでは、埋め込み層のパラメータがモデルボリュームの大部分を占めることが多いですが、埋め込み層の計算の複雑さはそれほど大きくありません。 NN パラメータの数はモデルボリュームのほんの一部を占めるだけですが、計算作業の大部分を占めます。これは、ハードウェアの観点で言えば、CPU はメモリが大きいが計算能力が低く、GPU はビデオメモリが小さいが計算能力が高い、ということに相当します。

既存のトレーニング フレームワークには GPU コンピューティング パワーが含まれていますが、各 GPU ワーカーは多数の PS を使用してデータとモデルを転送する必要があり、通信のボトルネックが発生することが多く、全体的な効率が低下します。

したがって、  ペルシャ システム設計では、NN は GPU メモリに配置され、勾配計算は GPU を通じて実行されます。 NN 部分については、PS をまったく経由せずに、GPU 間の効率的な集合通信を通じて直接同期されます。埋め込みはメモリに配置され、CPU によって計算されます。 Persia は、PS に 2 層アーキテクチャ設計 (Embedding PS、Embedding Worker、後ほど紹介) を使用します。これにより、ほとんどのシナリオで GPU ワーカーの帯域幅消費をさらに削減し、全体的なトレーニング効率を向上させることができます。

同期+非同期混合トレーニング

さらに、既存のシステムでは、完全に同期したトレーニング方法や完全に非同期のトレーニング方法を採用することがよくあります。完全に同期したトレーニングでは、すべての GPU ワーカーがデータのバッチに対してモデルをトレーニングおよび更新し、すべてが完了したら次のデータのバッチに進みます。モデルが大きくなり、マシンの数が増えると、マシン間の待機や同期にかかる時間が大幅に増加し、限られた時間内にトレーニングを完了することが難しくなります。この場合、システムのトレーニング プロセスは、下の図の最初の行 (Full Sync) に表示されます。完全に非同期のトレーニングでは、各ワーカーが PS パラメータを独立してトレーニングおよび更新します。作業員同士が待つ必要がなく、トレーニング効率は高いのですが、機械の台数が増えると、作業員ごとに使用するモデルの差が大きくなり、モデルのトレーニング効果が大幅に低下してしまいます。この場合、システムのトレーニング プロセスは、下の図の 2 行目 (Full Async) に表示されます。

これら2つの問題に対してペルシャは ハイブリッド トレーニング アーキテクチャは、トレーニング結果を確保しながら、ほぼ完全に非同期のトレーニング効率を実現するように設計されています。推奨シナリオのトレーニングにおける重要な観察は、埋め込みの更新が非常にまばらであり、2 つの更新間の交差が小さいことが多いことです。したがって、埋め込みが非同期的に更新されたとしても、最終的なトレーニング結果にはほとんど影響がありません。 NN 部分の更新はその逆で、毎回すべてのパラメータが更新されます。非同期トレーニングを実行すると、トレーニング結果に大きな差が生じます。 Persia が提案したハイブリッド トレーニング メソッドでは、NN 部分を同期的にトレーニングし、埋め込み部分を非同期的にトレーニングできます。最終的なトレーニング効率は純粋な非同期トレーニングの効率に近くなりますが、モデル効果は完全な同期トレーニングと一致します。両方の長所を活かそう。この場合、システムのトレーニング プロセスは、下の図の 3 行目 (Naive Hybrid) に示されています。 Persia は、通信とコンピューティング操作を重複させ、それらを並行して実行することで、システム効率をさらに向上させます。最終的なシステムトレーニングプロセスは、下の図の 4 行目 (ペルシャ) に示されています。

理論上の保証

既存のシステムとは異なり、Persia はハイブリッド アルゴリズムの設計に対して厳密な理論的保証を提供します。損失の期待値の最適化問題（各サンプルが損失に対応する推奨における最も一般的なシナリオなど）の場合：

ここで、f(w) はデータセット全体の平均損失を表し、ξ はサンプルを表し、w はモデルパラメータを表し、F(w; ξ) はサンプル ξ の損失を表します。モデルトレーニングの目標は、データセット全体の平均損失を最小限に抑えることです。使用  Persia Hybrid のトレーニング方法により、モデルの収束速度が次のようになることが証明できます。

ここで、σ はデータセットの分散、T は反復回数、τ は GPU ワーカーの数、α は ID クラス機能の衝突確率です。最初の 2 つの項目は完全同期トレーニングの収束速度であり、最後の項目はハイブリッド トレーニングによって導入されるエラーです。推奨シナリオでは、Embedding の更新が非常にスパースであり、衝突確率 α が 1 よりはるかに小さいため、Hybrid の収束速度は完全同期トレーニングとほぼ同じです。ただし、同期オーバーヘッドが削減されるため、トレーニングの各ステップの実行効率が大幅に向上します。具体的な理論的証明については、  ペルシャの論文[1]。

その他の最適化

アルゴリズムの革新に基づき、究極のパフォーマンスを実現します。 Persia は、実装レベルの最適化を多数提供します。例えば：

• すべてのPSサービス通信は、トレーニングシナリオに最適化されたゼロコピーを使用します。   Persia RPC システムは、トレーニング シナリオ (埋め込みや勾配などの大量のテンソル データを含む非常に大きなペイロードを特徴とする) において、従来の RPC フレームワーク (gRPC や bRPC など) よりも優れたパフォーマンスを発揮します。
• GPU 間の通信には、Kuaishou によってオープンソース化された Bagua トレーニング アクセラレーション フレームワークが使用されます。これにより、GPU 間の集合的な通信パフォーマンスが大幅に向上し、勾配圧縮などのアルゴリズムによって通信オーバーヘッドをさらに削減できます。
• 埋め込みは、特別に設計されたデータ構造 (Persia Embedding Array List) を通じて PS に保存され、PS の効率とモデル アクセスの効率が大幅に向上します。これには、操作中に動的に新しいメモリを要求する必要性を排除し、CPU キャッシュ メカニズムをより有効に活用することが含まれます。埋め込み削除ロジックをサポートします。モデルの保存および読み取りプロセスは、連続メモリの直接ダンプ/ロード プロセスに簡素化されます。

• 導入  Persia Embedding Worker コンポーネントは、合計プーリングの埋め込みや生データの処理などの操作を GPU に送信する前に実行し、GPU 帯域幅の使用量を大幅に削減します。

• 元のデータは、  ペルシャ  Compact Batch 形式には、ID 重複排除とデータ圧縮表現プロパティが組み込まれており、一般的な表現方法と比較してデータ量が 1/4 に削減され、システムのデータ処理効率が向上します。

検証と比較

テストでは、Alibaba-Ad、Avazu-Ad、Criteo-Ad などのさまざまなオープンソース データ セットが使用され、全体的なトレーニング効率が 8 倍以上向上しました。

ペルシャ 最大 100 兆のモデル トレーニングをサポートし、モデル サイズの増加に応じてトレーニング効率を維持できます。

資源が拡大するにつれ、ペルシャ ほぼ直線的に拡張できます (リソースの量を n 倍投資すると、トレーニング効率がほぼ n 倍向上します)。

ペルシャ 使用事例

使用 ペルシャ 非常にシンプルで、主にトレーニングの展開、モデルの定義、カスタム データ セットの部分に分かれています。

• 分散展開：  ペルシャ オペレーターはワンクリックでKubernetesクラスターにPERSIAタスクをデプロイできる
• モデル定義: PyTorchを直接使用する
• カスタムデータセット: カスタム前処理ロジック、結果を渡す ペルシャ 提供されるPythonツールキットは、  ペルシャ コンパクトバッチ

完全な例とその他のシナリオについては、以下を参照してください。  ペルシャ チュートリアルドキュメント (https://persiaml-tutorials.pages.dev/)。ペルシャ オンライン推論のモデル

ペルシャ トレーニング済みモデルの埋め込み部分は、Embedding PS と Embedding Worker をオンラインでデプロイすることで、サービスとして直接提供できます。 NN 部分はネイティブの PyTorch モデルです。  ペルシャ  TorchServe による推論の簡単な例がチュートリアルに記載されています。ユーザーは、TensorRT モデルへの変換など、さまざまなネイティブ PyTorch ツールを使用して、推論パフォーマンスをさらに向上させることもできます。​