[論文レビュー] Mixed-Precision Embedding Using a Cache

Jie Yang, Jianyu Huang+3 more|arXiv (Cornell University)|Oct 21, 2020

Topic Modeling参考文献 32被引用数 15

ひとこと要約

本稿では、頻繁または最近アクセスされた埋め込み行を格納する小規模で高精度のキャッシュを用いた混合精度埋め込み学習手法を提案する。多数の埋め込み行は低精度（INT8 や INT4）で学習されるが、キャッシュは頻繁にアクセスされる行の全精度勾配とパラメータを保持する。このアプローチにより、GPU からホストへの転送を最小限に抑え、効率的なキャッシュ置換ポリシーを活用することで、工業的スケールのモデルで最大 7× のメモリ削減と、完全な精度保持、16% の学習速度向上を達成した。

ABSTRACT

In recommendation systems, practitioners observed that increase in the number of embedding tables and their sizes often leads to significant improvement in model performances. Given this and the business importance of these models to major internet companies, embedding tables for personalization tasks have grown to terabyte scale and continue to grow at a significant rate. Meanwhile, these large-scale models are often trained with GPUs where high-performance memory is a scarce resource, thus motivating numerous work on embedding table compression during training. We propose a novel change to embedding tables using a cache memory architecture, where the majority of rows in an embedding is trained in low precision, and the most frequently or recently accessed rows cached and trained in full precision. The proposed architectural change works in conjunction with standard precision reduction and computer arithmetic techniques such as quantization and stochastic rounding. For an open source deep learning recommendation model (DLRM) running with Criteo-Kaggle dataset, we achieve 3x memory reduction with INT8 precision embedding tables and full-precision cache whose size are 5% of the embedding tables, while maintaining accuracy. For an industrial scale model and dataset, we achieve even higher >7x memory reduction with INT4 precision and cache size 1% of embedding tables, while maintaining accuracy, and 16% end-to-end training speedup by reducing GPU-to-host data transfers.

研究の動機と目的

• 深層学習推薦システムにおける大規模埋め込みテーブルの増大するメモリ使用量に対処すること。
• 大規模学習ワークロードにおいてモデルの精度を損なわずに GPU メモリ使用量を削減すること。
• 埋め込みテーブルに対して低精度学習と高精度キャッシュを組み合わせる有効性を調査すること。
• キャッシュサイズ、置換ポリシー、精度レベルがモデルの精度と学習効率に与える影響を評価すること。
• よりスマートなメモリアクセスパターンにより、PCIe バンド幅の使用量を減らし、学習速度を向上させること。

提案手法

• 本手法は、埋め込みテーブル全体の 1–5% の小規模で高精度のキャッシュを用い、頻繁または最近アクセスされた埋め込み行を格納する。
• 多数の埋め込み行は低精度（INT8 や INT4）で学習されるが、キャッシュはアクティブな行の全精度勾配とパラメータを保持する。
• LRU や LFU などのキャッシュ置換ポリシーを用い、アクセス頻度や再利用の新しさに基づいてどの行を高精度に昇格させるかを管理する。
• 低精度学習中に確率的丸めを適用することで、系統的な丸めバイアスを軽減し、精度を維持する。
• 既存のディープラーニングフレームワークと統合され、大規模なテーブルでは統一メモリ（UVM）が使用され、キャッシュにより PCIe トラフィックが削減される。
• 本手法は、オープンソース（DLRM と Criteo-Kaggle）および工業的スケールのモデルの両方で評価され、メモリ削減、精度、学習速度が測定された。

実験結果

リサーチクエスチョン

• RQ1 多数の埋め込み行が低精度（INT8/INT4）で学習される中で、小規模な高精度キャッシュがモデル精度を効果的に維持できるか？
• RQ2 異なるキャッシュ置換ポリシー（LRU、LFU）が、低精度埋め込み学習におけるモデル精度とキャッシュヒット率に与える影響は何か？
• RQ3 メモリ削減と精度回復のバランスを考慮した場合、埋め込みテーブルに対して最適なキャッシュサイズはどの程度か？
• RQ4 高精度キャッシュの有無にかかわらず、確率的丸めと四捨五入の違いが低精度埋め込み学習に与える影響は何か？
• RQ5 このキャッシュ機構により、GPU からホストへのデータ転送をどの程度削減でき、エンドツーエンドの学習速度向上に寄与できるか？

主な発見

• 工業的スケールのモデルにおいて、INT4 精度と 1% の高精度キャッシュを用いることで、最大 7× のメモリ削減を達成し、精度に影響を与えない。
• 5% の高精度キャッシュを用いることで、オープンソースの DLRM モデルで INT8 精度で 3× のメモリ削減を達成し、完全な精度を維持した。
• 確率的丸めは、高精度キャッシュと組み合わせた場合に特に顕著に、四捨五入よりも優れた性能を示し、系統的な丸めバイアスを低減した。
• キャッシュヒット率は学習速度と強く相関しており、高いヒット率では GPU メモリ帯域幅が飽和するが、低いヒット率でも統一メモリ（UVM）アクセスを上回る性能を示した。
• 工業的モデルでは、効果的なキャッシュにより GPU からホストへのデータ転送を削減することで、エンドツーエンドの学習速度が 16% 向上した。
• キャッシュサイズを拡大することで得られる精度向上の効果は、次第に逓減する傾向にあり、小規模なキャッシュ（1–5%）でほぼ最適な性能が達成可能であることが示された。