文｜沙剑蚂蚁集团高级技术专家专注分布式深度学习领域主要负责蚂蚁大规模分布式训练引擎的设计和开发
本文4491字阅读12分钟
2022年6月，蚂蚁集团决定全面引入ESG框架，启动并确立了“数字普惠”、“绿色低碳”、“科技创新”、“开放生态”四位一体的可持续发展战略。针对“绿色低碳”，设立了4个子议题，包括绿色运营、科技助力产业碳中和、生态保护与修复绿色低碳生活。
在此背景下，绿色AI也成为蚂蚁AIInfra团队的一个重要工作方向。作为绿色AI的重要板块，工程提效项目致力于打造高性能离在线AI工程体系，通过提升算力效率和资源利用率，最终达到节省资源降低碳排放的目的。
当前，用户提交分布式训练作业的工具有Yarn或者KubeFlow/Training-Operator。在提交作业时，用户需要在作业中指定作业资源，包括不同角色的节点数量和资源规格（CPU核数、内存、GPU等）。
在训练作业提交后，作业可能遇到如下问题：-集群资源不足以启动作业的所有节点，作业只能等待。-训练作业的节点可能会出错，比如被高优任务抢占、机器故障、IO故障等，导致作业失败。
出现这些问题后，用户只能修改作业资源来重新提交作业。
针对这两个问题，蚂蚁集团早期基于Kubernetes开源了ElasticDL项目来支持K8s上TF2.x分布式训练的弹性容错。在项目落地过程中我们又发现了如下问题：-用户配置的资源可能过少引起OOM和训练性能差。-用户为了保障作业成功率和速度，通常会配置超额资源导致利用率低。-越来越多的用户使用PyTorch或其他TF之外的框架来开发和训练模型。-越来越多的分布式集群开始支持AI作业，比如Ray、Spark集群，能否适配任意计算集群？-在线学习越来越被广泛采用的情况下，如何运用一套系统同时解决兼容离在线训练？
前两个问题使得集群CPU利用率通常只有20%上下，同时算法开发人员需要投入很多人工运维成本，为了解决训练端资源提效的需求，支持在不同集群上针对在离线多种训练模式，给不同框架的分布式训练作业自动地寻找最优资源配置。
蚂蚁AIInfra团队基于ElasticDL弹性容错的思路，升级扩展并开源了DLRover，其目标在于提升分布式模型训练的智能性，目前很多公司的训练作业都是跑在混部的集群中，运行环境复杂多变，正如其名，DLRover作为分布式训练领域的“路虎”，不管多么崎岖的地形，都可以轻松驾驭。
DLRover提出了“MLforSystem”的理念来提升分布式训练的智能性，那么这样的系统应该具备哪些能力呢？
我们认为主要体现在如下几个方面：-解耦：不和底层训练框架耦合在一起，只依赖接口抽象，遵循依赖倒置原则。（*i.e.ElasticRuntime*）-资源调度：具备上帝视角的资源调度管控能力。和建立在对作业精准画像的决策能力。-数据驱动：同时收集掌握集群资源数据，也掌握训练作业数据。以数据驱动智能。-作业交互：以对训练作业以及模型白盒化的理解，动态根据实际情况，对训练作业进行优化调整。超越简单机械的弹性容错！-智能：通过对集群以及作业信息的收集，结合算法模型+固定策略产出精准的作业优化策略。
我们希望设计并实现一个系统，让用户完全摆脱资源配置的束缚，专注于模型训练本身。在没有任何资源配置输入的情况下，DLRover仍然可以为每个训练作业提供最佳资源配置。考虑到用户可能会以不同的方式运行他们的训练作业，DLRover除了面向训练平台进行作业统一管理的ClusterMode，也提供Single-JobMode方便独立的算法开发者也能享受到弹性容错等基本特性。
DLRover由四个主要组件组成：ElasticJob、ElasticTrainer、Brain服务和ClusterMonitor。
上图显示了DLRover如何在K8s集群上管理深度学习训练作业。DLRover以ElasticJobCRD的形式将作业提交到集群。收到CRD后，ElasticJobOperator会拉起一个MasterPod作为ElasticTrainer。其从Brain服务中获取初始资源计划。ElasticTrainer用它来创建ScaleCRD，并应用ScaleCRD通知ElasticJobController启动所需的Pod，每个Pod将在其上启动一个ElasticAgent。
在训练过程中，ElasticTrainer的TrainingMaster将数据分片分发给Worker。同时，ClusterMonitor监控每个作业的运行状态（*i.e.每个节点的Workload*）和集群状态（*i.e.资源水位*）。这些数据将定期报告给Brain，Brain将数据持久化到数据库中。
然后DLRoverBrain根据作业的运行状态，选择合适的算法生成新的资源计划，并通知ElasticTrainer开始资源调整。
用户提交分布式作业时无需提供任何资源信息，DLRover会自动对作业进行画像，推导出最优的资源配置，同时运行时可以根据实际情况（*集群资源、样本流量、当前利用率、…*）自动对资源进行调整。下面展示了两种提交脚本的配置对比：
DLRover支持单点恢复ParameterServer和Worker角色的失败退出而不需要整体作业重启，对于非用户代码和数据类型的错误可以实现用户无感的重启。例如集群中，很常见的一类错误是由于用户配置了不足的内存，导致训练OOM。在DLRover的帮助下，我们可以自动拉起一个优化配置的节点来恢复失败的Node。在真实环境下，DLRover管理的训练作业，相比基线的KubeflowTF-Operator作业，训练成功率从84%提升到了95%以上。
DLRover针对ParameterServer和Worker级别都支持在训练运行时进行自动的调节训练资源以提升训练性能。通过监控作业节点的Workload，DLRover可以分析资源配置的瓶颈。常见的资源瓶颈有：节点抢占、Workload不平衡、CPU不足导致算力低下、节点数目不足。DLRover可以通过动态的资源热更新来持续优化训练性能。
通常不同的模型训练作业，需要不同的资源配置。然而用户倾向于超额配置作业的资源以保障作业的成功率。这通常会导致大量的资源浪费。DLRover的自动扩缩容能力，可以自动根据作业的真实需求配置资源，以最少的资源达到最优的训练性能，从而减少资源浪费。下图显示了自动资源对比手动资源的资源利用率曲线对比：
混部集群存在资源超卖和抢占的情况，部分节点消费数据慢，快节点需要等待慢节点，降低训练速度。DLRover可以通过数据动态分发给慢节点少分发一些数据，减少等待。此外DLRover应该保证训练任务尽可能按照用户配置参数消费数据，避免重复消费/丢失数据，这会给训练带来不确定性，影响模型性能。
当扩容或者缩容时，需要有个全局协调者知道记录节点当前消费数据详情。当节点失败重启后，全局协调者需要知道节点已经消费和尚未消费的数据。如果这些逻辑让训练节点来做，训练节点和训练节点之间需要交互，增加训练节点逻辑的复杂性。DLRoverMaster充当了这个全局协调者的角色。
总而言之，在我们看来，通过动态数据可以简化训练节点逻辑的复杂性，训练节点只管从DLRoverMaster获取Shard，然后读取数据，不需要处理其他的逻辑。
上述动态数据分片特性，实际上帮助我们将DataSource和训练作业进行了解耦，在此基础上DLRover可以同时支持离线训练，也可以支持消费实时样本流的在线学习作业。（*可以通过Dlrover.trainer直接对接样本流，也可以作为流计算引擎的训练Sink节点*）
在蚂蚁的实践中，DLRover可以作为一个理想的组件，来帮助我们构建出一个端到端的在线学习系统。DLRover可以提供数据源消费位点记录与恢复，在线学习长跑作业稳定性与性能保障，资源利用率保障等一系列实际问题。我们的开源仓库中也提供了简单的范例，后续我们也会开放更多周边组件。
训练集群中每天都运行着不同业务域性质各异的训练作业：推荐系统的大规模稀疏模型通常运行在PS/Worker架构的训练模式下进行异步参数更新，资源也多以CPU计算为主。CV/NLP领域的稠密模型则多以数据并行的方式在GPU服务器上进行同步训练，这时只有Worker一种角色。
DLRover在设计上，可以同时支持同步和异步更新模式，做到针对各种训练范式的统一。
DLRover支持用户使用任何自己的训练框架，底层训练代码通过提供约定的API接口以实现自动弹性扩缩等需要同底层分布式代码深度交互。集群中部署完成后，终端算法同学基本可以无感接入。
DLRover目前已经在蚂蚁大规模落地，集群资源利用率相对于基线稳定获得了15%以上的提升。同时也有效解决了由于资源配置不合理造成的训练吞吐不及预期的问题。我们希望通过DLRover的开源可以帮助更多同行一起推行低碳、绿色、AI的理念。同时也切实降低模型开发中的运维成本，释放更多的生产力去解决业务的问题。
当前DLRover的调优算法，以及资源，作业画像策略主要针对蚂蚁内部技术栈优化。考虑到不同机构实际技术栈的多样性，在设计上，DLRover在API层做了统一接口抽象，具体调优算法与作业画像策略则可灵活自定义。我们欢迎不同机构的开发者也能根据自身特点，同我们一起共建DLRover项目，将其发展壮大。