平台架构 - 分布式训练与加速

分布式集群

要做分布式处理/训练，首先要有分布式集群。在数据处理/训练前，先根据用户的需求创建需要的分布式集群，例如：

• tf 分布式训练集群
• pytorch 分布式训练集群
• spark 分布式数据处理集群
• ray 分布式超参搜索集群
• mpi 分布式训练集群
• horovod 分布式训练集群
• nni 分布式超参搜索集群
• mxnet 分布式集群
• volcano 分布式集群
• ……

以及在此衍生出来的分布式的数据下载、hdfs 拉取、cos 上传下载、视频采帧、音频抽取。

要实现这些分布式数据处理/训练，在云原生中定义对应的 crd，然后部署相应的 operator，来解析我们提交的 crd yaml。例如下面我们要实现一个 tf 的 ps 分布式训练集群的结构。

实现自动创建这样的分布式集群，需要四个层面的支持。

1、配置分布式训练集群的部署（CRD）

用户每次任务使用的集群框架和集群规模都是不一样的，需要一个简单的方式能够描述用户本次所需要的集群。在 k8s 中通过 crd 来实现这样的集群定义，可以理解为一个 yaml 文件。

2、控制器部署分布式训练集群（operator）

有了 crd 定义的这样一个集群，还需要一个 operator 来监听、解析、创建、监控 yaml 文件中定义的集群。对于没法直接搬移到 k8s 之上的分布式，要开发对应的 crd+operator：例如分布式多机多卡训练在 train-operator 框架上做开发，horovod 分布式训练在 mpi 的分布式上做开发。tf/pytorch/mxnet/ray 官方都有对应的 k8s 分布式解决方案，可以直接使用对应的方式加入到 k8s 中来。

3、代码识别分布式角色（有状态）

在 k8s 中分布式集群创建好以后，对于有状态的分布式集群，需要用户的代码能识别自己是哪个角色。如果单独为每个角色都写一份代码是很不友好的，所以用户代码需要按照分布式策略识别自己在集群中的角色和身份，进而加入集群、协调工作内容。对于无状态集群，代码只需要接收自己要处理的任务、处理返回即可。在 k8s 中，大部分场景是通过环境变量来识别身份的。

4、训练框架支持分布式（协议和策略）

当每个 pod 都进入各自的工作岗位以后，还需要该用户所使用的算法框架必须支持分布式的训练模式，支持分布式集群通信的协议和协商方法，以及不同角色的节点之间的分工，这样才能协调一致工作，不然就要用户自己手动编写分工逻辑。

在 k8s 中角色（pod）之间的区别主要通过环境变量或者名称来识别，例如 role:ps、role:worker；通信的方式主要通过域名:端口，或者挂载相同的分布式存储、使用文件来沟通。

通过上面的方法，我们就可以实现一个分布式任务在集群中的创建运行。

选对分布式框架

k8s 中分布式框架按照角色形式主要分为下面几种：

1、无状态的框架：优势在于任务 worker 可以横向扩展，任务的完成没有数据倾斜的问题。在集群启动前并不定义 worker 要做哪些数据的处理任务，运行起来以后才开始决定处理哪些任务（比如 kafka 的消费者、ray 的分布式）。这类分布式任务不会因为某个 worker 而严重拖慢整体任务的进度，但一定程度上受限于任务分发机制：比如 kafka 需要将任务发送到队列，ray 只能传输可序列化的函数、无法传递大内存变量。比较适合无状态的任务函数，每次处理之间不存在依赖或者不依赖基础对象。

2、有序分布式任务：提前告诉每个 worker 在集群中的角色和位置，但并不定义它们之间的通信逻辑。比如 volcano 分布式训练，这类任务可以在数据处理前先定义好每个 worker 要处理的任务，这种就不需要任务分发器；但这样用户的任务容易存在数据倾斜的问题（比如某个 worker 所要处理的数据因为某些原因被拖慢，那么整个 pipeline 的结束时间会因为这一个 worker 而拖慢，例如 volcano-job）。也可以在运行中有专门的角色 worker 来分发任务，这种任务如果官方没有分发器的话需要自己写一个专门的分发器。

3、分角色有序分布式任务：在上面的基础上，同时定义好了角色之间的通信逻辑。这种一般是因为除了做分工还要做协同，比如 ring allreduce 的分布式训练方案。这样每个 worker 不仅要知道自己是什么角色（ps-worker），还要知道访问其他 worker 的方法。这种分布式一般来源于特定分布式任务的框架，例如 tf/pytorch。

对于同一个分布式框架，也有多种分布式的方式。比如对于 tf 分布式，有 ps 模式和 ring allreduce 的方式。在单机模型可承受的情况下，ring allreduce 的效率一般比 ps 的方式高；对于单机无法承受模型大小的情况下，选择 ps 模式进行模型并行优化。同时还有 cpu/gpu 方式的选择、稀疏 embedding 大模型和稠密大模型的选择等。

分布式 worker 的均衡度

有了分布式的能力，在机器学习平台中 CubeStudio 加入了通过拖拉拽制作分布式 ml 任务流的能力（见 平台架构-ML任务流.md），每个 task 可能都是一个分布式任务，也可能是单机的任务。不同的用户会有很多分布式的任务流。

如果不做均衡，分布式任务很容易无法更好地利用多机的算力，会更早受到机器的性能瓶颈，进而影响任务运行的效率。为了更好地利用多机的能力、加快整体任务流的运行效率，平台在三个层次做优化。

亲和度 + 反亲和度

一个分布式任务的不同 worker

在一个分布式框架启动了一个分布式集群进行训练时，这些 worker 应该根据实际需求确定尽量分布在哪里，就像 impala 尽量分布在 hdfs datanode 所在的节点上。因为分布式集群中除了分工协作，还有沟通的过程，相当于既独立又合作，所以分布式的方式要能利于这两种操作。例如在 tf ring allreduce 的方式进行分布式训练时，不同 worker 节点之间需要通信传递参数，所以通信的时延也不可忽略：对于在 gpu 上训练时，如果训练中其他过程耗时很小（io、gpu 计算等耗时小），那么通信时延的性能影响就会比较突出，这个时候我们可能更偏向于将 gpu 的任务分布在相同的机器上；如果是 cpu 的训练任务，前向计算的时长比较大、通信时延可以忽略，我们为了能更多地利用 cpu 资源，可能会选择将 cpu 的任务分布在不同的机器上。

同一个 pipeline，不同的 task

在同一个 pipeline 下，我们有时候会并行很多 task，它们并不是完全串行的。比如我们同时跑 7 个 hdfs 数据拉取的任务，分别拉取最近 7 天的数据，后面用这 7 天的数据同时跑 3 个分布式的多机多卡训练。在 k8s 中内存、cpu 等资源是相互隔离的，但是对于部分资源无法做到很好的隔离，比如机器的网络 io、磁盘 io、三方平台的黑白名单限制限速等。对于这种作用相同的任务，我们倾向于将同一个 pipeline 下不同的 task 区分到不同的机器上。如果某个 task 是一个分布式任务，那么 task 下面的 worker 再由 task launcher 自行决定如何调度。

kube-scheduler 打分策略

不同的 pipeline 可能归属于不同的用户，可能会运行在同一个资源组。在 k8s 中 pod 调度依赖于算法和策略，策略部分主要由过滤和打分两个阶段组成，就像推荐中的召回和排序一样：过滤会将满足必备条件的机器筛选出来，打分阶段则选择其中最适合部署当前 pod 的机器。

过滤段主要是：

• PodFitsHostPorts：检查 Pod 请求的端口（网络协议类型）在节点上是否可用。
• PodFitsHost：检查 Pod 是否通过主机名指定了 Node。
• PodFitsResources：检查节点的空闲资源（例如 CPU 和内存）是否满足 Pod 的要求。
• MatchNodeSelector：检查 Pod 的节点选择算符和节点的标签是否匹配。
• NoVolumeZoneConflict：给定该存储的故障区域限制，评估 Pod 请求的卷在节点上是否可用。
• NoDiskConflict：根据 Pod 请求的卷是否在节点上已经挂载，评估 Pod 和节点是否匹配。
• MaxCSIVolumeCount：决定附加 CSI 卷的数量，判断是否超过配置的限制。
• PodToleratesNodeTaints：检查 Pod 的容忍是否能容忍节点的污点。
• CheckVolumeBinding：基于 Pod 的卷请求，评估 Pod 是否适合节点（包括绑定的和未绑定的 PVC）。

打分段主要包含如下接口：

• SelectorSpreadPriority：属于同一 Service、StatefulSet、ReplicaSet 的 Pod 跨主机部署。
• InterPodAffinityPriority：实现了 Pod 亲和性和反亲和性的优先级。
• LeastRequestedPriority：偏向最少请求资源的节点。换句话说，节点上的 Pod 越多、使用的资源就越多，此策略给出的排名就越低。
• MostRequestedPriority：支持最多请求资源的节点。该策略将 Pod 调度到整体工作负载所需的最少的一组节点上。
• RequestedToCapacityRatioPriority：使用默认的打分方法模型，创建基于 ResourceAllocationPriority 的 requestedToCapacity。
• BalancedResourceAllocation：偏向平衡资源使用的节点。
• NodePreferAvoidPodsPriority：根据节点的注解 scheduler.alpha.kubernetes.io/preferAvoidPods 对节点进行优先级排序。
• NodeAffinityPriority：根据节点亲和中 PreferredDuringSchedulingIgnoredDuringExecution 字段对节点进行优先级排序。
• TaintTolerationPriority：根据节点上无法忍受的污点数量，给所有节点进行优先级排序。
• ImageLocalityPriority：偏向已在本地缓存 Pod 所需容器镜像的节点。
• ServiceSpreadingPriority：对于给定的 Service，确保该 Service 关联的 Pod 在不同节点上运行，偏向把 Pod 调度到没有该服务的节点。
• EqualPriority：给予所有节点相等的权重。
• EvenPodsSpreadPriority：实现了 Pod 拓扑扩展约束的优先级排序。

其中在 RequestedToCapacityRatioPriority 打分阶段，k8s 默认的方式是在之前运行成功过的机器上调度任务，这个策略不太适用于分布式训练——我们需要在资源申请率（注意不是申请量）最小的机器上调度新的 pod，所以可以创建修正后的调度打分策略。

ℹ️ 核实：配置文件 scheduler-policy-config.json 真实存在于 install/kubernetes/scheduler-policy-config.json，其 RequestedToCapacityRatioPriority 的 shape 设为 utilization:0 -> score:10、utilization:100 -> score:0，即资源利用率越低分数越高，从而优先调度到申请率最低的机器，与上文描述一致。该文件的部署位置与用法见 02-部署安装/。

gang 调度

在一个算力资源池中运行多个分布式任务时，会存在资源死锁的问题。例如下图所示的 3 台 4 卡的 gpu 机器，如果每个分布式任务有 4 个 worker、每个 worker 占 1 张卡，那么最多能同时启动 3 个这样的分布式任务；但如果同时启动了 4 个这样的分布式任务，就有可能存在每个分布式任务只成功启动了 3 个 worker，而另一个 worker 因为占不到卡而启动不了的情况，进而所有分布式任务都在等待新的卡释放出来。这种情况经常出现在剩余资源不足够多的场景下。

应对这种情况，CubeStudio 加入了 kube-batch 工具，使用 gang 调度算法，配置新的 scheduler 组件，为分布式任务配置批量调度的能力：在调度队列中，先计算剩余资源是否满足所有 worker 的算力要求，完全满足并且全部占用以后，整个集群才算调度成功。

同时引入 volcano 实现批调度的同时，加入优先级调度的能力，在 scheduler 层面统筹批调度的能力。

分布式任务单 worker 的性能加速

通过分布式框架是为了能使用更多的算力帮助用户尽快完成任务，但能尽快完成的前提也是分布式任务中每个 worker 也能尽快完成自己所属的任务。

对于均匀式的任务分配

即每个 worker 所做的工作时长基本一致，那么就需要每个 worker 尽快完成自己所属的数据处理工作，并尽快完成不同 worker 之间的交流，尽快进入下一个迭代。

对于独立的单 worker 任务加速，我们可以像单机优化方案一样，使用多进程多线程的并发方案，或者使用 ray、异步 io 等并发方案，让单 worker 申请的资源能尽可能使用起来。虽然用户申请的 cpu/内存资源如果没有使用起来会被 CubeStudio 的智能算法自动调整到合理的范围，但这些方案只是提高了整体利用率，并没有加速任务的运行，所以一定程度上还是需要用户加速单 worker 的运行。

当然还有尽快完成每个 worker 之间的协商，例如 ring allreduce 之间的通信时长。对于不同 worker 之间的通信，在 k8s 中均是通过内部服务域名解析、在 pod 之间进行通信，主要优化点是在内核方面优化容器 tcp 通信时的效率：使用更高版本的内核，从 linux3 升级到 linux4 后，压测 pod 之间的通信效率从 1.3Gi 升级到 18Gi（主机带宽 25Gi）。还有 nvlink 等解决方案。

对于非均匀式任务分配

尽量减少数据倾斜，或者对未完成任务进行重分配、重新均衡。这种方式的不足是需要用户手动进行再均衡，优势是用户可以自由决定角色的定义，比如用户自行决定 rank=0 的 worker 进行最后结果数据的合并操作。

gpu 利用率提高

gpu 相对于内存和 cpu 比较特殊：内存和 cpu 可以通过算法进行智能修改，可以通过增大 worker 数量替代单 worker 性能不高的问题；但对于 gpu，因为在训练中是独占的形式（可以使用 vgpu 的方式，但并不推荐在训练中使用 vgpu，而是将 vgpu 的能力应用到推理中，详见 平台架构-推理服务.md），所以 gpu 使用率不高的话非常影响训练任务的效率。

因为 gpu 的计算能力比 cpu 强很多，所以之前在 cpu 下不凸显的问题在 gpu 下都变得非常明显。

gpu 利用率监控

通过监控按钮，可以进入查看任务运行的资源使用率，对于资源使用超标可以手动配置增加资源。

ℹ️ 核实：gpu 利用率采集自 NVIDIA dcgm-exporter 的指标 DCGM_FI_DEV_GPU_UTIL，平台侧查询见 myapp/utils/py/py_prometheus.py:104/202/333/458。

其中 gpu 由于是整卡占用，需要调整任务的部分参数和代码，提高 gpu 显存占用率和 gpu 使用率。

注意：很多框架对 gpu 的显存占用默认是独占，所以在容器内通过 nvidia-smi 等命令查看，容易看到 gpu 显存占用 100%。

gpu 利用率低的原因

主要思想是 cpu 操作慢，进而阻塞了 gpu 的计算。

常见的 CPU 计算操作如下：

• 数据加载
• 数据预处理
• 模型保存
• loss 计算
• 评估指标计算
• 日志打印
• 指标上报
• 进度上报

代码层面 gpu 利用率优化

1、数据加载相关

• 存储计算不在同一个城市：数据导入到集群存储。
• 磁盘 io 性能太差：对于临时数据可以将内存映射为磁盘。
• 小文件太多、频繁 io：合并为大文件处理。
• 未启用多进程并行读取数据：pytorch 提高 num_workers，tf 配置 num_parallel_calls/num_parallel_reads。
• 未启用提前加载机制来实现 CPU 和 GPU 的并行：pytorch 配置 prefetch_factor，tf 配置 Dataset.prefetch()。
• 未设置共享内存 pin_memory：设置为 true。
• 每次送入 gpu 的 batch_size 太少：模型固定后，调整 batch_size 尽量增大显存的利用率，然后再调节 num_workers 提高 gpu 利用率。

2、数据预处理相关

• 数据处理和训练耦合在一起：将数据处理和训练分成两个 task，训练中需要的配置之类全部提前加载到内存，让 gpu 只做训练任务。
• 使用 Nvidia DALI，在 gpu 中做数据处理。

3、频繁 IO 操作

• 模型保存太频繁：减少保存模型（checkpoint）的频率。
• tensorboard 文件保存太频繁。
• 日志打印太频繁、频繁 cpu/gpu 切换：不要打印训练迭代中个人日志。
• 存储 io 性能太低。
• 对于频繁使用的文件（库文件、lib 等）放入了性能不高的存储中。

vgpu + 共享 gpu

常规我们会先采用 gpu 利用率提高的方法，把申请的 gpu 尽量占满，但这部分优化一定程度上是代码层面的优化，需要用户的配合。但有时候用户的代码无法做到很完美，这时候就需要架构层面来解决利用率的问题。

1、使用 vgpu 的形式

现在有很多种方案使用 vgpu，将一张卡隔离成多张卡。但在训练中使用 vgpu 有个调度问题，如下图所示：分布式任务 t1 大量占用 1/3 卡，造成大量算力碎片化，进而对于正常的任务 t2，无法找到满足条件的机器 1 卡进行调度。这是因为 k8s 首先还是在机器上，pod 的算力申请上限无法突破单机。另外一个问题，就算通过零散化算力只是为了提高 gpu 的利用率，并没有加速任务的运行。这种 vgpu 更适合在推理阶段使用，因为推理阶段输入的流量不是训练开发者决定的，而是 C 端用户决定的。

虚拟化 gpu 以及虚拟内存+gpu 利用率提高在推理端的方案，在 平台架构-推理服务.md 中有介绍。

2、共享 gpu

所以我们可以采用共享 gpu 的方式来提高 gpu 利用率，同时加速任务的训练进度。在常规的多机多卡训练中，gpu 的占用是独享的，以 ring allreduce 为例，每个 worker 独享一张 gpu 卡。

但并不是每个开发者都能将 gpu 卡的利用率使用得非常高。比如用户的代码使用 gpu 训练但 gpu 利用率只有 10%，对于这种情况，一般开发者会采用使用更多的 gpu 卡来加速训练，这样 gpu 的浪费情况就更加严重。在架构层面我们调整为用户可以配置共享同一个 gpu 的进程数目，以此让在相同 gpu 卡上跑出更多分布式 worker 的效果，直到机器的某个资源达到瓶颈。比如 3 个进程共享一张卡、一台机器 4 张卡、同时 12 个进程，此时 cpu 使用率已经接近 100%，那么就要开始从其他地方进一步加速了。

这种方式重置了每个 pod 的角色，而是角色直接传递到进程，rank 和 world_size 都被重置。

任务资源配置动态调整

用户在运行任务的时候也不知道任务的资源该设置为多少。除了监控、通知用户任务的资源使用情况，还需要能够智能地调整用户的配置，因为 CubeStudio 平台是统筹算力的，一个资源池里面会有很多任务在同时运行。

但对于某些特殊的任务，也支持模板里面提前设置好固定资源配置 TASK_RESOURCE_CPU/TASK_RESOURCE_MEMORY/TASK_RESOURCE_GPU，屏蔽用户的配置。同时对于某些不合理的资源动态修正，也允许在模板中通过环境变量 NO_RESOURCE_CHECK 屏蔽动态资源配置。

说明：模板固定资源环境变量为 TASK_RESOURCE_CPU / TASK_RESOURCE_MEMORY / TASK_RESOURCE_GPU / TASK_RESOURCE_RDMA（见 myapp/init/init-job-template.json:72 等多处）；NO_RESOURCE_CHECK 屏蔽动态校验的判断见 myapp/tools/watch_workflow.py:180，模板侧固定资源读取见 myapp/views/view_task.py:641。

在动态资源配置中仅对内存、cpu 进行修正；对 gpu 资源，采用的仍然是上面 gpu 利用率提升的方法。因为大部分场景下，cpu 上的任务可以通过减少每个 worker 的 cpu 数目、增加 worker 数量来降下来，但在 gpu 训练场景我们采用独占方式，所以无法缩减配置。

资源池动态均衡

分布式集群的加速能力与可以占用的资源几乎可以成线性关系，所以保障分布式任务有资源能调度起来就变得异常重要，尤其在多集群、多资源组的平台下。平台在联邦控制器下，多个集群之间不进行均衡，但在一个集群下的多个项目组之间进行均衡配置：限制保障每个项目组的最低算力以保障基本任务的运行，当有临时性任务发起时，就会从公共池中将算力加入到指定项目组，运行结束后释放资源回公共池。