云原生分布式训练技术综述_nvswitch3总交换带宽

大模型训练必然绕不开大规模分布式训练，因此本文主要阐述基于云原生架构的大规模分布式训练涉及的技术，覆盖通信、调度、存储、通信等。
本文主要借鉴了云原生机器学习平台技术综述文章，对部分内容进行了扩充，来叙述分布式训练技术在云原生上涉及的一些技术。
本文借鉴了很多网上材料，也基本都在每一节后粘贴了参考的链接，如有侵权，请联系删除。

整体架构

上图为机器学习平台的普遍选型架构图，本文根据该架构图的层级进行分点阐述。从上到下，依次是：

• Pipeline Orchestrator（流程编排引擎）：模型训练流程中的各个环节（component）可以被定义为一个DAG，DAG主要描述了每个component的输入输出和拓扑关系。Orchestrator会根据DAG定义的顺序依次执行各个component。在这个方向的开源工具有：Kubeflow Pipeline、Airflow、MLflow等。
• Distributed Training Framework（分布式训练框架）：Pipeline Orchestrator不会提供分布式训练功能，分布式训练需要在训练框架层面实现并行算法、提供通信机制。Tensorflow, PyTorch, Horovod框架基于不同并行化方法提供了各自的分布式训练API。
• Training Job Operator：每个训练框架都有各自的operator，operator的主要作用是将一个包含多个worker的分布式训练作业拆分为不同的built-in workload，向pod中注入用来达成共识的参数，交给scheduler调度。Kubeflow社区提供了TFJob、PytorchJob、MPIJob等多种框架的operator。
• Batch Job Scheduler：scheduler监听集群内新的pod，为pod找到合适的node调度上去。深度学习训练场景的调度器应该支持批量调度、公平调度等高级调度策略。但kubernetes一开始主要管理在线服务，原生的scheduler无法满足批处理任务的调度需求。目前在模型训练场景业界常用的开源调度器是Volcano。
• NUMA感知调度（Topology Manager）：Topology Manager在v1.16作为alpha功能引入kubelet的组件，v1.18升级为beta版本。主要作用是在调度的最后阶段，当服务已经绑定node时，Topology Manager会从第三方Device Manager获取到亲和性调度建议，决定将服务绑定在哪个Numa节点上。

Kubernets基本原理（一些简单的普及，可跳过）

云原生就必然绕不过kubernets，因此本章将简单阐述kubernets的一些的基本概念知识。
如果你懂Kubernets可以跳过该章节。

上图是一个简单的架构图。

• Master节点
  • **etcd：**兼顾一致性和可用性的kv数据库，是k8s的数据库。
  • **kube-apiserver：**所有服务不会跟etcd直接建立连接，而是通过apiserver读写数据。
  • **kube-controller-manager：**负责编排。监听api server变更，发现待编排对象期望状态和实际状态的差异，将实际状态调整为期望状态。
  • **kube-scheduler：**负责调度。为新创建出来的pod寻找最合适的节点调度上去。
• Node节点
  • **CNI(Container Networking Interface)：**管理容器的网络。
  • **kubelet：**接受Scheduler调度过来的pod，保证pod中容器的监控运行。
  • **CRI(Container Runtime Interface)：**kubelet通过CRI跟容器交互。
  • CSI(Container Storage Interface)：管理容器持久化存储。
  • **Device Plugin：**管理硬件设备，向kubernetes上报设备信息，在调度时给容器绑定为其分配的设备。

控制器（Controller Manager）

控制器主要负责编排的工作，也就是kubernets最核心的功能之一。它的功能用一句话概括就是：检测期望状态和实际状态的一致性，并控制系统达到期望状态。
假设申请一个需要3副本的deployment，则处理流程如下：

1. informer组件会通过List&Watch监听到集群内出现了一个新的deployment；
2. 队列通知控制循环（也叫control loop、reconcile loop）处理这个请求；
3. control loop接受到请求后执行具体的diff操作，发现集群内现在pod数为0，而期望启动3个；
4. 向集群内创建3个pod。

kubernetes原生的kube-controller-manager由一系列的控制组成，包括EndpointController、ReplicationController、PodGCController、DaemonSetController、JobController等。因此，它只能控制原生对象，比如deployment、statefulset、service、configmap等。

自定义资源类型（Operator）

正如上文所述，kubernetes原生控制器支持的资源类型有限，但是kubernetes提供了很好的扩展方式，用于用户自定义资源类型和对应的编排控制器。CoreOS开源的Operator框架提供了很好的脚手架，提高了编码效率。
Operator主要做要两件事情：

• **定义资源类型（Custom Resource Definition）：**它其实只是往etcd数据库注册了一条数据，用于声明定义的资源类型、分组等信息，以下为TFJob的示例：

apiVersion: apiextensions.k8s.io/v1
kind: CustomResourceDefinition
metadata:
  generateName: tfjob.kubeflow.org
spec:
  # 资源分组
  group: kubeflow.org
  names:
    # 复数
    plural: tfjobs
    # 单数
    singular: tfjob
    # 资源名称
    kind: TFJob
    listKind: TFJobList
  scope: Namespaced
  versions:
    # 版本
    - name: v1
      served: true
      storage: true
      schema:
        ...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

自定义控制器（controller）：自定义controller一般以deployment的形式启动。它会监听CRD类型的资源，运行你实现的control loop逻辑，将资源拆分成pod、service、configmap等细粒度资源，应用到kubernets集群。

调度器（Scheduler）

调度器主要完成为pod找到最适合的节点（node），然后将nodeName更新到pod的yaml上。
在scheduler的control loop中，scheduler会监听etcd的变更，拿到未被调度的pod，为pod选择node。选择node时会执行预选和优选两个策略：

• **预选：**找到能够被调度的node，node需要满足一些必要条件：剩余资源大于pod所申请的资源、污点、亲和性等策略。
• **优选：**在满足要求的node里找到最适合的node，优选会考虑：选择调度后最空闲的节点、选择调度后资源使用最平均的节点、亲和性等策略。统一打分，选择分数最高的节点。

可以看到，kubernets原生的scheduler调度是以POD为基本单元，通过队列进行调度的。因此，加入一个job中包含多个pod的调度，它也是以pod为基础单元调度的，因此无法满足批处理任务的调度。

kubelet

每个node都会有一个kubelet，它的设计模式和控制器、调度器类似，主要工作流程如下：

1. kubelet的control loop会监听包含当前节点nodeName字段，但还未被调度的pod。
2. kubelet会再次判断当前节点的资源是否满足调度需求，如果满足，则调用容器运行时的接口创建容器，并指定资源申请相关的参数。

资源类型原生只支持cpu和memory，但是提供厂家也可以提供device plugin进行自定义，比如说NVIDIA Device Plugin用于管理和指定GPU资源，后续也会介绍NVIDIA Device Plugin。

底层存储

大规模集群训练时，容错恢复是绕不过去的话题，要做到容错恢复，那就需要checkpoint。在训练过程中，每隔一段时间保存一个checkpoint，在遇到训练中断的情况时，可以根据最近保存的检查点进行恢复。但是，大模型场景下的checkpoint的保存和恢复会对存储带来新的挑战。
以Llama2-7B为例， 使用AdamW优化器进行训练，保存出来的一个checkpoint大小为76GB。由此可见，若有集群中有128个训练实例进行70B模型训练恢复时，会同时读取存储服务中的760GB的checkpoint，可想而知存储的压力会很大，总数据读取量可能会有128×760GB≈97TB。如果对进行模型切分，每个读取的checkpoint会变小，但是总数据读取量的量纲也是在TB级别。
那么存储服务的性能会直接影响训练实例的恢复情况，影响GPU有效训练率，因此常用的基于CS架构开发的NFS是不够用的。

Minio+JuiceFS

Minio和JuiceFS是基于云原生的分布式训练平台一个不错的存储选型。这个选型可能不是最优的，但是对于云原生的场景这个选型是兼顾性能和成本的方案。

对象存储（Minio）

云原生+机器学习必然绕不过对象存储，当前华为云、阿里云、百度云存储都是基于对象存储完成的。**为什么对象存储会适合云原生+机器学习场景呢？**下图是一个简单的产品架构图:

• Bucket(桶)：对象存储的基本组成之一，使用来存储object的容器，桶可以将数据进行逻辑上的隔离；
• Object(对象)：对象存储的基本组成之一，数据存储的基本单元，一个对象实际是一个文件的数据与其相关属性信息的集合体，包括Key、Metadata、Data三部分。
• Key：键值，即对象的名称，一个桶里的每个对象必须拥有唯一的对象键值。
• Metadata：元数据，即对象的描述信息，包括系统元数据和用户元数据，这些元数据以键值对（Key-Value）的形式被上传到OBS中。
  • 系统元数据由OBS自动产生，在处理对象数据时使用，包括Date，Content-length，Last-modify，ETag等。
  • 用户元数据由用户在上传对象时指定，是用户自定义的对象描述信息。
• Data：数据，即文件的数据内容。
• API：对象存储的访问形式，一般都会支持亚马逊的s3协议，因为对象存储时亚马逊首先提出的，并制定了相关的业内标准。

Minio是一个开源的高性能对象存储项目，下面结合Minio来说明对象存储的为什么适合云原生+机器学习平台场景？
**高扩展性：**对象存储的特性之一就是可以几乎无限扩容，正如下图Minio的分布式集群架构所示，可以在server内增加磁盘空间，也可以增加server的数量。

• **高可靠性：**如果其中一个或多个磁盘或设备故障，所有集群中的其他磁盘可以进行替代并保证系统照常运行。数据恢复只需要数分钟便可以完成，而且数据可用性不会中断。相反，在传统存储中，当一个RAID磁盘出现故障，系统会慢如蜗牛需要数小时或数天来重建阵列。 Minio采用Reed-Solomon code将对象拆分成N/2数据和N/2 奇偶校验块。 这就意味着如果是12块盘，一个对象会被分成6个数据块、6个奇偶校验块，你可以丢失任意6块盘（不管其是存放的数据块还是奇偶校验块），你仍可以从剩下的盘中的数据进行恢复。
• **负载均衡：**对象存储集群的每个节点都是独立的，访问负载可以平均分配到集群中的所有节点上，避免出现NAS和集群文件系统中常见的资源利用不合理的问题。并且可以让数据读取自动选择合理的节点，保证系统性能最大化。可以从上图中Load Balancer看出这一点。

分布式文件系统（JuiceFS）

JuiceFS是一个开源的高性能、云原生分布式文件系统。有了Minio之后为什么还会需要JuiceFS？
下图为JuiceFS的技术架构：

从上述结构架构图，可以发现JuiceFS是对象存储和机器学习之间的一个桥梁，它有以下优势：

• **提供文件系统：**JuiceFS可以对接对象存储，为用户提供文件系统，机器学习场景下的算法很多都是基于文件系统开发，这就需要一个文件系统的入口，但是对象存储只提供了RestAPI；
• **弹性扩展吞吐：**JuiceFS提供缓存机制，对象存储的数据读取速度是瓶颈之一，大模型场景下的数据读取量是很大的。JuiceFS提供了缓存机制，可以将对象存储里的数据拉取到本地进行缓存，加速用户的频繁读取。而大模型训练任务中的数据集和checkpoint便是经常需要频繁读取的数据类型；
• **云原生：**可以支持云原生，通过kubernetes统一纳管。

由上图可知，写入数据时会将数据真实存放在对象存储中，JuiceFS会维护一套元信息来保证文件系统的持久性和稳定性。

小结

存储的选型没有最优，只有最适合。大模型分布式训练场景下，数据量猛增、对数据读写性能的要求猛增，那么用Minio+JuiceFS的方案可能就是为此而生的。Minio的高可靠性、易扩容特性解决了数据量猛增的问题，JuiceFS提供了弹性吞吐的特性。
参考：

• https://blog.csdn.net/qq_45704048/article/details/139046705
• https://aws.amazon.com/cn/what-is/object-storage/
• https://support.huaweicloud.com/productdesc-obs/zh-cn_topic_0045829060.html?utm_source=obs_Growth_map&utm_medium=display&utm_campaign=help_center&utm_content=Growth_map
• https://www.elibaron.com/arch/minio/minio-erasure-code.html

GPU硬件互联通信

主要讲述NVIDIA GPU硬件上如何进行互联和搭建通信渠道的。

单机内通信（NVLink+NVSwitch）

单机或者单节点的多GPU通信，往往会用到NVLink。NVLink是不同于PCIe的通信协议，它可以直接进行GPU间的通信，绕过CPU处理，减少与通信相关的网络开销，提高计算效率。截止2024年6月，NVLink已经伴随GB200升级到第五代，GB200支持18个100GB/s带宽的NVLink，总带宽可以达到1.8TB/s。

NVSwitch是一种物理芯片，类似于交换机，可通过高速的NVLink接口连接多个GPU，提高服务器内的通信和带宽。NVSWITCH3芯片可以搭载64个NVLink4端口，提供总共12.8 Tbps的单向带宽或3.2 TB/s的双向带宽。下图为NVIDIA官网的照片。

一般的组网形式如下图所示。

多机通信（RDMA）

分布式训练，往往会需要多机间通信，传统TCP/IP通信的效率低，因此需要RDMA（Remote Direct Memory Access，远程直接内存访问）技术来绕过CPU的频繁介入和数据拷贝，增强通信网络性能。

如上图所示，为传统的TCP/IP网络数据传输过程，可以发现数据会在各个层级内流转拷贝。

如上图所示，虚线为TCP/IP的数据传输流程，实线为RDMA传输流程。可以很直观地看出RDMA可以直接绕过CPU和多次数据拷贝，进行数据间的通信和传输。
RDMA技术的主要实现有三个:

• **IB（InfiniBand）：**NVIDIA基于 InfiniBand 架构的 RDMA 技术，需要专用的 IB 网卡和 IB 交换机。从性能上，很明显Infiniband网络最好，但网卡和交换机是价格也很高。
• **RoCE（RDMA over Converged Ethernet）：**基于以太网的 RDMA 技术，也是由 IBTA 提出，可以认为是IB的“低成本解决方案”。RoCE支持在标准以太网基础设施上使用RDMA技术，但是需要交换机支持无损以太网传输，网卡必须是支持RoCE的特殊的NIC，并且存在RoCEv1和RoCEv2两个版本。
• **iWARP（internet Wide Area RDMA Protocal）：**基于 TCP/IP 协议的 RDMA 技术(在现有TCP/IP协议栈基础上实现RDMA技术,在TCP协议上增加一层DDP)，由 IETF 标 准定义。iWARP 支持在标准以太网基础设施上使用 RDMA 技术，而不需要交换机支持无损以太网传输，但服务器需要使用支持iWARP 的网卡。与此同时，受 TCP 影响，性能稍差。

需要注意的是，上述几种协议都需要专门的硬件（网卡）支持。

小结

大模型背景下的分布式训练，多机多卡通信是绕不过的事情。因此为保证训练稳定性，对于硬件组网情况进行监控就必不可少。
参考：

• https://wiki.h3c.com/cn/detail.html?WikiName=RDMA
• https://blog.csdn.net/bandaoyu/article/details/112859853
• https://www.nvidia.cn/data-center/nvlink/

Kubernetes使用GPU

使用kubernets组件集群，进行统一调度和编排不是难事，只需要按照教程一步步安装即可。但是原生的kubernets只能支持CPU和内存的资源分配调度，缺少对于GPU等硬件的调度。
因此，上述本章主要描述kubernetes如何去调度GPU等第三方资源，下面均以GPU为例（其实国内的训练推理加速卡都做了类似的工作）。Kubernetes涉及的技术层级图如下所示：

安装完GPU硬件后，需要安装NVIDIA驱动，然后需要安装NVIDIA Container Toolkit使得docker能够调度GPU，然后再安装NVIDIA Deivce Plugin来支持kuberenetes调度。

NVIDIA Container Toolkit

Kubernets的最终管理的是容器，因此如果需要在容器中使用NVIDIA GPU设备，那么NVIDIA Container Toolkit是必不可少的组件。它的主要作用是将NVIDIA GPU设备挂载到容器中。
支持docker的NVIDIA Container Tookit由如下的组件构成：

• nvidia-docker2
• nvidia-container-runtime
• nvidia-container-toolkit
• libnvidia-container

如下图所示，就是一些NVIDIA Container Toolkit的软件安装包。

安装完NVIDIA Container Toolkit后，就可以用下面这条命令进行GPU的调度

docker run -it --privileged --gpus='"device=0"' {image-id} bash
1

通过NVIDIA官网的图可以了解其中的调用流程

1. 容器启动时， gpu-container-runtime 调用 gpu-containers-runtime-hook；
2. gpu-containers-runtime-hook 根据 --devices 参数，调用 nvidia-container-cli prestart；
3. nvidia-container-cli 根据 --devices ，将GPU设备映射到容器中。 并且将宿主机的Nvidia Driver Lib 的so文件也映射到容器中。 此时容器可以通过这些so文件，调用宿主机的Nvidia Driver。

上述第三步提到的挂载文件具体可以表示为下图

其中/usr/libX/libnvidia*.so是在驱动安装完后就可以看到的，如下图所示。

其中/dev/nvidia*是linux下的逻辑设备文件。

为什么需要挂载/usr/libX/libnvidia*so呢 ?

正如上图所示，CUDA应用在调度的时候，其实会依赖各个不同层级的CUDA。CUDA一般分为CUDA Runtime和CUDA Driver两个版本，而挂载进去的文件/usr/libX/libnvidia*so便是CUDA Driver中的库，而一般镜像中安装的CUDA版本（比如nvidia官网提供的GPU镜像）便是CUDA Runtime版本。

NVIDIA Device Plugin

Kubernets提供了device plugin，允许进行自定义设备资源类型，它的本质是一个gPRC server，Device 插件一般推荐使用 DaemonSet 的方式部署。
当 Device Plugin 向 kubelet 注册后，kubelet 就通过 RPC 与 Device Plugin 交互：

• ListAndWatch() ：让 kubelet 发现设备资源和对应属性，并且在设备资源发生变动的时候接收通
• Allocate() ：kubelet 在创建容器前通过 Allocate来申请相关设备资源。

它的具体工作流程如下：

1. NVIDIA Device plugin 部署到GPU节点上，通过 ListAndWatch 接口，上报注册节点的GPU信息和对应的DeviceID。
2. 当有声明 nvidia.com/gpu 的GPU Pod创建出现，调度器会综合考虑GPU设备的空闲情况，将Pod调度到有充足GPU设备的节点上。
3. 节点上的kubelet 启动Pod时，根据request中的声明调用各个Device plugin 的 allocate接口， 由于容器声明了GPU。 kubelet 根据之前 ListAndWatch接口收到的Device信息，选取合适的设备，DeviceID 作为参数，调用NVIDIA Device Plugin的 Allocate 接口。
4. NVIDIA Device Plugin接收到调用，将DeviceID 转换为 NVIDIA_VISIBLE_DEVICES 环境变量，返回kubelet。
5. kubelet将环境变量注入到Pod， 启动容器。（后面几步其实具体为NVDIA Container Toolkit的功能）
6. 容器启动时， gpu-container-runtime 调用 gpu-containers-runtime-hook。
7. gpu-containers-runtime-hook 根据容器的 NVIDIA_VISIBLE_DEVICES 环境变量，转换为 --devices 参数，调用 nvidia-container-cli prestart。
8. nvidia-container-cli 根据 --devices ，将GPU设备映射到容器中。 并且将宿主机的Nvidia Driver Lib 的so文件也映射到容器中。 此时容器可以通过这些so文件，调用宿主机的Nvidia Driver。

GPU Operator

根据上述可知，如果你做了以下工作：

• 节点上安装nvidia驱动；
• 节点上安装nvidia-docker；
• 集群部署gpu device plugin，用于为调度到该节点的pod分配GPU设备。

便可以调度GPU设备了，但是你可以还需要安装DCCM exporter结合Prometheus输出GPU资源监控信息。为了方便统一安装和维护上述组件，因此NVIDIA开源了一款叫NVIDIA GPU Operator的工具。它由以下组件构成：

• 安装nvidia driver的组件
• 安装nvidia container toolkit的组件
• 安装nvidia devcie plugin的组件
• 安装nvidia dcgm exporter组件
• 安装gpu feature discovery组件

其中nvidia container toolkit组件镜像会自动安装工具，并挂载到宿主机目录，具体目录为** **/usr/local/nvidia。除了上述目录，还会挂载一些别的目录，用于保证NVIDIA Container Toolkit有效性，比如：

• /etc/docker => /etc/docker（仅对docker有效，对containerd需要挂载/etc/containerd）：因为需要修改节点/etc/docker/daemon.json这个文件，然后重启docker服务，NVIDIA Container Toolkit才会生效。
• /run/nvidia => /run/nvidia: 前面我们介绍过，基于容器安装NVIDIA驱动后，节点的/run/nvidia/driver其实是整个driver的rootfs，NVIDIA Container Toolkit需要使用到该目录。
• /var/run => /var/run：重启docker的过程中，需要用到/var/run/docker.sock这个文件，所以该目录需要挂载。

以下为某集群安装的GPU Operator组件截图：

小结

本节主要讲述了如何让Kubernets能够调度到GPU，需要安装哪些组件。从底层驱动到docker层的NVIDIA Container Toolkit，再到kubernets层的device plugin，以及为了统一管理的开源组件GPU Operator。
因此，在分布式训练集群搭建的时候，只需要获取设备厂家的组件安装包进行安装即可。但是组件的服务状态是需要监控的，因为这会直接影响节点是否可以调度。
参考：

• https://developer.aliyun.com/article/784150
• https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/arch-overview.html
• https://github.com/NVIDIA/k8s-device-plugin
• https://developer.aliyun.com/article/784150#slide-10

NUMA感知调度（Topology Manager）

Topology Manager是kubelet的一个组件，在kubernetes 1.16加入，而kubernetes 1.18中该feature变为beta版。如果要使用其实只要在启动项中加入下面参数即可（当然也可以通过配置文件指定）：

--topology-manager-policy=
    [none | best-effort | restricted | single-numa-node]
1
2

如果你对Topology Manager的原理不感兴趣，只需要了解为什么需要Topology Manager。

为什么需要Topology Manager

现代计算机的CPU架构多采用NUMA（Non-Uniform Memory Access，非统一内存）架构。NUMA就是将cpu资源分开，以node 为单位进行分组，每个node都有着独有的cpu、memory等资源，当一个NUMA节点内的资源相交互时，性能将会有很大的提升；但是，如果是两个NUMA节点之间的资源交互将会变得很慢。
下面这幅图中有两个NUMA节点存在：

• NUMA0：由cpu0、cpu1、cpu2、cpu3以及gpu0、nic0和一块本地内存组成
• NUMA1：由cpu4、cpu5、cpu6、cpu7以及gpu1、nic1和一块本地内存组成

假设某个pod需要的资源清单如下：4个CPU、200MB内存、1个GPU和1个NIC。但是由于cpu和其他外围设备（比如GPU）的分配由不同的组件完成，cpu的分配由CPU Manager完成，外围设备由Device Manager完成（正如Kubernetes使用GPU章节所述）。
因此，它们在给pod分配设备时，都是独立工作的，不会有一个全局观念。这会使调度结果很有可能是CPU和内存分配在Node0上，但是GPU却分配在Node1上。

TopologyHint

TopologyHint是Topology Manager中一个概念。它的数据结构定义如下：

type TopologyHint struct {
    NUMANodeAffinity bitmask.BitMask
    Preferred bool
}
1
2
3
4

• bitmask：表示NUMA分配情况
• Preferred：表示这个NUMA节点组合对于某个pod而言是不是“优先考虑的”

举个例子，假设有两个NUMA节点（编号分别为0和1），那么可能出现的组合为：[0]、[1]、[0,1]，用bitmask表示为：01，10，11（从右往左开始，组合中有哪一个NUMA节点，那一位就是1）。
假设现在有两个NUMA节点（编号为0和1），每个NUMA节点上都有两个cpu，如果某个pod需要请求两个cpu，那么TopologyHint有如下几个：

• {01: True}代表从NUMA0上分配两个cpu给pod，这两个cpu都在一个NUMA节点上，涉及的NUMA节点个数最少（为1），所以是“优先考虑的”。
• {10: True}代表从NUMA1上分配两个cpu给pod，这两个cpu也在一个NUMA节点上，涉及的NUMA节点个数也最少（为1），所以是“优先考虑的”。
• {11: False}代表从NUMA0和NUMA1上各取一个cpu，涉及的NUMA节点个数为2，所以不是“优先考虑的”。

HintProviders

HintProvider 是 kubelet 内部的一个组件，用于协调与 TopologyManager 对齐的资源分配。 目前 Kubernetes 中唯一的 HintProvider 是 CPUManager 和 DeviceManager。
比如NVIDIA Device Plugin会注册自己的DeviceManager。
以下是一个多种资源的TopologyHints示例。

               cpu: [{01: True}, {10: True}, {11: False}]
gpu-vendor.com/gpu: [{01: True}, {10: True}]
nic-vendor.com/nic: [{01: True}, {10: True}]
1
2
3

上述每类资源类型生成的TopologyHints进行merge操作，并寻找最优的合并后TopologyHint。

Topology Manager策略

Topology Manager提供了四种策略供用户组合各个资源的TopologyHint。这四种策略是：

• **none：**什么也不做，与没有开启Topology Manager的效果一样。
• **best-effort: **允许Topology Manager通过组合各个资源提供的TopologyHint，而找到一个最优的TopologyHint，如果没有找到也没关系，节点也会接纳这个Pod。
• **restricted：**允许Topology Manager通过组合各个资源提供的TopologyHint，而找到一个最优的TopologyHint，如果没有找到，那么节点会拒绝接纳这个Pod，如果Pod遭到节点拒绝，其状态将变为Terminated。
• **single-numa-node：**允许Topology Manager通过组合各个资源提供的TopologyHint，而找到一个最优的TopologyHint，并且这个最优的TopologyHint所涉及的NUMA节点个数是1。如果没有找到，那么节点会拒绝接纳这个Pod，如果Pod遭到节点拒绝，其状态将变为Terminated。

小结

如果要做千卡训练，那么集群规模会上百，那么出现跨NUMA节点分配情况的可能性会增加，因此将该策略开启，应该是个不错的选择。但是能带来多大的提升，还需要通过测试进行量化比较。
借鉴这个逻辑，大规模集群其实除了NUMA感知还可以将网络拓扑的最优子图获取亲和度调度建议，例如调度前获取整个可用资源的RoCE网络拓扑获取最优调度子图亲和度建议。
参考：

• https://developer.aliyun.com/article/784148

批量调度引擎（Volcano）

分布式训练场景下的调度瓶颈

一个训练任务（job）包含了多个pod，但k8s的默认scheduler是以pod为粒度调度的，不是以job维度调度。这种简单的调度策略在深度学习模型训练场景存在一些缺陷：

• **死锁：**如下图所示，a4和b4会无法调度，造成死锁。

• **资源碎片：**如下图所示，由于job1会申请3个pod，每个pod需要1个GPU，而job2的b1需要2个GPU，但是没有满足的节点。但其实资源是够用的。

• **亲和性与反亲和性：**对于ps架构来说，希望ps和worker尽量调度到同一个node上。如果没有使用亲和性调度策略，会造成网络通信的压力，进而影响计算资源的利用率。因此，job统一调度，增加job内多个pod的亲和度，可以解决该问题。

• **调度公平性：**当2个job包含的pod数存在明显差异时，如果我们对所有用一视同仁的优先级调度，就会出现pod越多，占用资源也越多的问题。

Volcano

Volcano（以前叫kube-batch）是华为开源的、kubernetes生态中高性能计算场景最常用的调度器，能够解决上述调度问题。

PodGroup

PodGroup CRD就是volcano调度的最小单位，它代表了一组强关联的pod集合。所以 VolcanoJob 背后对应一个 K8s 里的自定义控制器（Operator 模式），这个控制器会根据 VolcanoJob 的具体配置去创建相应的 PodGroup 出来。而 PodGroup 最终会被当做一个整体被 Volcano Scheduler 调度。在调度的过程中，Volcano 还用到了 Queue 来实现 PodGroup 的排队、优先级控制等逻辑。

Plugin

volcano通过提供各种plugin能够解决上述提到的瓶颈：

• **Gang Plugin：**解决死锁问题。对PodGroup使用“All or nothing”策略；
• **Binpack Plugin：**解决资源碎片问题。根据资源使用率计算节点权重，在优选阶段为pod选择打分最高的节点；
• **Priority、DRF Plugin：**解决Job调度的公平性问题，支持以fair-share的形式共享资源；
• **Proportion Plugin：**为不同团队划分资源使用比例；
• **Task-topology Plugin：**提供亲和性和反亲和性配置策略。

更多可以查看https://volcano.sh/zh/docs/plugins/
volcano的强可扩展性也提供给我们很大的自由度，去定制符合业务要求的调度策略。
NUMA感知调度中提到的优化调度策略可以通过写Plugin来开发。

Scheduler架构体系

具体流程讲解可以参考https://bbs.huaweicloud.com/blogs/239645。
Scheduler支持动态配置和加载。左边为apiserver，右边为整个Scheduler,apiserver里有Job、Pod、Pod Group，其中Scheduler分为三部分：左边为Cache；中间层为整个调度的过程；右边是以插件形式存在的调度算法。
具体工作流程如下：

1. Cache会将apiserver里创建的Pod、Pod Group这些信息存储并加工为Jobinfors；
2. 中间层的OpenSession会从Cache里拉取Pod、Pod Group，同时将右边的算法插件一起获取；
3. 进行调度工作，依次执行enqueue（入队）、allocate（分配）、preempt（抢占）、reclaim（回收）、backfill（预留）等动作，为每个Job找到一个最合适的节点。

小结

Volcano是一个开源的高性能计算场景最常用的调度器。大模型分布式训练场景下，可以通过自定义plugin来定义最适合该业务的调度策略，提升训练效率。
参考：

• https://laiye.com/news/post/2627.html
• https://bbs.huaweicloud.com/blogs/239645
• https://www.cnblogs.com/daniel-hutao/p/17935624.html

Training Job Operator

为什么会需要自定义资源类型（Operator）？需要先介绍常用的训练框架是如何进行分布式训练，然后再介绍kubeflow提供的training operator。

SOTA框架的分布式训练

PyTorch

PyTorch的DDP使用了Ring AllReduce进行分布式训练，DDP依赖c10d（collective communication）库，提供了gloo、MPI、NCCL等多种通信方式。

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
def setup(rank, world_size):
 os.environ['MASTER_ADDR'] = 'localhost'
 os.environ['MASTER_PORT'] = '12355'
 # initialize the process group
 dist.init_process_group("gloo", rank=rank, world_size=world_size)
def cleanup():
 dist.destroy_process_group()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

正如上述代码所示，PyTorch框架进行分布式训练时，每个训练节点需要获取以下几个参数：

• MASTER_ADDR：Master节点地址
• MASTER_PORT：Master节点暴露的端口
• RANK：当前的节点的标识符
• WORD_SIZE：训练节点总数

Tensorflow

Tensorflow的parameter server分布式训练策略从tf1.X到tf2.x都是有的，需要parameter server和worker，不同的角色会执行不同的初始化逻辑。

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
        "ps": ["host4:port", "host5:port"]
    },
   "task": {"type": "worker", "index": 1}
})
strategy = tf.distribute.experimental.ParameterServerStrategy(
    tf.distribute.cluster_resolver.TFConfigClusterResolver(),
    variable_partitioner=variable_partitioner)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy)
1
2
3
4
5
6
7
8
9
10
11
12

正如上述代码所示，Tensorflow框架在进行分布式训练时，需要注入TF_CONFIG的环境变量，来告知每个训练节点Master的信息、Worker的信息、以及当前节点的类型和index。

Training Operator(Kubeflow)

正如上述提到的，一个分布式训练任务时，所有训练框架都要求使用者通过环境变量或命令行参数传入节点地址和标识符。这些信息会用来互相发现，建立拓扑关系，从而在各个节点内实现参数的传递和同步。
但是这种操作带来诸多不便，每次启动都要传很多参数，不同节点标识符也不一样，使用MPI通信的训练框架还需要配置节点间访问秘钥，如果同时用很多机器训练模型，启动过程会非常繁琐，而这个启动流程完全是可以被自动化的。
因此，Kubeflow开源了基于operator的自定义了一些CRD，用于支持不同框架实现的分布式训练任务的编排自动化。

Kubeflow

Kubeflow 是一种开源的 Kubernetes 原生框架，可用于开发、管理和运行机器学习（ML）工作负载。Kubeflow 是一个 AI/ML 平台，它汇集了多种工具:

• Kubeflow 中央控制面板为访问 Kubeflow 及其生态系统组件提供了一个经过验证的 Web 界面。作为一个集中式中心，它聚合了集群内各种工具和服务的用户界面，为管理机器学习平台提供了一个统一的接入点
• Kubeflow 与 Jupyter Notebooks 集成，为数据探索、实验和模型开发提供了一个交互式环境。Notebooks 支持各种编程语言，包括 Python、R 和 Scala，允许用户以协作且可再现的方式创建和执行 ML 工作流。
• Kubeflow Pipelines 让用户能够以有向无环图（DAG）的形式定义和执行复杂的 ML 工作流。Kubeflow Pipelines 提供了一种方法，可编排并自动执行数据预处理、模型训练、评估和部署的端到端流程，从而促进了 ML 项目的可重现性、可扩展性和协作性。Kubeflow Pipelines SDK 是一组 Python 软件包，允许用户精确而高效地定义和执行机器学习工作流。
• Kubeflow Training Operator 为大规模训练机器学习模型提供了工具。这包括支持使用 TensorFlow、PyTorch 和 XGBoost 等框架进行分布式训练。用户可以利用 Kubernetes 的可扩展性和资源管理功能，跨机器集群高效地训练模型。
• Kubeflow Serving 支持用户将经过训练的 ML 模型部署为可扩展的生产就绪型服务。它为使用 TensorFlow Serving、Seldon Core 等流行框架或自定义推理服务器部署模型提供了一致的界面。模型可在实时或批处理场景中部署，通过 HTTP 端点提供预测。

更多可以查看Kubeflow。

Training Operator

正如上述提到的，Kubeflow Training Operator提供了一系列CRD用于分布式训练任务的自动化。具体支持如下表所示。

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-simple
  namespace: kubeflow
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

从上述示例中，可以发现用户不需要关心master和pod一些信息，可以专注于训练本身相关的工作。

Training Operator with Volcano

Training Operator支持将任务给volcano调度器调度。如官网所描述：在k8s中安装volcano调度器后，并为Training Operator设置一个命令行参数即可，如下所示：

    spec:
      containers:
        - command:
            - /manager
+           - --gang-scheduler-name=volcano
          image: kubeflow/training-operator
          name: training-operator
...
1
2
3
4
5
6
7
8

配置完上述参数后，在集群中发起新的训练任务时会自动被volcano调度。
Training Operator会监听TFJob类型的CRD并创建pod。当设置了volcano调度器时，Training Operator还做了其他事情：

• 创建一个PodGroup类型的对象并应用到集群中。这个PodGroup CRD就是volcano调度的最小单位，它代表了一组强关联的pod集合。在PodGroup的yaml中，支持为其设置最小启动的pod数、最小资源申请量、关联使用的优先级队列。
• 将pod绑定podGroup、并将pod的schedulerName设置为volcano，保证volcano调度器的control loop可以发现并处理这个pod。

小结：

因此，云原生下的分布式训练已经离不开Training Job Operator了，那么在大模型背景下，可以根据自身需要定义自己的Training Job Operator。
参考：

• https://www.kubeflow.org/docs/components/training/user-guides/job-scheduling/
• https://www.kubeflow.org/

分布式训练策略

分布式训练策略中往往绕不过并行策略和分布式计算策略。
大模型场景下的模型规模越来越大，比如 GPT-3 模型的参数量达到1750亿。即使用1024张 80 GB 的 A100，那么完整训练 GPT-3 的时长都需要1个月，因此并行策略显得尤为重要。
知道了如何并行，但是如何去高效地实现这些并行策略，那么分布式计算便是为此而生的。

并行策略

主要借鉴OneFlow的文章进行说明。
简单的机器堆叠并不一定会带来算力的增长。因为神经网络的训练并不是单纯的“把原来一个设备做的事情，现在分给多个设备各自做”，它不仅需要多个设备进行计算，还涉及到设备之间的数据传输，只有协调好集群中的计算与通信，才能做高效的分布式训练。
我们将以矩阵乘法的例子，解释数据并行、模型并行的区别。
先了解以下逻辑上的矩阵乘法例子:
假设神经网络中某一层是做矩阵乘法，其中的输入