PyTorch に基づく WeChat の大規模レコメンデーション システム トレーニングの実践

#この記事では、PyTorch をベースにした WeChat の大規模レコメンデーション システム トレーニングについて紹介します。他のディープ ラーニング分野とは異なり、レコメンデーション システムは依然として Tensorflow をトレーニング フレームワークとして使用していますが、これは開発者の大多数によって批判されています。 PyTorchを活用したレコメンデーション研修などの実践もあるが、規模が小さく、実際のビジネス検証が行われていないため、ビジネスの早期導入を促進することが難しい。

2022 年 2 月、PyTorch チームは公式推奨ライブラリ TorchRec を開始しました。私たちのチームは5月に社内業務でTorchRecを試し始め、TorchRecチームとの一連の協力を開始しました。数か月間試用した結果、TorchRec の多くの利点を実感しましたが、非常に大規模なモデルでは、TorchRec にはまだいくつかの欠点があるとも感じています。これらの欠点に対応して、私たちはその問題を補う拡張機能を設計しました。 2022 年 9 月に、私たちが設計した拡張機能の動的埋め込みは、TorchRec のメイン ブランチに正式に統合され、現在も公式チームとともに継続的に最適化が行われています。

1.トーチライクが私たちにもたらすもの

まず、TorchRec が私たちに何をもたらすかについて話しましょう。レコメンデーション システムは企業のキャッシュ フローに直結することが多く、試行錯誤のコストが非常に高いことは誰もが知っています。そのため、誰もが必要としているのは、ビジネスでテストされたフレームワークです。これが、PyTorch に基づく以前のレコメンデーション フレームワークの一部が広く使用されていない理由です。公式の推奨フレームワークとして、TorchRec は 2022 年 1 月に開始されました。Meta はすでにこれを使用して、Instagram Reels ビジネスで 1,250 億のパラメーター モデルのトレーニングと立ち上げに成功し、ビジネスでテストされた PyTorch フレームワークになりました。 Instagram のような大企業のサポートにより、私たちはさらに自信を持ち、ついに PyTorch に基づくレコメンデーション フレームワークの利点を合理的に検討できるようになりました。

#TorchRec には、チームのメンバーごとに異なるメリットがあります。まず第一に、チーム内のアルゴリズム エンジニアの大多数にとって、PyTorch 推奨フレームワークにより、CV および NLP エンジニアが体験したよりユーザー フレンドリーなダイナミック グラフィックスとデバッグ エクスペリエンスを最終的に誰もが楽しむことができるようになります。

さらに、PyTorch は優れた互換性を備えており、PyTorch1.8 ベースのモデルは、コード行を変更することなく最新バージョン 1.13 で実行できます。アルゴリズム エンジニアは最終的に自信を持ってフレームワークをアップグレードできるため、最新のフレームワーク機能とより優れたパフォーマンスを享受できます。一方で、TensorFlow に基づく一部の推奨フレームワークは、TensorFlow の特定のバージョンに依存していることが多く、たとえば、多くのチームは依然として TensorFlow 1.x に基づく内部フレームワークを使用している可能性があります。 TensorFlow 1.x は 2021 年 1 月にメンテナンスを停止しました。つまり、過去 2 年間、すべての新しいバグと新機能は十分にサポートされなくなります。使用中に発生した問題は社内のメンテナンス チームのみが修理できるため、追加コストがかかります。タイムリーなフレームワークのアップグレードにより、無料の速度向上も実現できます。PyTorch の上位バージョンは、多くの場合、CUDA の上位バージョンと一致するだけでなく、トレーニング速度をさらに向上させ、トレーニング効率を向上させる CUDA グラフなどの新機能と一致します。

アルゴリズム エンジニアに加えて、フレームワーク チームも推奨チームの重要な部分です。オープンソースフレームワークの選定後、社内のフレームワークチームが社内ニーズに基づいて二次開発を行います。彼らにとって、PyTorch の推奨フレームワークは、より合理化された開発エクスペリエンスにつながります。従来の TensorFlow 推奨フレームワークの多くは、TF サービングを模倣して C セッションに基づいた拡張機能を作成します。この設計ソリューションは当時、非常に高度なソリューションと考えられていました。しかし、これには、コードを 1 行変更するだけで TensorFlow 全体を完全にコンパイルする必要があります。非常に時間がかかり、イントラネット上で外部の依存関係をダウンロードするなどの些細な問題も解決する必要があり、開発エクスペリエンスはあまり良くありません。

PyTorch を使用する場合、そのような問題は発生しません。PyTorch は Python の哲学を核としており、誰もがそれを自由に拡張できることを望んでいるからです。二次開発を行う場合は、pybind11 などの比較的成熟した Python ライブラリでカプセル化し、ライブラリをダイナミック リンク ライブラリにパッケージ化するだけでロードできます。このようにして、全体的なコンパイル速度は当然大幅に速くなり、学習コストは大幅に低くなります。

前述したように、PyTorch は非常に優れた下位互換性を備えたフレームワークです。これにより、保守チームは複数のバージョンを保守する必要がなくなり、多くの一般的な問題を解決できます。正式なソリューションを入手できれば、全員が特殊なニーズに集中できるようになり、チームの効率が大幅に向上します。

上記はすべて、PyTorch レコメンデーション フレームワークとしての TorchRec の利点です。私たちを非常に喜ばしく思うのは、TorchRec チームが PyTorch レコメンデーションに留まらなかったことです。フレームワーク。 。彼らは既存のレコメンデーション モデルとハードウェアの特性を観察し、多くの新機能をフレームワークに追加し、TorchRec に従来のレコメンデーション フレームワークに比べて明らかなパフォーマンス上の利点を与えました。次に、GPU エンベディング、TorchRec の優れた GPU カーネル、TorchRec がネットワーク通信に基づいて実行できるエンベディング部門をいくつか選んで紹介します。

1 つ目は GPU の組み込みです。まず、従来のレコメンデーション システムの GPU トレーニング プロセスを確認しましょう。特定のモデルを GPU ワーカーに配置し、エンベディングはリモート PS に保存されます。各反復ステップでは、まずリモート PS からパラメーターを取得し、次に GPU 上でモデルの順方向計算と逆方向計算を実行し、勾配を PS に転送して PS 上のパラメーターを更新します。

図の緑色の部分は GPU で実行される操作で、赤色の部分はネットワークまたは CPU で実行されます。 GPU はシステムの中で最も高価な部分ですが、多くの操作が GPU に配置されていないことがわかります。

従来のプロセスでは、GPU が十分に活用されていません。同時に、ハードウェアの観点から見ると、単一の GPU カードのメモリはますます大きくなり、一部の高密度モデルでは GPU を完全に活用することはできません。NVIDIA の継続的な最適化により、NV リンクと GPU ダイレクト RDMA の相互接続も可能になりました。 -カード通信がますます速くなります。

#GPU の埋め込みは非常にシンプルなソリューションです。彼はエンベディングを直接分割して GPU 上に配置します。たとえば、1 台のマシンに 8 枚のカードがある場合、エンベディングを直接 8 つの部分に分割し、各部分をカード上に配置します。これにより、すべての操作がカード上に残ることが保証されます。 . . GPUの利用効率が大幅に向上し、トレーニング速度も質的に向上します。 GPU 上のビデオ メモリ容量不足が心配な場合は、TorchRec は、ビデオ メモリの補足としてホスト上のメモリの一部を事前に分割できる UVM もサポートしており、これにより、内部に配置できるエンベディング サイズを増やすことができます。単一のマシン。

GPU の組み込みに加えて、TorchRec は非常に優れた GPU カーネルも実装しています。これらのカーネルは、最新のハードウェア機能と CUDA 機能を最大限に活用します。

#たとえば、エンベディング ルックアップ カーネルを実装したい場合は、大規模なエンベディングから ID に対応する一連のエンベディングを見つける必要があります。 .vector を使用すると、一般的な実装では、各 GPU スレッドに ID が割り当てられ、対応する埋め込みをそれぞれ検索できるようになります。このとき、GPU の最下層はワープに従ってスケジュールされ、ワー​​プ内の 32 スレッドが一緒にビデオ メモリの読み書きを行うことを考慮する必要があります。これは、上記の処理において、ID の読み取り時にはビデオメモリへのアクセスが継続されるものの、それ以降のコピーはランダムな読み書き状態となることを意味します。ハードウェアの場合、ランダムな読み書きではメモリ帯域を十分に活用できず、動作効率が十分ではありません。

#TorchRec は shuffle_sync のようなワープ プリミティブを使用して、ワープ内のすべてのスレッド、ラップ内の 32 のスレッドに ID をブロードキャストします。同じエンベディングを同時に処理できるため、連続的なメモリの読み取りと書き込みが実行できるため、ビデオ メモリの帯域幅利用効率が大幅に向上し、カーネルの速度が数倍向上します。

#この表は、埋め込みルックアップのパフォーマンス向上の公式テストです。ここでは Fused EBC が最適化されたカーネルであり、さまざまな設定の下で、TorchRec はネイティブの PyTorch と比較して数十倍のパフォーマンス向上があることがわかります。 TorchRec に基づいて、エンベディングが比較的小さい (128 未満) 場合、半分またはそれ以上のスレッドがアイドル状態になる可能性があることがわかりました。そこで、ワープ内のスレッドをさらにグループ化し、複数のエンベディングを同時に処理できるようにしました。

#私たちの改善により、小さな埋め込みカーネルがオンになります。 dim はさらに 10% ～ 30% 改善されました。この最適化は公式リポジトリにも組み込まれています。なお、TorchRec のカーネルは FBGEMM ライブラリに置かれているので、興味のある方は覗いてみてください。

#最後に、TorchRec の埋め込み分割の仕組みを紹介します。前述したように、GPUエンベディングはエンベディングを分割してカード上に載せることなので、それをどのように分割するかが考慮すべき問題となります。従来、行単位と列単位という 2 つの分割の考え方があります。行単位とは、特徴量が 20,000 個ある場合、カード 1 に 0 ～ 10000 の数字が配置され、カード 2 に 10000 ～ 20000 の数字が配置されることを意味します。このようにして、トレーニングするときに、ID がカード 1 に対応する場合、次から開始します。カード1から取るとカード2に対応するのでカード2から取ります。 Row wise の問題は、最初の 10,000 番号と最後の 10,000 番号のトラフィック量に大きな差があるかどうかがわからないため、通信が不均衡になり、ネットワーク ハードウェアを十分に活用できないことです。

#カラム単位は、埋め込み長の観点から分割されます。たとえば、埋め込みの全長は 128 です。最初の 64 次元と最後の 64 次元を異なる位置に配置できます。このようにすると、通信のバランスが良くなりますが、読み取り時には、すべてのカードまたは PS と通信する必要があります。 。

分割モードの違いにより、選択にトレードオフが生じます。従来のレコメンデーション フレームワークは、デザイン内の埋め込み分割方法を修正しますが、TorchRec は、行単位、列単位、さらにはテーブル単位、データ並列などの複数の分割方法をサポートし、Planner、Estimator、PerfModel などのモジュールやその他のモジュールを内部で自動的に提供します。帯域幅、ビデオメモリ、メモリ、モデルサイズ、および使用シナリオのその他のパラメータに基づいて分割方法を計算します。このようにして、実際のハードウェア条件に応じて最も効率的に埋め込みを分割することができ、ハードウェアを最も効率的に利用することができます。これらの関数のほとんどは Python で実装されています。これにより、社内環境に合わせてカスタマイズすることができ、社内環境に最適なレコメンドシステムを簡単に構築することができます。

2. 数百億のモデルでの実験結果

私たちの実験では、DeepFM や DCN などの標準モデルの場合、TorchRec は以前のベースライン推奨フレームワークと比較して、驚くべき 10 ～ 15 倍のパフォーマンス向上を示しました。このようなパフォーマンスの向上により、TorchRec をビジネスに導入する自信が得られます。

WeChat 読み取りの細かい配置モデルについては、位置合わせの精度に基づいて、実際のデータでは約 3 倍のパフォーマンス向上が見られ、偽データでは約 10 倍の向上も見られます。ここでの違いは、トレーニング用のデータの読み取りがボトルネックになっており、この点に関してはまださらなる最適化を行っていることです。

#03

##当初の計画は 1,000 億ドルまたはもっと見る 大規模モデルの欠点

以前に紹介したモデルは基本的に数百億以下のモデルです。 1台に収まるモデルです。 TorchRec をより大きなモデルにプッシュする際に、TorchRec のネイティブ デザインにいくつかの問題が確認されました。大規模なモデルの場合、TorchRec の純粋な GPU 埋め込みソリューションにはさらに多くのカードが必要です - おそらく 8 枚のカードの元のトレーニング速度ですべてのデータを吸収できますが、埋め込みを配置するには 16 枚のカードを使用する必要があるため、改善が困難になります。利用効率が再び低下しました。

#そして、大規模なモデル シナリオの場合、アルゴリズム チームは、1 週間アクセスされなかった ID の削除など、埋め込みのための動的な追加および削除の要件を提案することがよくあります。 。 TorchRec のソリューションはそのような機能をサポートしていません。さらに、非常に大規模なモデル ビジネスには一般に多くのチームが関与しており、基本的なフレームワークの移行には大きな抵抗が生じます。私たちに必要なのは、徐々に段階的に移行するためのサポートであり、全員が目の前の仕事を放棄することはできませんが、それはあまりにもコストがかかり、リスクが高くなります。

上記の要件に基づいて、非常に大規模なシナリオに適応できるように、TorchRec を変更する方法を検討しました。 -スケールモデル。リモート CPU PS は非常に成熟しており、埋め込みの動的な追加を簡単にサポートできるため、超大規模トレーニングではリモート PS 接続をサポートする必要があると考えています。同時に、チーム間の協力のために、PS を使用してトレーニングと推論を分離し、段階的な移行を実現できます。

#次に問題となるのは、PS をどのように導入するかです。 PS が GPU 組み込みに直接接続されている場合、各反復ステップでは引き続きリモート PS にアクセスする必要があるため、ネットワーク全体と CPU 操作の割合が増加し、GPU 使用率が低下します。

#04

WeChat チームでの動的埋め込みの問題を解決する方法

現時点で、単位時間あたりのデータ内の新しい ID は、実際には全体のほんの一部にすぎないことがわかりました。データ、HugeCTR が論文を発表 同様の結論についても言及されており、頻繁にアクセスされるのはごく一部の ID だけです。このことから、通常は GPU 埋め込みを使用してトレーニングを行い、表示メモリがいっぱいになったときに ID を PS にバッチで追い出すことを考えました。

このアイデアによると、GPU 組み込みが n 個の ID しか保存できず、合計 ID が N 個である場合、あるいは無限にたくさんあります。グローバル ID は 0、1、2、3... に順番にマッピングでき、マッピング関係は ID 変換と呼ばれる構造に保存され、GPU 埋め込みで通常のトレーニングにマッピング結果を使用できるようになります。 GPU の埋め込みがいっぱいの場合、つまり ID トランスフォーマーにマッピングのペアが n 個ある場合、ID はバッチで PS に削除されます。

#この設計では、PS がほとんど介入しない可能性があります。 , GPU ワーカーと PS の通信はエビクション中にのみ必要です。

さらに、この設計では、PS は KV として使用するだけでよく、KV として使用する必要はありません。パラメータの更新をサポートするため、オプティマイザ関連の操作を実装する必要がないため、PS チームはストレージ関連の作業に集中できます。また、任意の KV ストレージを実装するプラグインもサポートしており、オープンソース バージョンには Redis プラグインが組み込まれているため、Redis を PS としても使用できます。

#次に、動的埋め込みの設計の詳細をいくつか紹介します。私たちが実装した最も単純かつ基本的な ID Transformer は、実際には、PyTorch の高性能 ska:: flat_hash_map を使用してハッシュ テーブルを使用します。

ID トランスフォーマー プロセス内の唯一の CPU 操作として、パフォーマンス要件が比較的高くなる可能性があるため、また、メモリアクセス効率をさらに向上させるために、L1キャッシュライン単位に格納する高性能バージョンも実装しました。

また、立ち退きスキームについては、LRU と LFU を効率的に融合したいと考えています。 Redis の LFU ソリューションからインスピレーションを得て、ID アクセス頻度のインデックスのみが保存される確率的アルゴリズムを設計しました。たとえば、32 回アクセスされた場合、5 が格納されます。頻度を更新する際、再度このIDにアクセスすると5桁の乱数が生成され、5桁がすべて0、つまり1/32の確率でイベントが発生した場合、頻度指数を増加させます。 6まで。このような確率アルゴリズムにより、LRU と LFU の周波数を uint32 に入れることができ、メモリ アクセスの負荷を高めることなく LRU と LFU を統合できます。

最後に、マルチカード ソリューションを簡単に紹介しましょう。現在、すべてのカード データをカード 1 の ID Transformer に収集し、ブロードキャストで送り返しています。実装した ID Transformer は非常に高性能であり、GPU コンピューティング パイプラインと組み合わせることができるため、特定のパフォーマンスのボトルネックになることはありません。

上記は、動的埋め込みの設計アイデアの一部です。当社の社内の兆レベルのビジネスの 1 つでは、アライメント精度の場合、動的エンベディング ソリューションは社内のオリジナル GPU Tensorflow フレームワークと比較して約 3 倍のパフォーマンス向上を実現しています。 TF 最適化バージョンと比較すると、依然として 50% 以上のパフォーマンス上の利点があります。

最後に、皆さんにもトーチレックを試してみることをお勧めします。数百億未満の企業など、比較的小規模な企業の場合は、ネイティブの TorchRec を直接使用することをお勧めします。プラグ アンド プレイで、二次開発が不要で、パフォーマンスが 2 倍になります。非常に大規模なビジネスの場合は、TorchRec の動的エンベディングを統合してみることをお勧めします。一方で、内部 PS への接続が容易になる一方で、エンベディングの拡張と段階的な移行もサポートされます。同時に、一定のパフォーマンスの向上も得られます。

#興味のある方のために、私たちが調整したいくつかの高精度モデルと既存のアプリケーション シナリオを紹介します。友達も試してみることができます。

以上がPyTorch に基づく WeChat の大規模レコメンデーション システム トレーニングの実践の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。