[論文レビュー] DreamShard: Generalizable Embedding Table Placement for Recommender Systems

Daochen Zha, Louis Feng+7 more|arXiv (Cornell University)|Oct 5, 2022

Recommender Systems and Techniques被引用数 7

ひとこと要約

DreamShardは、分散推薦システムにおける一般化可能な埋め込みテーブル配置のための強化学習アプローチを提案する。コストネットワークを用いて統合演算のコストを予測し、推定されたMDP上で訓練された方策ネットワークにより、テーブル数やデバイス数が異なる未確認の設定に対してもゼロショット一般化を達成する。合成および本番ワークロードにおいて、最先端のベースラインを上回る最大19%のトレーニングスピードアップを達成する。

ABSTRACT

We study embedding table placement for distributed recommender systems, which aims to partition and place the tables on multiple hardware devices (e.g., GPUs) to balance the computation and communication costs. Although prior work has explored learning-based approaches for the device placement of computational graphs, embedding table placement remains to be a challenging problem because of 1) the operation fusion of embedding tables, and 2) the generalizability requirement on unseen placement tasks with different numbers of tables and/or devices. To this end, we present DreamShard, a reinforcement learning (RL) approach for embedding table placement. DreamShard achieves the reasoning of operation fusion and generalizability with 1) a cost network to directly predict the costs of the fused operation, and 2) a policy network that is efficiently trained on an estimated Markov decision process (MDP) without real GPU execution, where the states and the rewards are estimated with the cost network. Equipped with sum and max representation reductions, the two networks can directly generalize to any unseen tasks with different numbers of tables and/or devices without fine-tuning. Extensive experiments show that DreamShard substantially outperforms the existing human expert and RNN-based strategies with up to 19% speedup over the strongest baseline on large-scale synthetic tables and our production tables. The code is available at https://github.com/daochenzha/dreamshard

研究の動機と目的

• 分散推薦システムにおける埋め込みテーブル配置の課題に対処すること。ここでは、演算の統合と、未確認の構成に対する一般化能力が重要である。
• 従来の強化学習ベースのデバイス配置手法の限界を克服すること。これらの手法は、新しいテーブル数やデバイス数に一般化できず、統合演算のコストについての考察も行わない。
• 微調整なしに、任意の未確認の配置タスクに直接一般化できるシステムを設計すること。スケーラビリティを高めるために表現の低次元化を活用する。
• 知的なテーブル統合と配置により、計算コストと通信コストのバランスを取ることで、トレーニング遅延を低減すること。
• 合成環境および本番環境の両方で、人間が設計したヒューリスティクスおよびRNNベースのベースラインを上回ること。

提案手法

• テーブルの次元、ハッシュサイズ、分布などの特徴に基づいて、統合された埋め込み演算の計算コストと通信コストを予測するためのコストネットワークを訓練する。
• 状態と報酬がコストネットワークの予測から得られる推定されたマルコフ決定過程（MDP）上で強化学習により方策ネットワークを訓練する。これにより、実際のGPU実行の必要がなくなる。
• テーブル表現に対しては和集合を、デバイス表現に対しては最大値集合を適用することで、異なる数のテーブルやデバイスを持つ新しいタスクに対してもゼロショット一般化を実現する。
• 方策ネットワークは、デバイス間の最大遅延を最小化するようにテーブルを配置することを学習し、計算コストと通信コストの両方のバランスを取る。
• 本手法は合成データ上でエンドツーエンドに訓練され、オープンソースの合成ベンチマークおよび実際の産業界の本番テーブルの両方で評価される。
• 本手法は、128 GPUクラスタを用いて、ほぼ1,000の埋め込みテーブルを有する超大規模産業用推薦モデルを対象として評価されている。

実験結果

リサーチクエスチョン

• RQ1 強化学習ベースのアプローチは、微調整なしに、テーブル数やデバイス数が異なる未確認の埋め込みテーブル配置タスクに一般化できるか？
• RQ2 学習されたコストネットワークは、多様なテーブル組み合わせにおける統合された埋め込み演算の性能をどの程度正確に予測できるか？
• RQ3 DreamShardは、計算コストと通信コストのバランスを取る点で、人間が設計したヒューリスティクスおよびRNNベースのベースラインをどの程度上回るか？
• RQ4 テーブル特徴（例：次元、サイズ、分布）が、コスト予測の正確性および最終的な配置パフォーマンスに与える影響はどの程度か？
• RQ5 本システムは、マルチテラバイトの埋め込みテーブルを有する実世界の産業ワークロードで顕著なスピードアップを達成できるか？

主な発見

• DreamShardは、大規模な合成テーブルおよび本番ワークロードにおいて、最も強力なベースラインを上回る最大19%のトレーニングスピードアップを達成する。
• ほぼ1,000のテーブルとマルチテラバイトのメモリを有する超大規模産業用推薦モデルにおいて、DreamShardはトレーニングスループットを45.3%向上させ、次に良いベースラインを30ポイントも上回る。
• 本番データセットでは、コストネットワークがテストMSE 0.303を達成し、テーブル次元、プーリング要因、分布特徴が正確な予測に最も影響を与える。
• アブレーションスタディの結果、個々の特徴を削除するとテストMSEが上昇するため、すべての特徴がコスト予測において重要であることが確認された。
• バランスの取れた状況から不均衡な状況まで、複数のテストケースにおいて、DreamShardはランダム配置および最良の人間エキスパート戦略を一貫して上回る。
• 稀にDreamShardの配置が最良のヒューリスティクスよりわずかに劣るケースがあるが、全体として競争力のある、あるいはそれ以上のパフォーマンスを維持しており、ロバストネスと一般化能力が裏付けられる。