必要な埋め込みパラメーターは 1% のみで、ハードウェア コストは 10 分の 1 に削減され、オープン ソース ソリューションは単一の GPU を使用して大規模な推奨モデルをトレーニングします。

ディープ レコメンデーション モデル (DLRM) は、ビデオ レコメンデーション、ショッピング検索、広告プッシュ、その他のトラフィック収益化サービスなど、インターネット企業におけるディープ ラーニングの応用にとって最も重要な技術シナリオとなっており、ユーザーの利便性を大幅に向上させています。経験とビジネスの商品価値です。ただし、膨大なユーザー データとビジネス データ、頻繁な反復更新要件、および高額なトレーニング コストはすべて、DLRM トレーニングに深刻な課題をもたらします。

DLRM では、まず埋め込みテーブル (EmbeddingBags) でルックアップを実行してから、ダウンストリームの計算を完了する必要があります。埋め込みテーブルは、DLRM のメモリ要件の 99% 以上に寄与することがよくありますが、計算にはわずか 1% しか寄与しません。 GPU オンチップの高速メモリ (高帯域幅メモリ) と強力なコンピューティング能力のおかげで、GPU は DLRM トレーニングの主流のハードウェアになりました。しかし、推奨システムの研究が深まるにつれて、埋め込みテーブルのサイズが増大し、GPU メモリが限られているため、大きな矛盾が生じています。 GPU メモリの壁の制限を突破しながら、GPU を使用して非常に大規模な DLRM モデルを効率的にトレーニングする方法は、DLRM の分野で解決する必要がある重要な問題となっています。

Colossal-AI は、これまでに異種戦略を使用して、同じハードウェアでトレーニングされた NLP モデルのパラメーター容量を数百倍に増やすことに成功しており、最近拡張に成功しました。 レコメンデーション システムでは、埋め込まれたテーブルはソフトウェア キャッシュ (Cache) メソッドを通じて CPU および GPU メモリに動的に保存されます。ソフトウェア キャッシュの設計に基づいて、Colossal-AI はパイプライン プリフェッチも追加し、将来入力されるトレーニング データを観察することでソフトウェア キャッシュの取得とデータ移動のオーバーヘッドを削減します。同時に、同期更新方式で GPU 上で DLRM モデル全体をトレーニングします。これは、広く使用されているハイブリッド並列トレーニング方法と組み合わせることで、複数の GPU に拡張できます。実験の結果、Colossal-AI は GPU に埋め込みパラメーターの 1% を保持するだけで済み、優れたエンドツーエンドのトレーニング速度を維持できることがわかりました。他の PyTorch ソリューションと比較して、グラフィックス メモリ要件が 1 桁削減され、1 枚のグラフィックス カードでテラバイト レベルの推奨モデルをトレーニングできます。コストの利点は顕著です。たとえば、91GB の埋め込みバッグを占有する DLRM をトレーニングするのに必要なビデオ メモリは 5GB だけです。トレーニング ハードウェアのコストは、A100 2 台の約 200,000 元から、エントリー レベルのグラフィックス カードに比べて 10 分の 1 に削減されます。 RTX 3050としては約2,000元しかかかりません。

オープンソース アドレス: https://github.com/hpcaitech/ColossalAI

既存の埋め込みテーブル拡張テクノロジー

埋め込みテーブルは、離散整数特徴を連続浮動小数点特徴ベクトルにマッピングします。次の図は、DLRM での埋め込みテーブル トレーニング プロセスを示しています。まず、埋め込みテーブル内の各特徴の埋め込みテーブルに対応する行を見つけ、最大、平均、合計演算などのリダクション演算を使用して特徴ベクトルに変換し、後続の高密度ニューラル ネットワークに渡します。 。 DLRM の埋め込みテーブル トレーニング プロセスは主に不規則なメモリ アクセス操作であるため、ハードウェアのメモリ アクセス速度によって大幅に制限されることがわかります。

#産業グレードの DLRM の組み込みテーブルは、数百 GB、さらには TB レベルに達する可能性があり、最大ビデオ メモリ容量の数十をはるかに超えています。単一 GPU の GB。単一 GPU のメモリの壁を突破して、DLRM の組み込みテーブルのサイズを増やす方法は数多くあります。以下の図に示す GPU クラスターのメモリ階層図を例として、いくつかの一般的なソリューションの長所と短所を分析してみましょう。

GPU モデルの並列処理: 埋め込みテーブルは複数の GPU のメモリに分割され分散され、トレーニング中に GPU 間の相互接続ネットワークを通じて中間結果が同期されます。 。この方法の欠点は、埋め込みテーブルのセグメント化の負荷が均一ではなく、スケーラビリティの問題の解決が難しいことです。第 2 に、GPU を追加するための初期ハードウェア コストが高く、DLRM トレーニング中に GPU のコンピューティング能力が十分に活用されず、HBM 帯域幅の利点のみが使用されるため、GPU 使用率が低くなります。

CPU 部分トレーニング: 埋め込みテーブルを 2 つの部分に分割し、1 つの部分は GPU でトレーニングされ、もう 1 つの部分は CPU でトレーニングされます。データ分散のロングテール効果を利用することで、CPU 計算の割合をできるだけ小さくし、GPU 計算の割合をできるだけ大きくすることができます。ただし、バッチサイズが大きくなると、すべてのミニバッチデータが CPU または GPU にヒットすることが難しくなり、データが同時に CPU または GPU にヒットした場合、この方法は扱いにくくなります。さらに、DDR 帯域幅と HBM はデータの大きさが 1 つ異なるため、入力データの 10% が CPU 上でトレーニングされたとしても、システム全体の速度は少なくとも半分低下します。さらに、CPU と GPU は中間結果を送信する必要がありますが、これにもかなりの通信オーバーヘッドが発生し、トレーニング速度がさらに遅くなります。したがって、研究者はこれらのパフォーマンス欠陥を回避するために非同期更新などの方法を設計しましたが、非同期方法はトレーニング結果に不確実性を引き起こす可能性があり、実際にはアルゴリズム エンジニアの第一選択ではありません。

ソフトウェア キャッシュ: すべてのトレーニングが GPU 上で実行され、埋め込みテーブルが CPU と GPU で構成される異種空間に存在することを確認します。ソフトウェア キャッシュ方式では、必要なパーツを GPU にスワップします。この方法では、埋め込みテーブルの増大するニーズを満たすためにストレージ リソースを安価に拡張できます。さらに、計算に CPU を使用する場合と比較して、この方法によるトレーニング プロセス全体は GPU 上で完全に完了し、HBM 帯域幅の利点を最大限に活用します。ただし、キャッシュ クエリとデータの移動により、さらなるパフォーマンスの低下が発生します。

テーブルを埋め込むための優れたソフトウェア キャッシュ ソリューションがすでにいくつかありますが、それらは多くの場合、fbgemm などのカスタマイズされた EmbeddingBags カーネル実装を使用するか、サードパーティのディープ ラーニング フレームワークを使用します。 Colossal-AI は、ネイティブ PyTorch に基づいてカーネル レベルの変更を行いません。すぐに使えるソフトウェア Cache EmbeddingBags 実装のセットを提供します。また、DLRM トレーニングをさらに最適化します。キャッシュのオーバーヘッドをさらに削減するためのパイプライン化。

メモリ階層

Colossal-AI組み込みテーブル ソフトウェア Cache

Colossal-AI は、ソフトウェア キャッシュを実装し、ユーザーが独自のモデルで使用できるようにそれを nn.Module にカプセル化します。 DLRM の埋め込みテーブルは通常、複数の Embedding で構成される EmbeddingBags であり、CPU メモリに常駐します。メモリ空間のこの部分は CPU ウェイトと呼ばれます。 EmbeddingBags のデータのごく一部は GPU メモリに保存されており、これにはトレーニングに使用されるデータが含まれます。メモリ空間のこの部分は CUDA Cached Weight と呼ばれます。 DLRM トレーニング中、最初に、この反復のミニバッチに入力されたデータに対応する埋め込みテーブルの行を決定する必要があります。一部の行が GPU にない場合は、CPU ウェイトから CUDA に転送する必要があります。キャッシュされた重み。 GPU に十分なスペースがない場合は、LFU アルゴリズムを使用して、キャッシュへのアクセス頻度の履歴に基づいて最も使用されていないデータを削除します。

キャッシュの取得を実現するには、いくつかの補助データ構造が必要です。cached_idx_map は、CPU 重みと CUDA の行番号間の対応を格納する 1 次元配列です。キャッシュされた重み、および GPU がアクセスされる頻度に関する情報の対応する行。 CUDA キャッシュ ウェイト サイズと CPU ウェイト サイズの比率は、cache_ratio という名前で、デフォルトは 1.0% です。

キャッシュは、各反復の前に実行され、具体的には 3 つのステップで CUDA 重みのデータを調整します。

ステップ 1: CPU インデックス: CPU ウェイトにキャッシュする必要がある行番号を取得します

これは必要ですto input mini -batch の input_ids と cached_idx_map の交差部分を取得して、CPU から GPU に移動する必要がある CPU Weight の行番号を見つけます。

#Step2: GPU インデックス: 使用頻度に応じて CUDA 重み付けでエビクトできる行を見つける

This require us 頻度に応じて、cache_idx_mapとinput_idsの差分セット以降の部分に対して、低い方から高い方へtop-k(最大k個を取る)演算を行います。

ステップ 3: データ転送:

CUDA キャッシュされた重みの対応する行を CPU 重みに移動してから、移動します。 CPU Weight の対応する行を CUDA Weight にします。

データ送信モジュールは、CUDA キャッシュ ウェイトと CPU ウェイト間のデータの双方向送信を担当します。非効率的なラインごとの送信とは異なり、最初にキャッシュを使用し、次に集中送信を使用して PCI-e 帯域幅の使用率を向上させます。メモリ内に分散された埋め込み行は、ソース デバイスのローカル メモリ内の連続したデータ ブロックに集められ、ブロックは CPU と GPU の間で転送され、ターゲット メモリ内の対応する場所に分散されます。データをブロック単位で移動すると、PCI-e 帯域幅の使用率が向上し、マージおよびスキャッター操作には CPU と GPU のオンチップ メモリ アクセスのみが含まれるため、オーバーヘッドはそれほど大きくありません。

Colossal-AI は、サイズ制限のあるバッファを使用して CPU と GPU の間でデータを転送します。最悪の場合、すべての入力 ID がキャッシュに届かず、大量の要素を転送する必要があります。バッファーがメモリを過剰に消費しないようにするために、バッファー サイズは厳密に制限されています。転送データがバッファより大きい場合、複数回に分けて転送が完了します。

キャッシュされた EmbeddingBag ワークフロー

ソフトウェア キャッシュ パフォーマンス分析

キャッシュ ステップ 1 とステップ 2 の上記の操作はどちらもメモリ アクセスを大量に消費します。したがって、GPU HBM の帯域幅を活用するために、これらは GPU 上で実行され、深層学習フレームワークによってカプセル化された API を使用して実装されます。それにもかかわらず、キャッシュ操作のオーバーヘッドは、テーブルを埋め込むための GPU でのトレーニング操作と比較して特に顕著です。

たとえば、合計 199 秒のトレーニング タスクでは、キャッシュ操作のオーバーヘッドは 99 秒で、ほぼ 50% を占めます。合計計算時間のうち。分析後、キャッシュの主なオーバーヘッドは主にステップ 1 とステップ 2 によって引き起こされます。下図の基本位置は、この時点のキャッシュ オーバーヘッドの時間内訳を示しており、キャッシュ ステップ 1 と 2 の赤とオレンジの段階がキャッシュ オーバーヘッド全体の 70% を占めています。

#キャッシュ操作の時間の内訳

上記の問題の理由は、従来のキャッシュ戦略がやや「近視眼的」であり、現在のミニバッチの状況に応じてキャッシュを調整することしかできないため、ほとんどの時間がクエリ操作に無駄に費やされることです。

キャッシュ パイプライン プリフェッチ

キャッシュのオーバーヘッドを削減するために、Colossal-AI は一連の

「Foresight」を設計しました。 キャッシュメカニズム。 Colossal-AI は、最初のミニバッチに対してのみキャッシュ操作を実行するのではなく、後で使用する複数のミニバッチをプリフェッチし、統合された方法でキャッシュ クエリ操作を実行します。

以下の図に示すように、Colossal-AI はプリフェッチを使用して複数のミニバッチ データをマージして統合キャッシュ操作を行います。また、パイプライン アプローチを使用してデータの読み取りとオーバーヘッドをオーバーラップします。計算。この例では、プリフェッチ ミニバッチ数は 2 です。トレーニングを開始する前に、ミニバッチ 0,1 データがディスクから GPU メモリに読み取られ、キャッシュ操作が開始され、順方向および逆方向の伝播と 2 つのミニバッチのパラメーター更新が実行されます。同時に、ミニバッチ 2 および 3 の初期データ読み取りが実行される可能性があり、このオーバーヘッドが計算と重なる可能性があります。

ベースラインのキャッシュ実行方法と比較した図 [キャッシュ操作の時間の内訳] は、プリフェッチ 8 ミニバッチとベースラインのキャッシュ時間を比較しています。壊す。合計トレーニング時間は 201 秒から 120 秒に短縮され、図に示されているキャッシュ フェーズの操作時間の割合も大幅に減少しました。キャッシュ操作を個別に実行する各ミニバッチと比較して、各部分、特にキャッシュ操作の最初の 2 ステップの時間が短縮されていることがわかります。

要約すると、キャッシュ パイプラインのプリフェッチには 2 つの利点があります。

a. キャッシュ インデックスのオーバーヘッドの軽減

プリフェッチの最も明白な利点は、ステップ 1 とステップ 1 のオーバーヘッドを削減できることです。 Step2 、この 2 ステップの操作がトレーニング プロセス全体の 5% 未満を占めることになります。 [キャッシュ操作の時間内訳] に示すように、8 個のミニバッチ データをプリフェッチすることにより、プリフェッチなしのベースラインと比較して、キャッシュ クエリのオーバーヘッドが大幅に削減されます。

b. CPU-GPU 間のデータ移動帯域幅を増やす

より多くのデータを集中させることで、データ送信の粒度を向上させ、CPU を最大限に活用します。 GPU 送信帯域幅。上記の例では、CUDA->CPU の帯域幅が 860MB/s から 1477 MB/s に増加し、CPU->CUDA の帯域幅が 1257 MB/s から 2415 MB/s に増加し、パフォーマンスの向上がほぼ 2 倍になります。

使いやすい

使用法は Pytorch EmbeddingBag と一致しています。レコメンデーション モデルを構築する場合、初期化に必要なのは次のコード行だけです。埋め込みテーブルの容量 テラバイトレベルの推奨モデルトレーニングを低コストで実現します。

Bashfrom colossalai.nn.parallel.layers.cache_embedding import CachedEmbeddingBag
emb_module = CachedEmbeddingBag(num_embeddings=num_embeddings,embedding_dim=embedding_dim,mode="sum"include_last_offset=True,sparse=True,_weight=torch.randn(num_embeddings, embedding_dim),warmup_ratio=0.7,cache_ratio = 0.01,)

パフォーマンス テスト

NVIDIA A100 GPU (80GB) および AMD EPYC 7543 32 コア プロセッサ (512GB) ハードウェア プラットフォーム、Colossal - AI は、Meta の DLRM モデルをテストターゲットとして使用し、超大規模データセット Cretio 1TB と Meta の dlrm_datasets を使用してテストモデルとしてデータセットを生成します。実験では、すべての埋め込みテーブルを GPU に保存する PyTorch トレーニング速度がベースラインとして使用されます。

Cretio 1TB

Cretio 1TB 埋め込みテーブルには合計 177944275 行があり、埋め込み dim=128 に設定されています。およびその組み込みテーブル メモリ要件 91.10 GB。最もハイエンドの NVIDIA A100 80GB であっても、すべての EmbeddingBags を 1 つの GPU メモリに保存するためのメモリ要件を満たすことはできません。

しかし、Colossal-AI を使用すると、トレーニングは 1 つの GPU で完了し、キャッシュ率 = 0.05 の場合、メモリ消費量はわずか 5.01 GB となり、約 18 分の 1 に削減されます。 Zhang GPU 上でテラバイトレベルの推奨システム モデルのトレーニングを実装します。トレーニング速度に関しては、次の図に示すように、さまざまなバッチ サイズで 100M サンプルをトレーニングする場合の遅延を示しています。緑色の Prefetch1 はプリフェッチなしの遅延、青色の Prefetch8 はプリフェッチを使用した遅延 (プリフェッチ ミニバッチ = 8) で、プリフェッチ パイプラインの最適化が全体のパフォーマンスの向上に重要な役割を果たしていることがわかります。図の各列の暗い部分はキャッシュ オーバーヘッドであり、プリフェッチを使用した後、キャッシュ オーバーヘッドは総トレーニング時間の 15% 以内に制御されます。

#マルチ GPU のスケーラビリティ

グローバル バッチ サイズとして 8192 を使用する, 8 GPU カード上の EmbeddingBags としてテーブルワイズ シャーディングを使用して、DLRM を並列トレーニングし、100M サンプルをトレーニングします。このとき、プリフェッチサイズを4、ColossalAI-mem-cr0.05はキャッシュ率=0.05、ColossalAI-mem-cr0.5=0.5に設定します。以下の図は、さまざまな GPU シナリオのトレーニング レイテンシーを示しています。 1 GPU を使用する場合は PyTorch OOM をトレーニングできないことを除けば、PyTorch と Colossal-AI のトレーニング時間は他のケースでは同様です。 4 および 8 GPU を使用しても、パフォーマンスが大幅に向上しないことがわかります。これは、1. 結果の同期に多大な通信オーバーヘッドが必要となるためです。 2. テーブルごとのシャーディングにより、シャーディング負荷の不均衡が発生します。また、複数の GPU を使用して埋め込みテーブル トレーニングのスケーラビリティを拡張するのはあまり良くないことも示しています。

#次の図は、ビデオ メモリの使用量を示しています。ビデオ メモリの使用量は、カードによって異なります。ここでは、ビデオ メモリの最大値を示しています。

GPU を 1 つだけ使用する場合、Colossal-AI のソフトウェア キャッシュ メソッドのみをトレーニングでき、複数のカードが並行して占有するメモリも大幅に数分の 1 に削減されます。

#Meta Research の合成データセット dlrm_datasets は、業界の埋め込みテーブルのトレーニング アクセス動作を模倣しているため、レコメンデーション システムとしてよく使用されます。研究中 関連するソフトウェアおよびハードウェア設計のテスト リファレンス。 5 億行の埋め込みテーブル項目をサブデータセットとして選択し、テスト用に 256GB と 128GB の 2 つの EmbeddingBag を構築します。

ビデオ メモリが不足しているため、PyTorch を 1 枚のカード A100 でトレーニングすることはできません。対照的に、Colossal-AI のソフトウェア キャッシュは GPU メモリ要件を大幅に削減し、256 GB ものサイズの埋め込みテーブルをトレーニングするのに十分であり、さらに TB レベルまで拡張できます。さらに、パイプライン プリフェッチによる高速化効果もあり、プリフェッチ数が 32 の場合、プリフェッチなしの場合に比べて合計時間が 60% 削減され、GPU ストレージの需要も増加しません。

One More Thing

##Colossal、大規模時代の一般的な深層学習システムモデル - AI は、効率的な多次元自動並列処理、異種メモリ管理、大規模な最適化ライブラリ、適応型タスク スケジューリングなど、多数の自社開発の最先端テクノロジーを使用して、AI 大規模モデルのトレーニングと推論の効率的かつ迅速な展開を実現します。 、AI ラージ モデル アプリケーションのコストを削減します。

Colossal-AI 関連ソリューションは、自動運転、クラウド コンピューティング、小売、医薬品、チップ、その他の業界の有名メーカーによって導入され、広く賞賛されています。

Colossal-AI は、オープンソース コミュニティの構築に重点を置き、中国語のチュートリアルを提供し、ユーザー コミュニティとフォーラムを開き、効率的なコミュニケーションとユーザー フィードバックの反復更新を実施し、次のような最先端のテクノロジーを継続的に追加しています。 PaLM、AlphaFold、OPT.アプリケーションとして。

Colossal-AI は、自然なオープンソースであるため、GitHub および Papers With Code のホット リストで何度も世界 1 位にランクされており、多くのスター オープンソース プロジェクトとともに国際的な注目を集めています。何万もの星が輝く、内も外も注目してください！

プロジェクトのオープンソース アドレス: https://github.com/hpcaitech/ColossalAI

以上が必要な埋め込みパラメーターは 1% のみで、ハードウェア コストは 10 分の 1 に削減され、オープン ソース ソリューションは単一の GPU を使用して大規模な推奨モデルをトレーニングします。の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。