DeepRec Large Scale Sparse Model Training Reasoning Engine

ガイド：

本記事では、以下の 3 つの側面を紹介します。

- DeepRec の背景（なぜ DeepRec が必要なのか）

- DeepRec の機能（設計の動機と実装）

- DeepRec コミュニティ（最新リリース 2206 の主な機能）

DeepRec の背景

なぜスパースモデルエンジンが必要なのでしょうか。現在のコミュニティ版 TensorFlow はスパースシナリオをサポートできますが、以下の 3 つの側面で機能的な弱みがあります。

- モデル効果を向上させるスパーストレーニング機能

- モデル反復効率のトレーニングパフォーマンスの向上

- スパースモデルのデプロイメント

そこで、DeepRec を提案しました。その機能は、スパースシナリオでのディープラーニングを最適化することです。

DeepRec の注力分野は主に 4 つの側面です。スパース機能、トレーニングパフォーマンス、サービング、デプロイメントと ODL です。

DeepRec の Alibaba 内部でのアプリケーションは、主に以下のコアシナリオで使用されています。レコメンデーション（好みのアイテムを推測）、検索（メイン検索）、広告（ダイレクトトレインとターゲティング）です。クラウド上の顧客にもスパースシナリオ向けのソリューションを提供しており、モデル効果と反復効率の向上に大きく貢献しています。

DeepRec の機能紹介

DeepRec の機能は主に 5 つの側面に分かれます。スパース機能の埋め込み、トレーニングフレームワーク（非同期、同期）、ランタイム（エグゼキュータ、PRMalloc）、グラフ最適化（構造モデル、SmartStage）、サービスデプロイメント関連の機能です。

1. 埋め込み

埋め込みセクションでは、以下の 5 つのサブ機能を紹介します。

1.1 動的弾性特徴（EV）

上記の図の左側は、TensorFlow がスパース機能をサポートする主な方法です。ユーザーがまず固定形状の変数を定義し、スパース特徴は Hash+Mod によって定義された変数にマッピングされます。ここには 4 つの論理的な問題があります。

- スパース特徴のコンフリクト。Hash+Mod は特徴コンフリクトを引き起こしやすく、有効な特徴が失われ、効果に影響を及ぼします。

- ストレージ部分でメモリの無駄が生じ、一部のメモリ空間が使用されません。

- 固定形状。変数の形状が固定されると、後から変更できません。

- 非効率的な IO。ユーザーがこの方法で変数を定義する場合、すべてフル量でエクスポートする必要があります。変数のディメンションが大きい場合、エクスポートとロードの両方に時間がかかりますが、スパースシナリオでは変更が少ないです。

この状況において、DeepRec で定義される Embedding Variable の基本原理は、静的な変数を HashTable に似た動的ストレージに変換することです。各キーが新しい埋め込みを作成するため、特徴コンフリクトの問題を自然に解決します。この設計により、特徴が多くなりすぎると、Embedding Variable が無秩序に拡張され、メモリ消費が非常に大きくなります。そこで、DeepRec は以下の 2 つの機能を導入しました。特徴アドミッションと特徴エリミネーションです。これらは特徴が大規模なディメンションに拡張されるのを効果的に防ぎます。検索とレコメンデーションのようなスパースシナリオでは、一部のロングテール特徴はモデルによってほとんどトレーニングされません。そのため、特徴アドミッションでは CounterFilter または BloomFilter を通じて Embedding Variable に入る特徴のしきい値を設定できます。モデルがチェックポイントにエクスポートされる際、特徴エリミネーション機能も備えており、時間的に古い特徴も排除されます。これにより、Alibaba 内のあるレコメンデーションサービスの AUC が 5 パーミル向上し、クラウド上のあるレコメンデーションサービスの AUC も 5 パーミル向上、pvctr は 4% 増加しました。

1.2 特徴頻度に基づく動的弾性ディメンション特徴（FAE）

一般的に、同じ特徴に対応する Embedding Variable は同じディメンションに設定されます。Embedding Variable を高いディメンションに設定すると、低頻度特徴の内容が過学習を引き起こしやすく、大量のメモリを消費します。逆に、ディメンションを低く設定しすぎると、高頻度特徴の内容が表現力不足によりモデルの効果に影響を及ぼす可能性があります。FAE 機能は、同じ特徴に対しても特徴に応じて異なるディメンションを提供します。これにより、モデルの自動トレーニング時に、まずモデルの効果を確保し、次にトレーニングリソースの使用を最適化できます。これは FAE 機能の設計動機の紹介です。この機能の使用により、現在ユーザーはディメンションと統計アルゴリズムを渡すことができます。FAE は実装されたアルゴリズムに従って異なる Embedding Variable を自動生成します。今後、DeepRec はシステム内で特徴ディメンションを適応的に発見して割り当てることで、ユーザーの使いやすさを向上させる予定です。

1.3 Adaptive Embedding Variable

この機能は、高頻度と低頻度の関係の定義から始まる点で、2 番目の機能と似ています。前述の EV が特に大きい場合、メモリ消費が特に高くなることがわかります。Adaptive Embedding Variables では、右図に示すように 2 つの変数を使用します。一方の変数を静的なものとして定義し、低頻度特徴をできるだけこの変数にマッピングします。もう一方は動的弾性ディメンション特徴として定義され、高頻度部分の特徴に使用されます。この変数は低頻度特徴と高頻度特徴の動的変換をサポートし、システムのメモリ使用量を大幅に削減します。たとえば、ある特徴のトレーニング後、最初のディメンションは 10 億に近づく可能性がありますが、重要な特徴はわずか 20 〜 30% です。この適応的方法により、メモリ使用量を大幅に削減できます。実際のアプリケーションでは、モデルの精度への影響は非常に小さいことがわかっています。

1.4 Multi-Hash Variable

この機能は特徴コンフリクトの問題を解決するために使用されます。以前は 1 つの Hash+Mod で特徴コンフリクトを解決していましたが、現在は 2 つ以上の Hash+Mod を使用して埋め込みを取得し、結果の埋め込みを削減します。これにより、少ないメモリで特徴コンフリクトの問題を解決できるという利点があります。

1.5 埋め込みのマルチレベルハイブリッドストレージ

この機能の出発点も、特徴が多い際に EV のメモリコストが非常に高く、トレーニング中にワーカーが占有するメモリが数十 GB から数百 GB に達する可能性があるという発見からです。特徴は典型的なべき分布に従うことを考慮し、ホットスポット特徴を CPU などのより価値のあるリソースに配置し、比較的ロングテールの低頻度特徴を比較的安価なリソースに配置します。右図に示すように、DRAM、PMEM、SSD の 3 つの構造があります。PMEM は Intel が提供するメモリで、速度は DRAM と SSD の中間に位置し、大容量です。現在、DRAM-PMEM、DRAM-SSD、PMEM-SSD のハイブリッドをサポートしており、ビジネスで成果を上げています。クラウド上のあるサービスでは、以前は 200 台以上のマルチ CPU 分散トレーニングを使用していましたが、マルチレベルストレージの使用により、現在はシングル GPU トレーニングに変更されています。

以上が埋め込みのすべての機能の紹介です。これらの機能を開発する動機は、TensorFlow のいくつかの問題（主に特徴コンフリクト）に起因しています。解決策は動的弾性特徴とマルチハッシュ特徴です。動的弾性特徴の大きなメモリオーバーヘッドの問題に対して、特徴アドミッションと特徴エリミネーション機能を開発しました。特徴頻度に関しては、3 つの機能グループを開発しました。動的弾性ディメンションと適応的動的弾性特徴はディメンションの方向で解決し、マルチレベルハイブリッドストレージはソフトウェアとハードウェアの方向で解決しています。

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

2 番目に紹介する機能はトレーニングフレームワークで、非同期方向と同期方向に分かれます。

2.1 非同期トレーニングフレームワーク StarServer

大規模タスクの場合、数千のワーカーとネイティブ TensorFlow の問題は、スレッドスケジューリングが非常に非効率的で、クリティカルパスのオーバーヘッドが顕著であり、パケット通信が非常に頻繁であることで、これらが分散通信のボトルネックとなっています。

StarServer はグラフのスレッドスケジューリングとメモリ最適化に優れています。フレームワーク内の Send/Recv を Push/Pull セマンティクスに変更し、PS は実行時にロックフリー方式を使用することで、実行効率が大幅に向上します。ネイティブフレームワークと比較して数倍のパフォーマンス向上を実現し、内部的には 3,000 ワーカー規模でも線形スケーリングを達成できます。

2.2 同期トレーニングフレームワーク HybridBackend

これは同期トレーニング向けに開発したソリューションです。データ並列処理とモデル並列処理のハイブリッド分散トレーニングをサポートします。データの読み取りはデータ並列処理で完了します。モデル並列処理は多数のパラメータを持つトレーニングをサポートできます。最後に、密な計算にはデータ並列処理を使用します。複数の埋め込みルックアップのマージとグループ化に対して最適化を行い、GPU Direct RDMA の利点を活用して、ネットワークトポロジーを考慮した同期フレームワーク全体を設計しています。

3. ランタイム

3 番目の主要機能はランタイムで、主に PRMalloc とエグゼキュータの最適化について紹介します。

3.1 PRMalloc

まずメモリ割り当てについて説明します。メモリ割り当ては TensorFlow と DeepRec の両方で広く使用されています。まず、スパーストレーニングでは大きなメモリ割り当てが多数のマイナーページフォールトを引き起こすことを確認しました。さらに、マルチスレッド割り当てでは同時割り当ての問題も存在します。DeepRec では、スパーストレーニングの順方向と逆方向の特性に基づいて、ディープラーニング専用のメモリ割り当て方式 PRMalloc を設計しました。これによりメモリ使用量とシステムパフォーマンスを改善します。図に示すように、中心となるのは MemoryPlanner で、モデルトレーニングの最初の k ラウンドのミニバッチで現在のトレーニング特徴を統計的に分析し、各割り当てで必要なセンサー数を把握し、バイナリバッファーを通じてこれらの動作を記録して対応する最適化を行います。k ステップ以降にこれを適用することで、前述の問題を大幅に軽減できます。DeepRec での使用により、マイナーページフォールトの発生を大幅に削減し、メモリ使用量を削減し、トレーニング速度を 1.6 倍に加速できることが確認できました。

3.2 エグゼキュータの最適化

TensorFlow のネイティブエグゼキュータの実装は非常にシンプルです。まず DAG をトポロジカルソートし、ノードを実行キューに挿入してから、タスクを通じてエグゼキュータがスケジューリングを行います。この実装はビジネスを考慮しておらず、スレッドプールはデフォルトで Eigen スレッドプールを使用しています。スレッド負荷が不均一な場合、多数のスレッドが Steal を競合し、多大なオーバーヘッドが発生します。DeepRec ではスケジューリングをより均一に定義し、クリティカルパスを定義することで、Ops の実行時に優先順位付きのスケジューリングを実現しています。最後に、DeepRec は Task ベースと SimpleGraph ベースの複数のスケジューリング戦略も提供しています。

4. グラフ最適化関連の機能

4.1 構造的特徴

これはビジネスから着想を得た機能です。検索シナリオでは、トレーニングも推論も、通常 1 人のユーザーが複数のアイテムと複数のラベルに対応するサンプルであることを発見しました。従来の処理方法ではこれらを複数のサンプルとして扱うため、ユーザーのストレージが冗長になります。このコストを削減するために、ストレージ形式をカスタマイズしてこの部分を最適化しました。ミニバッチ内のこれらのサンプルが同じユーザーの場合、一部のユーザーネットワークとアイテムネットワークをそれぞれ個別に計算し、最後に対応する論理計算を実行することで、計算オーバーヘッドを削減できます。これにより、ストレージ側と計算側の両方から構造最適化を行いました。

4.2 SmartStage

スパースモデルのトレーニングには通常、サンプル読み取り、埋め込みルックアップ、MLP ネットワーク計算が含まれます。サンプル読み取りと埋め込みルックアップは計算集約的ではなく、計算リソースを効果的に使用できません。ネイティブフレームワークが提供するプリフェッチインターフェイスである程度非同期操作を実現できますが、埋め込みルックアッププロセス中に設計するいくつかの複雑なサブグラフは TensorFlow のプリフェッチを通じてパイプライン化できません。TensorFlow のパイプライン機能では、実際の使用時にユーザーがステージ境界を明示的に指定する必要があります。一方面では使用の難易度を上げ、もう一方面ではステージの精度が不十分で op レベルまで正確に指定できないため、高レベルの API ユーザーは手動で挿入できず、多くのステップが並列化できません。下図は SmartStage の具体的な動作を示しており、Ops を異なるステージに自動的に分類することで、並列パイプラインのパフォーマンスを向上させます。ModelZoo でのモデルテスト効果の最大高速化率は 1.1 〜 1.3 倍に達します。

5. サービング

5.1 モデルのインクリメンタルエクスポートとロード

冒頭で、埋め込みの導入における重要なポイントの 1 つは非効率的な IO でした。前述の動的弾性機能を使用すれば、自然にインクリメンタルエクスポートが可能になります。以前アクセスされたスパース ID がグラフに追加されている限り、インクリメンタルエクスポート時に必要な ID を正確にエクスポートできます。この機能には 2 つの動機があります。まず、モデルトレーニングの従来の方法では、各ステップでフル量のモデルエクスポートを行い、プログラムの中断時の復元チェックポイントとしても使用していました。最悪の場合、2 つのチェックポイント間隔の結果をすべて失う可能性がありました。インクリメンタルエクスポートにより、密な部分をフル量で、スパース部分をインクリメンタルにエクスポートすることで、実際のシナリオでは 10 分間のインクリメンタルエクスポートが復元による損失を大幅に削減できます。もう 1 つはオンラインサービングのシナリオです。スパースシナリオではモデルが非常に大きく、毎回フル量でロードすると長時間かかります。インクリメンタルエクスポートは、オンライン学習が困難な場合でも ODL シナリオでも使用されます。

5.2 ODL

左側がサンプル処理、上下がオフラインとオンラインのトレーニング、右側がサービングです。多数の PAI コンポーネントを使用してパイプラインを構築しています。

DeepRec コミュニティ