深度推荐模型（DLRMs）已经成为深度学习在互联网公司应用的最重要技术场景，如视频推荐、购物搜索、广告推送等流量变现业务，极大改善了用户体验和业务商业价值。但海量的用户和业务数据，频繁地迭代更新需求，以及高昂的训练成本，都对 DLRM 训练提出了严峻挑战。
在 DLRM 中，需要先在嵌入表（EmbeddingBags）中进行查表（lookup），再完成下游计算。嵌入表常常贡献 DLRM 中 99% 以上的内存需求，却只贡献 1% 的计算量。借助于 GPU 片上高速内存（High Bandwidth Memory）和强大算力的帮助，GPU 成为 DLRM 训练的主流硬件。但是，随着推荐系统研究的深入，日益增长的嵌入表大小和有限的 GPU 显存形成显著矛盾。如何让利用 GPU 高效训练超大 DLRM 模型，同时突破 GPU 内存墙的限制，已成为 DLRM 领域亟待解决的关键问题。
Colossal-AI此前已成功利用异构策略将相同硬件上训练NLP模型的参数容量提升上百倍，近期成功将其拓展到推荐系统中，通过软件缓存（Cache）方法在 CPU 和 GPU 内存中动态存储嵌入表。基于软件 Cache 设计，Colossal-AI 还添加流水预取，通过观察未来即将输入的训练数据，降低软件 Cache 检索和数据移动开销。同时，它以同步更新方式在 GPU 上训练整个 DLRM 模型，结合广泛使用的混合并行训练方法，可以扩展到多个 GPU。实验表明，Colossal-AI 仅需在 GPU 中保留 1% 的嵌入参数，仍能保持优秀的端到端训练速度。相比 PyTorch 其他方案，显存需求降低一个数量级，单块显卡即可训练 TB 级推荐模型。成本优势显著，例如仅需 5GB 显存即可训练占据 91GB 空间 Embedding Bag 的 DLRM，训练硬件成本从两张约 20 万元的 A100，降低十倍至仅需 2000 元左右的 RTX 3050 等入门级显卡。
开源地址：https://github.com/hpcaitech/ColossalAI
现有的嵌入表扩展技术
嵌入表将离散的整型特征映射成连续的浮点特征向量，下图展示了 DLRM 中的嵌入表训练过程。首先，在嵌入表中对每个特征查找 Embedding Table 对应的行，然后通过规约操作，比如 max，mean, sum 操作，变成一个特征向量，传递给后续的稠密神经网络。可见，DLRM 的嵌入表训练过程主要是不规则的内存访问操作，因此严重受限于硬件访存速度。
而工业级 DLRM 的嵌入表可能达到数百 GB 甚至 TB 级别，远超单 GPU 最高数十 GB 的显存容量。突破单 GPU 的内存墙来增大 DLRM 的嵌入表规模有很多方法。根据下图展示的 GPU 集群的内存层级图为例，让我们来分析几种常见方案的优劣。
GPU 模型并行：将嵌入表切分后分布在多个 GPU 的内存中，训练中通过 GPU 之间互联网络同步中间结果。这种方式的缺点首先是嵌入表切分负载并不均匀，扩展性问题难以解决。其次，增加 GPU 的前期硬件成本大，而且 DLRM 训练时 GPU 的计算能力并没有被充分利用，而是仅仅利用了它的 HBM 带宽优势，导致 GPU 使用率不高。
CPU 部分训练：将嵌入表分割成两部分，一部分在 GPU 上训练，另一部分在 CPU 上训练。通过利用数据分布的长尾效应，我们可以让 CPU 计算比例尽可能少，让 GPU 计算比例尽可能大。但是，随着 batch size 增大，让 mini-batch 的数据全部命中 CPU 或者 GPU 很困难，如果同时命中 CPU 或 GPU 这种方法很难处理。另外，由于 DDR 带宽和 HBM 相差一个数据量级，即使 10% 的输入数据在 CPU 上训练，整个系统也会有至少一半速度下降。此外，CPU 和 GPU 需要传输中间结果，这也有不小的通信开销，进一步拖慢训练速度。因此，研究人员设计了异步更新等方式来避免这些性能缺陷，但是异步方式会造成训练结果的不确定性，在实践中并不是算法工程师的首选方案。
软件 Cache：保证训练全部在 GPU 上进行，嵌入表存在 CPU 和 GPU 组成的异构空间中，每次通过软件 Cache 方式，将需要的部分换入 GPU。这种方式可以廉价扩展存储资源，满足嵌入表不断增大的需求。而且，相比使用 CPU 来计算，这种方式的整个训练过程完全在 GPU 上完成，充分利用 HBM 带宽优势。但 Cache 的查询、数据移动会带来额外性能损耗。
目前已经有一些针对嵌入表优秀的软件 Cache 方案实现，但是它们往往使用定制的 EmbeddingBags Kernel 实现，比如 fbgemm，或者借助第三方深度学习框架。而 Colossal-AI 在原生 PyTorch 基础上不做任何 Kernel 层次改动，提供了一套开箱用的软件 Cache EmbeddingBags 实现，还进一步针对 DLRM 训练流程进行优化，提出预取流水来进一步降低 Cache 开销。
Memory Hierarchy
Colossal-AI 的嵌入表软件 Cache
Colossal-AI 实现了一个软件 Cache 并封装成 nn.Module 提供给用户在自己模型中使用。DLRM 的嵌入表，一般是由多个 Embedding 组成的 EmbeddingBags，驻留在 CPU 内存中。这部分内存空间被命名为 CPU Weight。而 EmbeddingBags 一小部分数据存储在 GPU 内存中，它包括即将被训练用到的数据。这部分内存空间被命名为 CUDA Cached Weight。在 DLRM 训练期间，首先需要确定本次迭代输入 mini-batch 的数据所对应嵌入表的行，如果有的行不在 GPU 中，需要将它们从 CPU Weight 传输到 CUDA Cached Weight 中。如果 GPU 中没有足够的空间，它会使用 LFU 算法，根据访问缓存的历史频率来淘汰被使用最少数据。
为了实现 Cache 的检索，需要一些辅助数据结构帮忙：cached_idx_map 是一维数组，存储 CPU Weight 中行号和 CUDA Cached Weight 的行号对应关系，以及对应行在 GPU 被访问的频率信息。CUDA Cached Weight 大小与 CPU Weight 大小的比值命名为 cache_ratio，默认为 1.0%。
Cache 在每个迭代 forward 之前运行，以调整 CUDA Weight 中的数据，具体来说分三个步骤。
Step1：CPU 索引：检索 CPU Weight 中需要被 Cache 的行号
它需要对输入 mini-batch 的 input_ids 和 cached_idx_map 取交集，找到 CPU Weight 中需要从 CPU 移动到 GPU 的行号。
Step2：GPU 索引：根据使用频率找到 CUDA Weight 中可以被驱逐的行
这需要我们根据频率以从低到高顺序，对 cache_idx_map 和 input_ids 取差集合之后的部分进行 top-k（取最大值 k 个数）操作。
Step3：数据搬运：
将 CUDA Cached Weight 中的对应行移动到 CPU Weight 中，然后将 CPU Weight 中的对应行移动到 CUDA Weight 中。
数据传输模块负责 CUDA Cached Weight 和 CPU Weight 之间的数据双向传输。不同于低效的逐行传输，它采用先缓存再集中传输方式来提升 PCI-e 的带宽利用率。分散在内存中的嵌入行在源设备的本地内存中集中为连续的数据块，然后块在 CPU 和 GPU 之间传输，并分散到目标内存的相应位置。以块为单位移动数据可以提高 PCI-e 带宽利用率，merge 和 scatter 操作只涉及 CPU 和 GPU 的片上内存访问，因此开销并不是很大。
Colossal-AI 用一个尺寸受限的缓冲区来传输 CPU 和 GPU 之间数据。在最坏的情况下，所有输入 id 都未命中缓存 cache，那就需要需要传输大量元素。为了防止缓冲区占用过多内存，缓冲区大小被严格限制。如果传输的数据大于缓冲区，会分为多次完成传输。
Cached EmbeddingBag Workflow
软件 Cache 性能分析
上述 Cache Step1 和 Step2 的操作都是访存密集的。因此为了能利用 GPU 的 HBM 的带宽，它们是在 GPU 上运行的，并使用深度学习框架封装好的 API 来实现。尽管如此，与嵌入表在 GPU 上的训练操作相比，Cache 操作的开销尤为突出。
比如在一次总计 199 秒训练任务中，Cache 操作的开销为 99 秒，占比总计算时间接近 50%。经过分析，Cache 的主要开销主要是 Step1 和 Step2 引起。下图 base 位置展示了此时的 Cache 开销时间分解，Cache 的 step1,2 红色和橙色两阶段占 Cache 总开销的 70%。
Cache 操作的时间分解
而上述问题的原因，是因为传统的 Cache 策略有些“短视”，只能根据当前 mini-batch 情况调整 Cache，因此大部分时间浪费在查询操作上。
Cache 流水预取
为了缩减 Cache 的开销，Colossal-AI 设计了一套 “高瞻远瞩” 的 Cache 机制。与其只对前 mini-batch 进行 Cache 操作，Colossal-AI 预取后续将会被使用的若干 mini-batch，统一进行 Cache 查询操作。
如下图所示，Colossal-AI 使用预取来合并多个 mini-batch 数据统一进行 Cache 操作，同时采用流水线方式来重叠数据读取和计算的开销。例子中预取 mini-batch 数量是 2。在开始训练前，先从磁盘读取 mini-batch 0,1 数据到 GPU 内存，随后开始 Cache 操作，然后执行这两个 mini-batch 的正、反向传播和参数更新。与此同时，可以和对 mini-batch 2，3 的开始数据读取，这部分开销可以和计算重叠。
和 baseline Cache 执行方式相比，图【Cache 操作的时间分解】对比了 prefetch 8 个 mini-batch 和 baseline 的 Cache 时间分解。训练总时间从 201 秒下降到 120 秒，图中所示的 Cache 阶段操作时间占比也显著下降。可以看到和每个 mini-batch 独立进行 Cache 操作相比，各部分时间都减少了，尤其是 Cache 的前两步操作。
总结起来，Cache 流水预取带来两个好处。
a.摊薄 Cache 索引开销
预取最显而易见的好处是减少了 Step1 和 Step2 的开销，使这个两步操作在总的训练过程占比小于 5%。如【Cache 操作的时间分解】所示，通过预取 8 个 mini-batch 数据，和没有预取的 baseline 相比，Cache 查询的开销显著降低。
b.增加 CPU-GPU 数据移动带宽
通过集中更多数据，提升数据传输粒度，从而充分利用 CPU-GPU 传输带宽。对于上面例子，CUDA->CPU 带宽从 860MB/s 提升到 1477 MB/s，CPU->CUDA 带宽从 1257 MB/s 提升到 2415 MB/s，几乎带来了近一倍的性能增益。
便捷使用
和 Pytorch EmbeddingBag 用法一致，在构建推荐模型时，仅需如下数行代码进行初始化，即可大幅提升嵌入表容纳量，低成本实现 TB 级超大推荐模型训练。
性能测试
在 NVIDIA A100 GPU （80GB）和 AMD EPYC 7543 32-Core Processor (512GB)硬件平台上，Colossal-AI 以 Meta 的 DLRM 模型作为测试目标，用超大数据集 Cretio 1TB 和 Meta 的 dlrm_datasets 生成数据集作为测试模型。实验中采用将嵌入表全部存储 GPU 上的 PyTorch 训练速度作为 baseline。
Cretio 1TB
Cretio 1TB嵌入表总共 177944275 行，设置 embedding dim=128，其嵌入表内存需求 91.10 GB。想把 EmbeddingBags 全部存储在单个 GPU 内存中，即使是最高端的英伟达 A100 80GB 也无法满足其内存需求。
但使用 Colossal-AI 仍然在单 GPU 上完成训练，当 cache ratio=0.05，显存消耗仅为 5.01 GB，直接降低约 18 倍，可进一步扩展到在单张 GPU 上实现 TB 级推荐系统模型的训练。在训练速度上，如下图所示，展示了不同 batch size 下训练 100M 个样本的延迟。绿色 Prefetch1 是不使用预取，蓝色 Prefetch8 是使用预取（prefetch mini-batch=8）的延迟，可见预取流水优化对整体性能提升发挥了重要作用。图中每个柱子深色部分为 Cache 开销，使用预取后，Cache 开销控制在训练总时间的 15% 范围内。
多 GPU 扩展性
用 8192 作为全局 batch size，在 8 张 GPU 卡上使用 table-wise sharding 作为 EmbeddingBags 并行方式训练 DLRM，训练 100M samples。此时设置 Prefetch 大小为 4，ColossalAI-mem-cr0.05 是 cache ratio=0.05，ColossalAI-mem-cr0.5=0.5。下图展示了不同 GPU 情况下的训练延迟。除了 1 GPU 时 PyTorch OOM 无法训练之外，其余情况 PyTorch 和 Colossal-AI 训练时间类似。可以观察到使用 4 和 8 GPU 并没有带来明显性能提升，这是因为，1. 同步结果需要通信开销巨大。2. table-wise sharding 会导致切分负载不均衡。也说明使用多 GPU 来扩展 embedding table 训练扩展性并不是很好。
下图展示了显存使用，显存使用在不同卡上并不相同，这里展示最大显存数值。在仅使用一张 GPU 时，只有 Colossal-AI 的软件 Cache 方法可以训练，多卡并行的占用内存也显著减少数倍。
Meta Research 的合成数据集 dlrm_datasets 模仿了工业界嵌入表的训练访问行为，因此常在研究中作为推荐系统相关的软硬件设计的测试参考。选取其中的 5 亿行嵌入表项的作为子数据集，构造 256GB 和 128GB 大小的两个 EmbeddingBags 用于测试。
PyTorch 由于显存内存不足无法在单卡 A100 上训练。作为对比, Colossal-AI 的软件 cache 将显著降低 GPU 内存需求，足以训练大至 256GB 的嵌入表，并可进一步扩展至 TB 级别。而且，流水预取也能体现出加速效果，当预取数为 32 时，相比没有预取总时间下降 60%，而且对 GPU 的存储的需求却没有增大。
One More Thing
面向大模型时代的通用深度学习系统 Colossal-AI，通过多项自研领先技术如高效多维自动并行、异构内存管理、大规模优化库、自适应任务调度等实现高效快速部署 AI 大模型训练和推理，降低 AI 大模型应用成本。
Colossal-AI 相关解决方案已成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用，广受好评。
Colossal-AI 注重开源社区建设，提供中文教程，开放用户社群及论坛，对于用户反馈进行高效交流与迭代更新，不断添加 PaLM、AlphaFold、OPT 等前沿应用。
自然开源以来，Colossal-AI 已经多次在 GitHub 及 Papers With Code 热榜位列世界第一，与众多已有数万 star 的明星开源项目一起受到海内外关注！
项目开源地址：https://github.com/hpcaitech/ColossalAI
