创建并运行使用 TPU 的作业

本文档介绍了如何创建和运行使用具有张量处理单元 (TPU) 的虚拟机 (VM) 实例的批量作业。TPU 是 Google 定制开发的应用专用集成电路 (ASIC),旨在加速机器学习 (ML) 和人工智能 (AI) 工作负载。创建批量作业时,您可以选择指定 TPU 机器类型,以加速容器化或基于脚本的 AI 和机器学习推理、评估、微调和训练工作负载。

如需详细了解 TPU,请参阅 Compute Engine 文档中的 Google Cloud上的 TPU 简介和 TPU 机器类型。或者,您也可以使用 GPU 来加速数据处理、AI 和机器学习工作负载。

准备工作

  1. 如果您之前未使用过 Batch,请参阅 Batch 使用入门,并通过完成项目和用户的前提条件来启用 Batch。
  2. 如需获得创建作业所需的权限,请让您的管理员为您授予以下 IAM 角色:

    如需详细了解如何授予角色,请参阅管理对项目、文件夹和组织的访问权限。

    您也可以通过自定义角色或其他预定义角色来获取所需的权限。

限制

当作业使用 TPU 时,除了 TPU 的 Compute Engine 限制之外,还需遵循以下所有限制。

  • 仅限单主机 TPU 切片:Batch 仅支持独立的 TPU 虚拟机,这些虚拟机称为单主机 TPU 切片。不支持多主机 TPU 切片,每个切片由多个互连的虚拟机组成。

    作业可以在一个或多个单主机 TPU 切片上运行。例如,如果作业将 taskCount 设置为 4,将 taskCountPerNode 设置为 1,则运行作业的虚拟机数量不超过 4 个。与多主机 TPU 切片不同,单主机 TPU 切片各自作为独立的虚拟机运行,不提供芯片间互连 (ICI) 网络来实现高带宽、低延迟通信。

    如需详细了解单主机和多主机 TPU 切片,请参阅 Compute Engine 文档。

  • 支持的 TPU 版本和机器类型:Batch 仅支持以下 TPU 版本和机器类型的单主机 TPU 切片(独立 TPU 虚拟机):

    • TPU7x (Ironwood):tpu7x-standard-4t
    • v6e (Trillium):ct6e-standard-1t、ct6e-standard-4t 和 ct6e-standard-8t

    • v5p:ct5p-hightpu-4t

    如需详细了解每种 TPU 机器类型(例如支持的磁盘),请参阅 Compute Engine 文档中的 TPU 机器类型。

  • 支持的使用选项:对于使用 TPU 的作业,Batch 不支持灵活启动型虚拟机(FLEX_START 预配模型)。除灵活启动型虚拟机外,您可以根据 Compute Engine 文档中的按 TPU 版本划分的消费选项可用性,指定作业的 TPU 版本支持的任何消费选项。

  • 需要单可用区位置:创建使用 TPU 的作业时,您必须将作业的允许位置(allowedLocations[] 字段)设置为包含 TPU 版本或预留(如果适用)的可用区,从而将作业限制在特定可用区。

  • 不支持的功能:使用 TPU 的作业不支持以下功能:

创建使用 TPU 的作业

如需创建使用 TPU 的作业,请执行以下操作:

  1. 规划使用 TPU 的作业的配置。
  2. 使用您确定的配置创建作业。 如需查看有关如何创建使用 TPU 的作业的示例,请参阅本文档中的创建使用 TPU 的作业。

规划使用 TPU 的作业的配置

在继续创建使用 TPU 的作业之前,您必须执行以下操作:

  1. 选择 TPU 版本:选择适合您的工作负载的 TPU 版本。如需查看按工作负载类型列出的 TPU 版本列表,请参阅按工作负载类型推荐的 TPU 版本。

    确保选择受支持的 TPU 版本之一。

  2. 选择 TPU 使用情况选项:选择最符合您的工作负载、其时长和费用需求的消耗选项。 确保选择支持的使用选项之一。下表总结了每个使用选项的用例和作业要求。

    使用选项及其作业要求

    预留绑定(日历模式下的未来预留或容量块的未来预留)

    • 应用场景:对于需要在特定时间段内高度确保 TPU 资源可用性的作业,我们建议使用预留绑定使用方案。在 Compute Engine 中,这些预留被描述为日历模式下的未来预留(1 到 90 天)和容量块的未来预留(1 年或更长时间)。

    • 详细信息:在创建作业之前,您必须获取或确定所选 TPU 版本和 RESERVATION_BOUND 预配模型的预留。预留创建后,会产生与运行虚拟机相同的指定虚拟机费用,直到您删除预留为止。使用预留的虚拟机不会产生单独的费用,但无论是否使用,预留都会产生费用。

      如需了解详情,请参阅 Compute Engine 文档中的选择预留类型。

    Batch 会为执行以下所有操作的作业使用绑定到预留的预留:

    • 将 provisioningModel 字段设置为 RESERVATION_BOUND。
    • 明确指向兼容的预留。将 reservation 字段设置为预留、预留块或预留子块的简称(如果预留位于同一项目和可用区中)或完整资源路径。

    按需

    • 应用场景:按需模式快速且易于试用,因此是大多数作业的默认使用选项。但是,如果其他使用选项更符合您的优先事项,或者您反复遇到资源可获取性问题,我们不建议您使用按需选项。

    • 详细信息:按需通常是访问 Compute Engine 虚拟机的默认方式。按需预留可让您在请求时访问资源(如果可用)。在作业完成或达到运行时长限制之前,系统不会删除这些资源。

    默认情况下,Batch 会对大多数作业使用按需模式。如果您想确保作业使用按需价格,请执行以下所有操作:

    • 将 provisioningModel 字段设置为 STANDARD。
    • 阻止使用预留。如需防止作业自动使用兼容的预留,请将 reservation 字段设置为 NO_RESERVATION。如需了解详情,请参阅创建和运行无法使用预留虚拟机的作业。

    Spot 虚拟机

    • 用例:我们建议您使用 Spot 虚拟机来降低容错型工作负载的费用。

    • 详细信息:如果创建使用 Spot 虚拟机的作业的请求得到满足,则虚拟机将按 Spot 虚拟机价格计费,该价格可能会经常变化,并且可能提供最大折扣,最高可优惠按需价格的 91%。但是,Spot 虚拟机可能并非始终可用,并且随时可能被抢占。

      详细了解 Spot 虚拟机和针对使用 Spot 虚拟机的作业的建议。

    Batch 会将 Spot 虚拟机用于执行以下所有操作的作业:

    • 将 provisioningModel 字段设置为 SPOT。
    • (可选)防止使用预留。您无法将 Spot 虚拟机与预留搭配使用。具体而言,作业可以选择将 reservation 字段设置为 NO_RESERVATION。如需了解详情,请参阅创建和运行无法使用预留虚拟机的作业。

  3. 选择您所选 TPU 可用的可用区:如需使用 TPU 版本,作业的 allowedLocations[] 字段必须设置为包含您所选 TPU 的可用区,如下所示:

    • 如果您选择了使用预留的消耗选项,则必须使用预留所在的可用区。
    • 否则,请根据 TPU 可用性选择包含您的 TPU 版本的可用区。

    无论是哪种情况,可用区都必须位于支持 Batch 的区域之一。

  4. 确保您的项目有足够的 TPU 配额:如果您要创建使用 TPU 搭配按需或 Spot 虚拟机的作业,则必须在要使用的区域中有足够的 TPU 配额。创建使用 TPU 预留的 TPU 作业不需要任何 TPU 配额,因为配额在创建预留时已使用。如需查看 TPU 配额名称列表,请参阅 TPU 配额;如需查看有关如何查看配额的说明,请参阅查看和管理配额

  5. 选择 TPU 机器类型:为您选择的 TPU 版本选择 TPU 机器类型。如果您选择了使用预留的使用选项,则必须使用预留的机器类型。如需查看每个 TPU 版本可用的机器类型列表,请参阅 TPU 机器类型。

    由于存在限制,您必须为 TPU 机器类型使用单主机拓扑。

创建使用 TPU 的作业

以下部分介绍了如何为每个使用选项创建使用 TPU 的示例作业。

使用选项 作业说明示例
预留绑定(日历模式下的未来预留或容量块的未来预留)

将 TPU 与受预留约束的使用方案搭配使用

按需虚拟机或 Spot 虚拟机

将 TPU 与按需虚拟机或 Spot 虚拟机搭配使用

将 TPU 与按需虚拟机或 Spot 虚拟机搭配使用

您可以使用 gcloud CLI 或 Batch API 创建使用 TPU 且采用按需虚拟机或 Spot 虚拟机的作业。

gcloud

  1. 创建一个 JSON 配置文件,用于指定受支持的单主机 TPU 机器类型,将预配模型设置为 STANDARD 或 SPOT,并在提供该 TPU 机器类型的位置运行。

    例如,如需创建基本脚本作业,请创建一个包含以下内容的 JSON 文件:

    {
      "taskGroups": [
        {
          "taskSpec": {
            "runnables": [
              {
                "script": {
                  "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
                }
              }
            ]
          },
          "taskCount": 3,
          "parallelism": 1
        }
      ],
      "allocationPolicy": {
        "instances": [
          {
            "policy": {
              "provisioningModel": "PROVISIONING_MODEL",
              "machineType": "MACHINE_TYPE"
            }
          }
        ],
        "location": {
          "allowedLocations": [
            "ALLOWED_LOCATIONS"
          ]
        }
      },
      "logsPolicy": {
        "destination": "CLOUD_LOGGING"
      }
    }
    

    替换以下内容:

    • PROVISIONING_MODEL:用量选项的预配模型:STANDARD(按需)或 SPOT(Spot 虚拟机)。

    • MACHINE_TYPE:受支持的单主机 TPU 机器类型,例如 ct5p-hightpu-4t。

    • ALLOWED_LOCATIONS:作业的允许位置。对于使用 TPU 的作业,您必须将 allowedLocations[] 字段设置为提供 TPU 机器类型的可用区,例如,zones/us-central1-a 提供 TPU 机器类型 ct5p-hightpu-4t。

  2. 如需创建并运行作业,请使用 gcloud batch jobs submit 命令:

    gcloud batch jobs submit JOB_NAME \
        --location LOCATION \
        --config JSON_CONFIGURATION_FILE
    

    替换以下内容:

    • JOB_NAME:作业的名称。

    • LOCATION:作业的位置。

    • JSON_CONFIGURATION_FILE:包含作业配置详细信息的 JSON 文件的路径。

API

向 jobs.create 方法发出 POST 请求,指定受支持的单主机 TPU 机器类型,将预配模型设置为 STANDARD 或 SPOT,并在提供该 TPU 机器类型的位置运行。

例如,如需创建基本脚本作业,请发出以下请求:

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "PROVISIONING_MODEL",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

替换以下内容:

  • PROJECT_ID:项目的项目 ID。

  • LOCATION:作业的位置。

  • JOB_NAME:作业的名称。

  • PROVISIONING_MODEL:用量选项的预配模型:STANDARD(按需)或 SPOT(Spot 虚拟机)。

  • MACHINE_TYPE:受支持的单主机 TPU 机器类型,例如 ct5p-hightpu-4t。

  • ALLOWED_LOCATIONS:作业的允许位置。对于使用 TPU 的作业,您必须将 allowedLocations[] 字段设置为提供 TPU 机器类型的可用区,例如,zones/us-central1-a 提供 TPU 机器类型 ct5p-hightpu-4t。

将 TPU 与受预留约束的使用方案搭配使用

您可以使用 gcloud CLI 或 Batch API 创建一个作业,该作业使用 TPU,并为预留绑定 (RESERVATION_BOUND) 预配模型提供使用选项,即容量块的未来预留或日历模式下的未来预留。

gcloud

  1. 创建一个 JSON 配置文件,该文件为单主机 TPU 机器类型指定预留,将预配模型设置为 RESERVATION_BOUND,并在预留的可用区中运行。

    例如,如需创建基本脚本作业,请创建一个包含以下内容的 JSON 文件:

    {
      "taskGroups": [
        {
          "taskSpec": {
            "runnables": [
              {
                "script": {
                  "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
                }
              }
            ]
          },
          "taskCount": 3,
          "parallelism": 1
        }
      ],
      "allocationPolicy": {
        "instances": [
          {
            "policy": {
              "provisioningModel": "RESERVATION_BOUND",
              "reservation": "RESERVATION",
              "machineType": "MACHINE_TYPE"
            }
          }
        ],
        "location": {
          "allowedLocations": [
            "ALLOWED_LOCATIONS"
          ]
        }
      },
      "logsPolicy": {
        "destination": "CLOUD_LOGGING"
      }
    }
    

    替换以下内容:

    • RESERVATION:替换为以下预留路径之一:

      • 如果您想使用日历模式下的未来预留或容量块的未来预留的任何部分,请指定预留:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME
        
      • 如果您想使用容量块的未来预留中的特定块,请指定预留块:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
        
      • 如果您想使用容量块的未来预留中的特定子块,请指定预留子块:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME
        
    • MACHINE_TYPE:预留的 TPU 机器类型,例如 ct5p-hightpu-4t。

    • ALLOWED_LOCATIONS:作业的允许位置。对于使用预留 TPU 的作业,您必须将 allowedLocations[] 字段设置为预留的可用区,例如,zones/us-central1-a 提供 TPU 机器类型 ct5p-hightpu-4t。

  2. 如需创建并运行作业,请使用 gcloud batch jobs submit 命令:

    gcloud batch jobs submit JOB_NAME \
        --location LOCATION \
        --config JSON_CONFIGURATION_FILE
    

    替换以下内容:

    • JOB_NAME:作业的名称。

    • LOCATION:作业的位置。

    • JSON_CONFIGURATION_FILE:包含作业配置详细信息的 JSON 文件的路径。

API

向 jobs.create 方法发出 POST 请求,该请求指定单主机 TPU 机器类型的预留,将预配模型设置为 RESERVATION_BOUND,并在预留的可用区中运行。

例如,如需创建基本脚本作业,请发出以下请求:

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "RESERVATION_BOUND",
          "reservation": "RESERVATION",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

替换以下内容:

  • PROJECT_ID:项目的项目 ID。

  • LOCATION:作业的位置。

  • JOB_NAME:作业的名称。

  • RESERVATION:替换为以下预留路径之一:

    • 如果您想使用日历模式下的未来预留或容量块的未来预留的任何部分,请指定预留:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME
      
    • 如果您想使用容量块的未来预留中的特定块,请指定预留块:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
      
    • 如果您想使用容量块的未来预留中的特定子块,请指定预留子块:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME
      
  • MACHINE_TYPE:预留的 TPU 机器类型,例如 ct5p-hightpu-4t。

  • ALLOWED_LOCATIONS:作业的允许位置。对于使用预留 TPU 的作业,您必须将 allowedLocations[] 字段设置为预留的可用区,例如,zones/us-central1-a 提供 TPU 机器类型 ct5p-hightpu-4t。

如需详细了解如何创建使用预留的作业,请参阅使用虚拟机预留确保资源可用性。

后续步骤