一块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和有代码热榜的论文上排名世界第一，与众多上万颗星的明星开源项目一起引起了国内外的关注！

项目的开放源地址:

参考链接:

这篇文章引用自: