DreamShard：面向推荐系统的可泛化嵌入表放置
机器学习
2022-10-06 v1
摘要
我们研究分布式推荐系统的嵌入表放置问题，旨在将表划分并放置到多个硬件设备（如GPU）上以平衡计算与通信成本。尽管先前工作已探索基于学习的方法用于计算图的设备放置，嵌入表放置仍是一个具有挑战性的问题，原因在于：1）嵌入表的操作融合，以及2）对具有不同表数和/或设备数的未知放置任务的泛化性要求。为此，我们提出DreamShard，一种用于嵌入表放置的强化学习（RL）方法。DreamShard通过以下方式实现操作融合推理与泛化性：1）一个成本网络直接预测融合操作的成本；2）一个在估计的马尔可夫决策过程（MDP）上高效训练的策略网络，无需真实GPU执行，其中状态和奖励由成本网络估计。借助sum和max表示归约，这两个网络可直接泛化到任意具有不同表数和/或设备数的未知任务而无需微调。大量实验表明，DreamShard在大规模合成表和我们的生产表上大幅优于现有人类专家和基于RNN的策略，相较最强基线最高加速19%。代码见 https://github.com/daochenzha/dreamshard
引用
@article{arxiv.2210.02023,
title = {DreamShard: Generalizable Embedding Table Placement for Recommender Systems},
author = {Daochen Zha and Louis Feng and Qiaoyu Tan and Zirui Liu and Kwei-Herng Lai and Bhargav Bhushanam and Yuandong Tian and Arun Kejariwal and Xia Hu},
journal= {arXiv preprint arXiv:2210.02023},
year = {2022}
}
备注
Accepted by NeurIPS 2022