深度推荐模型（DLRMs）已经成为互联网公司深度学习最重要的技术场景，如视频推荐、购物搜索、广告推送等流量变现服务，大大提升了业务的用户体验和商业价值。然而海量的用户和业务数据、频繁的迭代更新需求、高昂的训练成本对DLRM训练提出了严峻挑战。在DLRM中，需要先在嵌入表（EmbeddingBags）中进行一次查找，然后完成下游的计算。嵌入式表通常占DLRM中99%以上的内存需求，但只占计算量的1%。借助GPU片上高速内存（HighBandwidthMemory）和强大的计算能力，GPU已经成为DLRM训练的主流硬件。然而，随着推荐系统研究的深入，嵌入式表的规模越来越大与GPU内存有限之间存在着显着的矛盾。如何利用GPU高效训练超大型DLRM模型，同时突破GPU内存墙的限制，成为DLRM领域亟待解决的关键问题。Colossal-AI此前曾成功利用异构策略将在同一硬件上训练的NLP模型参数容量提升数百倍，近期又成功扩展到推荐系统，采用软件缓存（Cache）的方式动态存储并嵌入到CPU和GPU表面的内存中。在软件缓存设计的基础上，Colossal-AI还增加了流水线预取，通过观察未来输入的训练数据，减少软件缓存检索和数据移动的开销。同时，它以同步更新的方式在GPU上训练整个DLRM模型，结合广泛使用的混合并行训练方法，可以扩展到多个GPU。实验表明，Colossal-AI只需在GPU中保留1%的嵌入参数，仍能保持出色的端到端训练速度。与其他PyTorch解决方案相比，内存需求降低一个数量级，单显卡即可训练TB级推荐模型。成本优势显着。例如，训练DLRM只需要5GB的显存，占用91GB的EmbeddingBag空间。训练硬件的成本从两块20万元左右的A100到RTX3050等入门级显卡只要2000元左右，降低了十倍。开源地址：https://github.com/hpcaitech/ColossalAI现有的embeddingtable扩展技术embeddingtable将离散的整型特征映射为连续的浮点型特征向量。下图展示了DLRM中embeddingtable的训练过程。首先在embedding表中为每个特征查找EmbeddingTable对应的行，然后通过归约操作，如max、mean、sum等操作，成为一个特征向量，传递给后续的dense神经网络。可以看出，DLRM的嵌入表训练过程主要是一次不规则的内存访问操作，因此受到硬件访问速度的严重限制。工业级DLRM的内嵌表可能达到数百GB甚至TB级别，远超单个GPU几十GB的最大显存容量。有很多方法可以通过突破单个GPU的内存墙来增加DLRM的嵌入表的大小。我们以下图所示的GPU集群内存层次图为例，分析一下几种常见方案的优缺点。GPU模型并行：embeddingtable被划分分布在多个GPU的内存中，训练时通过GPU之间的互连网络同步中间结果。这种方式的缺点是内嵌表分担负载不均匀，扩展性问题难以解决。其次，增加GPU的前期硬件成本高，DLRM训练时没有充分利用GPU的算力，仅利用其HBM带宽优势，导致GPU使用率低。CPUPartialTraining：将embeddingtable分成两部分，一部分在GPU上训练，一部分在CPU上训练。利用数据分布的长尾效应，我们可以让CPU计算的比例尽可能小，GPU计算的比例尽可能大。但是随着batchsize的增大，很难让所有的mini-batch数据都打到CPU或者GPU上，同时打到CPU或者GPU上的方法很难处理。另外，由于DDR的带宽与HBM相差一个数量级，即使有10%的输入数据在CPU上训练，整个系统的速度也会下降至少一半。另外，CPU和GPU需要传输中间结果，也有较大的通信开销，进一步拖慢了训练速度。因此，研究人员设计了异步更新等方法来避免这些性能缺陷，但异步方法会导致训练结果的不确定性，在实践中并不是算法工程师的首选方案。SoftwareCache：确保所有训练都在GPU上进行，嵌入表存储在CPU和GPU组成的异构空间中。每次通过软件Cache的方式将需要的部分替换到GPU中。通过这种方式，可以廉价地扩展存储资源以满足嵌入式表不断增长的需求。而且，相比使用CPU计算，这种方式整个训练过程完全在GPU上完成，充分利用了HBM的带宽优势。但是Cache查询和数据移动会带来额外的性能损失。已经有一些优秀的嵌入式表软件缓存实现，但它们通常使用自定义的EmbeddingBagsKernel实现，例如fbgemm，或者使用第三方深度学习框架。另一方面，Colossal-AI在原生PyTorch的基础上，并没有对Kernel层面做任何改动。它提供了一套开箱即用的软件CacheEmbeddingBags实现，并进一步优化了DLRM训练过程，并提出了prefetchpipeline以进一步降低Cache开销。MemoryHierarchyColossal-AI的嵌入式表软件CacheColossal-AI实现了一个软件Cache，封装成nn.Module，供用户在自己的模型中使用。DLRM的embedding表一般是由多个Embedding组成的EmbeddingBags，驻留在CPU内存中。这部分内存空间被命名为CPUWeight。而EmbeddingBags的一小部分数据存储在GPU内存中，其中包括将用于训练的数据。这部分内存空间被命名为CUDACachedWeight。在DLRM训练时，首先需要确定本次迭代输入mini-batch的数据对应的embedding表的行。如果某些行不在GPU中，则需要将它们从CPUWeight转移到CUDACachedWeight。如果GPU中没有足够的空间，它会使用LFU算法根据访问缓存的历史频率驱逐最少使用的数据。为了实现对Cache的检索，需要一些辅助数据结构：cached_idx_map是一个一维数组，存储了CPUWeight中行号与CUDACachedWeight中行号的对应关系，以及GPU上正在访问的相应行的频率信息。CUDACachedWeight与CPUWeight的比值命名为cache_ratio，默认为1.0%。Cache在每次迭代前运行，调整CUDA中的数据Weight，具体分为三步。Step1：CPUindex：检索CPUWeight中需要缓存的行号。它需要将inputmini-batch的input_ids和cached_idx_map相交，找到CPUWeight中需要从CPU移动到GPU的行号。Step2：GPU索引：根据使用频率在CUDAWeight中找到可以被逐出的行。这就需要我们对cache_idx_map和input_ids的差集之后的部分按照频率从低到高的顺序（取最大的k数）进行top-k操作。Step3：数据处理：将CUDACachedWeight中对应行移动到CPUWeight，再将CPUWeight中对应行移动到CUDAWeight。数据传输模块负责CUDACachedWeight和CPUWeight之间的双向数据传输。不同于低效的逐行传输，它采用先缓冲后集中传输的方式，提高PCI-e的带宽利用率。分散在内存中的嵌入式行在源设备的本地内存中聚集成连续的数据块，这些块在CPU和GPU之间传输并分散到目标内存中的相应位置。以块为单位移动数据提高了PCI-e带宽利用率，合并和分散操作仅涉及CPU和GPU的片上内存访问，因此开销并不大。Colossal-AI使用大小受限的缓冲区在CPU和GPU之间传输数据。在最坏的情况下，所有输入id都会错过缓存，需要传输大量元素。为防止缓冲区占用过多内存，缓冲区大小受到严格限制。如果传输的数据大于缓冲区，则会分成多次完成传输。CachedEmbeddingBagWorkflow软件Cache性能分析上面CacheStep1和Step2的操作都是内存密集型的。所以为了利用GPU的HBM的带宽，他们运行在GPU上，使用深度学习框架封装的API来实现。尽管如此，与GPU上嵌入表的训练操作相比，缓存操作的开销尤为突出。例如，在总计199秒的训练任务中，Cache操作的开销为99秒，占总计算时间的近50%。经过分析，Cache的主要开销主要是Step1和Step2造成的。下图中的base位置展示了此时Cache开销的时间细目。Cache的第1步和第2步的红色和橙色阶段占总Cache开销的70%。出现上述问题的原因是传统的缓存策略有些“短视”，只能根据当前的mini-batch情况调整缓存，所以大部分时间都浪费在了查询操作上。CachePipelinePrefetch为了减少Cache的开销，Colossal-AI设计了一套“前瞻性”的Cache机制。Colossal-AI不是只对前面的mini-batch进行Cache操作，而是预取几个mini-batch后面会用到，统一进行Cache查询操作。如下图所示，Colossal-AI通过prefetching将多个mini-batch数据合并进行统一的缓存操作，通过pipeline重叠数据读取和计算的开销。例子中prefetchedmini-batch的个数为2，在开始训练之前，先从磁盘读取mini-batch0,1数据到GPU内存，然后开始Cache操作，然后进行前向和反向传播和两个小批量的参数更新。同时可以和mini-batch2、3的起始数据一起读取，这部分开销可以和计算重叠。与baselineCache执行方式对比，图【CacheOperationTimeDecomposition】比较了prefetch8mini-batch和baselineCache时间分解。总训练时间从201秒下降到120秒，图中Cache阶段的运行时间占比也明显下降。可以看到，相比于每个mini-batch的独立Cache操作，减少了各部分的时间，尤其是Cache的前两步。总结一下，Cachepipelineprefetching带来了两个好处。A。DilutedCacheindexoverhead预取最明显的好处是减少了Step1和Step2的开销，使得这两个步骤的操作占总训练过程的不到5%。如【缓存操作的时间分解】所示，通过预取8个mini-batch数据，与没有预取的baseline相比，Cache查询的开销明显降低。b.通过集中更多的数据和提高数据传输的粒度来增加CPU-GPU的数据移动带宽，从而充分利用CPU-GPU的传输带宽。对于上面的例子，CUDA->CPU带宽从860MB/s增加到1477MB/s，CPU->CUDA带宽从1257MB/s增加到2415MB/s，性能增益几乎翻倍。方便的使用与PytorchEmbeddingBag的用法一致。在构建推荐模型时，只需以下几行代码初始化，即可大幅提升embedding表的容量，低成本实现TB级超大型推荐模型训练。Bashfromcolossalai.nn.parallel.layers.cache_embeddingimportCachedEmbeddingBagemb_module=CachedEmbeddingBag(num_embeddings=num_embeddings,embedding_dim=embedding_dim,mode="sum"include_last_offset=True,sparse=True,_weight=torch.randn(num_embeddings,embedding_dim),0.7比率,cache_ratio=0.01,)NVIDIAA100GPU（80GB）和AMDEPYC754332核处理器（512GB）硬件平台上的性能测试，Colossal-AI以Meta的DLRM模型为测试目标，使用超大数据集Cretio1TB和Metadlrm_datasets生成数据集作为测试模型。在实验中，以在GPU上存储所有嵌入表的PyTorch训练速度作为基线。Cretio1TBCretio1TBembeddingtable共有177944275行，设置embeddingdim=128，嵌入表内存要求为91.10GB。如果要将所有的EmbeddingBags存储在单个GPU内存中，即使是最高端的NvidiaA10080GB也无法满足其内存要求。然而，Colossal-AI仍然在单个GPU上进行训练。当cacheratio=0.05时，内存消耗仅为5.01GB，直接减少了约18倍。可进一步扩展，实现单GPU上TB级推荐系统模型训练。在训练速度方面，如下图所示，展示了在不同batchsize下训练100M样本的延迟。绿色Prefetch1不使用预取，蓝色Prefetch8使用预取（prefetchmini-batch=8）延时。由此可见，预取流水线优化对整体性能的提升起到了重要作用。图中每一列的深色部分是缓存开销。使用预取后，缓存开销控制在总训练时间的15%以内。对于多GPU可扩展性，使用8192作为全局批大小，使用table-wisesharding作为8个GPU卡上的EmbeddingBags并行训练DLRM，并训练100M样本。此时设置Prefetchsize为4，ColossalAI-mem-cr0.05为cacheratio=0.05，ColossalAI-mem-cr0.5=0.5。下图显示了不同GPU的训练延迟。除了PyTorchOOM不能用1个GPU训练外，PyTorch和Colossal-AI的训练时间在其他情况下是相似的。可以观察到，使用4个和8个GPU并没有带来显着的性能提升，因为，1.同步结果需要巨大的通信开销。2.Table-wisesharding会导致负载不均衡。也说明使用多个GPU来扩展embeddingtable训练的可扩展性不是很好。下图显示了显存使用情况。不同卡上的视频内存使用情况不同。此处显示最大显存值。只用一块GPU时，只能训练Colossal-AI的软件缓存方式，多卡并行占用内存也明显减少数倍。MetaResearch的合成数据集dlrm_datasets模仿了业界嵌入式表的训练访问行为，因此在研究中常被用作推荐系统相关软硬件设计的测试参考。选取5亿行嵌入表项作为子数据集，构建256GB和128GB两个EmbeddingBags进行测试。由于显存不足，无法在单卡A100上训练PyTorch。相比之下，Colossal-AI的软件缓存将显着降低GPU内存需求，足以训练高达256GB的嵌入表，并可以进一步扩展到TB级别。此外，流水线预取还可以显示加速效果。当预取次数为32次时，总时间比没有预取减少60%，对GPU存储的需求没有增加。OneMoreThing是面向大规模模型时代的通用深度学习系统Colossal-AI。通过高效多维自动并行、异构内存管理、大规模优化库、自适应任务调度等多项自主研发的领先技术，模型训练和推理降低AI大模型应用成本。Colossal-AI相关解决方案已在自动驾驶、云计算、零售、医药、芯片等领域的知名厂商成功落地，获得广泛好评。Colossal-AI专注于开源社区建设，提供中文教程，开通用户社区和论坛，针对用户反馈进行高效沟通和迭代更新，不断添加PaLM、AlphaFold、OPT等前沿应用。自开源以来，Colossal-AI多次荣登GitHub和PapersWithCode热榜全球第一，与众多拥有数万star的明星开源项目一起受到国内外关注！项目开源地址：https://github.com/hpcaitech/ColossalAI
