腾讯北大合作的稀疏大模型训练加速方案HET入选国际顶会VLDB

近日，腾讯TEG数据平台部机器学习团队与北京大学-腾讯协同创新实验室，合作研发了全新的稀疏大模型训练加速解决方案HET，其研究成果《HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework》已被国际顶会VLDB 2022录用。HET提出了一种新颖的基于Embedding缓存的训练方法，能够显著降低稀疏大模型分布式训练时通信开销，提升模型训练整体效率。

HET目前已正式开源：https://github.com/PKU-DAIR/Hetu

稀疏大模型日益多见，通信瓶颈或成训练效率“致命”问题

图1 深度学习模型规模发展情况

稀疏大模型，是目前重要的深度学习模型类型之一，广泛应用在搜索广告推荐、图表示学习等场景。近年来，随着数据规模的逐步增长，工业界稀疏大模型的规模日益庞大，参数量可以达到万亿规模。如图1所示，Facebook今年提出的ZionEX[详见注解1]系统所支持的推荐模型（DLRM）大小已经超过10万亿规模，远远超过了Google之前发布的1.6万亿参数的Switch Transformer[详见注解2]模型。

稀疏模型的参数，即Embedding参数，可以达到总模型参数量的99%以上。相比于其他模型，这类模型有着更低的计算密度和更大的模型规模，这也对分布式深度学习系统带来了严峻的挑战。近年来，如何提升稀疏大模型的训练效率逐渐成为了学术界和工业界都在关注的热点问题。

对于万亿规模的模型，仅模型参数就需要3.7TB的内存空间。由于稀疏大模型中的稀疏参数规模极大，因此工业界目前普遍采用基于参数服务器（Parameter Server）的解决方案，将Embedding均匀地切分到不同服务器上。在训练过程中，计算节点采用稀疏通信的形式，动态地从参数服务器上拉取所需的Embedding向量，完成当前轮次的计算后，再将Embedding的梯度提交回参数服务器。尽管这种方式可以灵活地扩展模型规模，但是也面临着严重的通信瓶颈。以主流深度学习框架TensorFlow为例，在实际数据测试中，通信时间甚至会占到总训练时间的80%以上。目前大多数改进方向是在参数服务器的工程实现上进行优化，例如充分挖掘硬件性能来提高整个系统的吞吐率。然而并没有从根本上解决稀疏参数通信量大的问题，通信仍然是系统的核心痛点。因此需要一种从源头上解决通信问题的方案。

HET：基于Embedding缓存的稀疏大模型训练系统

核心思路