仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

深度推荐模型(DLRMs)已经成为深度学习在互联网公司应用的最重要技术场景,如视频推荐、购物搜索、广告推送等流量变现业务,极大改善了用户体验和业务商业价值。但海量的用户和业务数据,频繁地迭代更新需求,以及高昂的训练成本,都对 DLRM 训练提出了严峻挑战。

在 DLRM 中,需要先在嵌入表(EmbeddingBags)中进行查表(lookup),再完成下游计算。嵌入表常常贡献 DLRM 中 99% 以上的内存需求,却只贡献 1% 的计算量。借助于 GPU 片上高速内存(High Bandwidth Memory)和强大算力的帮助,GPU 成为 DLRM 训练的主流硬件。但是,随着推荐系统研究的深入,日益增长的嵌入表大小和有限的 GPU 显存形成显著矛盾。如何让利用 GPU 高效训练超大 DLRM 模型,同时突破 GPU 内存墙的限制,已成为 DLRM 领域亟待解决的关键问题。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

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 的嵌入表训练过程主要是不规则的内存访问操作,因此严重受限于硬件访存速度。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

而工业级 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 开销。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

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,那就需要需要传输大量元素。为了防止缓冲区占用过多内存,缓冲区大小被严格限制。如果传输的数据大于缓冲区,会分为多次完成传输。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

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%。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

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 的开始数据读取,这部分开销可以和计算重叠。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

和 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 级超大推荐模型训练。

Bashfrom colossalai.nn.parallel.layers.cache_embedding import CachedEmbeddingBag
emb_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),warmup_ratio=0.7,cache_ratio = 0.01,)

性能测试

在 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% 范围内。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

多 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 训练扩展性并不是很好。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

下图展示了显存使用,显存使用在不同卡上并不相同,这里展示最大显存数值。在仅使用一张 GPU 时,只有 Colossal-AI 的软件 Cache 方法可以训练,多卡并行的占用内存也显著减少数倍。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

Meta Research 的合成数据集 dlrm_datasets 模仿了工业界嵌入表的训练访问行为,因此常在研究中作为推荐系统相关的软硬件设计的测试参考。选取其中的 5 亿行嵌入表项的作为子数据集,构造 256GB 和 128GB 大小的两个 EmbeddingBags 用于测试。

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

PyTorch 由于显存内存不足无法在单卡 A100 上训练。作为对比, Colossal-AI 的软件 cache 将显著降低 GPU 内存需求,足以训练大至 256GB 的嵌入表,并可进一步扩展至 TB 级别。而且,流水预取也能体现出加速效果,当预取数为 32 时,相比没有预取总时间下降 60%,而且对 GPU 的存储的需求却没有增大。

One More Thing

仅需1% Embedding参数,硬件成本降低十倍,开源方案单GPU训练超大推荐模型

面向大模型时代的通用深度学习系统 Colossal-AI,通过多项自研领先技术如高效多维自动并行、异构内存管理、大规模优化库、自适应任务调度等实现高效快速部署 AI 大模型训练和推理,降低 AI 大模型应用成本。

Colossal-AI 相关解决方案已成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用,广受好评。

Colossal-AI 注重开源社区建设,提供中文教程,开放用户社群及论坛,对于用户反馈进行高效交流与迭代更新,不断添加 PaLM、AlphaFold、OPT 等前沿应用。

自然开源以来,Colossal-AI 已经多次在 GitHub 及 Papers With Code 热榜位列世界第一,与众多已有数万 star 的明星开源项目一起受到海内外关注!

项目开源地址:https://github.com/hpcaitech/ColossalAI

文章版权声明

 1 原创文章作者:2159,如若转载,请注明出处: https://www.52hwl.com/46412.html

 2 温馨提示:软件侵权请联系469472785#qq.com(三天内删除相关链接)资源失效请留言反馈

 3 下载提示:如遇蓝奏云无法访问,请修改lanzous(把s修改成x)

 免责声明:本站为个人博客,所有软件信息均来自网络 修改版软件,加群广告提示为修改者自留,非本站信息,注意鉴别

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2023年7月15日 上午9:00
下一篇 2023年7月15日 上午9:01