为 Vertex AI 无服务器训练配置计算资源

执行无服务器训练时，您的训练代码会在一个或多个虚拟机 (VM) 实例上运行。您可以配置用于训练的虚拟机类型：使用计算资源较多的虚拟机可加快训练速度，并允许您使用更大的数据集，不过，它们也可能会产生更多训练费用。

在某些情况下，您还可以使用 GPU 加速训练。GPU 会产生额外费用。

您也可以选择自定义训练作业的启动磁盘的类型和大小。

本文档介绍了可用于无服务器训练的不同计算资源及其配置方法。

管理费用和可用性

为帮助管理费用或确保虚拟机资源的可用性，Vertex AI 提供了以下功能：

为确保虚拟机资源在训练作业需要时可用，您可以使用 Compute Engine 预留。预留为获取 Compute Engine 资源的容量提供了较高的保障。如需了解详情，请参阅将预留与训练搭配使用。
如需降低训练作业运行费用，您可以使用 Spot 虚拟机。Spot 虚拟机是过剩的 Compute Engine 容量的虚拟机 (VM) 实例。Spot 虚拟机的折扣力度很大，但 Compute Engine 可能会提前停止或删除 Spot 虚拟机来随时收回容量。如需了解详情，请参阅将 Spot 虚拟机与训练搭配使用。
对于请求 GPU 资源的无服务器训练作业，动态工作负载调度器可让您根据请求的 GPU 资源何时可用来调度作业。如需了解详情，请参阅根据资源可用性安排训练作业。

指定计算资源的位置

在 WorkerPoolSpec 中指定配置详细信息。根据您执行无服务器训练的方式，请将此 WorkerPoolSpec 放入以下某个 API 字段中：

如果您要创建一个CustomJob资源，在 CustomJob.jobSpec.workerPoolSpecs 中指定 WorkerPoolSpec。

如果您使用的是 Google Cloud CLI，则可以在 gcloud ai custom-jobs create 命令上使用 --worker-pool-spec 标志或 --config 标志来指定工作器池选项。

详细了解如何创建 CustomJob。
如果您要创建一个HyperparameterTuningJob资源，在 HyperparameterTuningJob.trialJobSpec.workerPoolSpecs 中指定 WorkerPoolSpec。

如果您使用的是 gcloud CLI，则可以使用 gcloud ai hpt-tuning-jobs create 命令上的 --config 标志来指定工作器池选项。

详细了解如何创建 HyperparameterTuningJob。
如果您要创建一个未进行超参数调节的TrainingPipeline资源，在 TrainingPipeline.trainingTaskInputs.workerPoolSpecs 中指定 WorkerPoolSpec。

详细了解如何创建自定义TrainingPipeline。
如果要使用超参数调节创建 TrainingPipeline，请在 TrainingPipeline.trainingTaskInputs.trialJobSpec.workerPoolSpecs 中指定 WorkerPoolSpec。

如果您执行分布式训练，则可以对每个工作器池使用不同的设置。

机器类型

在 WorkerPoolSpec 中，您必须在 machineSpec.machineType 字段中指定以下机器类型之一。工作器池中的每个副本都在具有指定机器类型的不同虚拟机上运行。

a4x-highgpu-4g^*
a4-highgpu-8g^*
a3-ultragpu-8g^*
a3-megagpu-8g^*
a3-highgpu-1g^*
a3-highgpu-2g^*
a3-highgpu-4g^*
a3-highgpu-8g^*
a2-ultragpu-1g^*
a2-ultragpu-2g^*
a2-ultragpu-4g^*
a2-ultragpu-8g^*
a2-highgpu-1g^*
a2-highgpu-2g^*
a2-highgpu-4g^*
a2-highgpu-8g^*
a2-megagpu-16g^*
e2-standard-4
e2-standard-8
e2-standard-16
e2-standard-32
e2-highmem-2
e2-highmem-4
e2-highmem-8
e2-highmem-16
e2-highcpu-16
e2-highcpu-32
n2-standard-4
n2-standard-8
n2-standard-16
n2-standard-32
n2-standard-48
n2-standard-64
n2-standard-80
n2-highmem-2
n2-highmem-4
n2-highmem-8
n2-highmem-16
n2-highmem-32
n2-highmem-48
n2-highmem-64
n2-highmem-80
n2-highcpu-16
n2-highcpu-32
n2-highcpu-48
n2-highcpu-64
n2-highcpu-80
n1-standard-4
n1-standard-8
n1-standard-16
n1-standard-32
n1-standard-64
n1-standard-96
n1-highmem-2
n1-highmem-4
n1-highmem-8
n1-highmem-16
n1-highmem-32
n1-highmem-64
n1-highmem-96
n1-highcpu-16
n1-highcpu-32
n1-highcpu-64
n1-highcpu-96
c2-standard-4
c2-standard-8
c2-standard-16
c2-standard-30
c2-standard-60
ct5lp-hightpu-1t^*
ct5lp-hightpu-4t^*
ct5lp-hightpu-8t^*
m1-ultramem-40
m1-ultramem-80
m1-ultramem-160
m1-megamem-96
g2-standard-4^*
g2-standard-8^*
g2-standard-12^*
g2-standard-16^*
g2-standard-24^*
g2-standard-32^*
g2-standard-48^*
g2-standard-96^*
g4-standard-48^*
g4-standard-96^*
g4-standard-192^*
g4-standard-384^*
cloud-tpu^*

* 上述列表中标有星号的机器类型必须与特定 GPU 或 TPU 搭配使用。请参阅本指南的以下部分。

如需了解每种机器类型的技术规范，请参阅有关机器类型的 Compute Engine 文档。如需了解每种机器类型进行无服务器训练的费用，请参阅价格。

以下示例演示了创建 CustomJob 时指定机器类型的位置：

控制台

在 Google Cloud 控制台中，您无法直接创建 CustomJob。但是，您可以创建一个创建 CustomJob 的 TrainingPipeline。在 Google Cloud 控制台中创建 TrainingPipeline 时，请在计算和价格步骤的机器类型中为每个工作器池指定机器类型。

gcloud

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --worker-pool-spec=machine-type=MACHINE_TYPE,replica-count=REPLICA_COUNT,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI

Java

在尝试此示例之前，请按照《Vertex AI 快速入门：使用客户端库》中的 Java 设置说明执行操作。如需了解详情，请参阅 Vertex AI Java API 参考文档。

如需向 Vertex AI 进行身份验证，请设置应用默认凭证。如需了解详情，请参阅为本地开发环境设置身份验证。


import com.google.cloud.aiplatform.v1.AcceleratorType;
import com.google.cloud.aiplatform.v1.ContainerSpec;
import com.google.cloud.aiplatform.v1.CustomJob;
import com.google.cloud.aiplatform.v1.CustomJobSpec;
import com.google.cloud.aiplatform.v1.JobServiceClient;
import com.google.cloud.aiplatform.v1.JobServiceSettings;
import com.google.cloud.aiplatform.v1.LocationName;
import com.google.cloud.aiplatform.v1.MachineSpec;
import com.google.cloud.aiplatform.v1.WorkerPoolSpec;
import java.io.IOException;

// Create a custom job to run machine learning training code in Vertex AI
public class CreateCustomJobSample {

  public static void main(String[] args) throws IOException {
    // TODO(developer): Replace these variables before running the sample.
    String project = "PROJECT";
    String displayName = "DISPLAY_NAME";

    // Vertex AI runs your training application in a Docker container image. A Docker container
    // image is a self-contained software package that includes code and all dependencies. Learn
    // more about preparing your training application at
    // https://cloud.google.com/vertex-ai/docs/training/overview#prepare_your_training_application
    String containerImageUri = "CONTAINER_IMAGE_URI";
    createCustomJobSample(project, displayName, containerImageUri);
  }

  static void createCustomJobSample(String project, String displayName, String containerImageUri)
      throws IOException {
    JobServiceSettings settings =
        JobServiceSettings.newBuilder()
            .setEndpoint("us-central1-aiplatform.googleapis.com:443")
            .build();
    String location = "us-central1";

    // Initialize client that will be used to send requests. This client only needs to be created
    // once, and can be reused for multiple requests.
    try (JobServiceClient client = JobServiceClient.create(settings)) {
      MachineSpec machineSpec =
          MachineSpec.newBuilder()
              .setMachineType("n1-standard-4")
              .setAcceleratorType(AcceleratorType.NVIDIA_TESLA_T4)
              .setAcceleratorCount(1)
              .build();

      ContainerSpec containerSpec =
          ContainerSpec.newBuilder().setImageUri(containerImageUri).build();

      WorkerPoolSpec workerPoolSpec =
          WorkerPoolSpec.newBuilder()
              .setMachineSpec(machineSpec)
              .setReplicaCount(1)
              .setContainerSpec(containerSpec)
              .build();

      CustomJobSpec customJobSpecJobSpec =
          CustomJobSpec.newBuilder().addWorkerPoolSpecs(workerPoolSpec).build();

      CustomJob customJob =
          CustomJob.newBuilder()
              .setDisplayName(displayName)
              .setJobSpec(customJobSpecJobSpec)
              .build();
      LocationName parent = LocationName.of(project, location);
      CustomJob response = client.createCustomJob(parent, customJob);
      System.out.format("response: %s\n", response);
      System.out.format("Name: %s\n", response.getName());
    }
  }
}

Node.js

在尝试此示例之前，请按照《Vertex AI 快速入门：使用客户端库》中的 Node.js 设置说明执行操作。如需了解详情，请参阅 Vertex AI Node.js API 参考文档。

如需向 Vertex AI 进行身份验证，请设置应用默认凭证。如需了解详情，请参阅为本地开发环境设置身份验证。

/**
 * TODO(developer): Uncomment these variables before running the sample.\
 * (Not necessary if passing values as arguments)
 */

// const customJobDisplayName = 'YOUR_CUSTOM_JOB_DISPLAY_NAME';
// const containerImageUri = 'YOUR_CONTAINER_IMAGE_URI';
// const project = 'YOUR_PROJECT_ID';
// const location = 'YOUR_PROJECT_LOCATION';

// Imports the Google Cloud Job Service Client library
const {JobServiceClient} = require('@google-cloud/aiplatform');

// Specifies the location of the api endpoint
const clientOptions = {
  apiEndpoint: 'us-central1-aiplatform.googleapis.com',
};

// Instantiates a client
const jobServiceClient = new JobServiceClient(clientOptions);

async function createCustomJob() {
  // Configure the parent resource
  const parent = `projects/${project}/locations/${location}`;
  const customJob = {
    displayName: customJobDisplayName,
    jobSpec: {
      workerPoolSpecs: [
        {
          machineSpec: {
            machineType: 'n1-standard-4',
            acceleratorType: 'NVIDIA_TESLA_T4',
            acceleratorCount: 1,
          },
          replicaCount: 1,
          containerSpec: {
            imageUri: containerImageUri,
            command: [],
            args: [],
          },
        },
      ],
    },
  };
  const request = {parent, customJob};

  // Create custom job request
  const [response] = await jobServiceClient.createCustomJob(request);

  console.log('Create custom job response:\n', JSON.stringify(response));
}
createCustomJob();

Python

如需了解如何安装或更新 Vertex AI SDK for Python，请参阅安装 Vertex AI SDK for Python。如需了解详情，请参阅 Python API 参考文档。

from google.cloud import aiplatform


def create_custom_job_sample(
    project: str,
    display_name: str,
    container_image_uri: str,
    location: str = "us-central1",
    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
):
    # The AI Platform services require regional API endpoints.
    client_options = {"api_endpoint": api_endpoint}
    # Initialize client that will be used to create and send requests.
    # This client only needs to be created once, and can be reused for multiple requests.
    client = aiplatform.gapic.JobServiceClient(client_options=client_options)
    custom_job = {
        "display_name": display_name,
        "job_spec": {
            "worker_pool_specs": [
                {
                    "machine_spec": {
                        "machine_type": "n1-standard-4",
                        "accelerator_type": aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,
                        "accelerator_count": 1,
                    },
                    "replica_count": 1,
                    "container_spec": {
                        "image_uri": container_image_uri,
                        "command": [],
                        "args": [],
                    },
                }
            ]
        },
    }
    parent = f"projects/{project}/locations/{location}"
    response = client.create_custom_job(parent=parent, custom_job=custom_job)
    print("response:", response)

如需了解详情，请参阅创建 CustomJob 指南。

GPU

如果您已编写训练代码以使用 GPU，则可以将工作器池配置为在每个虚拟机上使用一个或多个 GPU。为了使用 GPU，您必须使用 A2、N1 或 G2 机器类型。此外，将较小的机器类型（例如 n1-highmem-2）与 GPU 搭配使用可能会因 CPU 限制而导致某些工作负载的记录失败。如果训练作业停止返回日志，请考虑选择更大的机器类型。

Vertex AI 支持以下类型的可用于无服务器训练的 GPU：

NVIDIA_GB200⁺（包括 GPUDirect-RDMA）
NVIDIA_B200^*（包括 GPUDirect-RDMA）
NVIDIA_H100_MEGA_80GB^*（包括 GPUDirect-TCPXO）
NVIDIA_H100_80GB
NVIDIA_H200_141GB^*（包括 GPUDirect-RDMA）
NVIDIA_A100_80GB
NVIDIA_TESLA_A100 (NVIDIA A100 40 GB)
NVIDIA_TESLA_P4
NVIDIA_TESLA_P100
NVIDIA_TESLA_T4
NVIDIA_TESLA_V100
NVIDIA_L4
NVIDIA_RTX_PRO_6000

^* 建议您使用共享预留或 Spot 虚拟机来获取容量。

⁺ 需要使用共享预留获取容量。

如需详细了解各种 GPU 的技术规范，请参阅有关计算工作负载 GPU 的 Compute Engine 简短文档。如需了解每种机器类型进行无服务器训练的费用，请参阅价格。

在 WorkerPoolSpec 中，指定要在machineSpec.acceleratorType 字段中使用的 GPU 类型，以及工作器池中每个虚拟机要在 machineSpec.acceleratorCount 字段中使用的 GPU 数量。但是，您对这些字段的选择必须符合以下要求：

您选择的 GPU 类型必须在执行无服务器训练的位置可用。并非所有区域都提供所有类型的 GPU。了解区域可用性。
您的配置只能使用特定数量的 GPU。例如，您可以在虚拟机上使用 2 或 4 个 NVIDIA_TESLA_T4 GPU，但不能使用 3 个 GPU。如需了解每种 GPU 有效的 acceleratorCount 值，请参阅以下兼容性表。
您必须确保 GPU 配置能够为要使用的 GPU 机器类型提供足够的虚拟 CPU 和内存。例如，如果您在工作器池中使用 n1-standard-32 机器类型，则每个虚拟机都有 32 个虚拟 CPU 和 120 GB 内存。由于每个 NVIDIA_TESLA_V100 GPU 最多可提供 12 个虚拟 CPU 和 76 GB 内存，因此每个 n1-standard-32 虚拟机必须至少使用 4 个 GPU 才能满足其需求。（2 个 GPU 提供的资源就已足够，您不能指定 3 个 GPU。）

以下兼容性表考虑到了此要求。

请注意，对使用 GPU 进行自定义训练与在 Compute Engine 中使用 GPU 的区别，还具有以下额外限制：
- 配备了 4 个 NVIDIA_TESLA_P100 GPU 的配置在所有区域和地区最多只提供 64 个虚拟 CPU 和 208 GB 的内存。
对于使用动态工作负载调度器或 Spot 虚拟机的作业，请将 CustomJob 的 scheduling.strategy 字段更新为所选策略。

以下兼容性表列出了 machineSpec.acceleratorCount 的有效值，具体取决于您对 machineSpec.machineType 和 machineSpec.acceleratorType 的选择：

每种机器类型的有效 GPU 数量
机器类型	`NVIDIA_H100_MEGA_80GB`	`NVIDIA_H100_80GB`	`NVIDIA_A100_80GB`	`NVIDIA_TESLA_A100`	`NVIDIA_TESLA_P4`	`NVIDIA_TESLA_P100`	`NVIDIA_TESLA_T4`	`NVIDIA_TESLA_V100`	`NVIDIA_L4`	`NVIDIA_H200_141GB`	`NVIDIA_B200`	`NVIDIA_GB200`	`NVIDIA_RTX_PRO_6000`
`a3-megagpu-8g`	8
`a3-highgpu-1g`		1^*
`a3-highgpu-2g`		2^*
`a3-highgpu-4g`		4^*
`a3-highgpu-8g`		8
`a3-ultragpu-8g`										8
`a4-highgpu-8g`											8
`a4x-highgpu-4g`												4
`a2-ultragpu-1g`			1
`a2-ultragpu-2g`			2
`a2-ultragpu-4g`			4
`a2-ultragpu-8g`			8
`a2-highgpu-1g`				1
`a2-highgpu-2g`				2
`a2-highgpu-4g`				4
`a2-highgpu-8g`				8
`a2-megagpu-16g`				16
`n1-standard-4`					1、2、4	1、2、4	1、2、4	1、2、4、8
`n1-standard-8`					1、2、4	1、2、4	1、2、4	1、2、4、8
`n1-standard-16`					1、2、4	1、2、4	1、2、4	2、4、8
`n1-standard-32`					2、4	2、4	2、4	4、8
`n1-standard-64`					4		4	8
`n1-standard-96`					4		4	8
`n1-highmem-2`					1、2、4	1、2、4	1、2、4	1、2、4、8
`n1-highmem-4`					1、2、4	1、2、4	1、2、4	1、2、4、8
`n1-highmem-8`					1、2、4	1、2、4	1、2、4	1、2、4、8
`n1-highmem-16`					1、2、4	1、2、4	1、2、4	2、4、8
`n1-highmem-32`					2、4	2、4	2、4	4、8
`n1-highmem-64`					4		4	8
`n1-highmem-96`					4		4	8
`n1-highcpu-16`					1、2、4	1、2、4	1、2、4	2、4、8
`n1-highcpu-32`					2、4	2、4	2、4	4、8
`n1-highcpu-64`					4	4	4	8
`n1-highcpu-96`					4		4	8
`g2-standard-4`									1
`g2-standard-8`									1
`g2-standard-12`									1
`g2-standard-16`									1
`g2-standard-24`									2
`g2-standard-32`									1
`g2-standard-48`									4
`g2-standard-96`									8
`g4-standard-48`													1
`g4-standard-96`													2
`g4-standard-192`													4
`g4-standard-384`													8

^* 只有在使用动态工作负载调度器或 Spot 虚拟机时，才能使用指定的机器类型。

以下示例演示了创建 CustomJob 时在何处指定 GPU：

控制台

在 Google Cloud 控制台中，您无法直接创建 CustomJob。但是，您可以创建一个创建 CustomJob 的 TrainingPipeline。在 Google Cloud 控制台中创建 TrainingPipeline 时，您可以在计算和价格步骤中为每个工作器池指定 GPU。首先指定机器类型。然后，您可以在加速器类型和加速器数量字段中指定 GPU 详细信息。

gcloud

如需使用 Google Cloud CLI 工具指定 GPU，您必须使用 config.yaml 文件。例如：

`config.yaml`

workerPoolSpecs:
  machineSpec:
    machineType: MACHINE_TYPE
    acceleratorType: ACCELERATOR_TYPE
    acceleratorCount: ACCELERATOR_COUNT
  replicaCount: REPLICA_COUNT
  containerSpec:
    imageUri: CUSTOM_CONTAINER_IMAGE_URI

然后，运行如下命令：

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --config=config.yaml

Node.js

在尝试此示例之前，请按照《Vertex AI 快速入门：使用客户端库》中的 Node.js 设置说明执行操作。如需了解详情，请参阅 Vertex AI Node.js API 参考文档。

如需向 Vertex AI 进行身份验证，请设置应用默认凭证。如需了解详情，请参阅为本地开发环境设置身份验证。

/**
 * TODO(developer): Uncomment these variables before running the sample.\
 * (Not necessary if passing values as arguments)
 */

// const customJobDisplayName = 'YOUR_CUSTOM_JOB_DISPLAY_NAME';
// const containerImageUri = 'YOUR_CONTAINER_IMAGE_URI';
// const project = 'YOUR_PROJECT_ID';
// const location = 'YOUR_PROJECT_LOCATION';

// Imports the Google Cloud Job Service Client library
const {JobServiceClient} = require('@google-cloud/aiplatform');

// Specifies the location of the api endpoint
const clientOptions = {
  apiEndpoint: 'us-central1-aiplatform.googleapis.com',
};

// Instantiates a client
const jobServiceClient = new JobServiceClient(clientOptions);

async function createCustomJob() {
  // Configure the parent resource
  const parent = `projects/${project}/locations/${location}`;
  const customJob = {
    displayName: customJobDisplayName,
    jobSpec: {
      workerPoolSpecs: [
        {
          machineSpec: {
            machineType: 'n1-standard-4',
            acceleratorType: 'NVIDIA_TESLA_T4',
            acceleratorCount: 1,
          },
          replicaCount: 1,
          containerSpec: {
            imageUri: containerImageUri,
            command: [],
            args: [],
          },
        },
      ],
    },
  };
  const request = {parent, customJob};

  // Create custom job request
  const [response] = await jobServiceClient.createCustomJob(request);

  console.log('Create custom job response:\n', JSON.stringify(response));
}
createCustomJob();

Python

如需了解如何安装或更新 Vertex AI SDK for Python，请参阅安装 Vertex AI SDK for Python。如需了解详情，请参阅 Python API 参考文档。

from google.cloud import aiplatform


def create_custom_job_sample(
    project: str,
    display_name: str,
    container_image_uri: str,
    location: str = "us-central1",
    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
):
    # The AI Platform services require regional API endpoints.
    client_options = {"api_endpoint": api_endpoint}
    # Initialize client that will be used to create and send requests.
    # This client only needs to be created once, and can be reused for multiple requests.
    client = aiplatform.gapic.JobServiceClient(client_options=client_options)
    custom_job = {
        "display_name": display_name,
        "job_spec": {
            "worker_pool_specs": [
                {
                    "machine_spec": {
                        "machine_type": "n1-standard-4",
                        "accelerator_type": aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,
                        "accelerator_count": 1,
                    },
                    "replica_count": 1,
                    "container_spec": {
                        "image_uri": container_image_uri,
                        "command": [],
                        "args": [],
                    },
                }
            ]
        },
    }
    parent = f"projects/{project}/locations/{location}"
    response = client.create_custom_job(parent=parent, custom_job=custom_job)
    print("response:", response)

如需了解详情，请参阅创建 CustomJob 指南。

GPUDirect 网络

在 Vertex Training 上，部分 H100、H200、B200 和 GB200 系列机器预配置了 GPUDirect 网络栈。与没有 GPUDirect 的 GPU 相比，GPUDirect 可以将 GPU 间联网速度提高多达 2 倍。

GPUDirect 通过减少在 GPU 之间传输数据包载荷所需的开销来实现这一点，从而大幅提高大规模吞吐量。

GPUDirect-TCPXO

a3-megagpu-8g 机器类型具有以下特性：

每个机器 8 个 NVIDIA H100 GPU
主 NIC 可支持最高 200 Gbps 的带宽
8 个次要 NIC，每个 NIC 支持最高 200 Gbps 的 GPU 数据传输
GPUDirect-TCPXO，可进一步改进 GPU 到虚拟机的通信

支持 GPUDirect 的 GPU 特别适合对大型模型进行分布式训练。

GPUDirect-RDMA

a4x-highgpu-4g 机器类型具有以下特性：

每个机器 4 个 GB200 GPU
2 个主机 NIC，提供 400 Gbps 的带宽
6 个 NIC，可提供最高 2,400 Gbps 的 GPU 数据传输
GPUDirect-RDMA，可通过 RoCE (RDMA over Converged Ethernet) 为大规模机器学习训练工作负载实现更高的 GPU 通信网络性能

a3-ultragpu-8g 和 a4-highgpu-8g 机器类型具有以下特性：

每个机器 8 个 NVIDIA H200/B200 GPU
2 个主机 NIC，提供 400 Gbps 的带宽
8 个 NIC，可提供最高 3,200 Gbps 的 GPU 数据传输
GPUDirect-RDMA，可通过 RoCE (RDMA over Converged Ethernet) 为大规模机器学习训练工作负载实现更高的 GPU 通信网络性能

TPU

如需使用张量处理单元 (TPU) 在 Vertex AI 上进行自定义训练，您可以将工作器池配置为使用 TPU 虚拟机。

在 Vertex AI 中使用 TPU 虚拟机时，您只能使用一个工作器池进行自定义训练，并且必须将此工作器池配置为仅使用一个副本。

TPU v2 和 v3

如需在工作器池中使用 TPU v2 或 v3 虚拟机，您必须使用以下配置之一：

如需使用 TPU v2 配置 TPU 虚拟机，请在 WorkerPoolSpec 中指定以下字段：
- 将 machineSpec.machineType 设置为 cloud-tpu。
- 将 machineSpec.acceleratorType 设置为 TPU_V2。
- 对于单个 TPU，将 machineSpec.acceleratorCount 设置为 8；对于 TPU Pod，则设置为 32 or multiple of 32。
- 将 replicaCount 设置为 1。
如需使用 TPU v3 配置 TPU 虚拟机，请在 WorkerPoolSpec 中指定以下字段：
- 将 machineSpec.machineType 设置为 cloud-tpu。
- 将 machineSpec.acceleratorType 设置为 TPU_V3。
- 对于单个 TPU，将 machineSpec.acceleratorCount 设置为 8；对于 TPU Pod，则设置为 32+。
- 将 replicaCount 设置为 1。

如需了解 TPU 的区域级可用性，请参阅使用加速器。

TPU v5e

TPU v5e 需要 JAX 0.4.6+、TensorFlow 2.15+ 或 PyTorch 2.1+。如需使用 TPU v5e 配置 TPU 虚拟机，请在 WorkerPoolSpec 中指定以下字段：

将 machineSpec.machineType 设置为 ct5lp-hightpu-1t、ct5lp-hightpu-4t 或 ct5lp-hightpu-8t。
将 machineSpec.tpuTopology 设置为相应机器类型支持的拓扑。有关详情，请参阅下表。
将 replicaCount 设置为 1。

下表列出了自定义训练支持的 TPU v5e 机器类型和拓扑：

机器类型	拓扑	TPU 芯片数量	虚拟机数量	推荐的使用场景
`ct5lp-hightpu-1t`	1x1	1	1	小到中等规模的训练
`ct5lp-hightpu-4t`	2x2	4	1	小到中等规模的训练
`ct5lp-hightpu-8t`	2x4	8	1	小到中等规模的训练
`ct5lp-hightpu-4t`	2x4	8	2	小到中等规模的训练
`ct5lp-hightpu-4t`	4x4	16	4	大规模训练
`ct5lp-hightpu-4t`	4x8	32	8	大规模训练
`ct5lp-hightpu-4t`	8x8	64	16	大规模训练
`ct5lp-hightpu-4t`	8x16	128	32	大规模训练
`ct5lp-hightpu-4t`	16x16	256	64	大规模训练

针对吞吐量和可用性对在 TPU v5e 虚拟机上运行的自定义训练作业进行了优化。如需了解详情，请参阅 v5e 训练加速器类型。

如需了解 TPU 的区域级可用性，请参阅使用加速器。如需详细了解 TPU v5e，请参阅 Cloud TPU v5e 训练。

机器类型比较：

机器类型	ct5lp-hightpu-1t	ct5lp-hightpu-4t	ct5lp-hightpu-8t
v5e 芯片数量	1	4	8
vCPU 数量	24	112	224
RAM (GB)	48	192	384
NUMA 节点的数量	1	1	2
抢占的可能性	高	中	低

TPU v6e

TPU v6e 需要 Python 3.10+、JAX 0.4.37+、使用 PJRT 作为默认运行时的 PyTorch 2.1+，或者仅使用 tf-nightly 运行时版本 2.18+ 的 TensorFlow。如需使用 TPU v6e 配置 TPU 虚拟机，请在 WorkerPoolSpec 中指定以下字段：

将 machineSpec.machineType 设置为 ct6e。
将 machineSpec.tpuTopology 设置为相应机器类型支持的拓扑。有关详情，请参阅下表。
将 replicaCount 设置为 1。

下表显示了支持自定义训练的 TPU v6e 机器类型和拓扑：

机器类型	拓扑	TPU 芯片数量	虚拟机数量	推荐的使用场景
`ct6e-standard-1t`	1x1	1	1	小到中等规模的训练
`ct6e-standard-8t`	2x4	8	1	小到中等规模的训练
`ct6e-standard-4t`	2x2	4	1	小到中等规模的训练
`ct6e-standard-4t`	2x4	8	2	小到中等规模的训练
`ct6e-standard-4t`	4x4	16	4	大规模训练
`ct6e-standard-4t`	4x8	32	8	大规模训练
`ct6e-standard-4t`	8x8	64	16	大规模训练
`ct6e-standard-4t`	8x16	128	32	大规模训练
`ct6e-standard-4t`	16x16	256	64	大规模训练

如需了解 TPU 的区域级可用性，请参阅使用加速器。如需详细了解 TPU v6e，请参阅 Cloud TPU v6e 训练。

机器类型比较：

机器类型	ct6e-standard-1t	ct6e-standard-4t	ct6e-standard-8t
v6e 芯片数量	1	4	8
vCPU 数量	44	180	180
RAM (GB)	48	720	1440
NUMA 节点的数量	2	1	2
抢占的可能性	高	中	低

TPU 7x（预览版）

TPU7x 需要 Python 3.12 及更高版本。

我们建议采用以下稳定组合进行功能测试和工作负载迁移：

JAX + JAX Lib：jax-0.8.1.dev20251104、jaxlib-0.8.1.dev2025104
稳定的 libtpu：libtpu-0.0.27

如需使用 TPU 7x 配置 TPU 虚拟机，请在 WorkerPoolSpec 中指定以下字段：

将 machineSpec.machineType 设置为 tpu7x-standard-4t。
将 machineSpec.tpuTopology 设置为相应机器类型支持的拓扑。有关详情，请参阅下表。
将 replicaCount 设置为 1。

下表显示了支持自定义训练的 TPU 7x 拓扑。所有拓扑都使用 tpu7x-standard-4t 机器类型。

拓扑	TPU 芯片数量	虚拟机数量	范围
2x2x1	4	1	单个主机
2x2x2	8	2	多主机
2x2x4	16	4	多主机
2x4x4	32	8	多主机
4x4x4	64	16	多主机
4x4x8	128	32	多主机
4x8x8	256	64	多主机
8x8x8	512	128	多主机
8x8x16	1024	256	多主机

如需了解 TPU 的区域级可用性，请参阅使用加速器。如需详细了解 TPU7x，请参阅 Cloud TPU7x 训练。

示例 `CustomJob`：指定 TPU 虚拟机

以下示例突出展示了如何在创建 CustomJob 时指定 TPU 虚拟机：

gcloud

如需使用 gcloud CLI 工具指定 TPU 虚拟机，您必须使用 config.yaml 文件。选择以下任一标签页以查看示例：

TPU v2/v3

workerPoolSpecs:
  machineSpec:
    machineType: cloud-tpu
    acceleratorType: TPU_V2
    acceleratorCount: 8
  replicaCount: 1
  containerSpec:
    imageUri: CUSTOM_CONTAINER_IMAGE_URI

TPU v5e

workerPoolSpecs:
  machineSpec:
    machineType: ct5lp-hightpu-4t
    tpuTopology: 4x4
  replicaCount: 1
  containerSpec:
    imageUri: CUSTOM_CONTAINER_IMAGE_URI

然后，运行如下命令：

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --config=config.yaml

Python

在尝试此示例之前，请按照《Vertex AI 快速入门：使用客户端库》中的 Python 设置说明执行操作。如需了解详情，请参阅 Vertex AI Python API 参考文档。

如需向 Vertex AI 进行身份验证，请设置应用默认凭证。如需了解详情，请参阅为本地开发环境设置身份验证。

如需使用 Python 版 Vertex AI SDK 指定 TPU VM，请参阅以下示例：

from google.cloud.aiplatform import aiplatform

job = aiplatform.CustomContainerTrainingJob(
    display_name='DISPLAY_NAME',
    location='us-west1',
    project='PROJECT_ID',
    staging_bucket="gs://CLOUD_STORAGE_URI",
    container_uri='CONTAINER_URI')

job.run(machine_type='ct5lp-hightpu-4t', tpu_topology='2x2')

如需详细了解如何创建自定义训练作业，请参阅创建自定义训练作业。

启动磁盘选项

您可以选择为训练虚拟机自定义启动磁盘。工作器池中的所有虚拟机均使用相同类型和大小的启动磁盘。

如需自定义每个训练虚拟机使用的启动磁盘类型，请在 WorkerPoolSpec 中指定 diskSpec.bootDiskType 字段。

您可以将此字段设置为以下任一项：
- pd-standard，以使用由标准硬盘支持的标准永久性磁盘
- pd-ssd，以使用由固态硬盘支持的 SSD 永久性磁盘
- hyperdisk-balanced，以获得更高的 IOPS 和吞吐率。
默认值为 pd-ssd（hyperdisk-balanced 是 a3-ultragpu-8g 和 a4-highgpu-8g 的默认值）。

如果训练代码对磁盘执行读写操作，则使用 pd-ssd 或 hyperdisk-balanced 可能会提升性能。了解磁盘类型。另请参阅 Hyperdisk 支持的机器。
如需自定义每个训练虚拟机使用的启动磁盘的大小（以 GB 为单位），在 WorkerPoolSpec 中指定 diskSpec.bootDiskSizeGb 字段。

您可将此字段设为一个介于 100 到 64000（含）之间的整数。默认值为 100。

如果您的训练代码将大量临时数据写入磁盘，则可能需要增加启动磁盘大小。请注意，您写入启动磁盘的任何数据都是临时的，在训练完成后，您便无法检索这些数据。

更改启动磁盘的类型和大小会影响自定义训练价格。

以下示例突出显示您在创建 CustomJob 时指定启动磁盘选项的位置：

控制台

在 Google Cloud 控制台中，您无法直接创建 CustomJob。但是，您可以创建一个创建 CustomJob 的 TrainingPipeline。在 Google Cloud 控制台中创建 TrainingPipeline 时，您可以在计算和价格步骤的磁盘类型下拉列表和磁盘大小 (GB)字段中为每个工作器池指定启动磁盘选项。

gcloud

如需使用 Google Cloud CLI 工具指定启动磁盘选项，您必须使用 config.yaml 文件。例如：

`config.yaml`

workerPoolSpecs:
  machineSpec:
    machineType: MACHINE_TYPE
  diskSpec:
    bootDiskType: DISK_TYPE
    bootDiskSizeGb: DISK_SIZE
  replicaCount: REPLICA_COUNT
  containerSpec:
    imageUri: CUSTOM_CONTAINER_IMAGE_URI

然后，运行如下命令：

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --config=config.yaml

如需了解详情，请参阅创建 CustomJob 指南。

后续步骤

了解如何创建永久性资源以运行自定义训练作业。
了解如何通过创建 CustomJob 执行自定义训练。

为 Vertex AI 无服务器训练配置计算资源 使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

管理费用和可用性

指定计算资源的位置

机器类型

控制台

gcloud

Java

Node.js

Python

GPU

控制台

gcloud

config.yaml

Node.js

Python

GPUDirect 网络

GPUDirect-TCPXO

GPUDirect-RDMA

TPU

TPU v2 和 v3

TPU v5e

TPU v6e

TPU 7x（预览版）

示例 CustomJob：指定 TPU 虚拟机

gcloud

TPU v2/v3

TPU v5e

Python

启动磁盘选项

控制台

gcloud

config.yaml

后续步骤

为 Vertex AI 无服务器训练配置计算资源

`config.yaml`

示例 `CustomJob`：指定 TPU 虚拟机

`config.yaml`