Pre-train and Search: Efficient Embedding Table Sharding with Pre-trained Neural Cost Models

要約

タイトル：予め学習されたニューラルコストモデルによる効率的な埋め込みテーブルシャーディング

要約：
– 大規模な機械学習モデルを複数のデバイスに分散してバランスをとるためには、シャーディングが重要です。
– しかし、パーティションはNP困難であり、コストを正確かつ効率的に推定することは難しいため、これは課題となっています。
– 本論文では、効率的なシャーディングのための「pre-train and search」というパラダイムを探求しています。
– このアイデアは、可能なすべてのシャードのコストを予測するためのユニバーサルで一度だけのニューラルネットワークを事前に学習することで、効率的なシャーディングシミュレーターとして機能するものです。
– この事前学習されたコストモデルをベースに、任意のシャーディングタスクに対して最適なシャーディングプランを特定するためにオンライン検索を実行します。
– このアイデアは、Deep Learning Recommendation Models（DLRMs）における埋め込みテーブルシャーディングに関して、NeuroShardという提案手法を示します。
– NeuroShardは、さまざまなシャーディングシナリオをカバーするために拡張されたテーブルに対してニューラルコストモデルを事前に学習し、ビームサーチと貪欲グリッドサーチを使用して最適な列単位およびテーブル単位のシャーディングプランを特定します。
– 実験結果は、NeuroShardがベンチマークシャーディングデータセットにおいて、従来の手法を大幅に上回り、最大23.8%の改善を達成していることを示しています。
– 超大規模プロダクションDLRMにデプロイされた場合、NeuroShardは最先端技術に比べて埋め込みコストで11.6%の改善を達成し、エンドツーエンドトレーニングスループットの改善に繋がります。
– ML for Systemsにおいて、「pre-train and search」パラダイムの将来の研究を容易にするために、本研究はコードをオープンソースで公開します。

要約(オリジナル)

Sharding a large machine learning model across multiple devices to balance the costs is important in distributed training. This is challenging because partitioning is NP-hard, and estimating the costs accurately and efficiently is difficult. In this work, we explore a ‘pre-train, and search’ paradigm for efficient sharding. The idea is to pre-train a universal and once-for-all neural network to predict the costs of all the possible shards, which serves as an efficient sharding simulator. Built upon this pre-trained cost model, we then perform an online search to identify the best sharding plans given any specific sharding task. We instantiate this idea in deep learning recommendation models (DLRMs) and propose NeuroShard for embedding table sharding. NeuroShard pre-trains neural cost models on augmented tables to cover various sharding scenarios. Then it identifies the best column-wise and table-wise sharding plans with beam search and greedy grid search, respectively. Experiments show that NeuroShard significantly and consistently outperforms the state-of-the-art on the benchmark sharding dataset, achieving up to 23.8% improvement. When deployed in an ultra-large production DLRM with multi-terabyte embedding tables, NeuroShard achieves 11.6% improvement in embedding costs over the state-of-the-art, which translates to 6.6% end-to-end training throughput improvement. To facilitate future research of the ‘pre-train, and search’ paradigm in ML for Systems, we open-source our code at https://github.com/daochenzha/neuroshard

arxiv情報

提供元, 利用サービス

arxiv.jp, OpenAI