介绍 JobSet

作者：Daniel Vega-Myhre (Google)，Abdullah Gharaibeh (Google)，Kevin Hannon (Red Hat)

在本文中，我们介绍 JobSet，这是一个用于表示分布式作业的开源 API。JobSet 的目标是为 Kubernetes 上的分布式机器学习训练和 HPC 负载提供统一的 API。

为什么需要 JobSet？

Kubernetes 社区最近对 Kubernetes 批处理生态系统的增强吸引了机器学习工程师，他们发现这非常适合运行分布式训练负载的需求。

大型机器学习模型（特别是 LLM）无法装入单个主机上的 GPU 或 TPU 芯片内存中，通常会分布在数万个加速器芯片上，而这些芯片又可能跨越数千个主机。

因此，模型训练代码通常被容器化并在所有这些主机上同时执行，执行分布式计算，通常将模型参数和/或训练数据集分片到目标加速器芯片上，使用 all-gather 和 all-reduce 等通信集合原语来执行分布式计算并在主机之间同步梯度。

这些工作负载特性使 Kubernetes 非常适合此类工作负载，因为它擅长在计算资源集群中高效地调度和管理容器化应用程序的生命周期。

它还具有很强的可扩展性，允许开发人员定义自己的 Kubernetes API、对象和控制器来管理这些对象的行为和生命周期，从而使工程师能够开发定制的分布式训练编排解决方案以满足其需求。

然而，随着分布式机器学习训练技术的不断发展，现有的 Kubernetes 原语已不足以单独对其进行建模。

此外，Kubernetes 分布式训练编排 API 的格局变得分散，而且这个分散格局中的每个现有解决方案都存在某些局限性，使其不适用于分布式机器学习训练。

例如，KubeFlow 训练 Operator 为不同的框架（例如 PyTorchJob、TFJob、MPIJob 等）定义了自定义 API；然而，这些作业类型实际上是专门针对目标框架的解决方案，每种都有不同的语义和行为。

另一方面，Job API 修复了运行批处理工作负载的许多差距，包括 Indexed 完成模式、更高的可扩展性、Pod 失败策略和 Pod 退避策略等等，这只是最近的一些增强功能。然而，使用上游 Job API 运行机器学习训练和 HPC 工作负载需要额外的编排来填补以下差距

多模板 Pod：大多数 HPC 或机器学习训练作业包含不止一种类型的 Pod。不同的 Pod 是同一工作负载的一部分，但它们需要运行不同的容器、请求不同的资源或具有不同的失败策略。一个常见的例子是驱动程序-工作程序模式。

作业组：大规模训练工作负载跨越多个网络拓扑，例如跨多个机架运行。此类工作负载对网络延迟敏感，旨在将通信本地化并最小化跨越高延迟网络链路的流量。为了实现这一点，需要将工作负载拆分为多个 Pod 组，每个组分配给一个网络拓扑。

Pod 间通信：创建和管理建立作业 Pod 之间通信所需的资源（例如 无头服务）。

启动顺序：一些作业需要特定的 Pod 启动顺序；有时驱动程序需要先启动（如 Ray 或 Spark），在其他情况下，工作程序需要在驱动程序启动前准备好（如 MPI）。

JobSet 旨在利用 Job API 作为构建块，为大规模分布式 HPC 和机器学习用例构建更丰富的 API 来解决这些差距。

JobSet 的工作原理

JobSet 将分布式批处理工作负载建模为一组 Kubernetes Job。这使用户可以轻松地为不同的 Pod 组（例如 leader、worker、参数服务器等）指定不同的 Pod 模板。

它使用 ReplicatedJob 的抽象来管理子 Job，其中 ReplicatedJob 本质上是一个指定了所需 Job 副本数量的 Job 模板。这提供了一种声明性的方式来轻松创建相同的子 Job 以在不同的加速器孤岛上运行，而无需借助脚本或 Helm chart 来生成具有不同名称的同一 Job 的多个版本。

JobSet 的其他一些解决上述问题的关键特性包括

Replicated Jobs：在现代数据中心中，像 GPU 和 TPU 这样的硬件加速器分配在同构加速器的孤岛中，通过专门的高带宽网络链接连接。例如，用户可能会配置包含一组位于同一机架上的主机的节点，每个主机都配备 H100 GPU，其中每个主机内的 GPU 芯片通过 NVLink 连接，并由 NVLink Switch 连接多个 NVLink。TPU Pod 是另一个例子：TPU ViperLitePods 由 64 个主机组成，每个主机连接 4 个 TPU v5e 芯片，所有芯片通过 ICI 网格连接。当在多个这样的孤岛上运行分布式训练作业时，我们通常希望将工作负载划分为一组较小的相同作业，每个孤岛一个，其中每个 Pod 主要与同一孤岛内的 Pod 通信以进行分布式计算的分段，并将 DCN（数据中心网络，其带宽低于 ICI）上的梯度同步保持在最低限度。

自动创建、配置和生命周期管理无头服务：默认启用通过 Pod 主机名进行 Pod 间通信，并自动配置和管理实现此功能的无头服务的生命周期。

可配置的成功策略：JobSet 具有可配置的成功策略，这些策略针对特定的 ReplicatedJob，并带有运算符以针对其子作业的“Any”或“All”。例如，您可以配置 JobSet，使其仅在“worker” ReplicatedJob 的所有 Pod 都完成时才标记为完成。

可配置的失败策略：JobSet 具有可配置的失败策略，允许用户指定在发生故障时 JobSet 应重新启动的最大次数。如果任何作业被标记为失败，整个 JobSet 将被重新创建，从而允许工作负载从最后一个检查点恢复。当未指定失败策略时，如果任何作业失败，JobSet 将直接失败。

每个拓扑域的独占放置：JobSet 允许用户表示子作业与拓扑域（通常是像机架一样的加速器孤岛）具有 1:1 的独占分配。例如，如果 JobSet 创建两个子作业，那么此功能将强制每个子作业的 Pod 位于同一孤岛上，并且每个孤岛只允许调度一个子作业。这对于我们希望使用分布式数据并行（DDP）训练策略来使用多个计算资源孤岛（GPU 机架或 TPU 切片）训练模型的场景很有用，在每个加速器孤岛中运行 1 个模型副本，确保前向和后向传播本身在单个模型副本内通过连接孤岛内加速器芯片的高带宽互连进行，并且只有模型副本之间的梯度同步通过较低带宽的数据中心网络跨加速器孤岛进行。

与 Kueue 集成：用户可以通过 Kueue 提交 JobSet 来超额订阅其集群、将工作负载排队以在容量可用时运行、防止部分调度和死锁、启用多租户等。

用例示例

使用 Jax 在多个 TPU 切片上进行分布式机器学习训练

以下示例是一个 JobSet 规范，用于在 4 个 TPU v5e 切片上运行 TPU 多切片工作负载。要了解有关 TPU 概念和术语的更多信息，请参阅这些文档。

本示例使用 Jax，这是一个通过 OpenXLA 原生支持针对 TPU 芯片进行即时（JIT）编译的机器学习框架。不过，您也可以使用 PyTorch/XLA 在 TPU 上进行机器学习训练。

此示例利用了多个 JobSet 特性（显式和隐式）来支持 TPU 多切片训练的独特调度要求，开箱即用，用户只需很少的配置。

# Run a simple Jax workload on 
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice
  annotations:
    # Give each child Job exclusive usage of a TPU slice 
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 3
  replicatedJobs:
  - name: workers
    replicas: 4 # Set to number of TPU slices
    template:
      spec:
        parallelism: 2 # Set to number of VMs per TPU slice
        completions: 2 # Set to number of VMs per TPU slice
        backoffLimit: 0
        template:
          spec:
            hostNetwork: true
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
              cloud.google.com/gke-tpu-topology: 2x4
            containers:
            - name: jax-tpu
              image: python:3.8
              ports:
              - containerPort: 8471
              - containerPort: 8080
              securityContext:
                privileged: true
              command:
              - bash
              - -c
              - |
                pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                python -c 'import jax; print("Global device count:", jax.device_count())'
                sleep 60                
              resources:
                limits:
                  google.com/tpu: 4

未来工作和参与

我们在 JobSet 路线图上计划了今年要开发的许多功能，可以在 JobSet 路线图中找到。

欢迎随时提供任何形式的反馈。我们也欢迎更多的贡献者，无论是修复或报告错误，还是帮助添加新功能或编写文档。

您可以通过我们的 仓库、邮件列表或在 Slack 上与我们联系。

最后但同样重要的是，感谢所有使这个项目成为可能的贡献者！

• ←上一篇
• 下一篇→