[译] 对于 512 个以上的 GPU 作业，Kubernetes 还是 Slurm 是最佳编排工具？

, , ,

写在前面

之前整理了 nebius 的一篇AI HPC 场景下的 Slurm 和 K8S，最近重拾 k8s 调度的脏活任务，遂把 stas00 在 Machine Learning Engineering Open Book 推荐的这篇相关文章的翻译也搬运顺便回顾下。

介绍

更多的参数和更大的训练集意味着需要更多的 GPU。但更多的 GPU 也意味着在管理、扩展和编排工作负载方面需要更多的开销。请考虑以下几点：

  1. 资源分配 ：您必须高效地将计算资源分配给各种作业和任务。这不仅关乎确保任务能够运行，更关乎确保 GPU 得到充分利用，避免资源过度分配，并防止不同作业之间发生资源争用。那么，哪些作业将在何时何地运行呢？您需要保持资源的繁忙状态以减少空闲时间，处理作业优先级、抢占和公平性问题，并管理作业依赖关系和工作流程。
  2. 容错性 ：GPU 会发生故障。编排器可以帮助工作负载在组件发生故障时继续正常运行。这包括作业恢复、处理节点故障以及确保高可用性。
  3. 可扩展性 ：我们希望从个位数的 GPU 集群扩展到数百甚至数千个。编排器通过管理节点和 GPU，并根据需要动态地扩展或缩减资源来应对这种不断增长的工作负载。精心设计的系统还能最大限度地减少这些大规模操作中的延迟和开销。

但编排的最佳方案是什么？两大平台始终名列前茅。Slurm 是高性能计算 (HPC) 作业的传统编排方式，而 Kubernetes 则是支撑现代 Web 的云原生容器化平台。本文将展示如何使用这两个平台设置大规模 GPU 作业，并阐述我们为何认为 Kubernetes 是构建人工智能的理想平台。

Slurm：传统，但依然强劲

如果你从事过高性能计算方面的工作，你就会接触到 Slurm 。它被大多数超级计算机用于调度工作负载，是基于 Linux 的高性能非分布式计算的核心实用程序。

但这一切听起来好像是为上世纪 80 年代计算机科学系的研究人员准备的。其实不然。Slurm 完全能够运行 21 世纪的大规模人工智能工作负载。它的可扩展性非常出色。这个缩写甚至从“Simple Linux Utility for Resource Management”（简单 Linux 资源管理工具）改成了“Scalable Linux Utility for Resource Management”（可扩展 Linux 资源管理工具）。正如一位 Reddit 用户所说：

我认为，Kubernetes 并不适合人工智能/机器学习训练，尤其是在分布式训练等方面。而 SLURM 在我们拥有 6000 个 GPU、近 400 个用户的集群上运行良好。

那么，是什么让 usnus 如此确信 Slurm 是 AI/ML 训练的最佳选择呢？Slurm 的开发者们在高性能计算领域深耕多年，不断优化工具，使其能够精准满足大型工作负载的需求：可扩展性、资源管理和作业调度。

Slurm 的核心是 Slurm 控制器 (slurmctld)。这是一个中央管理守护进程，负责协调资源分配。而每个节点都有自己的守护进程 slurmd，负责执行作业。

然后，您可以通过配置文件、脚本和客户端命令的组合来控制这些。假设我们要构建一个包含 512 个 GPU 的集群，我们需要做什么？

首先，我们需要在 slurm.conf 中定义我们的节点。假设每个节点有 8 个 GPU，那么我们需要 64 个节点才能达到 512 个 GPU。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# slurm.conf

# Basic Slurm configuration
ClusterName=ai_cluster
ControlMachine=slurm-controller
SlurmUser=slurm
SlurmctldPort=6817
SlurmdPort=6818

# Node definitions
NodeName=gpu-node[001-064] Gres=gpu:8 CPUs=32 RealMemory=128000 Sockets=2 CoresPerSocket=16 ThreadsPerCore=2 State=UNKNOWN

# Partition definitions
PartitionName=gpu_partition Nodes=gpu-node[001-064] Default=YES MaxTime=INFINITE State=U

在此，我们定义了节点 gpu-node001 至 gpu-node064，每个节点包含 8 个 GPU、32 个 CPU 和 128 GB 内存。我们还创建了一个名为 gpu_partition 的分区，其中包含所有 GPU 节点。

这里的关键在于 Gres =gpu:8。Gres 是 Slurm 处理通用资源（例如 GPU）的方式。我们可以在 gres.conf 文件中配置这些：

1
2
3
# /etc/slurm/gres.conf

Name=gpu Type=h100 File=/dev/nvidia[0-7]

这告诉 Slurm 我们的 GPU 在哪里，以及它们是 H100 系列显卡。配置完成后，我们就可以设置运行了。为此，我们编写了一个作业脚本，并使用 srun 命令运行它。在脚本中，我们使用 SBATCH 命令来请求资源并定义作业参数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#!/bin/bash
#SBATCH --job-name=AI_Training
#SBATCH --partition=gpu_partition
#SBATCH --nodes=64                # 64 nodes x 8 GPUs = 512 GPUs
#SBATCH --ntasks-per-node=8       # One task per GPU
#SBATCH --gres=gpu:8              # Request 8 GPUs per node
#SBATCH --cpus-per-task=4         # CPUs per task (adjust as needed)
#SBATCH --mem=120G                # Memory per node
#SBATCH --time=48:00:00           # Maximum job runtime
#SBATCH --output=logs/output_%j.txt
#SBATCH --error=logs/error_%j.txt

# Load modules, environments, variables
...

# Run the training script using srun
srun --mpi=pmix_v3 python train.py

该脚本请求必要的资源并配置训练任务的作业参数。让我们来详细了解一下每个部分的作用：

  1. 资源分配 ：该脚本请求 64 个节点（–nodes=64），每个节点配备 8 个 GPU（–gres=gpu:8），共计 512 个 GPU。它为每个节点分配 8 个任务（–ntasks-per-node=8），每个 GPU 对应一个任务，每个任务使用 4 个 CPU（–cpus-per-task=4）。每个节点分配 120GB 内存（–mem=120G）。
  2. 作业配置 ：设置作业名称（–job-name=AI_Training），指定要使用的分区（–partition=gpu_partition），并设置最大运行时间为 48 小时（–time=48:00:00）。
  3. 作业执行 ：srun 命令用于在所有已分配的节点上启动训练脚本 (train.py)。

在最后一部分中，--mpi=pmix_v3 选项指定使用 PMIx 进行进程间通信。这在训练需要跨多个 GPU 和节点同步参数的深度学习模型时尤为重要。

通过这种配置，Slurm 将：

  1. 预留集群中所需的资源。
  2. 将工作负载均匀分配到所有 512 个 GPU 上。
  3. 管理分布式训练期间的节点间通信以实现同步。
  4. 处理任何节点故障或作业中断，如果已配置，则可能从上一个检查点重新启动作业。
  5. 工作完成后，提供详细的会计和使用统计数据。

很好。现在你只需要运行它即可：

1
sbatch submit_job.sh

Slurm 随后会将作业排队，并在请求的资源可用时执行该作业。

Slurm 提供了大量的配置和微调选项，可用于优化集群性能。您可以调整作业优先级、公平共享设置和抢占策略等参数，以确保不同用户组和工作负载都能高效利用资源。

例如，为了确保高效利用 512 个 GPU，您可以启用回填功能，以便将空闲资源用于较小的作业。您只需修改 slurm.conf 文件即可：

1
2
3
4
# slurm.conf continued

SchedulerType=sched/backfill
SchedulerParameters=bf_continue,bf_max_job_test=1000,bf_interval=30

其中：

  • bf_continue：在优先级更高的作业开始后继续回填。
  • bf_max_job_test：考虑用于回填的职位数量。
  • bf_interval：回填调度频率

这足以说明 Slurm 功能强大，但它也存在一些问题。

  1. 学习曲线 ：如果您不习惯使用 Linux 命令行，那么 Slurm 丰富的配置选项和命令行界面可能会让您感到不知所措。配置错误的可能性很大，如果您在扩展工作负载的同时也扩展了团队，那么让他们熟悉 Slurm 将会是一个难题。
  2. 对容器的支持有限 ：Slurm 可以与容器配合使用，但它对容器的原生支持程度不如 Kubernetes 等平台。如果您只处理本地模型，这倒不是什么大问题。但如果您要将 AI 产品化，则需要容器化工作负载所具备的可移植性和可复现性。
  3. 缺乏自动扩缩容 ：上述配置文件承担了 Slurm 的所有繁重工作。您在其中设置的内容就是 Slurm 实际使用的内容。Slurm 本身并不提供基于需求自动扩缩容的默认支持，而这对于波动较大的 AI 工作负载尤为重要。这可能导致在低需求时期资源利用率低下，或在高峰时段出现资源短缺。
  4. 与现代工具集成 ：如果您正在构建生产环境应用，那么您使用的工具和平台都是为云原生环境设计的。将这些工具与基于 Slurm 的高性能计算 (HPC) 环境集成需要投入大量精力或采用定制解决方案，这会减慢您的开发周期。

如果人工智能仍然局限于实验室超级计算机上构建的本地机器学习模型，那么 Slurm 就足够了。但人工智能已经走出实验室，进入互联网，因此我们需要一个网络规模的编排平台。

Kubernetes：未来发展方向及支持增长点

Kubernetes 的优势在于 Web 规模。在人工智能领域，或许证明 Kubernetes 实用性的最佳例证是 OpenAI 已将其在 Kubernetes 上的运营规模扩展到单个集群中的 7,500 个节点 （至少，因为该案例研究来自 2021 年）。

如果我们回到“为什么要进行编排”这个问题，我们就可以在 Kubernetes 中找到答案。

  1. 资源分配 。与 Slurm 配置文件类似，Kubernetes 使用 Pod 规范中的资源请求和限制来声明所需的 CPU、内存和 GPU 资源量。借助 Kubernetes，您可以通过设置配额，利用命名空间和资源配额来管理不同团队或项目之间的资源。
  2. 作业调度 。这可能是 Kubernetes 落后于 Slurm 等高性能计算框架的一个方面。Kubernetes 可以配置为批量作业，但需要自定义调度器或扩展。您可以使用 Kubernetes Jobs 和 CronJobs 进行批量处理。
  3. 容错性是 Kubernetes 的一大优势。它能够检测节点故障并重启故障的 Pod，从而维持系统状态。 检查点机制自 v1.30 版本起就已可用。
  4. 可扩展性是 Kubernetes 的另一项核心优势。通过水平 Pod 自动扩缩容 ，您可以根据资源利用率自动调整 Pod 副本数。通过集群自动扩缩容 ，您可以根据工作负载需求添加或移除节点。正如 OpenAI 所展示的，Kubernetes 也支持包含数千个节点的集群。

Kubernetes 并非具备构建完美高 GPU 任务所需的一切。Kubernetes 生态系统正积极朝着这个目标努力。随着 Kubernetes 融入 AI 生态系统，各个团队正在开发各种工具，帮助企业构建更完善的 AI 工作流程。

NVIDIA 就是一个很好的例子。过去几年，NVIDIA 见证了 AI 开发人员如何使用 Kubernetes 以及他们必须克服的种种困难：

  • 异构节点软件栈 ：管理集群中不同的 GPU 类型、驱动程序版本和操作系统版本是一项挑战。
  • 驱动程序配置和管理 ：在每个节点上安装、配置和更新 GPU 驱动程序是一个手动且容易出错的过程。
  • GPU 分区和共享 ：为不同的共享策略（如 MIG、时间切片或 MPS）配置 GPU 非常复杂，需要节点级配置。
  • 容器运行时配置 ：设置容器运行时以与 GPU 一起工作需要手动配置和钩子。
  • GPU 健康监控 ：没有可靠的解决方案来监控 GPU 健康并在出现问题时采取措施。
  • 互操作性问题 ：随着集群的增长和组件的更新，GPU 软件堆栈不同部分之间的互操作性问题变得越来越常见。

因此，NVIDIA 开发了 GPU Operator 来帮助解决这些问题。接下来，我们将介绍如何使用它来运行一个包含 512 个以上 GPU 的作业。首先，我们可以使用 Helm 安装 NVIDIA GPU Operator：

1
2
3
4
helm repo add nvidia https://nvidia.github.io/gpu-operator
helm repo update
kubectl create namespace gpu-operator
helm install gpu-operator nvidia/gpu-operator -n gpu-operator

大多数情况下， gpu-operator 的默认值就足够了。如果需要自定义设置，可以使用 custom-values.yaml 文件来优化性能和资源管理。

例如，您可以利用 NVIDIA 在 A100 及更新的 GPU 中引入的 MIG（ 多实例 GPU） 功能。它允许将单个物理 GPU 划分为多个独立的 GPU 实例，每个实例都拥有专用的内存、缓存和计算核心。这意味着单个 GPU 可以同时服务于多个用户或进程，从而有效提高 GPU 资源的利用率和效率。

定义 MIG 配置允许我们指定如何对每个 GPU 进行分区。例如，我们可以将每个 GPU 分区为七个大小相等的实例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
devicePlugin:
  enabled: true
  migStrategy: mixed

migManager:
  enabled: true
  config:
    create: true
    name: "custom-mig-config"
    data: |-
      config.yaml: |-
        version: v1
        mig-configs:
          all-disabled:
            - devices: all
              mig-enabled: false
          custom-mig:
            - devices: [0]
              mig-enabled: true
              mig-devices:
                "1g.5gb": 7

通过使用 MIG 对 GPU 进行分区，我们可以有效地增加可用 GPU 资源的数量。我们可以并行运行更多进程，从而显著加快受益于分布式计算的模型的训练速度。我们还可以使用资源匹配，为每个工作负载分配适量的 GPU 资源，防止 GPU 资源利用率不足。

然后我们可以应用这些值：

1
helm upgrade gpu-operator nvidia/gpu-operator -n gpu-operator --values custom-values.yaml

（您也可以将 MIG 与 Slurm 结合使用）

然后我们可以使用 kubeflow 训练操作符来运行我们的分布式作业，该作业包含 64 个 pod，每个 pod 请求 8 个 GPU（这里使用的是 PyTorch ，但您也可以使用 Tensorflow ）：

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
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: gpu-training-job
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: your-registry/your-gpu-app:latest
            resources:
              limits:
                nvidia.com/gpu: 8
    Worker:
      replicas: 63
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: your-registry/your-gpu-app:latest
            resources:
              limits:
                nvidia.com/gpu: 8

这看起来可能比 Slurm 的设置更复杂，在某些方面也确实如此。但这种模式能带来灵活性、可扩展性和可移植性。

通过修改容器规格和资源请求，此配置可以轻松适应不同的 GPU 类型，甚至异构集群。Kubernetes 配置的声明式特性使得整个基础架构的复制变得轻而易举。您可以将此配置迁移到不同的云提供商，并与团队成员轻松共享和协作开发代码。

Kubernetes 上 AI 工作负载的未来

在 Fluidstack，我们认为 Kubernetes 是高级 AI 工作负载的理想编排平台。Kubernetes 的声明式特性让您可以灵活地切换不同的服务商，并根据需求轻松地扩展或缩减工作负载。这也意味着您可以在同一基础设施上运行训练和推理。

如果您已经在使用 Slurm，那很好。它拥有非常坚实的基础，在可预见的未来仍能继续提供高性能的作业。但它并非未来发展的方向。Kubernetes 才是最适合扩展 AI 工作负载的平台，尤其是在您需要数百甚至数千个 GPU 的情况下。

随着人工智能的发展，Kubernetes 也将随之发展，为云原生人工智能应用提供基础架构，正如 Slurm 为高性能计算应用奠定基础一样。NVIDIA、OpenAI 和 Fluidstack 等组织将不断壮大生态系统和工具，为人工智能开发者提供构建最佳模型所需的全部资源。