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

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

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

图 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 缓存的稀疏大模型训练系统

核心思路

图 2 三个常用公开数据集上的 Embedding 访问频率分布情况

根据来自在业务场景中的观察,高维稀疏大模型的输入数据特征往往具有倾斜分布的特性,具有幂律分布(如图 2 所示),从而导致模型在训练过程中对 Embedding 向量的不均衡访问。以推荐数据集 Criteo 为例,约有 10%的 Embedding 向量引发了整个数据集 90%的 Embedding 访问。训练过程中,这些高频 Embedding 会被频繁地拉取和推送,成为了通信的主要负载。

我们利用这一特性,提出了 Embedding 缓存的思想:如果能够在计算节点利用有限的内存空间缓存这些高频 Embedding,那就有机会避免大量的远程 Embedding 访问,从而缓解通信瓶颈。根据这一思想,我们提出了基于 Embedding 缓存的新一代稀疏大模型训练框架 HET。

技术点 1:支持 Embedding 参数缓存的混合通信架构

图 3 HET 系统架构

针对稀疏大模型参数中同时存在稀疏以及稠密部分的特点,HET 整体上采用参数服务器(Parameter Server,PS)加全局规约(AllReduce)的混合通信架构,以充分发挥两者优势,如图 3 所示。其中 AllReduce 适合于稠密参数的同步,可借助 NCCL 等通信库充分发挥 GPU 间带宽,而参数服务器则天然支持稀疏通信,并且在同步协议上也具有较高的灵活性。同时,我们还在计算节点上设计了 Cache Embedding Table 结构,用于缓存高频访问的 Embedding 参数。

每个计算节点上采用 Cache Embedding Table 可节省大量的通信量,但是也带来了一个新的问题,即对于某个特定的 Embedding 来说,其副本可能同时存在于多个不同的计算节点缓存当中。如果不考虑副本间的一致性,可能会导致模型训练发散,无法收敛。为此,我们进一步提出了一种基于细粒度 Embedding 时钟的有限异步协议,来解决如何在不同的节点间同步这些 Embedding 副本的问题。

技术点 2:基于细粒度 Embedding 时钟的有限异步协议

图 4 HET 中的 Cache Embedding Table 结构

一般来说,Embedding 参数采用表的方式进行组织以支持稀疏访问。为了衡量 Embedding 副本间的一致性,我们对于每个 Embedding 向量,在常规的 key-value 数据结构基础之上,引入了一个重要的 Lamport 时钟,用来记录 Embedding 向量的状态。在模型训练过程中,通过比较 Embedding 的时钟,就可以知道该副本的延迟或超前程度。

图 5 HET 中的 Cache 读写操作

对于 Embedding 缓存表,我们既允许读取较为陈旧的 Embedding,也允许延迟写回缓存上的梯度更新。为了在充分发挥缓存加速效果的同时保证模型的训练质量,我们限制了每个 Embedding 副本和全局 Embedding 间的时钟相差不超过一个预先设定的阈值。在这种情况下,Embedding 的每个副本都不会过于超前或落后于它的其他副本。

从全局视角来看,整个模型的稀疏和稠密部分分别采用不同的同步模式,稠密参数采用全同步协议进行通信,稀疏参数采用基于细粒度 Embedding 时钟的有限异步协议进行通信。经过理论分析,我们进一步证明了,这种基于细粒度 Embedding 时钟的有限异步协议可以保证和全同步协议相似的收敛性。(详情见论文链接)

实验结果

我们将 HET 和基于传统参数服务器架构的 TensorFlow 以及同样是参数服务器加全局规约的混合通信架构的 Parallax[详见注解 3]进行了对比,选取的数据集和模型包括:推荐模型 Wide&Deep(WDL)、DeepFM(DFM)、Deep&Cross(DCN)和数据集 Criteo,拥有三千多万稀疏特征,当 Embedding 维度扩大到 4K 时,模型参数可以达到万亿级别;以及图学习模型 GraphSAGE 和数据集 Reddit、Amazon、ogbn-mag(OGB 也是目前最权威的图学习基准数据集之一,Open Graph Benchmark)。

端到端对比

图 6 收敛效果对比

图 7 端到端收敛速度对比

结合图 6 和图 7,我们可以看出,在时钟相差阈值上界设置为 100 时,相比于 TensorFlow 和 Parallax,HET 可以实现 6.37-20.68 倍的加速,并且不会对模型收敛性造成显著影响。对于 HET 本身,细粒度 Embedding 缓存带来了 4.36-5.14 倍的加速,最多可以减少 88%的稀疏参数通信。

缓存效果对比:

图 8 在不同 Cache 空间大小情况下 Cache 失效率情况

从图 8 可以看出,只需要很少的 Cache 空间,比如 15%的总参数量大小,就可以实现几乎 97%的缓存命中率,即 97%的 Embedding 访问都可以通过本地缓存访问,而无需通信。另外我们也注意到,不同的 Cache 实现策略在效果上也稍有不同,LFU 可以捕获长期访问倾向性,从而比 LRU 失效率更低。

可扩展性:

图 9 在不同参数规模情况下的收敛效果

我们将模型扩展到 32 节点,Embedding 维度设置到 4096,此时总参数量已经达到了万亿规模,从图 9 可以看到 HET 执行时间仍然显著优于其他基线方案,从而说明了 HET 的有效性。

腾讯 TEG 数据平台部机器学习团队:

该团队致力于研发腾讯分布式机器学习平台 Angel,解决高维模型和稀疏数据训练问题。Angel 诞生于腾讯大数据生态中,通过融合大数据、传统机器学习和深度学习生态,建立起一个端到端机器学习平台,功能涵盖传统机器学习、图挖掘、图学习、深度学习和隐私计算等。在腾讯公司内部, Angel 已经广泛应用于广告推荐、金融风控、用户画像和短视频推荐等业务。除了服务于公司内部业务外,Angel 于 2017 年对外开源,是国内第一个 LF AI 基金会顶级项目。

针对模型规模扩大后对性能和扩展性带来的挑战。Angel 平台工程师联合北京大学-腾讯协同创新实验室,推出了稀疏大模型训练框架 HEAP,并在广告推荐全链路中,依托腾讯内部业务需求对各种规模模型的训练提速增效,开展了包含基于 Embedding 缓存的新一代稀疏大模型训练,基于层次化参数服务器的万亿 Embedding 模型训练、多 GPU 分布式训练性能优化等在内的多项前瞻研究,并落地到广告精排、粗排、预排序和召回等多个业务模型的训练中,在腾讯各个业务线条累计获取 GMV 提升约 4%。本次发表的研究成果 HET 也是在该框架下进行的新的探索。

北京大学-腾讯协同创新实验室:

北京大学-腾讯协同创新实验室成立于 2017 年,主要在人工智能、大数据等领域展开前沿探索和人才培养,打造国际领先的校企合作科研平台和产学研用基地。

实验室通过合作研究,在理论和技术创新、系统研发和产业应用方面取得重要成果和进展,已在国际顶级学术会议和期刊发表学术论文 20 余篇,除合作研发 Angel 外,实验室还自主开发了多个开源系统。

今年 8 月,实验室已经宣布,将自研深度学习框架“河图”融入 Angel 生态,北京大学与腾讯团队将联合共建 Angel4.0——新一代分布式深度学习平台,面向拥有海量训练数据、超大模型参数的深度学习训练场景,为产业界带来新的大规模深度学习破局之策。

参考资料:

[1] Mudigere D, Hao Y, Huang J, et al. High-performance, distributed training of large-scale deep learning recommendation models[J]. arXiv preprint arXiv:2104.05158, 2021.

[2] Fedus W, Zoph B, Shazeer N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity[J]. arXiv preprint arXiv:2101.03961, 2021.

[3] Kim S, Yu G I, Park H, et al. Parallax: Sparsity-aware data parallel training of deep neural networks[C]//Proceedings of the Fourteenth EuroSys Conference 2019. 2019: 1-15.