不够大!还不够大?
在NLP领域,从BERT的亿级参数,到OpenAI的1700亿参数,再到Google Switch Transformer的1.6万亿参数,研究人员对参数量增长的渴望从来没有停止过,而我等吃瓜群众对参数量也早已经麻了。
那,100000000000000够不够?不用数了,后面14个零!一百万亿!
最近来自快手和苏黎世联邦理工学院的研究人员提出了一个新的推荐系统Persia,最高支持100万亿级参数的模型训练,比目前最大的FB-ZionEX推荐系统还高出一个数量级。
论文地址:https://arxiv.org/pdf/2111.05897.pdf
开源地址:https://github.com/PersiaML/Persia
为了支持如此庞大规模模型的训练,同时保证训练效率和训练精度,文中提出了一种全新的混合训练算法:通过混合异步、同步机制对embedding层和dense层进行分别训练。
用了这套机制的推荐系统Persia(parallel recommendation training system with hybrid acceleration,即混合加速的并行推荐系统)在一百万亿级参数量的模型上进行了理论和实验验证,证实了这种设计和实现的合理性。
作者表示,使用这套机制,任何人都可以很容易训练一个百万亿级参数量的模型。(只有3060Ti可以吗?)
值得一提的是,除了2016年Alibaba-XDL和2018年Baidu-AIBox以外,其他大规模推荐系统模型都没有开源。
在推荐系统中,深度学习模型目前也成了主流,同样遵循着「参数量大一级压死人」原则。
是不是感觉购物app、短视频app越来越容易猜到了你的想法了?
实际上,现代推荐系统取得的进展很大程度上就是来源于不断增大的模型规模,Google旗下的Youtube在2016年首次将推荐系统的模型规模推向十亿参数,从此模型参数量开始一路狂奔,Facebook(Meta)最新的模型将规模拉到12万亿,比Youtube的参数量高出12000倍!
模型参数量的规模每次上一个新台阶,都能带来性能上的明显改进,所以剑指100万亿也是有实际意义的。
一般来说,推荐系统模型首先需要将不同的ID特征(如用户ID和session ID)映射到一个固定长度的低维向量,而系统中的用户ID、交叉特征数量都特别多,就需要更大规模的模型来捕获特征和映射。但更大规模的embedding layer也需要更大的内存来载入,不得不说大模型太费钱了!
有了embedding后,剩下的工作就简单了,设计后续layer来适配不同的任务。通常只占据整个模型的0.1%,无需大内存,主要是一些计算密集型的工作。
虽然性能好了,但随着模型规模的不断扩大,模型的训练也是越来越难!
而想要训练好一个推荐系统,主要看这5个关键步骤:
由于embedding层和推荐任务是相关联的,所以上述五个任务必须顺序执行,而这种串行机制也导致了模型训练的硬件效率很低。
如果不严格顺序执行这五步,而是采用异步分布式训练(Asynchronous distributed training)方式的话,就能显著提升训练速度。
目前业界对神经网络的异步分布式训练主要有两个观察结论:
在不考虑模型结构的时候,肯定全异步的硬件效率最高,但在实际情况中,全异步优化会降低模型的性能,对于商用推荐模型来说是不能接受的,所以「同步、异步两手抓」才是最终解决方案!
文中提出的混合训练算法(Hybrid training)对模型的embedding层和dense层分别使用不同的优化机制来提升硬件效率。
因为embedding层更大、更稀疏,采用异步训练的方式。
而dense层更稠密,就采用同步训练的方式。
有了思路,下一步就是系统实现了!
Persia系统的设计上主要有两个难点:
对于第一个难题,研究人员为了支持基于深度学习的推荐模型的分布式训练,就直接使用了Tensforflow框架提供的PS范式,同时将embedding和神经网络参数的存储和更新放在一组PS节点(即CPU机器)中,将前向和后向传播的计算放在一组工作节点(即GPU机器)中。
但只使用这种方式还远远不够高效,甚至在实际情况中可能根本没法完成部署。
例如推荐模型太大,GPU的RAM很容易就不够用了,而且将embedding和神经网络模块进行统一视图管理可能还会引入大量不必要的网络流量。
虽然之前的研究也提出过一些PS范式优化、带缓存的colocated PS framework来减少通信开销等等,但这些方法根本无法支撑百万亿级模型的训练。
Persia的研究人员分解出4大模块来解决这个问题:
还没完!上面还提到了第二个难点就是分布式的训练过程,研究人员又来了一个「七步走」解决了这个难题:
最后的训练结果也可以看到,Persia的训练速度相比其他大规模模型来说要快了很多,从而能够支撑百万亿级模型的训练。
参考资料:
https://arxiv.org/pdf/2111.05897.pdf