Xiaohongshu プッシュ検索シナリオにおける計算能力のボトルネックを打破するために機械学習の異種ハードウェア推論を最適化する方法!

この記事では、Xiaohongshu プッシュ検索シナリオの完全な GPU 構築プロセスにおけるモデル サービス、GPU 最適化、およびその他の関連作業について説明します。

1. はじめに

近年、動画、画像、テキスト、検索など機械学習分野のアプリケーションは発展を続けており、モデル計算量やパラメータ量はCPUムーアの法則の成長率をはるかに上回っています。この文脈では、GPU コンピューティング能力の開発と大規模モデルの開発は同時に進行しています。多くの企業が GPU コンピューティング能力の開発を組み合わせて、自社の機械学習の問題に対するソリューションを模索しています。

他の企業と同様に、Xiaohongshu のさまざまなシナリオのモデルも大きくなっていますが、CPU の開発は必要な計算能力に追いついていません。そこで、2021年に、推論性能と効率性を向上させるために、一般的な検索モデルのGPUベースの変換を開始しました。移行プロセスでは、以前の CPU アーキテクチャの作業を GPU アーキテクチャにスムーズに移行する方法や、Xiaohongshu のビジネス シナリオとオンライン アーキテクチャに基づいて独自のソリューションを開発する方法など、いくつかの困難にも直面しました。同時に、モデルの反復を継続するためにコストを削減し、効率を高める必要があります。

II. 背景

具体的な実践を紹介する前に、まずは小紅書の応用シナリオを紹介したいと思います。 Xiaohongshu APP のホームページには、推奨ページと検索ページが含まれており、広告コンポーネントも含まれています。一部の GPU は、洗練された CTR モデル、CVR モデル、相関モデル、および一部のソート シナリオとリコール シナリオで使用されます。現在、すべてのファインランキングシナリオは GPU 推論に移行されており、ファインランキングの主な目的は CTR、CVR またはその他の複数のターゲットを正確に推定することです。

計算パラメータの量から見ると、小紅書の計算規模は2021年の初めから2022年の終わりにかけて大幅に拡大しました。推奨シナリオを例にとると、各リクエストには 400 億フロップスかかり、パラメータの総数は数千億に達します。

3. モデルサービス

1. モデルの特徴

2022年末にChatGpt型モデルが提案される前は、業界内や検索プッシュ企業においてDenseモデルの特に大規模な応用シナリオはありませんでした。したがって、2022年末までに、Xiaohongshuのメインモデルの大きなパラメータは完全にスパースになります。具体的には、レコメンデーションメインモデルを例にとると、多数のパラメータをIDタイプと交差させる必要があります。たとえば、Xiaohongshuノートをユーザーの都市と交差させ、XiaohongshuノートをユーザーIDと交差させるなどです。特徴埋め込みの構築は、パラメータスパース化プロセスになります。デカルト積問題により、スパース化プロセス全体のパラメータの数は TB、数千億、数兆に達する可能性があります。多くの大手企業では、数兆、さらには数十兆のパラメータに達しています。ただし、コストと CTR モデルの利点の制限を考慮して、Dense 部分の計算量はそれほど大きくしませんでした。高密度部分の計算は基本的にグラフィック カードが収容できる 10 GB 以内で制御されます。ただし、シナリオが異なるため、他の違いもあります。たとえば、オンライン シナリオでは、コンピューティング能力の制限に加えて、高い同時実行性やレイテンシなどの制限もあります。

2. トレーニングと推論フレームワーク

次に、Xiaohongshu のトレーニングと推論のフレームワークを紹介します。

GPU を使用する場合、モデルは主に 2 つのカテゴリに分けられます。1 つは、先ほど紹介したスパース CTR モデルを含む検索および予測モデルです。これらには独自の推論フレームワークとトレーニング フレームワークがあります。もう 1 つのモデルには、CV モデルと NLP モデルが含まれます。これらは PyTorch テクノロジー スタックに基づいて構築されており、独立したトレーニングおよび推論フレームワークの進化を備えています。

CTR型大規模テクニカルモデルはTBレベルのモデルです。推論部分とトレーニング部分をそれぞれ以下に紹介します。

推論フレームワークに関しては、2020年以前、XiaohongshuはTensorFlowテクノロジースタックを使用し、TensorFlow Servingを基本的な推論フレームワークとして採用し、このオープンソースフレームワークに基づいて独自の研究を行っていました。

TensorFlow Serving の特徴は、多くのモデルライフサイクル管理をホストし、多数の周辺コンポーネントを備えていることです。ただし、TensorFlow Serving には独自のオーバーヘッドがあり、特に無駄なメモリコピーが発生します。オリジナルの TensorFlow の基盤となる API は、TensorFlow のグラフ エンジンである CTensor に基づいています。ただし、インターフェースを抽象化するために、TensorFlow Serving は、リクエストに応答するために、外側に Proto を持つ TensorProto 構造のレイヤーをカプセル化します。そのため、2020 年に最適化を行い、TensorProto の外側のレイヤーを削除しました。TensorFlow Serving に直接依存するのではなく、基盤となる CTensor API、つまり CPI に直接基づいた独自の Lambda Service 推論フレームワークを構築しました。その後の反復では、グラフ スケジューラやコンパイル最適化などのテクノロジがさらに統合されました。

トレーニングフレームワークに関しては、20年前まで、Xiaohongshuには大規模なスパーストレーニングフレームワークがありませんでした。2020年初頭に構築が完了した後、検索と昇格シナリオのトレーニングフレームワークは、TensorFlowに基づいて独自に開発したWorker＆Psトレーニングフレームワークに統合されました。同時に、独自のオペレーターも開発しました。その後、GPU のアップグレード後に CPU、IO などのブロックの問題を解決するために、この自社開発のトレーニング フレームワークに対して複数回の技術的なアップグレードを実施しました。

3. 機械の特性

マシンの特性上、Xiaohongshu には独自のコンピュータ ルームはなく、すべてのマシンはクラウド ベンダーから購入されています。ただし、購入できるマシンモデルは限られており、帯域幅のカスタマイズ、メモリのカスタマイズ、カード数のカスタマイズなど、カスタマイズすることはできません。したがって、テクノロジを選択する際には、購入可能なマシン モデルに基づいて、最も最適化されたアーキテクチャを選択する必要があります。

4. GPUの機能

GPU 機能の問題に関しては、Xiaohongshu は他の企業と同じ問題に直面しています。GPU カーネルの実行は、データ転送、カーネルの起動、カーネルの計算、結果の転送といういくつかの段階に分かれています。まず、データをホストメモリから GPU メモリに転送する必要があります。カーネルの起動には、カーネルコードをホストから GPU に転送し、GPU 上でカーネルを起動する必要があります。カーネルが計算結果を実行した後、結果が GPU からホストに転送されます。全体のプロセスは、転送、計算、転送の 3 つの段階で構成されています。データ転送とカーネルの起動に多くの時間が費やされると、GPU の使用率が低下したり、GPU が空になったりすることもあります。

4. GPU最適化の実践

上記の問題に対して、Xiaohongshu は以下の具体的な実用的な最適化ソリューションを提案しました。

1. システムの最適化

（１）物理マシンの最適化

まず、システム最適化の物理マシン最適化の側面に焦点を当てましょう。先ほども述べたように、私たちはクラウドベンダーからマシンを購入しました。ハードウェアの複雑さを隠すために、クラウドベンダーのマシンでは多くの仮想化が行われています。仮想化自体は優れています。K8S エコシステムでは、仮想化によってアプリケーションの柔軟性が高まり、ハードウェアの詳細を保護できます。しかし、仮想化ではハードウェアの速度が GPU の速度に追いつけないため、パフォーマンスが低下する可能性があります。そこで、仮想化の中間工程を削減するために、以下の点についてクラウドベンダーと連携しました。

• 割り込み分離: エージェント分離とその他の割り込み分離など、GPU 割り込みを分離します。主な目的は、仮想マシンのジッターを減らすことです。
• カーネル バージョンのアップグレード: GPU ドライバー システムの互換性とパフォーマンスを向上させるために使用されます。
• 命令パススルー: GPU 命令を物理デバイスに直接渡し、GPU の計算速度を高速化します。上記の物理マシンに基づくシステム最適化により、パフォーマンスが 1% ～ 2% 向上します。

（２）マルチカード最適化

パート2では、マルチカードの最適化を実行しました。

購入した CPU は AMD 製であるため、NUMA の問題があります。オンライン フェーズ中に NUMA を介してメモリにアクセスするのにかかる時間は、CPU がローカル メモリに直接アクセスするのにかかる時間よりもはるかに長くなります。そこで、使用中にPodにNUMAバインディングを実施しました。バインドされたPodは、NUMAをまたいでメモリを申請したりノードを呼び出したりする必要がないため、CPUとGPU間のデータ転送速度が向上し、CPUステージでのメモリ消費が削減されます。 CPU ステージの速度が速くなり、後続の GPU 推論に多くの時間が確保されます。

（３）コンパイル最適化

システム最適化の 3 番目の部分も非常に一般的です。 CPU の発展に伴い、異なるモデルは異なる命令セットをサポートするようになりました。そのため、カーネルを実装すると、異なるモデルに対して多数の異なる定義が存在することになります。この問題を解決するために、異なるモデルの命令セットをクロスコンパイルしました。対応するモデルに最適なカーネルの実装を見つけ、対応するコンパイル命令を開く必要がありました。

Alibaba Cloud から購入した Intel 8163 と 2 つの A10 モデルを例にとり、このモデルの特性に基づいてコンパイル命令セットをカスタマイズしました。そのため、このモデルでは、実行可能なランタイム バイナリをランタイムに置き換えました。

システム段階でコンパイル命令を最適化すると、パフォーマンスが 10% 向上します。

2. 計算最適化

コンピューティングの最適化に関しては、依然として前述の中核的な問題、つまり GPU は高速に実行されているが CPU のメモリ アクセスが追いつかない場合に問題をどのように解決するかという問題を中心に展開しています。

CPU はメモリに頻繁にアクセスし、メモリ ページ フォールトの頻度が高くなるため、CPU リソースの浪費やレイテンシの増大の問題が発生します。レイテンシが高い場合、オンライン推論フェーズでタイムアウトの問題が発生する可能性があります。同時に、オンライン推論サービスでは、単一リクエストのバッチサイズが小さく、単一サービスの同時実行規模が大きく、数千 QPS の単一リクエストの小さなバッチでは GPU の計算能力を十分に活用できないため、元の TensorFlow の単一の Cuda Stream 起動カーネルがボトルネックになります。この時点で、推論シナリオでの GPU 使用率はわずか 50% です。さらに、小規模なモデル シナリオでは、Dense 部分は比較的単純であり、推論計算のためにデータを GPU に転送してから CPU に転送するコストは、直接 CPU 計算するよりも高くなり、コスト効率が良くありません。

したがって、さまざまなシナリオに直面して、ページ フォールト、高レイテンシ、バッチ制約の問題という 2 つの問題を解決する必要があります。

（１）メモリページフォールトの高頻度化への対応

まず、高メモリ ページ フォールトの問題には 2 つの解決策があります。

1つは、エンジン側のデータ構造を最適化することです。 Tensor データ構造をグラフ スケジューリング エンジンに渡す前に、追加のオーバーヘッドを削減してゼロ コピーを実現する、つまり、ユーザー側のリクエスト データをグラフ エンジンの入力に直接移動する必要があります。そのため、TensorFlow の基盤となる CPI のデータ構造をカスタマイズする必要があります。このカスタマイズされたデータ構造はリクエスト側に接続され、基盤となるグラフ エンジンに直接送信されるため、コストのオーバーヘッドが削減されます。

2 つ目は、オペレーティング システム レベルで透過的なラージ ページ機能を有効にして、ページ フォールトを減らすことです。同時に、jemalloc ライブラリを使用して、ガベージ コレクション時間を長く設定し、継続的に最適な状態に調整することで、追加のオーバーヘッドを削減し、ページ フォールトが多い問題を解決しました。

（２）TensorFlowの単一Cudaストリームの問題

さらに、TensorFlow の単一 Cuda ストリーム問題に対処するために、マルチ Cuda ストリームとマルチ コンテキストの機能をサポートしており、これにより、ミューテックス ロックのパフォーマンス ボトルネックを回避し、同時実行性と GPU 使用率を向上させることができます。同時に、複数の GPU カードを同時に使用する場合に存在する問題を解決するために、Nivida が提供する Cuda MPS 機能を使用して GPU の空間分割多重化を実現しました。もともと、Nvidia を Cuda コンピューティングに使用する場合、通常の状況では、一度に実行できるコンテキストは 1 つだけです (CPU プロセスで Cuda プログラムが呼び出されたときに、GPU 側でリソースとデータを要求するための唯一の単位)。ただし、一部のカードが Hyper-Q 機能をサポートしている場合は、Nvidia の MPS サービスを有効にして GPU 空間分割多重化 (複数の Cuda 実行を同時にサポート) を実現し、GPU の使用率をさらに向上させることができます。ただし、空間分割多重化により、プロセス障害分離の問題が発生する可能性があります。そのため、プッシュ検索などの特定のシナリオや、コンピューティング タスクが固定されているなど、比較的固定されたシナリオでは、GPU 空間分割多重化を有効にすることをお勧めします。

（３）計算と併合

計算最適化の 3 番目の部分は、計算プロセス全体にわたって計算をマージする可能性を見つけることです。

カーネルレベルの融合に加えて、アップストリームの Op レベルの融合も実行できます。手書きまたはグラフコンパイル最適化ツールを使用して、より高性能な Tensorflow 演算子を生成します。たとえば、Mat と Bn および Relu の計算は、最適化された FusedMatmul 演算子を使用して実行できます。コンパイルレベルの最適化を統合した後、演算子を再コンパイルして新しい演算子で実行し、計算オーバーヘッドを削減できます。

右の図に示すシナリオの 1 つでは、コア コンピューティング コストを 50% 削減できます。

（4）冗長コンピューティングの最適化

計算の最適化のためのより賢明なアイデアは、冗長な計算を見つけて最適化することです。

プッシュ検索 CTR モデルには多くの冗長な計算があります。プッシュ検索シナリオの最も顕著な特徴は、1 ユーザー + N アイテムの推定であり、元の計算入力段階で Tensor 構造を平坦化する必要があり、つまり、ユーザー側の計算はデータ側から N 個の部分に分割され、この部分の計算は冗長です。 2 番目の冗長性は、元の TensorFlow API を使用してグラフ全体を構築する場合、データ要求が GPU 上で複数回処理され、一部の計算が CPU が終了した後に GPU に転送されるため、効率が十分ではないことです。

そのため、演算子を統合して導出を最適化し、CUDA の API を使用して実装を統一することで、各リクエストが GPU に 1 回だけ入力されるようにし、GPU に複数回入力される冗長なコンピューティング機能を削減しました。

さらに、事前に計算を実行することで冗長な計算を削減できます。たとえば、モデルを出力するときに、変数演算子を凍結し、変数演算子を定数演算子に変換する必要があります。たとえば、グラフ凍結プロセス中に +1+2 の計算がありますが、これは静的に直接 +3 に変換できます。

グラフのフリーズにより、全体的なコンピューティング使用量が出力で 12% 削減され、非常に効果的です。計算後に外部データを GPU にコピーする問題を軽減し、より高いパフォーマンスの向上を実現するために、GPU の CUDA を使用してマージ計算部分を実装するよう最善を尽くす必要があります。

（５）より良いハードウェアへの変更

コンピューティングの最適化における非常に直接的な最適化は、より優れたハードウェアを交換することです。

オンライン推論段階では、A100 のような特に高価なカードを使用する必要はありません。T4 を A10 に置き換えると、パフォーマンスが 1.5 倍向上しますが、コストは 1.2 倍しかかかりません。将来的には、オンライン推論に A30 やその他のモデルを使用することも検討します。同時に、ハードウェア メーカーがバージョンを繰り返すにつれて、オンラインでのコスト メリットを実現するために、より優れたハードウェアに直接切り替えることになります。

しかし、ハードウェアを変更すると、CPU に配置された計算部分はさらに役に立たなくなります。Vlookup 機能やマージクロスハッシュ計算など、GPU での処理に適さないものが常に存在します。この時点で、3 番目のステップに進む必要があります。つまり、いくらかのレイテンシを無駄にして、前の CPU 計算部分を分割することになります。推論アーキテクチャは、最初に CPU によって計算され、次に GPU オンライン推論で実行される要求に変更されます。その後、ハードウェア GPU がアップグレードされ続け、CPU が追いつかなくなると、分割され続けることになります。

（６）DLスタック自動コンパイル最適化

DL スタックの自動コンパイル最適化もコア最適化ポイントです。

当社のコアテクノロジースタックは、Alibaba のオープンソース BladeDISC と TensorFlow の公式 XLA テクノロジースタックです。 Alibaba の Blade フレームワークは、MRI に基づく動的 Shark のディープラーニングコンパイル機能をサポートしています。従来のコンパイラ、たとえば C++ は、コンパイル最適化を通じて特定の C++ コードをマシン命令にコンパイルします。つまり、マシン コードはマシンの対応するシステムによって実行されます。ディープラーニング ML コンパイラーでは、グラフを DL テクノロジー スタックにコンパイルし、中間実装を見つける必要があります。たとえば、MLIR はこの中間層の表現に基づいてハードウェア上の最終的な実行状態をコンパイルします。このハードウェア状態は、GPU 上の実際のカーネル実行計算です。この部分では、ブレードと XLA テクノロジー スタックを統合し、大きなメリットを実現しました。ここでは特に Blade フレームワークをお勧めします。これは非常に良い結果をもたらしました。また、Apache 2.0 オープン ソースであり、すべての企業で使用できます。

3. トレーニングの最適化

（１）トレーニング最適化フェーズ１および２

3 番目の部分では、主にトレーニングの初期段階でのいくつかの最適化について説明します。

トレーニングと推論で発生する問題は若干異なります。トレーニングにはレイテンシの制約がなく、特定のミリ秒数以内に戻る必要はありませんが、スループットが高く、より多くのデータを取得することが期待されます。データ埋め込み層のプロセスは、後方勾配バックプロパゲーションのコストを可能な限り低く抑えながら、IO ステージで完了できます。 GPU コンピューティングにさらに多くの時間とデータを使用します。このプロセス中に、次の問題が発生しました。

1 つ目は、データ IO の問題です。元の TF トレーニング データは TFRecord 行に基づいて保存されます。つまり、データはフラット化されて 1 行に配置されるため、非常に大きな IO が発生します。したがって、まず入力データ行を列に変換して、IO データ取得次元を削減します。都市、ID などの列挙可能な機能入力の場合、行を列に変換すると、データ ディメンションを取得する IO のスケールが大幅に縮小されます。データ行が列ストレージに変換された後、列ストレージのシナリオに適応するためにトレーニング フレームワークを適宜アップグレードする必要があります。同時に、列ストレージと行ストレージではデータ分布に大きな違いがあり、行ストレージの方が分散に適しています。そのため、データを列に格納する場合は、データを列に格納してからデータ分布を復元し、ワーカーを入力してトレーニングし、GPUを入力して処理するようにしてください。

データを取得した後、データの前方参照を行う必要があります。 GPU の使用率を向上させるために、フル プリフェッチを積極的に実装しました。これにより、トレーニング中に CPU 側のルックアップ ブロッキングによって発生するトレーニング待機を解決できます。ただし、このタイプの更新は完全に非同期の更新となり、更新の遅れが生じ、効果の安定性に影響を与えるため、メカニズム設計による強力な保護が必要になります。また、トレーニングの逆更新自体への依存を解除したり、ルックアップ逆更新非同期戦略をアップグレードしたり、データ勾配を累積的に更新したりするなど、逆方向のアップグレードも行いました。次のバッチトレーニング勾配が更新されると、現在の更新を書き込むことができます。

（２）トレーニングと推論の最適化

このセクションでは、全体的なトレーニングと推論の最適化について説明します。

当社の機械学習全体は、トレーニングと推論のアプリケーション層に基づいて実装されており、主にアプリケーション層、ソフトウェア層、ハードウェア コンテナ層が含まれます。この図は、2021 年から 2022 年にかけて行った作業を示しています。各タスクには、独自の独立した制約とアプリケーション シナリオがあります。同様の大規模プッシュ検索モデルを作成したことがある方で、これまでどの方向も試したことがない方は、こちらの内容を参考にしてください。たとえば、精度が圧縮されている場合は、Float 16 や Int 8 などのより低いレベルの精度を使用できます。しかし、実際の使用シナリオでは、エフェクトの影響を制御するという問題があります。したがって、精度圧縮を使用する場合は、ホワイトボックス作業を行う可能性が高くなります。つまり、モデルをエクスポートするときに、すべてのレイヤーの精度を下げるのではなく、GPUに入る部分を選択して精度を下げます。同時に、後続の最適化ロジックを通じて精度の低下の程度を保証し、AUC全体または実際の指標がロスレスであることを保証する必要があります。つまり、ホワイトボックスの場合は、精度圧縮を合理的に選択します。

4. 未来

最後に、Xiaohongshuの機械学習エンジン全体の今後の進化の方向性を紹介したいと思います。

疎な大規模モデルのトレーニングでは、いくつかのことを行います。まず、トレーニング側でさまざまなパラメータ スケールを区別します。特に小さいパラメータ スケールについては、HPC 化して、追加のオーバーヘッドを発生させずに 1 台のマシンで完了するようにします。 1TB レベルに達するパラメータの場合、A100+ フル高速インターコネクトを使用します。数十 TB または数百 TB に達する非常に大きなパラメータ ボリュームの場合は、これを実装するために小型 GPU (A10) + インフラストラクチャ通信機能を使用してインフラストラクチャ通信を実行します。

推論に関しては、ハッシュ、マルチレベル キャッシュ、および軽量モデルのテクノロジを将来的にアップグレードする予定です。同時に、Nvidia のアップグレード リズムに従い、ドライバーとハードウェア バージョン全体を更新します。

最後に、会社全体の反復効率を最適化するために、ドラッグアンドドロップ、キャンバススタイルのアプローチを備えたワンストップ機械学習プラットフォームをその後開発します。同時に、DSL ベースの機能管理も実行し、機能管理を抽象オペレータ処理にアップグレードします。