一块RTX3050搞定DLRM训练!仅需1%Embedding参数,硬件
推荐模型已经成为互联网公司深度学习应用最重要的技术场景,如视频推荐,购物搜索,广告推送等流量变现服务,大大提升了用户体验和商业价值。
可是,海量的用户和业务数据,频繁的迭代更新需求以及高昂的培训成本给DLRM培训带来了严峻的挑战。
在DLRM,在完成下游计算之前,需要查找嵌入式表。
在DLRM,嵌入式表通常贡献了99%以上的内存需求,但只贡献了1%的计算量。
借助GPU的片上高速内存和强大的计算能力,GPU已经成为DLRM培训的主流硬件。
可是,伴随着推荐系统的深入研究,不断增加的嵌入式表大小和有限的GPU内存形成了显著的矛盾如何利用GPU高效训练超大型DLRM模型,突破GPU内存墙的限制,成为DLRM领域亟待解决的关键问题
庞—AI已经成功使用异构策略,在相同硬件上提高NLP模型训练的参数容量数百倍于之前最近,spool—AI成功将其扩展到推荐系统,通过软件缓存的方式,将嵌入式表格动态存储在CPU和GPU内存中
在软件缓存的设计基础上,庞—AI还增加了流水线预取,通过观察未来要输入的训练数据,降低软件缓存检索和数据移动的成本。
同时通过同步更新的方式在GPU上训练整个DLRM模型,并结合目前广泛使用的混合并行训练方法,可以扩展到多个GPU上。
实验表明,庞氏人工智能只需要在GPU中保留1%的嵌入参数,仍然可以保持优秀的端到端训练速度。
与其他PyTorch方案相比,对显存的需求降低了一个数量级,单块显卡就可以训练出TB级推荐模型。
成本是显而易见的,比如训练占用91GB嵌入包的DLRM只需要5GB显存,训练硬件的成本从两个大约20万元的A100降低到RTX 3050等入门级显卡的十分之一,只需要2000元左右。
开放源地址:
基于现有嵌入式表格扩展技术
嵌入式表将离散的整数特征映射为连续的浮点特征向量下图显示了DLRM嵌入式表的培训过程
首先为嵌入表中的每个特征搜索嵌入表的对应行,然后通过常规运算,如最大,均值,和运算,将特征向量转化为特征向量,再传递给后续的密集神经网络。
可以看出,DLRM的嵌入式表训练过程主要是不规则的内存访问操作,因此受到硬件内存访问速度的严重限制。
而工业DLRM的嵌入式表可能会达到几百GB甚至TB级别,远远超过单个GPU高达几十GB的内存容量。
有许多方法可以突破单个GPU的内存墙来增加DLRM的嵌入式表大小。
GPU模型并行性:
嵌入式表被分割分布在多个GPU的内存中,训练时通过GPU之间的互联网络同步中间结果。
这种方法的缺点首先是嵌入式表分段的负载不均匀,扩展性问题难以解决。
其次,增加GPU的前期硬件成本较高,DLRM训练时GPU的计算能力没有得到充分利用,只利用了其HBM带宽优势,导致GPU利用率较低。
部分CPU训练:
嵌入式表分为两部分,一部分在GPU上训练,一部分在CPU上训练。
利用数据分布的长尾效应,可以让CPU计算比尽可能小,GPU计算比尽可能大但伴随着批量的增加,很难让mini—batch的所有数据都命中CPU或GPU,如果同时命中CPU或GPU,这种方法很难处理
此外,由于DDR带宽与HBM相差一个数据量级,所以即使在CPU上训练10%的输入数据,整个系统的速度也会下降至少一半。
此外,CPU和GPU需要传输中间结果,这也有很大的通信开销,进一步拖慢了训练速度。
因此,研究人员设计了异步更新方法来避免这些性能缺陷,但异步方法会造成训练结果的不确定性,在实践中并不是算法工程师的首选。
软件缓存:
保证训练全部在GPU上进行,嵌入式表存在于CPU和GPU组成的异构空间中每次都是通过软件缓存把需要的部分换成GPU
该方法可以低成本地扩展存储资源,满足日益增长的嵌入式表格需求。
而且与CPU相比,该方法的整个训练过程完全在GPU上完成,充分利用了HBM的带宽优势但是,缓存的查询和数据移动会带来额外的性能损失
目前有一些优秀的嵌入式表的软件缓存方案,但往往是通过自定义的EmbeddingBags内核来实现,比如fbgemm,或者借助第三方深度学习框架来实现。
庞—AI在原生PyTorch的基础上,不做任何内核级的改动,提供了一套开箱即用的软件Cache EmbeddingBags实现,进一步优化了DLRM训练过程,并提出预取流水线,进一步降低缓存开销。
内存层次无损AI嵌入式表格软件缓存
庞—AI实现了一个软件缓存,并将其封装为nn模块,供用户在自己的模型中使用
DLRM的嵌入式表,通常由多个嵌入式包组成,驻留在CPU内存中。
这部分内存空间称为CPU权重但是,嵌入包的少量数据存储在GPU内存中,其中包括用于训练的数据
这部分内存空间称为CUDA缓存权重。
在DLRM训练期间,首先需要确定表中嵌入的行,这些行对应于在该迭代中输入到小批量中的数据如果有些行不在GPU中,需要从CPU权重转移到CUDA缓存权重
如果GPU空间不够,它会根据访问缓存的历史频率,使用LFU算法剔除最少使用的数据。
为了实现缓存检索,需要一些辅助的数据结构:cached_idx_map是一个一维数组,存储CPU权重中的行号和CUDA缓存权重的行号的对应关系,以及对应行在GPU中被访问的频率信息。
dacachedweight大小与CPU权重大小的比值命名为cache_ratio,默认值为1.0%。
每次迭代前运行缓存来调整CUDA中的数据权重,具体分三步。
第一步:CPU索引
检索需要缓存的CPU权重中的行号。
它需要对输入小批的input_ids和cached_idx_map求交集,找到CPU权重中需要从CPU移到GPU的行号。
第二步:GPU索引
根据使用频率找到CUDA权重中可以被驱逐的行。
这就要求我们按照频率从低到高,对cache_idx_map和input_ids的差集之后的部分进行top—k。
第三步:数据处理:
将CUDA缓存权重中的对应行移到CPU权重中,然后将CPU权重中的对应行移到CUDA权重中。
传输模块负责CUDA缓存权重和CPU权重之间的数据双向传输。
与低效的逐行传输不同,它采用先缓存再集中传输的方法,提高PCI—E的带宽利用率。
分散内存中的嵌入行在源设备的本地内存中集中成连续的数据块,然后在CPU和GPU之间传输,分散到目标内存中相应的位置分块移动数据可以提高PCI—E的带宽利用率,合并和分散操作只涉及CPU和GPU的片内内存访问,所以开销不是很高
庞—AI使用有限大小的缓冲区在CPU和GPU之间传输数据。
在最坏的情况下,所有的输入id都没有命中缓存,因此需要传输大量的元素为了防止缓冲区占用过多内存,缓冲区大小受到严格限制如果传输的数据大于缓冲区,传输将分多次完成
缓存嵌入包工作流软件的缓存性能分析
高速缓存步骤1和步骤2的上述操作是存储器访问密集型的。
因此,为了利用GPU的HBM的带宽,它们运行在GPU上,由深度学习框架封装的API来实现但是,相对于GPU上嵌入式表的训练操作,缓存操作的开销尤为突出
例如,在一个总共199秒的训练任务中,缓存操作的开销为99秒,占总计算时间的近50%。
经过分析,缓存的主要开销主要是由Step1和Step2造成的下图中的基准位置显示了此时的缓存开销时间偏差缓存步骤1和2的红色和橙色阶段占总缓存开销的70%
高速缓存操作的时间偏差
产生上述问题的原因是传统的缓存策略短视,只能根据当前的小批量情况调整缓存,因此大部分时间浪费在查询操作上。
缓存预取
为了降低缓存的开销,庞—AI设计了一种前瞻性的缓存机制spool—AI不是只缓存以前的小批量,而是预取几个以后要用的小批量,统一进行缓存查询操作
如下图所示,庞—AI利用预取合并多个小批量数据进行统一缓存操作,同时利用流水线技术重叠数据读取和计算的开销。
在本例中,预取小批量的数量是2在训练之前,将mini—batch 0,1数据从磁盘读入GPU内存,然后启动缓存操作,再对这两个mini—batch进行正反向传播和参数更新
同时可以用来读取Mini—Batch 2,3的数据,这部分开销可以和计算重叠。
与基线缓存的执行模式相比,该图比较了八个小批量预取和基线的缓存时间偏差。
总训练时间从201秒下降到120秒,图中显示的缓存阶段的操作时间比例也明显下降可以看出,与每个小批量独立执行缓存操作相比,每个部分的时间都有所减少,尤其是缓存操作的前两步
综上所述,缓存流水线预取带来了两个好处。
1.稀释缓存索引开销
预取最明显的优点是减少了Step1和Step2的开销,使得这两步操作占整个训练过程的比例不到5%如所示,通过预取8个小批量数据,与没有预取的基线相比,缓存查询的开销显著降低
2.增加CPU—GPU数据移动带宽
通过集中更多的数据,提高数据传输的粒度,从而充分利用CPU—GPU的传输带宽上面的例子,CUDA—gt,CPU带宽从860MB/s提升到1477 MB/s,CPU—gt,CUDA的带宽从1257 MB/s提升到2415 MB/s,性能增益几乎翻倍
使用方便
这与Pytorch嵌入包的用法是一致的在构建推荐模型时,只需要初始化以下几行代码,可以大幅增加嵌入式表的容量,低成本实现TB级超大推荐模型训练
fromcolosalai . nn . parallel . layers . cache _ embedding import cached embedding bag _ module = cached embedding bag,warmup _ ratio = 0.7,cache _ ratio = 0.01。)
在NVIDIA A100 GPU和AMD EPYC 7543 32核处理器(512 GB)硬件平台上,庞—AI以Meta DLRM模型为测试对象,使用超大数据集Cretio 1TB和Meta dlrm_datasets作为测试模型。
实验中,以存储所有嵌入式表的GPU上的PyTorch训练速度为基线。
Cretio 1TB
Cretio 1TB嵌入式表共有177,944,275行设置嵌入dim=128需要91.10 GB的嵌入式表内存
如果想在单个GPU内存中存储所有的EmbeddingBags,即使是最高端的NVIDIA A100 80GB也无法满足其内存需求。
可是,庞—AI仍然用于在单个GPU上完成训练当缓存比=0.05时,内存消耗仅为5.01 GB,直接降低了18倍左右还可以进一步扩展,在单个GPU上实现TB级推荐系统模型训练
在训练速度上,如下图所示,显示了不同批量下训练100M样本的延迟。
绿色预取1不使用预取,蓝色预取8是预取的延迟可见预取流水线优化对整体性能提升的重要作用
图中每列的黑色部分是缓存开销预取后,缓存开销控制在总训练时间的15%以内
多GPU可扩展性
8192被用作全局批量大小,表式分片被用作在8个GPU卡上嵌入标签的并行模式,以训练DLRM和100M样本。
此时预取大小设置为4,ColossalAI—mem—cr0.05的缓存比是0.05,庞AI—mem—cr0.5是0.5。
下图显示了不同GPU情况下的训练延迟除了PyTorch OOM不能在1个GPU上训练之外,PyTorch和庞—AI的训练时间差不多
可以看出,使用4个和8个GPU不会带来显著的性能提升,因为:
因此,同步需要巨大的通信开销。
表方向分片将导致不平衡的拆分负载也说明用多个GPU来扩展嵌入表的训练扩展性不是很好
下图显示了显存的使用,不同的卡上使用的显存是不同的此处显示了最大视频内存值
在只使用一个GPU的情况下,只需要训练庞—AI的软件缓存方法,多张卡并行占用的内存就可以显著降低数倍。
Meta Research的合成数据集dlrm_datasets模仿了业内嵌入式表的训练访问行为,因此在研究中经常作为推荐系统相关软硬件设计的测试参考。
其中,选择表项中嵌入的5亿行作为子数据集,构建256GB和128GB两个嵌入包进行测试。
PyTorch由于显存不足无法在单卡A100上训练相比之下,庞—AI的软件缓存将大幅降低GPU内存需求,足以训练高达256GB的嵌入式表格,并可以进一步扩展到TB级别
而且流水线预取也能体现加速效果预取次数为32时,总时间比不预取时减少了60%,且不增加GPU的存储需求
还有一点
面向大模型时代的通用深度学习系统panol—AI,通过高效多维自动并行,异构内存管理,大规模优化库,自适应任务调度等多项自主研发的领先技术,实现AI大模型训练和推理的高效快速部署,降低了AI大模型的应用成本
庞氏人工智能相关解决方案已成功应用于自动驾驶,云计算,零售,医药,芯片等行业的知名厂商,并获得好评。
庞—AI注重开源社区建设,提供中文教程,开放用户社区和论坛,高效交流和迭代更新用户反馈,不断增加PaLM,AlphaFold,OPT等前沿应用。
自自然开源以来,庞—AI多次在GitHub和有代码热榜的论文上排名世界第一,与众多上万颗星的明星开源项目一起引起了国内外的关注!
项目的开放源地址:
参考链接:
这篇文章引用自:
。声明:以上内容为本网站转自其它媒体,相关信息仅为传递更多企业信息之目的,不代表本网观点,亦不代表本网站赞同其观点或证实其内容的真实性。投资有风险,需谨慎。
猜你喜欢