当前位置: 首页 > 科技观察

突破万亿参数规模,追求极致效率和性价比:中国团队开源首个异构并行推荐系统训练框架Persia

时间:2023-03-18 09:52:10 科技观察

个性化推荐是互联网行业提升DAU的核心技术手段(每日活跃用户)和收入。随着深度学习的广泛应用,现代推荐系统通过神经网络变相“记住”用户的行为习惯,从而准确预测用户的偏好。移动互联网普及后,用户行为数据呈指数级增长,单位时间内产生和收集的用户行为数据极其庞大,因此需要更大的模型对用户兴趣进行编码。更大的数据规模意味着需要更大的模型容量。模型参数数量从五年前的十亿迅速增长到前段时间Facebook公开的十万亿参数规模。在这样的趋势下,更大规模的训练需求无疑成为下一个需要攻克的里程碑。近日,由两个中国团队联合开发的开源训练框架Persia,通过设计混合架构,成功将模型规模推向新的高度,成功将模型规模推向新的高度——万亿级参数(需要上百个T存储),并且可以兼顾效率和精度。目前,该框架已受邀集成到Pytorch生态系统PytorchLightning中。GitHub:https://github.com/PersiaML/PERSIA论文:https://arxiv.org/abs/2111.05897随着模型参数的数量呈指数级增长,对高性能训练框架的需求也越来越迫切。传统架构面对越来越多的参数也越来越无能为力。传统架构采用CPU的同构并行机制,相应的参数分配采用模型并行。它最大的优点是方便机器的水平扩展,以支持相对更大的机型,所以在很多公司仍然被广泛使用(虽然每个公司的命名方式不同,本文统称为mio架构).当推荐模型从传统的Logistic回归升级到基于DeepLearning的模型,参数数量急剧增加时,传统方案变得捉襟见肘——效率低下,精度难以平衡。后面的进化版本通过引入GPU来承担深度网络部分的计算(本文称为mio+架构),因为仍然采用同构设计思想(只是将CPU机器换成CPU带GPU的机器),虽然可以取得了一定的效率提升,部分缓解了效率和准确率之间的矛盾。然而,当需要处理不同规模的网络结构时,大量昂贵的GPU资源往往处于闲置状态,导致性价比严重下降。为了解决模型规模不断扩大带来的这两个问题,Persia的核心设计思路是:采用异构架构设计来解决GPU资源利用问题。当CPU和GPU的比例绑定时,任何框架都很难同时保证任何模型结构下的资源利用率。因此,Persia设计了灵活的异构架构,根据模型需求分配资源,在保证效率的前提下充分利用资源,大幅提升性价比。在兼顾效率和准确性的同时,采用了同步和异步混合训练模式。在传统的解决方案中,要么使用纯同步训练,要么使用纯异步训练。当模型越来越大,机器数量越来越多时,同步训练会导致机器互相等待,很容易严重损害训练效率。异步训练方式虽然避免了机器之间的相互等待,显着提高了训练效率,但是随着机器数量的增加,模型的准确率(Accuracy)会明显下降。针对超大模型场景下的此类挑战,Persia设计了同步异步Hybrid训练架构,集两者之长,扬长避短。并且从理论和实践两个维度都验证了Persia可以同时达到同步训练的准确性和异步训练的效率。下面是Persia的几个特点:原生支持PyTorch生态:由于PyTorch极大地降低了研究人员定义模型的门槛,因此它在整个深度学习领域中的占比往往会越来越大,这与现有的推荐有所不同。对于训练框架(如XDL、PaddlePaddle等),Persia决定将其基于PyTorch生态系统。用户模型定义等操作可以直接用PyTorch实现,因此即使是研究领域最新最前沿的模型(如Transformer等)也可以直接调用,实现最大的灵活性和易用性。高性能:在Criteo标准数据集上,与其他流行的开源推荐模型训练框架相比,Persia在同等资源条件下可以实现两倍以上的性能提升。Persia支持CPU-GPU异构训练,支持GPU-GPU直接通信,显着降低训练成本。可扩展性:Persia在高达100万亿的模型参数训练规模上保持着较高的训练效率。同时在大部分场景下可以接近线性加速(投入n倍的资源量,训练效率提升近n倍)。工业级场景大规模验证:Persia为Kubernetes实现定制算子,支持云原生部署。实现了各种容错机制,经过两年多线上生产环境的稳定运行验证。Persia经过数亿DAU的多个核心业务场景测试,取得了显着的性能和业务指标提升。安全易排错:Persia采用Rust语言实现,注重内存安全、速度和并发,在编译时消除了大量内存安全问题。原生提供大量打点监控,与Grafana完美结合,可自定义各种告警条件。同时,基于tracing实现了模块式、分级式的日志输出,在实际场景中更容易排查问题。灵活的特征处理:支持交叉特征等多种常见的特征处理方式,用户可以通过Python脚本定义多种自定义特征处理方式。将灵活性与易用性相结合。线上线下一致性:线下训练和线上训练代码统一,解决了工程师经常需要花费大量时间排查模型线上效果不一致的痛点。Persia设计思路的整体结构在推荐模型中,模型往往由下图几个部分组成:EmbeddingLayer:Embedding层,由userid、itemid等ID特征对应的Embedding组成。每个id对应一个预设大小的向量(称为Embedding)。由于id的数量往往非常多,这些向量往往占据了整个模型体积的99%以上。Non-IDTypeFeatures:图像信息、LDA等真实向量特征。这部分会和id对应的Embedding向量结合,输入到DNN中预测点击率等。DenseNeuralNetwork(以下简称NN):这部分是接受Embedding向量和真实numbervectorfeatures,并输出你想要预测的数量,比如点击率。在这个推荐模型中,EmbeddingLayer的参数往往占模型体积的绝大部分,但是EmbeddingLayer不需要太多的计算量。神经网络的参数量只占模型量的一小部分,却占了大部分的计算量。这正好对应:硬件上的CPU内存大但计算能力低,而GPU显存小但计算能力高。现有的训练框架虽然包含了GPU算力,但是每个GPUworker都需要和大量的PS进行数据和模型的传输,往往会引发通信瓶颈,从而拖累整个效率。因此在Persia系统设计中,将NN放在GPU内存中,通过GPU进行梯度计算。对于NN部分,直接通过GPU与GPU之间高效的setcommunication进行同步,完全不用PS。Embedding放在内存中,由CPU计算。Persia针对PS(EmbeddingPS,EmbeddingWorker,稍后介绍)实现了两层架构设计,可以进一步降低大部分场景下的GPUworker带宽消耗,提升整体训练效率。同步+异步混合训练此外,现有系统往往采用全同步训练或全异步训练。在全同步训练中,所有GPUworker在一批数据上训练和更新模型,全部完成后再输入下一批数据。当模型越来越大,机器数量越来越多时,会导致机器之间的等待和同步时间显着增加,难以在有限的时间内完成训练。在这种情况下,系统的训练过程如下图第一行(FullSync)所示。在完全异步训练中,每个工作人员独立训练和更新PS参数。虽然worker之间不需要互相等待,训练效率高,但是随着机器数量的增加,每个worker上使用的模型差异会变大,导致模型的训练效果明显下降.在这种情况下,系统的训练过程如下图第二行(FullAsync)所示。针对这两种方式存在的问题,Persia设计了Hybrid训练架构,在保证训练效果的同时,可以达到近乎完全异步的训练效率。推荐场景训练中的一个核心观察是Embedding的更新非常稀疏,往往两次更新的交集很小,所以即使Embedding异步更新,对最终的训练结果影响也很小。NN部分的更新则相反,每次都会更新所有参数。如果进行异步训练,会导致训练结果的巨大差异。Persia提出的Hybrid训练方法可以同步训练NN部分,异步训练Embedding部分。最终训练效率接近于纯异步训练的效率,而模型效果与全同步训练保持一致。充分利用两者。在这种情况下,系统的训练过程如下图第三行(NaiveHybrid)所示。除此之外,Persia还重叠了可以并行执行的通信和计算操作,以进一步提高系统效率。最终系统的训练过程如下图第四行(Persia)所示:TheoreticalGuarantees与现有系统不同,Persia为Hybrid算法的设计提供了严格的理论保证。对于ExpectationofLoss的优化问题(比如推荐场景中最常见的每个样本对应一个loss的场景):其中f(w)表示整个数据集上的平均loss,ξ表示一个样本,w表示模型参数,F(w;ξ)表示样本ξ上的损失。模型训练的目标是最小化整个数据集的平均损失。使用PersiaHybrid训练方法,可以证明模型的收敛速度为:其中σ为数据集的方差,T为迭代次数,τ为GPUworker数,α为碰撞概率ID类特征。前两项是全同步训练的收敛速度,最后一项是Hybrid训练引入的误差。在推荐场景下,由于Embedding的更新很稀疏,碰撞概率α远小于1,因此Hybrid的收敛速度与全同步训练的收敛速度相差无几,但由于减少了同步开销,大大提高了每一步的训练执行效率。具体的理论证明可以参考Persia的论文[1]。其他优化都是基于算法创新,以发挥极致性能。Persia提供了很多实现层的优化。例如,所有PS服务通信都使用针对训练场景优化的零拷贝PersiaRPC系统。在训练场景(特点是payload非常大,包括大量的Embedding、gradients等张量数据),性能远超传统RPC框架(如gRPC、bRPC);GPU之间的通信采用相同的开源八卦训练加速框架,显着提升了GPU之间的集体通信性能,并且可以通过梯度压缩等算法进一步降低通信开销;通过专门设计的数据结构(PersiaEmbeddingArrayList)存储嵌入到PS中,大大提高了PS效率和模型访问效率。这包括在运行过程中不需要动态申请新的内存,更好地利用CPU的缓存机制。支持嵌入驱逐逻辑。模型保存和读取过程简化为直接Dump/Load过程,连续内存;引入PersiaEmbeddingWorker组件进行EmbeddingSumPooling等操作,对原始数据进行处理后发送给GPU,大大降低GPU带宽占用;原始数据处理采用PersiaCompactBatch格式,自带ID去重、数据压缩表征等特性,与一般表示方式相比,数据量减少1/4,提高了系统数据处理效率。验证对比测试使用Alibaba-Ad、Avazu-Ad、Criteo-Ad等开源数据集,整体训练效率提升8倍以上:Persia可支持高达100万亿模型训练,并保持当模型尺寸变大时保持相同的水平。训练效率:当资源量膨胀时,Persia可以接近线性扩展(投入n倍的资源量,训练效率会提高近n倍):Persiausecases使用Persia很简单,主要分为训练部署,模型定义和自定义数据集部分。分布式部署:通过Persiaoperator在Kubernetes集群上一键部署PERSIA任务模型定义:直接使用PyTorch自定义数据集:自定义预处理逻辑,通过Persia提供的Python工具包将结果转化为PersiaCompactBatch是一个完整的例子等等场景,欢迎参考PersiaTutorial文档(https://persiaml-tutorials.pages.dev/)。Persia模型在线推理Persia训练的模型的Embedding部分可以通过在线部署EmbeddingPS和EmbeddingWorker直接对外提供服务。NN部分是原生PyTorch模型,PersiaTutorial中提供了一个通过TorchServe进行推理的简单示例。用户还可以使用各种原生PyTorch工具,例如转换为TensorRT模型,进一步提高推理性能。