关于人工智能:DLRover-云上自动扩缩容-DeepRec-分布式训练作业

背景

现在，深度学习已广泛应用在搜寻、广告、举荐等业务中，这类业务场景广泛有两个特点：1）训练样本量大，须要分布式训练晋升训练速度；2）模型稠密，即模型构造中离散特色计算逻辑占比拟高。

DeepRec 是阿里云机器学习平台 PAI 开源的面向举荐场景的高性能深度学习框架，针对稠密模型在分布式、图优化、算子、Runtime 等方面进行了深度的性能优化，同时提供了搜寻、举荐、广告场景下特有的动静弹性特色，动静弹性维度，自适应 EmbeddingVariable、增量模型导出及加载等一系列性能。

越来越多的公司为了节省成本开始在云上训练 AI 模型。通常，用户在云上提交一个分布式训练作业须要给作业配置资源，包含节点的数量和每个节点的资源（CPU & memory 等）。资源是影响训练性能的一个重要因素。在察看蚂蚁的 kubernetes 集群上几千个训练作业后，咱们发现大量的训练作业存在资源配置不合理的景象。比方：

1. PS 或 worker 节点的资源规格太少。CPU 不够会导致训练慢，memory 不够会导致 OOM 并导致作业失败。
2. PS 很少，worker 很多。PS 资源成为了训练瓶颈，worker 很多也没法晋升训练速度。
3. PS 和 worker 的资源太多。用户给每个 PS 和 worker 都配置了大量的 CPU 和 memory，利用率低下。
4. 每个节点的 PS 资源规格一样，然而存在热点 PS。PS 上的负载不平衡可能导致局部 PS 的 CPU 负载显著高于其余 PS，热点 PS 成了训练瓶颈。

因为模型构造千差万别，用户很难在作业提交时精确地配置适合资源。想要寻找合理配置，用户须要重复地调整资源和启动作业以便察看性能。这个过程是非常费时的。为此，DLRover 主动地为分布式训练作业配置资源，并在训练过程中针对资源瓶颈进行主动扩缩容，达成充分利用资源来晋升训练性能的目标。应用 DLRover 的 auto-scale 性能，用户提交训练作业时无需关怀资源配置。

基于运行时优化的主动扩缩容

DLRover 采纳运行时优化思维来对分布式训练作业进行主动扩缩容。DLRover 首先让训练作业跑起来，而后监控训练作业所有节点的负载和整体训练性能，并依据监控指标来动静调整作业的资源，从而晋升训练性能和资源利用率。为了实现主动扩缩容，DLRover 须要反对如下性能：

1. 资源主动优化：监控训练性能和节点负载，并依据以后性能和负载给出最优的资源配置，包含节点数量和资源规格。
2. 弹性调度：在训练运行时动静地给作业减少节点、缩小节点或者用新资源规格的节点来替换原有节点。
3. 训练框架反对节点容错和弹性：在训练运行时，新的节点能够退出训练；缩减局部节点时，其余的节点不受影响，能够继续训练。

为此，咱们设计的 DLRover 的整体设计架构如下：

• Brain Service: 负责资源主动优化。基于实时采集的训练速度和各个节点负载来主动优化作业的资源配置。
• DLRover Master：负责弹性调度。每个训练作业领有一个 master 节点，master 节点负责训练速度采集、节点负载收集、训练样本治理和弹性调度。
• Elastic Agent：负责与训练框架协调来反对训练的容错和弹性。每个节点上都有一个 Elastic Agent，Agent 从 master 上获取作业以后运行的节点信息，告诉训练框架更新分布式训练状态。Agent 还负责从 master 获取训练样本信息来供训练框架迭代模型，从而使训练样本分片反对 worker 的弹性。

零碎模块介绍

负载感知的主动资源配置策略

训练作业主动资源配置次要分为 3 个阶段：启动阶段、资源预估阶段、动静调优阶段。

启动阶段

DLRover 的 Brain Service 会将所有运行过的训练作业的信息存入数据库中，包含运行时各节点的负载、训练速度、用户信息、数据集信息等。在作业启动阶段，DLRover Brain 会依据新作业的用户和数据集信息寻找相干的历史作业，并依据历史作业来配置 PS 资源（PS 节点数量和节点的 CPU & memory）和第一个 worker（chief）的资源。

PS 启动资源预估：

• 预估 PS CPU 和 memory 总量：从历史作业中，计算每个相干作业的所有 PS 节点的 CPU 和 memory 耗费总量。取 CPU 总量的均值作为新作业的 PS CPU 总量（total_PSCPU），取内存总量的最大值作业新作业 PS 内存的总量（total_PSMEM）。
• PS 节点的 CPU 预估：取历史作业中所有 PS 节点的 CPU 应用的最大值作业 PSCPU。
• PS 数量预估：PSNUM = total_PSCPU / PSCPU。
• PS 节点的 memory 预估：PSMEM = total_PSMEM / PSNUM。

Chief 资源预估：取历史作业中所有 worker 的 CPU 耗费和内存耗费的最大值。

启动阶段的非凡状况：

1. 没有历史相干作业，比方新用户应用新数据集提交的训练作业。此时，DLRover 采纳默认资源配置来启动 PS 和 chief。比方 PSNUM=1，PSCPU=8，PSMEM=8G，chiefCPU=8，chiefMEM=8G。
2. 初始资源导致节点失败，最常见的是，内存不足导致 OOM。当节点产生 OOM 时，DLRover 通过容错机制会主动减少节点的内存来重启节点，直到训练失常运行。

资源预估阶段

该阶段，训练曾经失常开始。当模型迭代一小段时间后（比方 3min），DLRover 开始资源预估。此时，DLRover 曾经采集到了 chief 训练模型时的各个节点的负载，包含 chief 的 CPU 使用量（chiefUsedCPU）和内存使用量（chiefUsedMem），PS 总的 CPU 使用量（total_PSUsedCPU）和内存使用量（total_PSUsedMem）。该阶段，咱们假如 PS 的负载是均匀分布的，worker 的机器性能齐全一样，来计算训练作业的最优资源配置。

• worker 数量预估：在 PS 异步训练中，各个 worker 相互独立且角色一样，所以一个 worker 迭代模型对 CPU 的需求量为 trainingCPU =chiefUsedCPU + total_PSUsedCPU。这样咱们依据 Job 配置的总 CPU 量 limitCPU 能够计算 worker 的数量 workerNUM = limitCPU/trianingCPU。
• worker CPU 和内存预估：因为 worker 的模型完全相同，所以 CPU 和 内存耗费也是类似的，新的 workerCPU = chiefUsedCPU * factor, workerMem = chiefUsedMem * factor ，factor 为冗余因子，比方 1.2。
• PS 数量预估：异步训练中，PS 的 CPU 负载与 worker 数量成正比，所以 PSNUM = total_PSUsedCPU * workerNUM / PSCPU。

• PS 内存预估：PS 存储模型参数，内存使用量并不会随 worker 的减少而减少。如果模型蕴含稠密 embedding，PS 的内存会随着训练的进行而减少，为此 PS 的内存预估分为两种状况：

  1. PS 的内存在训练开始后保持稳定，PSMEM= (total_PSUsedMem / PSNUM)* factor，factor 为冗余因子，个别要大于 1。
  2. PS 的内存持续增长，那么 DLRover Brain 会计算 PS 内存随迭代步数的增长率 memRate，而后计算总的 totalPSMEM = memRate * totalStep，则每个 PS 的内存 PSMEM = totalPSMEM / PSNUM。

动静调优阶段

资源预估阶段假如：PS 负载均匀分布和 worker 性能完全一致，并不一定齐全成立。PS 的负载可能存在不平衡的状况，比方 TensorFlow 应用 RoundRobin 策略给 PS 调配 Variable，如果 variable 的大小差别很大，PS 的负载就可能不平衡。同时咱们的节点负载采样也可能存在误差，导致预估不准。所以，动静调优阶段会实时监控训练速度和节点负载，并采纳启发式策略来调整作业资源配置。

• PS 热点查看：如果发现某些 PS 节点的 CPU 负载很高，DLRover 会新拉起一个 CPU 更多的节点来替换热点 PS 节点，从而缓解热点 PS 的瓶颈。
• OOM 危险预查：如果发现节点的内存使用量达到危险门限，DLRover 会新拉起一个内存更多的节点来替换有 OOM 危险的节点。因为在稠密模型中，PS 的 embedding 表中的 item 数量可能越来越多，导致 PS 内存增长。
• worker 数量动静调整：如果发现 PS 的 CPU 负载还有较大的余量，DLRover 会减少肯定数量的 worker，而后察看 PS 的 CPU 负载以及训练速度是否晋升，直到 PS CPU 负载大于阈值或者新增 worker 没有晋升训练速度。

弹性调度

训练的弹性调度包含：1）能动静地拉起或者删除节点，2）新拉起的节点可能获取训练样本训练模型，3）被删除的节点未生产的数据能给其余节点应用。为此，DLRover 实现了反对在 k8s 上动静调整作业节点的 ElasticJob controller 和动静样本散发服务。

弹性调度 ElasticJob Controller

在 DLRover 的 Job master 拿到新的资源优化后果 plan 后，master 会生成一个 ScalePlan CRD 来告诉 ElasticJob Controller 更改训练作业的节点规模。例如，上面的 ScalePlan 指定了 worker 为 4 个。

apiVersion: elastic.iml.github.io/v1alpha1
kind: ScalePlan
metadata:
  labels:
    scale-type: auto
  name: deepctr-auto-job-scaleplan-2
spec:
  createPods: []
  ownerJob: deepctr-auto-job
  psHosts:
  - deepctr-auto-job-edljob-ps-0.dlrover.svc:2222
  removePods: []
  replicaResourceSpecs:
    worker:
      replicas: 4
      resource:
        cpu: "1"
        memory: 4170Mi

当 ElasticJob Controller 拿到 ScalePlan 后，会查看 deepctr-auto-job 这个作业以后的 worker 数量，如果 worker 数量少于 4 个，就会新加 worker；如果大于 4 个，就会缩小 worker。

动静样本散发服务

在 PS 异步数据并行训练中，咱们须要将训练样本分给不同的 worker。如果 worker 数量固定，咱们能够将样本等分给 worker。然而在弹性过程中，worker 数量会变动，所以等分策略并不适宜弹性训练。为此，DLRover 的 master 实现了一个动静样本散发服务。

master 会依据样本索引在数据集切分为很多小的 shard，放入一个 TODO 队列，每个 shard 并不蕴含样本数据，而是蕴含样本索引。worker 启动后，须要从 Job master 节点申请 shard，master 每次给 worker 分一个 shard，并将 shard 放入 DOING 队列。worker 拿到 shard 后就能够依据索引读取相应样本来训练模型。这样新 worker 启动后只须要向 master 申请 shard 即可。worker 生产完 shard 的数据后，向 master 汇报 shard 已实现，master 会将 shard 从 DOING 队列移除。如果 worker 挂了，master 会将对应的 shard 会从新放入 TODO 队列。

基于 DeepRec 的弹性 PS 训练

DeepRec EmbeddingVariable 反对 PS 扩缩容

对于搜寻广告和举荐模型，模型网络参数次要有两局部：一部分是 user 或者 item 的稠密 embedding，另一部分是 dense 层。随着模型一直遇到新的 user 或者 item，embedding 表会一直减少，导致 PS 的内存一直减少。正是因为这种起因，算法工程师无奈精确预估 ps 资源规格（并发度和内存）。并且经常出现以下状况：在提交作业的时候资源规格是足够的，随着工作运行，PS 内存逐步减少，直到最终 OOM。此外 PS 数量不合理也有可能影响训练速度。这就要求咱们能在训练过程中裁减 PS 的数量。裁减 PS 后，训练框架须要对 embedding 从新 partition，从而让 embedding 参数从新均分散布到新的 PS 节点汇合上。

为解决上述问题，DeepRec 新设计了一套反对动静 Embedding 语义的 EmbeddingVariable，在特色无损训练的同时以最经济的形式应用内存资源。具体能够参考 DeepRec。

基于 checkpoint 的 PS 弹性扩缩容

PS 架构中，模型的权重是存储在 PS 内存的。如果 PS 变动，模型训练须要将权重重新分配到 PS 上。为了保障训练的连续性和一致性，DLRover 采纳 checkpoint 的计划实现 PS 弹性扩缩容。在扩缩容开始前，DLRover 会告诉 PS 将模型权重放弃到 checkpoint 中。扩缩容实现后，PS 会从 checkpoint 中复原模型权重。

DLrover 在 PS 扩容 / 缩容时被动保留 checkpoint，流程如下：worker 在每一个 step 之后会运行相干 hook，在 hook 中会向 DLRover master 轮询 PS 汇合。当 worker-0 发现新的 PS 汇合与现有的 PS 联合不统一时，worker-0 会告诉所有 worker 进行训练，并告诉 PS 放弃 checkpoint。在 PS checkpoint 放弃完后，worker-0 会依据新的 PS 汇合来结构计算图，更新 session，从新组网，而后告诉新的 PS 加载 checkpoint。最初 worker-0 告诉所有的 worker 连贯新的 PS 开始训练。

阿里云 ACK 上验证 DLRover 主动扩缩容

为了验证主动扩缩容的可行性，咱们在阿里云的 ACK 上创立了一个小的 Kubernetes 集群。而后咱们针对 CRITEO 广告点击预测数据集开发了一个 DeepFM 模型，将训练作业 的 CPU 限度配置为 15 核。在作业起始阶段，DLRover 给作业启动了一个 PS 和一个 worker，如下所示：

NAME                                          READY   STATUS    RESTARTS   AGE
dlrover-auto-scale-edljob-chief-0             1/1     Running   0          32s
dlrover-auto-scale-edljob-ps-0                1/1     Running   0          32s
elasticjob-torch-mnist-dlrover-master         1/1     Running   0          39s

此时的训练速度约为 30 step/s。大概 3 min 后，DLRover 主动给作业新增了 3 个 worker，速度晋升到 100 steps/s 如下所示：

NAME                                          READY   STATUS    RESTARTS   AGE
dlrover-auto-scale-edljob-chief-0             1/1     Running   0          6m17s
dlrover-auto-scale-edljob-ps-0                1/1     Running   0          6m17s
dlrover-auto-scale-edljob-worker-0            1/1     Running   0          3m19s
dlrover-auto-scale-edljob-worker-1            1/1     Running   0          3m19s
dlrover-auto-scale-edljob-worker-2            1/1     Running   0          3m19s
elasticjob-torch-mnist-dlrover-master         1/1     Running   0          6m24s

后续打算

DLRover 反对了 PS 异步训练的主动扩缩容来晋升训练速度。下一步咱们将针对 DeepRec 的同步训练提供主动扩缩容性能。除了搜推广场景，DLRover 也将摸索 foundation model 的分布式训练的主动扩缩容，以晋升预训练大模型的训练效率和升高训练老本。