预训练与搜索：基于预训练神经代价模型的高效嵌入表分片
机器学习
2023-05-04 v1 分布式、并行与集群计算
信息检索
性能
摘要
将大型机器学习模型分片到多个设备以平衡开销，在分布式训练中十分重要。这具有挑战性，因为划分是 NP 难的，且准确高效地估计开销很困难。本工作中，我们探索一种用于高效分片的“预训练，然后搜索”范式。其思想是预训练一个通用的、一次性的神经网络来预测所有可能分片的开销，它充当高效的分片模拟器。基于该预训练代价模型，我们随后执行在线搜索以针对任意特定分片任务确定最佳分片方案。我们在深度学习推荐模型（DLRMs）中实例化该思想，并提出用于嵌入表分片的 NeuroShard。NeuroShard 在增广表上预训练神经代价模型以覆盖各种分片场景。然后它分别通过束搜索和贪心网格搜索确定最佳的列级和表级分片方案。实验表明，NeuroShard 在基准分片数据集上显著且持续地优于最先进水平（state-of-the-art），实现了高达 23.8% 的提升。当部署于具有多万亿字节嵌入表的超大规模生产型 DLRM 时，NeuroShard 在嵌入开销上比最先进水平提升 11.6%，相当于端到端训练吞吐率提升 6.6%。为促进未来对 ML for Systems 中“预训练，然后搜索”范式的研究，我们在 https://github.com/daochenzha/neuroshard 开源了代码。
引用
@article{arxiv.2305.01868,
title = {Pre-train and Search: Efficient Embedding Table Sharding with Pre-trained Neural Cost Models},
author = {Daochen Zha and Louis Feng and Liang Luo and Bhargav Bhushanam and Zirui Liu and Yusuo Hu and Jade Nie and Yuzhen Huang and Yuandong Tian and Arun Kejariwal and Xia Hu},
journal= {arXiv preprint arXiv:2305.01868},
year = {2023}
}
备注
Accepted by MLSys 2023. Code available at https://github.com/daochenzha/neuroshard