通过 Pathways 进行弹性训练

Pathways 通过以下方式提供弹性优势：

  • 暂停-恢复：在面对抢占通知等计划内中断时具有容错能力，无需用户编写任何自定义抢占处理代码。
  • 弹性训练：在发生意外硬件故障时具有容错能力，不会导致客户端崩溃，但需要用户编写特定于模型的恢复代码。

准备工作

请确保您已备妥：

暂停-恢复

通常，GKE 会在抢占加速器 Pod 之前向其发送抢占通知。默认情况下，所有 Cloud Deploy 都会启用 Pathways 抢占容忍度，并且 Pathways 加速器作业会监听这些通知。

当收到抢占通知时，Pathways 首先会确定当前工作负载是否可恢复，即 Pathways 能否以透明方式保存和恢复工作负载。如果是，则它会尝试在 GKE 逐出加速器作业之前，通过将其当前状态写入永久性存储（例如 Cloud Storage）来透明地暂停机器学习工作负载。当 GKE 稍后重新安排作业时，Pathways 会通过读回其持久化状态来恢复机器学习工作负载。

如果工作负载无法恢复，Pathways 会关闭加速器作业，并在配置了弹性训练的情况下将故障转发到您的作业。如果未配置弹性训练，GKE 会根据 JobSet 重启政策重启整个工作负载。

使用 JAX 定义的典型机器学习工作负载依赖于无状态 Pathways XLA 组件，这些组件可以使用高带宽内存 (HBM) 快照进行恢复。某些机器学习工作负载（例如使用 JAX 共置 Python API 定义的工作负载）依赖于有状态 Pathways 组件；这些组件无法恢复。

弹性训练

弹性训练使您的训练作业即使在发生硬件故障时也能继续。这是通过结合使用 Pathways 系统功能和用户定义的模型恢复逻辑来实现的：

  • 故障检测：发生硬件故障（例如 TPU 工作器崩溃）时，Pathways 系统会检测到此情况，并在下次访问该硬件上的数据时通过异常通知用户的训练作业。此通知不会导致工作负载崩溃；它允许您的代码处理通知并重新配置资源，以继续处理或正常退出。
  • 用户定义的弹性处理程序：用户的模型代码需要能够处理此异常。这就是“特定于模型的恢复”。
    • 快照：最常见的方法是定期保存模型状态的快照。发生故障时，您可以从最近的快照加载以恢复训练。
    • 重新配置：您可能需要重新配置训练作业，以根据可用切片数量进行调整。例如，如果一个切片停止工作，您可能会将活跃切片的数量减少一个，直到有可用的替换切片。如需了解详情，请参阅 Elastic 处理程序。
    • 数据/计算图更新：您的代码需要根据需要重新创建计算图，以处理可用于计算的设备数量的任何变化。这可能涉及重新划分数据或重新编译模型。
  • Pathways 在恢复中的作用：Pathways 提供支持用户定义的重新配置的原语：
    • 切片替换：如果故障切片被替换，则可以在新切片可用时通知客户端。然后，您的代码可以重新配置为使用此新切片。
    • 透明恢复：Pathways 会处理恢复的底层细节，例如重新建立与集群中健康部分的连接。
  • pathwaysutils 中的实用程序：pathways-utils 中定义的一组 Pathways 实用程序。

实现弹性处理程序

您必须编写的大部分代码都将位于用户定义的弹性处理程序中。此处理程序通过重新创建网格并重新初始化训练循环来响应弹性事件（例如 TPU 切片不可用）。

每个工作负载都是独一无二的。弹性处理程序的复杂性可能会随着工作负载的复杂性而增加。处理程序的输入和输出应为重新初始化训练循环所需的最小参数和返回值。

def elastic_handler(elastic_utils, *args, **kwargs):
  mesh = initialize_mesh(**kwargs["mesh_kwargs"])
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **kwargs["initialize_training_loop_kwargs"])

  step, snapshot = elastic_utils.get_next_snapshot()
  state = initial_state.replace(**snapshot)

  return state, step, mesh, jitted_train_step, other_variables

更新训练循环

您需要对训练循环进行以下更改：

  1. 创建弹性管理器
  2. 将训练循环封装在处理 jax.errors.JaxRuntimeError 的 try-except 块中
  3. 在 jax.errors.JaxRuntimeError 处理程序中，调用 maybe_reshard_down。 如果错误与弹性事件相关，弹性管理器将重新分片，否则会重新引发错误。
  4. 在训练循环结束时调用 maybe_snapshot 和 maybe_reshard_up
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()

def initialize_mesh(**kwargs):
  ...


def initialize_training_loop(**kwargs):
  ...


def train_loop(
    final_step,
    elastic_manager,
    mesh_kwargs,
    initialize_training_loop_kwargs,
):
  mesh = initialize_mesh(**mesh_kwargs)
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **initialize_training_loop_kwargs)

  step = initial_step
  while step < final_step:
    try:
      state = jitted_train_step(state)

      elastic_manager.maybe_snapshot(step=step, snapshot=state)
      handler_returns = elastic_manager.maybe_reshard_up(
          step=step,
          snapshot=state,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns
      step += 1
    except jax.errors.JaxRuntimeError as error:
      handler_returns = elastic_manager.maybe_reshard_down(
          error=error,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns

  return state


def main():
  elastic_manager = manager.Manager(
      devices=jax.devices(),
      snapshot_period=10,
      snapshot_buffer_size=1,
      reshard_check_period=5,
      max_elastic_down_event_count=10,
      max_reshard_retry_count=3,
  )

  train_loop(100, elastic_manager, {}, {})

配置弹性管理器

弹性管理器可以通过几种不同的方式进行配置。快照频率由快照周期决定。快照周期会影响因弹性事件而丢失的平均步数。重新分片检查周期决定了训练循环轮询切片可用性的频率。借助 max_elastic_down_event_count，您可以设置训练循环将支持的因切片丢失而导致的弹性事件数量。max_reshard_retry_count 指定弹性管理器应重试重新分片的次数。管理器是一个单例对象，应仅创建一次。

快照

根据弹性管理器配置，该函数可能会将数据快照到主机内存中，供弹性处理程序在弹性事件期间使用。

减少分片

捕获 jax.errors.JaxRuntimeError 后，Pathways 将检查错误是否是由切片丢失导致的弹性事件引起的。如果是，它会在循环中调用弹性处理程序，直到成功或达到最大重试次数。如果错误不是由弹性事件引起的，则会再次引发错误。弹性处理程序的返回值会传递给调用方。

增加分片

根据弹性管理器配置，如果存在不可用的切片，Pathways 将检查是否有其他切片可用。如果是，它会立即保存快照（如果尚未为当前步骤创建预先存在的快照），并循环调用弹性处理程序，直到成功或达到重试次数上限。如果发生重新分片，弹性处理程序的返回值会传递给调用方。否则，系统会返回 None。

热插拔

热插拔是指 GKE JobSet API 的一项功能，其中优先级较高的作业可以快速接管优先级较低的作业中的资源，从而最大限度地减少停机时间并确保更快地恢复。

创建 JobSet 后，GKE 会按照 JobSet 配置中的指定，在多个切片上调度工作负载。如果一个或多个切片发生硬件故障，受影响的 Pod 会被标记为失败。重新调度此 JobSet 时，如果您选择在 GKE 集群中保留一个可用于低优先级作业的备用切片，JobSet 系统会将高优先级作业的故障切片的工作负载重新映射到同一 GKE 集群中低优先级作业正在使用的备用切片上。此重新映射通常用时不超过 1 分钟。

JobSet 重启后，在以下情况下可能会发生热插拔：

  1. 默认模式：如果同一集群中有空闲的备用 TPU 切片，Kubernetes 调度器会优先将重启的作业调度到这些切片上，而不是等待故障切片修复。这样可以加快恢复速度。
  2. 异构工作负载：在运行多个工作负载且配置了 Kubernetes PriorityClass 的集群中，重启的 JobSet 可以触发热插拔。如果重启的作业的亲和性与优先级较低的作业的资源匹配，Kubernetes 会抢占优先级较低的作业，从而使优先级较高的作业立即启动。例如，您可以使用 PriorityClass 为 Pathways 工作器 Pod 配置不同的优先级。

如需在集群中使用优先级，请定义优先级类，例如：

kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

将此 YAML 应用到您的 GKE 集群：

kubectl apply -f high-prior-job.yaml

接下来，通过将以下文本添加到 pathways-worker Pod 的 podspec，将新的优先级类附加到 Pathways 工作器作业。

priorityClassName: high-prior-job

后续步骤