目录
1. 引言
在深度学习模型训练过程中，随着模型的复杂度增加和数据量的增多，单机单卡的计算能力逐渐无法满足需求。多机多卡分布式训练可以显著缩短训练时间，并提高训练效率。本文将介绍如何构建一个适合的多机多卡训练环境，利用Kubernetes（K8s）来进行分布式训练资源的管理与调度。
2. 构建镜像
为了进行分布式训练，我们首先需要准备好一个合适的Docker镜像。该镜像应包含训练所需的依赖库和深度学习框架，比如PyTorch、CUDA、NCCL等。
Dockerfile 样例
DockerfileFROM nvidia/cuda:12.1-cudnn8-devel-ubuntu20.04 # 安装必要的依赖项 RUN apt-get update && apt-get install -y \ python3 \ python3-pip \ libnccl2 \ libnccl-dev \ && rm -rf /var/lib/apt/lists/* # 安装PyTorch RUN pip3 install torch==2.3.0 torchvision torchaudio # 复制训练脚本 COPY ./scripts /app/scripts WORKDIR /app/scripts # 设置训练的环境变量 ENV PATH="/app/scripts:${PATH}"
3. Kubernetes 训练任务配置
在Kubernetes中，可以通过PyTorchJob或自定义的Job来调度训练任务。本文以PyTorchJob为例，展示如何在多机多卡环境下启动训练。
3.1 PyTorchJob YAML 文件
yamlapiVersion: "kubeflow.org/v1" # 定义该资源使用的API版本，这里是Kubeflow的PyTorchJob API版本
kind: PyTorchJob # 指定资源类型为PyTorchJob，用于管理分布式PyTorch训练任务。PyTorchJob 是 Kubeflow 中专门用于运行 PyTorch 分布式训练任务的自定义资源类型。是由 Kubeflow 提供的一个标准 API 资源，用于简化在 Kubernetes 上的分布式 PyTorch 任务部署。
metadata: # Kubernetes 资源的元数据信息，包括资源的名称、命名空间以及额外的注释信息
name: pytorch-multinode-job # 该任务的名称，用于唯一标识PyTorchJob资源。在 Kubernetes 中，名称用于唯一标识一个资源对象。
namespace: ai-namespace # 命名空间，用于隔离资源，这里指定该任务在ai-namespace命名空间下运行。Kubernetes 命名空间（namespace）用于在集群中隔离资源对象。
annotations:
scheduling.volcano.sh/queue-name: default # 这是给调度器的注解信息，scheduling.volcano.sh/queue-name 指定了使用 Volcano 调度器，并且任务将放入名为 default 的队列中。Volcano 是一种高性能的 Kubernetes 调度器，特别适合 AI 和深度学习任务的资源调度。这个参数指定了任务的优先级队列，帮助调度器进行任务排序和资源分配。
aijob.cce.baidubce.com/bccl.tracehang.enabled: "true" # 这个注解与百度云的 AIJob 相关，开启了 bccl.tracehang 功能，具体是用来监控和跟踪任务在网络通信中是否发生“hang”或死锁问题。当设置为 "true" 时，这个功能被启用，用于防止分布式训练中的 NCCL 网络通信死锁问题。【这是百度服务器自己的东西】
spec: # 定义该资源的具体规范。spec 是 Kubernetes 中的一个标准字段，用于定义资源的具体规范和配置。在每个 Kubernetes 资源中，spec 部分用于详细说明资源应该如何行为和运行。它决定了所创建资源的期望状态，Kubernetes 会根据这个规范来确保资源达到所定义的状态。在不同的资源类型中，spec 有不同的含义：Pod：在 Pod 的 spec 中，你可以定义容器列表、镜像、环境变量、卷挂载等。PyTorchJob（如你提到的）是一个自定义资源，spec 用来定义 PyTorch 分布式训练任务的配置，例如 Master 和 Worker 节点的数量、镜像、环境变量等。spec 是 Kubernetes 的一个关键字段，用于描述资源的期望行为。
pytorchReplicaSpecs: # 定义不同角色的副本规格，包括Master和Worker。pytorchReplicaSpecs 是 Kubeflow 中 PyTorchJob 的一个特有字段，用于定义分布式 PyTorch 训练任务中各个角色（如 Master 和 Worker）的副本数量和具体配置。
Master: # 定义主节点的配置
replicas: 1 # 主节点的副本数量，这里指定只有一个主节点
restartPolicy: Never # restartPolicy: Never 是 Kubernetes 中 Pod 的重启策略配置。它定义了 Pod 或容器在终止（异常或正常终止）后的重启行为。Always: Pod 中的容器如果终止，无论是正常还是异常，都会被重新启动。这是 Deployment 等资源的默认策略。OnFailure: 只有在容器非正常终止（退出码非0）时才会重新启动。这种策略适合那些任务失败时需要重试的场景。Never: 无论容器是否正常或异常终止，都不会重新启动容器。这通常用于一次性任务，例如批处理任务或训练任务，在任务完成后不需要重新启动。
template: # 模板部分，定义Master节点的Pod配置。template: 是 Kubernetes 中用于定义 Pod 模板的关键字段。
metadata:
spec: # Pod的具体规范
schedulerName: volcano
hostNetwork: true # hostNetwork: true 是 Kubernetes 中 Pod 配置的一个选项，它决定了 Pod 是否使用宿主机（节点）的网络栈。当设置为 true 时，Pod 将直接使用宿主机的网络，而不是创建一个独立的网络命名空间。这意味着：Pod 中的容器将与宿主机共享相同的 IP 地址。Pod 中的容器可以直接访问宿主机上的网络接口。容器内部的端口将与宿主机上的端口一致，因此如果多个 Pod 在同一节点上使用相同的端口号，会导致端口冲突。跟docker run --net host 参数类似。
dnsPolicy: ClusterFirstWithHostNet # dnsPolicy: ClusterFirstWithHostNet 是 Kubernetes 中 Pod 的 DNS 配置策略之一，它与 hostNetwork: true 搭配使用，用于控制 Pod 如何解析 DNS。ClusterFirst：Pod 首先使用集群内的 DNS 服务解析域名，也就是集群中的 Kubernetes DNS 服务（如 CoreDNS 或 kube-dns）。WithHostNet：由于 hostNetwork: true，Pod 直接使用宿主机网络。此策略确保即使 Pod 使用宿主机的网络栈，它依然可以优先使用 Kubernetes 集群内的 DNS 服务，而不会默认回退到宿主机的 DNS 设置。
containers: # 定义运行在Pod中的容器列表
- name: pytorch # 容器名称，这里命名为pytorch。在 Kubernetes 中，一个 Pod 可以包含一个或多个容器，每个容器需要一个名称来唯一标识它。
image: registry.example.com/ai/pytorch-training:latest # 使用的Docker镜像
imagePullPolicy: IfNotPresent # 是 Kubernetes 中用来控制容器镜像拉取策略的一个字段。它决定了 Kubernetes 在启动 Pod 时是否从镜像仓库拉取新的镜像。Always：每次创建 Pod 时，无论本地是否存在镜像，都会尝试从远程镜像仓库拉取最新版本的镜像。IfNotPresent：只有在本地不存在该镜像时，才会从远程仓库拉取镜像。如果本地已经有镜像，则直接使用本地镜像，不再重新拉取。
command: [ "/bin/bash", "-c", "/app/scripts/multi_train_run.sh" ] # 启动容器时执行的命令，运行训练脚本。这个脚本是要自己写的。这个 command 字段用于定义容器启动时运行的指令，它会在容器成功启动后立即执行。当容器启动时，Kubernetes 会按照 command 字段的配置启动容器内部的 Bash shell。Bash shell 执行 multi_train_run.sh 这个脚本，通常是你的训练任务脚本。
env: # 定义容器的环境变量
- name: NUM_PROCESSES # 环境变量NUM_PROCESSES，定义实例数量【需要设置！】
value: "32" # 设定为32个进程，表示节点数*每节点卡数【需要设置！】在 Kubernetes 中，环境变量的设置方式就是通过类似你提供的这种 YAML 结构来指定的。在 Kubernetes 的 Pod 或容器配置中，env 字段用于定义环境变量，name 指定环境变量的名称，value 指定对应的值。
- name: TRAIN_DATA_PATH # 环境变量TRAIN_DATA_PATH，定义训练数据的路径【需要设置！】
value: "/mnt/data" # 训练数据的挂载路径【需要设置！】
- name: MODEL_PATH # 环境变量MODEL_PATH，定义预训练模型的路径【需要设置！】
value: "/mnt/model" # 模型文件的挂载路径【需要设置！】
- name: OUTPUT_PATH # 环境变量OUTPUT_PATH，定义模型输出的路径【需要设置！】
value: "/mnt/output" # 输出结果的挂载路径【需要设置！】
- name: NCCL_DEBUG
value: INFO # 环境变量NCCL_DEBUG，用于控制NCCL库的调试信息，设置为INFO级别，表示输出一般的调试信息
- name: NCCL_IB_DISABLE
value: "0" # 环境变量NCCL_IB_DISABLE，用于控制是否禁用InfiniBand通信，设置为"0"表示不禁用InfiniBand通信，即启用InfiniBand。"0"：表示不禁用 InfiniBand，允许 NCCL 使用 InfiniBand 进行通信。如果集群的硬件支持 InfiniBand（如 RDMA 网络），NCCL 将利用它来加速节点之间的通信。"1"：表示禁用 InfiniBand，强制 NCCL 使用其他通信协议（例如 TCP/IP）来进行 GPU 之间的数据交换。
- name: CUDA_DEVICE_MAX_CONNECTIONS
value: "1" # 用于控制每个GPU的最大连接数，设置为1，表示每个GPU只有一个连接。这个环境变量用于控制每个 GPU 设备与其他设备（例如主机或其他 GPU）之间的最大并行连接数。在 CUDA 中，GPU 可以与主机 CPU 或其他 GPU 进行并行通信。默认值：通常情况下，这个值的默认设置是较高的，用于允许更多的并行连接，以提升数据传输的带宽和效率。"1"：将此值设置为 "1" 意味着每个 GPU 设备只能保持一个连接。这样可以减少 GPU 与其他设备的连接数，限制并行通信。这通常用于特定的调试场景或需要减少并发以防止通信瓶颈的场景。
- name: MASTER
value: "1" # 环境变量MASTER，用于标识主节点，设置为"1"，表示该节点是主节点。分布式训练中通常有一个 Master 节点和多个 Worker 节点：Master 节点：负责管理任务、协调训练过程和处理节点间的通信。Worker 节点：负责实际的计算工作。
resources:
limits: # 在 Kubernetes 的 resources 字段中，limits 用于为容器分配资源上限，确保容器使用的计算资源不会超出设定的范围。
baidu.com/a800_80g_cgpu: 8 # 限制为最多使用 8 个 Baidu AI 加速卡（如 GPU）。这是一个自定义资源名称，表示百度云中的 A800 80GB 计算资源，通常指代一种高性能 GPU 资源（如 A800 GPU，具有 80GB 显存）。
rdma/hca: 1 # 1就是使用，0就是不使用。
securityContext: # securityContext是 Kubernetes 中用于定义容器或 Pod 的安全配置的字段，允许你设置用户 ID、组 ID、权限、Linux 内核功能等。在 Kubernetes 中，securityContext 用于配置 Pod 或容器的安全相关设置。capabilities 是其中一个用于管理容器权限的字段，允许你向容器添加或移除特定的 Linux 功能（capabilities）。IPC_LOCK 是 Linux 内核中一个特定的权限，允许进程锁定内存。
capabilities: # 在 Linux 系统中，capabilities 是一组权限，可以将这些权限细化并授予给容器。相比于给容器授予 root 权限，通过 capabilities 可以更精细地控制容器所具有的特定权限。add 字段允许你向容器中添加特定的权限功能。
add: [ "IPC_LOCK" ] # 这是 Linux 内核中的一个能力，表示允许进程锁定内存。具体来说，IPC_LOCK 允许进程调用 mlock 或 mlockall 系统调用，把进程的某些内存区域锁定在物理内存中，避免它们被交换（swap）到磁盘中。这对于一些高性能、实时应用很重要，因为锁定内存可以避免因交换导致的性能损耗或延迟。
volumeMounts: # 在 Kubernetes 中，volumeMounts 用于将定义的卷（volume）挂载到容器的指定路径。通过 volumeMounts，容器可以访问持久化存储或共享内存区域。
- mountPath: /dev/shm # 这是 Linux 系统中的共享内存（Shared Memory）路径。共享内存是多个进程可以同时访问的内存区域，通常用于在同一节点上的进程之间共享数据，减少内存复制的开销。在深度学习任务或高性能计算中，使用共享内存可以提高不同进程之间的数据传输效率，特别是在多 GPU 或多进程情况下。
name: cache-volume # 这个卷的名称 cache-volume 对应一个具体的卷定义，通常是 emptyDir 类型的卷，它在节点内存中创建。挂载到 /dev/shm 意味着使用这个卷作为容器的共享内存空间，提供更快的 I/O 访问速度。
- name: models # 这个卷的名称 models，通常代表的是一个用于存放模型文件或其他数据的持久化卷。这个卷可能是某种持久化存储（如 PersistentVolume），以便训练任务可以从此路径读取预训练模型或保存模型的输出。
mountPath: /mnt/ # /mnt/ 是 Linux 系统中常见的挂载路径，通常用于临时存储设备的挂载。在这里，容器将通过这个路径访问某个数据卷。
volumes: # volumes 的定义属于 Kubernetes 中的存储资源配置，但它不同于我们通常所指的 CPU、内存等计算资源配置。volumes 主要用于定义 存储资源，即如何在 Pod 内使用不同类型的卷（Volume）来存储数据。它是 Kubernetes 中用于容器数据存储的一种资源定义方式。
- emptyDir: # emptyDir 卷是 Kubernetes 中的一种临时卷。当 Pod 启动时，Kubernetes 在节点上为 Pod 分配一个空的存储目录（或在内存中），当 Pod 停止运行时，emptyDir 卷中的数据会被删除。
medium: Memory # 这是 emptyDir 卷的一个特殊设置，指定该卷是存储在节点的内存中，而不是磁盘。这意味着它会使用系统的 RAM，提供非常快速的数据存取，但容量受限于可用的内存。
sizeLimit: 20G # 该设置限制了这个内存卷的最大大小为 20GB。如果这个卷存储的数据超过 20GB，可能会导致数据存储失败。这个限制用于防止该内存卷过度占用系统的 RAM。
name: cache-volume # cache-volume 被定义为一个存储在内存中的卷，主要用于缓存数据。使用内存作为存储介质可以提高读取和写入的速度，特别适合高性能应用或需要频繁读写的缓存场景。由于它是临时卷，当 Pod 停止运行时，数据会被清除，因此不适合存储重要的持久化数据。
- name: models
persistentVolumeClaim: # persistentVolumeClaim 是 Kubernetes 中的一种持久化存储方式，用于请求持久存储卷（PersistentVolume，PV）。PVC 会根据指定的存储需求，绑定到集群中可用的持久卷上，这样即使 Pod 停止或迁移，数据也不会丢失。
claimName: cloud-ai-test-pvc-1 # claimName 用于引用已经定义的 PVC，cloud-ai-test-pvc-1 是 PVC 的名称。该 PVC 绑定了一个持久化卷（PV），用于存储模型或其他需要在容器重启或迁移后仍然保留的数据。
Worker: # 定义Worker节点的配置
replicas: 3 # Worker节点的副本数量，这里指定3个Worker节点
restartPolicy: Never
template: # 模板部分，定义Worker节点的Pod配置
metadata:
spec: # Pod的具体规范
schedulerName: volcano
hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet
containers: # 定义运行在Pod中的容器列表
- name: pytorch # 容器名称，同样命名为pytorch
image: registry.example.com/ai/pytorch-training:latest # 使用与Master节点相同的Docker镜像
imagePullPolicy: IfNotPresent
command: [ "/bin/bash", "-c", "/app/scripts/multi_train_run.sh" ] # 启动容器时执行的命令，运行训练脚本
env: # 定义容器的环境变量
- name: NUM_PROCESSES # 环境变量NUM_PROCESSES，定义实例数量
value: "32" # 设定为32个进程，和Master节点相同
- name: TRAIN_DATA_PATH # 环境变量TRAIN_DATA_PATH，定义训练数据的路径
value: "/mnt/data" # 训练数据的挂载路径，与Master节点相同
- name: MODEL_PATH # 环境变量MODEL_PATH，定义预训练模型的路径
value: "/mnt/model" # 模型文件的挂载路径，与Master节点相同
- name: OUTPUT_PATH # 环境变量OUTPUT_PATH，定义模型输出的路径
value: "/mnt/output" # 输出结果的挂载路径，与Master节点相同
- name: NCCL_DEBUG
value: INFO
- name: NCCL_IB_DISABLE
value: "0"
- name: CUDA_DEVICE_MAX_CONNECTIONS
value: "1"
resources:
limits:
baidu.com/a800_80g_cgpu: 8
rdma/hca: 1
securityContext:
capabilities:
add: [ "IPC_LOCK" ]
volumeMounts:
- mountPath: /dev/shm
name: cache-volume
- name: models
mountPath: /mnt/
volumes:
- emptyDir:
medium: Memory
sizeLimit: 20G
name: cache-volume
- name: models
persistentVolumeClaim:
claimName: cloud-ai-test-pvc-1
Volcano 调度器是 Kubernetes 的一个高性能批处理调度器，专门用于处理大规模、高密度的计算任务，特别适合用于 AI、深度学习、机器学习、HPC（高性能计算）等场景。它为这些计算密集型任务提供了比 Kubernetes 默认调度器更优化的资源分配和管理能力。
Volcano 调度器的主要特点：
-
批处理任务优化：
Volcano 专为批处理任务设计，能够有效地调度和管理多个大规模作业的运行，确保高效使用集群资源。相比 Kubernetes 默认的调度器，它更加适合处理大量短期、高负载的任务。 -
队列管理：
Volcano 支持任务队列化，可以根据队列优先级来调度不同的任务。通过队列管理，任务可以按优先级依次调度，帮助提升资源的利用率和任务执行的效率。 -
任务依赖和任务拓扑结构：
Volcano 支持复杂任务的调度，包括那些存在任务依赖（例如某些任务必须在另一个任务完成后执行）或者需要特定拓扑结构（例如某些任务需要在特定节点或位置运行）的任务。 -
弹性调度和抢占式调度：
Volcano 可以根据任务的优先级进行资源的动态调整，支持抢占式调度。当更高优先级的任务需要资源时，它可以暂停或终止低优先级任务，并将资源重新分配给更重要的任务。 -
资源公平性：
Volcano 提供了资源公平分配的策略，确保多个任务或用户之间的资源分配是公平的，避免某个任务或用户占用过多资源。 -
作业监控和故障恢复：
Volcano 支持对作业的实时监控，并能够处理任务失败的情况，比如通过重新调度或重启任务，确保作业能够成功完成。
Volcano 的应用场景：
- 深度学习和机器学习： 当多个深度学习作业并行运行时，Volcano 可以有效调度 GPU 资源，优化资源利用率。
- 高性能计算（HPC）： Volcano 适用于需要大量计算资源的高性能计算任务，支持复杂的任务依赖。
- 批处理作业： 大规模的数据处理、计算任务可以通过 Volcano 实现批量调度，提升集群计算效率。
Volcano 是 Kubernetes 上专为计算密集型场景设计的调度器，可以为深度学习模型训练等场景中的资源调度和管理带来显著的性能提升。
3.2 启动训练任务
将上述PyTorchJob YAML文件提交给Kubernetes：
bashkubectl apply -f pytorch-job.yaml
以下是一些你可能会用到的基本命令：
-
- 查看 Pod 状态 查看集群中正在运行的所有 Pod 信息（包括名称、状态、节点等）：
bashkubectl get pods
如果你想查看特定命名空间下的 Pod：
bashkubectl get pods -n <namespace>
-
- 查看 Pod 的详细信息 获取某个特定 Pod 的详细信息，包括资源使用、容器状态、事件等：
bashkubectl describe pod <pod-name>
如果 Pod 位于特定的命名空间：
bashkubectl describe pod <pod-name> -n <namespace>
-
- 查看容器日志 查看某个 Pod 内容器的日志：
bashkubectl logs <pod-name>
如果该 Pod 内有多个容器，可以指定容器名查看该容器的日志：
bashkubectl logs <pod-name> -c <container-name>
-
- 进入某个容器
通过
kubectl exec命令进入容器中的 shell，会话类似于 SSH 登录：
- 进入某个容器
通过
bashkubectl exec -it <pod-name> -- /bin/bash
如果 Pod 内有多个容器，指定要进入的容器：
bashkubectl exec -it <pod-name> -c <container-name> -- /bin/bash
如果容器内没有 bash，可以尝试 sh：
bashkubectl exec -it <pod-name> -- /bin/sh
-
- 查看资源的 YAML 配置 如果想查看某个 Pod 或其他资源的详细配置，可以使用以下命令：
bashkubectl get pod <pod-name> -o yaml
也可以查看其他资源的 YAML，例如 deployment 或 service：
bashkubectl get deployment <deployment-name> -o yaml
-
- 删除某个 Pod 如果需要删除某个 Pod：
bashkubectl delete pod <pod-name>
-
- 查看 Pod 的事件或故障 如果 Pod 出现异常，查看事件可以帮助你了解具体问题：
bashkubectl describe pod <pod-name>
-
- 监控 Pod 状态的变化
你可以使用
watch命令实时查看 Pod 的状态变化：
- 监控 Pod 状态的变化
你可以使用
bashkubectl get pods --watch
4. 训练启动脚本解读
4.1 参数解释
NUM_PROCESSES：实例数量，通常等于节点数量×每节点GPU数量。WORLD_SIZE：分布式训练的总节点数，系统自动配置。MASTER_ADDR：主节点地址，由系统自动配置。MASTER_PORT：主节点端口，由系统自动配置。RANK：当前节点的序列号，主节点默认为0，系统自动配置。
4.2 脚本示例
bash#/bin/bash
accelerate launch \
--num_processes ${NUM_PROCESSES} \
--num_machines ${WORLD_SIZE} \
--rdzv_backend c10d \
--main_process_ip ${MASTER_ADDR} \
--main_process_port ${MASTER_PORT} \
--machine_rank ${RANK} \
/app/lora-scripts/sd-scripts/sdxl_train.py \
--pretrained_model_name_or_path ${MODEL_PATH} \
--train_data_dir ${TRAIN_DATA_PATH} \
--output_dir ${OUTPUT_PATH}
5. 结语
通过Kubernetes管理多机多卡训练任务，可以轻松实现模型的分布式训练，从而充分利用集群资源，提升训练效率。结合容器化和Kubernetes的自动调度能力，大大简化了训练环境的搭建和管理。
如果对你有用的话，可以打赏哦
打赏
本文作者:Dong
本文链接:
版权声明:本博客所有文章除特别声明外，均采用 CC BY-NC。本作品采用《知识共享署名-非商业性使用 4.0 国际许可协议》进行许可。您可以在非商业用途下自由转载和修改，但必须注明出处并提供原作者链接。 许可协议。转载请注明出处！
目录