【SIGMOD 2023】深度学习弹性数据流水线系统 GoldMiner，大幅提升任务和集群效率

第一板块：开篇

近日，阿里云机器学习平台 PAI 和北京大学杨智老师团队合作的论文《GoldMiner: Elastic Scaling of Training Data Pre-Processing Pipelines for Deep Learning》被数据库领域顶会 SIGMOD 2023 接收。

GoldMiner 观察到深度学习任务中的数据预处理流水线具有无状态的特点，具有内在的资源弹性。基于此，GoldMiner 将数据预处理流水线和模型训练部分分离执行，通过自动计算图分析来识别无状态的数据预处理计算，并对其实现高效的并行加速和弹性伸缩，从而缓解数据预处理瓶颈，提升训练性能。通过与集群调度器的协同设计，GoldMiner 进一步发挥了数据预处理计算的资源弹性，大幅提升集群调度效率。实验显示 GoldMiner 可提升训练性能达 12.1 倍，提升 GPU 集群利用率达 2.5 倍。

第二板块：背景

近年来，随着 GPU 加速器的不断进化，以及各类软件优化技术的层出不穷，深度学习训练的计算效率正不断被提升到新的层次。但与此同时，深度学习本质上仍是一种多阶段、多资源的任务类型：不仅需要在 GPU 上进行大量的训练计算，同时往往也需要 CPU 端的数据预处理流水线（如数据增强、特征转换等），这类预处理计算是训练出高质量模型的必要步骤。因此，GPU 端训练性能的提升也带来了更大的数据预处理压力，使后者成为新的性能瓶颈。

我们观察到数据预处理瓶颈对任务训练性能和集群资源利用效率都有深刻的影响。一方面，对于单个训练任务而言，数据预处理瓶颈意味着训练性能的受损。我们在一台配备 8 块 V100 GPU 和 64 个 vCPU 核的虚拟机上使用一块 GPU 和不同的 vCPU 数进行了性能测试，观察不同模型需要多少 vCPU 才能达到最优性能。结果显示（下图）大部分模型都需要超过 8 个 vCPU（即一块 GPU 平均能得到的 vCPU 数）才能达到最优性能，甚至部分模型需要消耗掉整机八卡的 64 个 vCPU。这就意味着这类模型在共享集群中很可能无法得到足够的 CPU 资源，导致数据预处理部分性能下降，最终影响训练效率（下图右侧纵轴表示仅使用 8 个 vCPU 时能达到的相对性能）。

而另一方面，上述的问题在云上场景中会更为严重，影响共享集群中的资源分配效率。目前企业一般都会构建或购买共享 GPU 集群来运行训练任务，在这种 GPU 集群中 GPU 利用率至关重要。而用户为了避免自己任务 CPU 不足，可能主动提高任务的 CPU-GPU 配比，然而这些用户自定义的 CPU-GPU 配比很容易导致集群资源碎片。比如某台机器上由于跑了一些 CPU-GPU 配比较高的任务，最终 CPU 比 GPU 先被分配完，这样机器上空闲的 GPU 会无法被分配出去，不仅导致昂贵的 GPU 资源被浪费，也会提高任务的等待时间。我们在阿里巴巴内部 GPU 集群中的观察发现有近 40%的任务等待时间浪费在了这种“GPU 足够但 CPU 不足”的情况上。

解决上述的两方面问题的一种手段是将 GPU 端训练和 CPU 端数据预处理分离，使这两部分计算的资源分配不必捆绑在同一台机器上。这样当机器 CPU 资源不足时可以使用其他机器的资源，一来可以为单个任务分配更多的 CPU 资源达到加速效果，同时也缓解了碎片 GPU 无法分配的问题。实际上这一思路并不是首次被提出，然而要用这种方式提高任务和集群效率仍存在一系列技术挑战。

第三板块：挑战

尽管已有一些方案（如 tf.data service、PyTorch DPP）支持数据预处理计算的分离执行，现有技术仍有以下几方面挑战需要解决：

1. 计算切分效率：现有技术简单地通过深度学习框架提供的 Dataset/DataLoader API，将这类 API 内部封装的计算作为数据预处理计算分离出去。然而我们发现即便在这类 API 之外，仍可能存在一部分计算是可以分离执行的，简单的切分方式错失了这部分并行加速的机会。

2. 用户代码侵入性：tf.data service [1]、PyTorch DPP [2]等技术要达到数据预处理的分离执行，均需要用户重构这部分代码逻辑，有较大的代码侵入性。我们希望达到以用户透明的方式实现这一分离的效果。

3. 与集群调度的结合：在与训练分离后，数据预处理计算实际上蕴含了内在的资源弹性，而现有技术均没有在集群调度的层次挖掘这部分弹性，以提升集群整体资源利用效率。

第四板块：破局

GoldMiner 是一个自动、弹性的数据预处理服务。如图所示，GoldMiner 使用 data worker（DW）和 training worker（TW）两类角色来分别进行数据预处理和训练两部分计算。GoldMiner 可从原始用户代码中自动化地识别出 data worker 部分的计算（也包含没有被封装在 Dataset/DataLoader API 中的计算）。同时 GoldMiner 实现了数据预处理计算的弹性伸缩，并通过与集群调度器的协同设计，进一步提升了集群效率。

GoldMiner 实现这一效果的关键在于利用了数据预处理计算的无状态性。这里的无状态是指数据预处理不依赖于模型参数，而模型参数需要在训练的每一次迭代中被反复更新，因此不依赖于模型参数的计算即可与训练部分异步化地执行。我们知道深度学习计算可以被表达为数据流图（DFG），GoldMiner 通过对于用户 DFG 的分析，自动找出其中无状态的子图。下图是一个典型推荐模型的 DFG，不同于直接仅切分 Dataset 的做法（图中的 Simple partition），GoldMiner 通过识别与模型参数的依赖关系，将切分范围自动扩展到后续的一些特征转换操作（Expected partition）。实验显示通过这一拓展我们可以将 data worker 并行加速的效果再提升 1.6 倍。

基于自动化的图切分，GoldMiner 进一步实现了 data worker 部分子图在多个 data workers 间的并行加速以及 data workers 和 training workers 间的数据传递。利用 data worker 的无状态性，这一分布式执行实现了 data workers 数量的动态伸缩，从而在集群空闲资源不断变化的过程中更高效地利用资源。

为了充分发挥 data worker 的资源弹性，GoldMiner 提供了一个 data worker 调度器，调度器会在任务和集群两个层次进行动态资源调整。对于每个任务，GoldMiner 调整其 DW 和 TW 的尺寸以搜索到一个扩展效率最高的配置；在集群维度上，GoldMiner 在不同任务之间动态调整 data workers 数量分配，以优化某些全局调度目标（如最小化任务完成时间）。这两层调整利用了一个统一的性能指标，即在 DW 和 TW 之间传递数据的队列，该队列的状态体现了 DW 和 TW 的相对速度，以及增加 DW 的潜在收益。在一个 64-GPU 的集群中进行的实验显示 GoldMiner 调度器可以缩短平均任务完成时间达 2.5 倍，提升 GPU 分配率达 2.1 倍。

第五板块：应用

我们此前已经在客户的真实推荐模型上评测了 GoldMiner 的效果，结果显示 GoldMiner 可为用户模型加速 1.43 倍，并可削减 13%的训练成本。目前在落地部署中。

同时我们也开发了 PyTorch 版本实现，并即将与 PAI-DLC 集成，向用户提供数据预处理加速的能力。

第六板块：

• 论文名字：GoldMiner: Elastic Scaling of Training Data Pre-Processing Pipelines for Deep Learning

• 论文作者：赵汉宇，杨智，程羽，田超，任仕儒，肖文聪，袁满，陈浪石，刘恺博，张杨，李永，林伟

• 论文 pdf 链接：https://dl.acm.org/doi/pdf/10.1145/3589773
  

• 参考文献：

[1] Andrew Audibert, Yang Chen, Dan Graur, Ana Klimovic, Jiri Simsa, Chandramohan A. Thekkath. A Case for Disaggregation of ML Data Processing. https://arxiv.org/abs/2210.14826

[2] Mark Zhao, Niket Agarwal, Aarti Basant, Bugra Gedik, Satadru Pan, Mustafa Ozdal, Rakesh Komuravelli, Jerry Pan, Tianshu Bao, Haowei Lu, Sundaram Narayanan, Jack Langman, Kevin Wilfong, Harsha Rastogi, Carole-Jean Wu, Christos Kozyrakis, Parik Pol. Understanding Data Storage and Ingestion for Large-Scale Deep Recommendation Model Training. ISCA'22

https://free.aliyun.com/?crowd=personal&pipCode=learn&tlog=0621pailwifq