Configurar recursos de computação para treinamento personalizado

Ao fazer o treinamento personalizado, o código de treinamento é executado em uma ou mais instâncias de máquina virtual (VM). Configurar os tipos de VM que serão usados no treinamento: usar VMs com mais recursos de computação pode acelerar o treinamento e permitir que você trabalhe com conjuntos de dados maiores, mas eles também podem gerar maiscustos de treinamento.

Em alguns casos, você também pode usar GPUs para acelerar o treinamento. As GPUs geram outros custos.

Também é possível personalizar o tipo e o tamanho dos discos de inicialização das suas VMs de treinamento.

Este documento descreve os diferentes recursos de computação que podem ser usados no treinamento personalizado e como configurá-los.

Gerenciar custo e disponibilidade

Para ajudar a gerenciar custos ou garantir a disponibilidade de recursos de VM, a Vertex AI oferece o seguinte:

  • Para que os recursos da VM estejam disponíveis quando os jobs de treinamento precisarem deles, use as reservas do Compute Engine. As reservas fornecem um nível alto de garantia da capacidade dos recursos do Compute Engine. Para mais informações, consulte Usar reservas com treinamento.

  • Para reduzir o custo da execução dos jobs de treinamento, use VMs spot. As VMs spot são instâncias de máquina virtual (VM) que excedem a capacidade do Compute Engine. Elas têm descontos significativos, mas o Compute Engine pode forçar a interrupção ou excluir essas VMs para recuperar a capacidade a qualquer momento. Para mais informações, consulte Usar VMs spot com o treinamento.

  • Para jobs de treinamento personalizados que solicitam recursos de GPU, o Dynamic Workload Scheduler permite programar os jobs com base no momento em que os recursos de GPU solicitados ficam disponíveis. Para mais informações, consulte Programar jobs de treinamento com base na disponibilidade de recursos.

Onde especificar recursos de computação

Especifique os detalhes de configuração em um WorkerPoolSpec. Dependendo de como você realiza o treinamento personalizado, coloque esse WorkerPoolSpec em um dos seguintes campos da API:

Se você estiver realizando o treinamento distribuído, use configurações diferentes para cada pool de workers.

Tipos de máquina

Em WorkerPoolSpec, é preciso especificar um dos seguintes tipos de máquina no campo machineSpec.machineType. Cada réplica no pool de workers é executada em uma VM separada que tem o tipo de máquina especificado.

  • 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*
  • cloud-tpu*

* Os tipos de máquina marcados com asteriscos na lista anterior precisam ser usados com algumas GPUs ou TPUs. Consulte as seções a seguir deste guia.

Para saber mais sobre as especificações técnicas de cada tipo de máquina, leia a documentação do Compute Engine sobre tipos de máquinas. Para saber mais sobre o custo de uso de cada tipo de máquina para treinamento personalizado, leia Preços.

Os exemplos a seguir destacam onde você especifica um tipo de máquina ao criar um CustomJob:

Console

No console do Google Cloud, não é possível criar um CustomJob diretamente. No entanto, é possível criar um TrainingPipeline que crie um CustomJob. Ao criar um TrainingPipeline no Console do Google Cloud, especifique um tipo de máquina para cada pool de workers na etapa Compute and pricing, no campo Machine type.

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

Antes de testar esse exemplo, siga as instruções de configuração para Java no Guia de início rápido da Vertex AI sobre como usar bibliotecas de cliente. Para mais informações, consulte a documentação de referência da API Vertex AI para Java.

Para autenticar na Vertex AI, configure o Application Default Credentials. Para mais informações, consulte Configurar a autenticação para um ambiente de desenvolvimento local.


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

Antes de testar esse exemplo, siga as instruções de configuração para Node.js no Guia de início rápido da Vertex AI sobre como usar bibliotecas de cliente. Para mais informações, consulte a documentação de referência da API Vertex AI para Node.js.

Para autenticar na Vertex AI, configure o Application Default Credentials. Para mais informações, consulte Configurar a autenticação para um ambiente de desenvolvimento local.

/**
 * 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_K80',
            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

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python.

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)

Para mais contexto, leia o guia para criar um CustomJob.

GPUs

Se você escrever seu código de treinamento para usar GPUs, poderá configurar o pool de workers para usar uma ou mais GPUs em cada VM. Para usar as GPUs, é preciso ter um tipo de máquina A2, N1 ou G2. Além disso, o uso de tipos de máquinas menores, como n1-highmem-2 com GPUs, pode causar falha na geração de registros para algumas cargas de trabalho devido a restrições de CPU. Se o job de treinamento parar de retornar registros, selecione um tipo de máquina maior.

O Vertex AI é compatível com os seguintes tipos de GPU para treinamento personalizado:

  • NVIDIA_H100_80GB
  • NVIDIA_A100_80GB
  • NVIDIA_TESLA_A100 (NVIDIA A100 40GB)
  • NVIDIA_TESLA_P4
  • NVIDIA_TESLA_P100
  • NVIDIA_TESLA_T4
  • NVIDIA_TESLA_V100
  • NVIDIA_L4

Para saber mais sobre a especificação técnica para cada tipo de GPU, leia a documentação breve do Compute Engine sobre GPUs para cargas de trabalho de computação. Para saber mais sobre o custo de usar cada tipo de máquina para treinamento personalizado, leia Preços.

Em WorkerPoolSpec, especifique o tipo de GPU que você quer usar no campo machineSpec.acceleratorType e o número de GPUs que você quer que cada VM no pool de workers para usar no campo machineSpec.acceleratorCount. No entanto, suas escolhas para esses campos precisam atender às seguintes restrições:

  • O tipo de GPU escolhido precisa estar disponível no local em que você está executando o treinamento personalizado. Nem todos os tipos de GPU estão disponíveis em todas as regiões. Saiba mais sobre a disponibilidade regional.

  • Só é possível usar determinados números de GPUs na configuração. Por exemplo, é possível usar duas ou quatro GPUs NVIDIA_TESLA_T4 em uma VM, mas não 3. Para ver quais valores acceleratorCount são válidos para cada tipo de GPU, consulte a tabela de compatibilidade a seguir.

  • Verifique se a configuração da GPU fornece CPUs virtuais e memória suficientes para o tipo de máquina com que você a usa. Por exemplo, se você usar o tipo de máquina n1-standard-32 no seu pool de workers, cada VM terá 32 CPUs virtuais e 120 GB de memória. Como cada GPU NVIDIA_TESLA_V100 pode fornecer até 12 CPUs virtuais e 76 GB de memória, é preciso usar pelo menos quatro GPUs para cada VM n1-standard-32 a fim de atender aos requisitos. Duas GPUs fornecem recursos insuficientes, e não é possível especificar três GPUs.

    A tabela de compatibilidade a seguir leva em conta esse requisito.

    Observe as outras limitações a seguir sobre o uso de GPUs para treinamento personalizado que difere do uso de GPUs com o Compute Engine:

    • Uma configuração com quatro GPUs NVIDIA_TESLA_P100 oferece até 64 CPUs virtuais e até 208 GB de memória em todas as regiões e zonas.
  • Para jobs que usam o Dynamic Workload Scheduler ou as VMs do Spot, atualize o campo scheduling.strategy do CustomJob para a estratégia escolhida.

A tabela de compatibilidade a seguir lista os valores válidos para machineSpec.acceleratorCount dependendo das suas escolhas para machineSpec.machineType e machineSpec.acceleratorType:

Números válidos de GPUs para cada tipo de máquina
Tipo de máquina NVIDIA_H100_80GB NVIDIA_A100_80GB NVIDIA_TESLA_A100 NVIDIA_TESLA_P4 NVIDIA_TESLA_P100 NVIDIA_TESLA_T4 NVIDIA_TESLA_V100 NVIDIA_L4
a3-highgpu-1g 1*
a3-highgpu-2g 2*
a3-highgpu-4g 4*
a3-highgpu-8g 8
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

* O tipo de máquina especificado só está disponível ao usar o Dynamic Workload Scheduler ou VMs Spot.

Os exemplos a seguir destacam onde você pode especificar GPUs ao criar um CustomJob:

Console

No console do Google Cloud, não é possível criar um CustomJob diretamente. No entanto, é possível criar um TrainingPipeline que crie um CustomJob. Ao criar um TrainingPipeline no Console do Google Cloud, é possível especificar GPUs para cada pool de workers na etapa Compute and pricing Primeiro, especifique um Tipo de máquina. Em seguida, é possível especificar detalhes da GPU nos campos Tipo de acelerador e Contagem de aceleradores.

gcloud

Para especificar GPUs usando a ferramenta CLI do Google Cloud, use um arquivo config.yaml. Por exemplo:

config.yaml

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

Em seguida, execute um comando como este:

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

Node.js

Antes de testar esse exemplo, siga as instruções de configuração para Node.js no Guia de início rápido da Vertex AI sobre como usar bibliotecas de cliente. Para mais informações, consulte a documentação de referência da API Vertex AI para Node.js.

Para autenticar na Vertex AI, configure o Application Default Credentials. Para mais informações, consulte Configurar a autenticação para um ambiente de desenvolvimento local.

/**
 * 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_K80',
            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

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python.

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)

Para mais contexto, leia o guia para criar um CustomJob.

TPUs

Para usar as Unidades de Processamento de Tensor (TPUs) no treinamento personalizado sobre a Vertex AI, configure um pool de workers para usar uma VM de TPU.

Ao usar uma VM da TPU na Vertex AI, use apenas um pool de workers para treinamento personalizado e configure esse pool de workers para usar apenas uma réplica.

TPU v2 e v3

Para usar VMs de TPU v2 ou v3 no seu pool de workers, utilize uma das seguintes configurações:

  • Para configurar uma VM de TPU com TPU v2, especifique os seguintes campos em WorkerPoolSpec:

    • Defina machineSpec.machineType como cloud-tpu.
    • Defina machineSpec.acceleratorType como TPU_V2.
    • Defina machineSpec.acceleratorCount como 8 para uma única TPU ou 32 or multiple of 32 para pods de TPU.
    • Defina replicaCount como 1.
  • Para configurar uma VM de TPU com TPU v3, especifique os seguintes campos em WorkerPoolSpec:

    • Defina machineSpec.machineType como cloud-tpu.
    • Defina machineSpec.acceleratorType como TPU_V3.
    • Defina machineSpec.acceleratorCount como 8 para uma única TPU ou 32+ para pods de TPU.
    • Defina replicaCount como 1.

TPU v5e

A TPU v5e requer o JAX 0.4.6+, o TensorFlow 2.15+ ou o PyTorch 2.1+. Para configurar uma VM de TPU com a TPU v5e, especifique os seguintes campos em WorkerPoolSpec:

  • Defina machineSpec.machineType como ct5lp-hightpu-1t, ct5lp-hightpu-4t ou ct5lp-hightpu-8t.
  • Defina machineSpec.tpuTopology como uma topologia compatível com o tipo de máquina. Veja mais detalhes na tabela a seguir.
  • Defina replicaCount como 1.

A tabela a seguir mostra os tipos de máquina e as topologias da TPU v5e compatíveis com o treinamento personalizado:

Tipo de máquina topologia Número de chips do TPU Número de VMs Caso de uso recomendado
ct5lp-hightpu-1t 1x1 1 1 Treinamento em pequena e média escala
ct5lp-hightpu-4t 2x2 4 1 Treinamento em pequena e média escala
ct5lp-hightpu-8t 2x4 8 1 Treinamento em pequena e média escala
ct5lp-hightpu-4t 2x4 8 2 Treinamento em pequena e média escala
ct5lp-hightpu-4t 4x4 16 4 Treinamento em grande escala
ct5lp-hightpu-4t 4x8 32 8 Treinamento em grande escala
ct5lp-hightpu-4t 8x8 64 16 Treinamento em grande escala
ct5lp-hightpu-4t 8x16 128 32 Treinamento em grande escala
ct5lp-hightpu-4t 16x16 256 64 Treinamento em grande escala

Os jobs de treinamento personalizados em execução nas VMs da TPU v5e são otimizados para capacidade e disponibilidade. Para mais informações, consulte Tipos de aceleradores de treinamento v5e.

As máquinas TPU v5e estão disponíveis em us-west1 e us-west4 para o treinamento personalizado da Vertex AI. Para mais informações sobre a TPU v5e, consulte Treinamento do Cloud TPU v5e.

Comparação entre tipos de máquina:

Tipo de máquina ct5lp-hightpu-1t ct5lp-hightpu-4t ct5lp-hightpu-8t
Número de chips v5e 1 4 8
Número de vCPUs 24 112 224
RAM (GB) 48 192 384
Número de nós NUMA 1 1 2
Probabilidade de preempção Alta Média Baixa

Exemplo de CustomJob especificando uma VM da TPU

O exemplo a seguir destaca como especificar uma VM de TPU ao criar um CustomJob:

gcloud

Para especificar uma VM de TPU usando a ferramenta de CLI gcloud, use um arquivo config.yaml. Selecione uma das seguintes guias para ver um exemplo:

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

Em seguida, execute um comando como este:

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

Python

Antes de testar esse exemplo, siga as instruções de configuração para Python no Guia de início rápido da Vertex AI sobre como usar bibliotecas de cliente. Para mais informações, consulte a documentação de referência da API Vertex AI para Python.

Para autenticar na Vertex AI, configure o Application Default Credentials. Para mais informações, consulte Configurar a autenticação para um ambiente de desenvolvimento local.

Para especificar uma VM TPU usando o SDK da Vertex AI para Python, consulte este exemplo:

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')

Para mais informações sobre como criar um job de treinamento personalizado, consulte Criar jobs de treinamento personalizados.

Opções de disco de inicialização

Opcionalmente, é possível personalizar os discos de inicialização das VMs de treinamento. Todas as VMs em um pool de workers usam o mesmo tipo e tamanho do disco de inicialização.

  • Para personalizar o tipo de disco de inicialização que cada VM de treinamento usa, especifique o campo diskSpec.bootDiskType no WorkerPoolSpec.

    Defina esse campo como pd-standard para usar um disco permanente padrão com suporte a um disco rígido padrão ou configure-o como pd-ssd para usar um disco permanente SSD com suporte a uma unidade de estado sólido. O valor padrão é pd-ssd.

    O uso de pd-ssd pode melhorar o desempenho se o código de treinamento for lido e gravado no disco. Saiba mais sobre tipos de disco.

  • Para personalizar o tamanho (em GB) do disco de inicialização que cada VM de treinamento usa, especifique o campo diskSpec.bootDiskSizeGb na WorkerPoolSpec.

    Você pode definir esse campo como um número inteiro entre 100 e 64.000. O valor padrão é 100.

    É recomendável aumentar o tamanho do disco de inicialização se o código de treinamento gravar muitos dados temporários no disco. Observe que os dados gravados no disco de inicialização são temporários, e não é possível recuperá-los após a conclusão do treinamento.

Alterar o tipo e o tamanho dos discos de inicialização afeta os preços de treinamento personalizados.

Os exemplos a seguir destacam onde você pode especificar opções de disco de inicialização ao criar um CustomJob:

Console

No console do Google Cloud, não é possível criar um CustomJob diretamente. No entanto, é possível criar um TrainingPipeline que crie um CustomJob. Ao criar um TrainingPipeline no Console do Google Cloud, é possível especificar opções de disco de inicialização para cada pool de workers na etapa Compute and pricing, na lista suspensa Tipo de disco e no campo Tamanho do disco (GB).

gcloud

Para especificar opções de disco de inicialização usando a ferramenta CLI do Google Cloud, use um arquivo config.yaml. Por exemplo:

config.yaml

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

Em seguida, execute um comando como este:

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

Para mais contexto, leia o guia para criar um CustomJob.

A seguir