[论文解读] 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 该系统是否能在包含多 TB 嵌入表的真实工业工作负载中实现显著的加速？

主要发现

• DreamShard 在大规模合成表与真实生产工作负载上，相较于最强基线模型，最高实现19%的训练加速。
• 在一个包含近1,000张表、内存达多 TB 的超大规模工业推荐模型上，DreamShard 将训练吞吐量提升了45.3%，比次优基线高出30个百分点。
• 成本网络在生产数据集上的测试均方误差（MSE）为0.303，其中表维度、池化因子与分布特征对准确预测影响最大。
• 消融实验表明，移除任意单一特征均会导致测试MSE上升，证实所有特征在成本预测中的重要性。
• DreamShard 在多个测试用例中，包括平衡与非平衡场景，始终优于随机放置与最佳人工专家策略。
• 即使在极少数情况下 DreamShard 的放置结果略逊于最优启发式方法，其整体性能仍保持竞争力或更优，展现出强大的鲁棒性与泛化能力。