第一板块:开篇
近日,阿里云机器学习平台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)支持数据预处理计算的分离执行,现有技术仍有以下几方面挑战需要解决:
- 计算切分效率:现有技术简单地通过深度学习框架提供的Dataset/DataLoader API,将这类API内部封装的计算作为数据预处理计算分离出去。然而我们发现即便在这类API之外,仍可能存在一部分计算是可以分离执行的,简单的切分方式错失了这部分并行加速的机会。
- 用户代码侵入性:tf.data service [1]、PyTorch DPP [2]等技术要达到数据预处理的分离执行,均需要用户重构这部分代码逻辑,有较大的代码侵入性。我们希望达到以用户透明的方式实现这一分离的效果。
- 与集群调度的结合:在与训练分离后,数据预处理计算实际上蕴含了内在的资源弹性,而现有技术均没有在集群调度的层次挖掘这部分弹性,以提升集群整体资源利用效率。
第四板块:破局
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